mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
Pearu Peterson
717d7a288d
This adds the `torch._native` override for the two chunked ops together with its eligibility condition, while the implementation still delegates to eager. Keeping the routing in its own commit leaves the kernel commit about the kernel. > `linear_cross_entropy` puts a whole fused loss behind one dispatcher symbol: an > `(N, F) x (F, C)` logits matmul, a row-shifted softmax, the cross-entropy > reduction, and three gradient matmuls, chunked over the batch. Per chunk the > eager loop walks the `(B, C)` logits buffer about seven times -- matmul write, > row max, subtract, gather, `exp_`, row sum, scale -- before the gradient > matmuls read it again. Collapsing those passes is what a DSL kernel is for. > > This installs overrides for both chunked ops -- the scalar-reduction > `_linear_cross_entropy_batch_chunked` and its `..._no_reduction` sibling -- and > nothing else. Both impls delegate to the op's own accumulator, so routing > through them is behaviourally indistinguishable from not routing through them, > and both conds gate nothing: they are the hooks a kernel will narrow to the > shapes it implements. No performance change is expected or claimed. Keeping the > wiring in its own commit means a later numeric or performance regression is > attributable to the kernel rather than to the dispatch. > > Both ops are registered from one `_OVERRIDES` table, which the registration > loop and the tests both read, so adding a third override is a row rather than an > edit in three places. > > Two things are worth knowing to review it. > > Registration has an ordering constraint. Overrides are installed at the end of > `import torch`, but the module defining the `torch_nn::` ops cannot be imported > from `torch/nn/modules/__init__.py` -- `torch.library` does not exist that > early -- so it is imported lazily on first use and the ops are absent when the > registry runs. `register_linear_cross_entropy_overrides()` therefore imports it > before registering, after asking the same question `cu.register_op_override` > asks -- runtime present, not disabled, and at a known-good version -- so the > import is only paid where an override can actually be installed. > > Installing at a backend dispatch key means the op's body never executes, so the > checks at the top of it would be skipped. `_batch_chunked_impl` calls > `_check_batch_chunked_grad_flags` for that reason -- the function the parent > commit extracted for exactly this caller. The no_reduction op needs no such > call: it takes no `compute_*_grad` flags, and its probability-target guard sits > in its backward, registered above this dispatch key and so not displaced. The > checks *inside* the accumulator are inherited only because the impls delegate to > it; a kernel will have to re-enforce those too, which a comment at each call > site records. > > `torch_nn` joins `_ALLOWED_LIB_SYMBOLS` here rather than in the registry > commit, since this is the change that registers on it and tests it. > > Unlike `aten`, whose ops the build guarantees, these ops exist only after > `torch/nn/modules/linear_cross_entropy.py` has executed, so the binding between > a registration's string and its referent is a runtime property whose failures > appear only where the DSL is installed -- never on CPU or standard GPU CI. The > module therefore declares `_NAMESPACE`, `_DEFINING_MODULE` and `_OVERRIDES`, > the registrar loops over that same table, and > test/python_native/test_override_declarations.py imports each defining module > and resolves every declared symbol. It needs neither a GPU nor the DSL, so it > runs on every shard: a one-token typo in `_OVERRIDES` fails it. Loud failure at > import remains the contract for anything that slips past. > > Test Plan: > > Two tests, because there are two things to test. The impls are > behaviour-preserving, so comparing their output to the ops' compares the ops to > themselves; what is not self-evident is that the overrides are installed at all, > given the lazily-imported module they target. One test asserts `import torch` > left every entry in `_OVERRIDES` registered with a live cutedsl node and a live > router. The other exports an `nn.LinearCrossEntropyLoss` call per reduction and > applies the registry's decomposition table: the override's own > `_native::<node_id>` op appears in the graph exactly when the call routes to it. > `mean` reaches the scalar-reduction override and `none` the no_reduction one, so > those two cover both. Nodes are matched on `target.namespace` and compared > against the node ids the registry minted, since one op symbol is a prefix of the > other. Exactly one such node is expected: there is one call, and the `_native` > node is opaque, so overrides of the aten ops inside it cannot add nodes. > > An eager call leaves no trace, since the impls are behaviour-preserving and the > router logs nothing per call, so the graph is the only place the choice is > observable. > > Both tests skip on those same three conditions. An installed CuTeDSL whose > version is not known-good registers nothing, and a test that looked for the > overrides anyway would fail where it should skip. > > ``` > python -m pytest test/python_native/test_linear_cross_entropy_override.py -q > python -m pytest test/python_native/ -q > python -m pytest test/test_nn.py -q -k linear_cross_entropy > CUDA_VISIBLE_DEVICES=0 python -m pytest test/test_cuda.py::TestCuda::test_lazy_init -q > python scripts/lintrunner.py > ``` > > `test_lazy_init` covers the added import: registration must not initialize CUDA, > since that poisons `fork` for anything importing torch. --- _🤖 Drafted by Claude Code (an AI agent) and reviewed & approved by pearu._ Pull Request resolved: https://github.com/pytorch/pytorch/pull/195332 Approved by: https://github.com/slayton58, https://github.com/Neilblaze ghstack dependencies: #194509, #194510, #195331
54 lines
2.2 KiB
Python
54 lines
2.2 KiB
Python
# Owner(s): ["module: dsl-native-ops"]
|
|
#
|
|
# Drift guard for overrides registered on namespaces other than `aten`.
|
|
#
|
|
# An `aten` op exists by construction: libtorch defines it before any Python
|
|
# runs, so a registration naming a symbol that does not exist dies on every
|
|
# `import torch`. An op defined by executing Python -- `torch_nn::` ops, for
|
|
# one -- exists only after its defining module has run, which turns the
|
|
# binding between a registration's string and its referent into a runtime
|
|
# property. Worse, the failure surfaces only where the DSL is installed, so
|
|
# ordinary CPU and GPU CI would never see it.
|
|
#
|
|
# This restores the build-grade invariant: a declaring module names its
|
|
# namespace, the module that defines its ops, and the overrides it registers,
|
|
# and every symbol must resolve. `import torch` has already imported these
|
|
# modules and their DSL imports are lazy, so this test needs neither a GPU
|
|
# nor any DSL runtime and runs on every shard.
|
|
|
|
import importlib
|
|
import sys
|
|
|
|
import torch # noqa: F401 # imports the registrars that carry the declarations
|
|
from torch._native.registry import _resolve_overload
|
|
from torch.testing._internal.common_utils import run_tests, TestCase
|
|
|
|
|
|
def _declaring_modules():
|
|
for name, mod in sorted(sys.modules.items()):
|
|
if name.startswith("torch._native.ops.") and getattr(mod, "_NAMESPACE", None):
|
|
yield name, mod
|
|
|
|
|
|
class TestOverrideDeclarations(TestCase):
|
|
def test_declared_overrides_resolve(self):
|
|
declaring = list(_declaring_modules())
|
|
self.assertTrue(
|
|
declaring,
|
|
"no module under torch._native.ops declares `_NAMESPACE`; the first "
|
|
"non-aten registrar must declare one so drift is caught here",
|
|
)
|
|
for name, mod in declaring:
|
|
importlib.import_module(mod._DEFINING_MODULE)
|
|
for op_symbol, *_ in mod._OVERRIDES:
|
|
op = _resolve_overload(mod._NAMESPACE, op_symbol)
|
|
self.assertIsNotNone(
|
|
op,
|
|
f"{name} registers {mod._NAMESPACE}::{op_symbol}, which does "
|
|
f"not resolve after importing {mod._DEFINING_MODULE}",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|