[PyTorch] Declare fusible ops as custom ops + HOP proof of concept - #28
Closed
pggPL wants to merge 11 commits into
Closed
[PyTorch] Declare fusible ops as custom ops + HOP proof of concept#28pggPL wants to merge 11 commits into
pggPL wants to merge 11 commits into
Conversation
… glue Registers an op's forward and backward as two independent two-tier custom ops and returns callables for both, leaving autograd to the caller. This lets a pipeline-level autograd.Function (which Dynamo traces as a higher-order op) group the forward and backward passes differently, as ops.OperationFuser does. Reuses the existing schema/adapter/TensorSpec machinery; register_custom_op is untouched. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Splits Bias into config resolution, pure compute and ctx saving, then registers the compute halves via register_op_halves. resolve_fwd_args reads module config and global FP8 state, so it stays in the traced region where Dynamo guards those reads; the impls take everything as arguments and never touch self. op_forward/op_backward keep their signatures and drive the same impls, so the eager path is unchanged. Adds tests/pytorch/test_ops_custom_ops.py with the fake-vs-real conformance harness every subsequent op will reuse. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
A custom op may not return one of its own inputs, which Bias did whenever the grad input is grad_output unchanged. Cloning would cost a full-size copy on the common unquantized path, so the impl returns None for that slot and the caller substitutes grad_output. Pass-through grads are common (Identity, Reshape, ConstantScale, Quantize), so this is the convention those ops will follow too. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Each activation subclass registers its own pair of compute halves, keyed by the class: the kernel pair is fixed by the class and a callable has no place in an op schema, so per-class registration keeps the op registry bounded by the op zoo rather than by the model. The subclass dispatch methods become staticmethods, which is what 'the compute half must not depend on an instance' means in practice; existing self._activation_forward_impl(...) call sites are unaffected. Activations are the first ops here that hand back a quantized tensor: with a next-operation input quantizer the kernel writes FP8 directly, so an FP8 tensor crosses the custom-op boundary instead of being dequantized at it. The forward's saved tensor is the input itself whenever dequantize and contiguous are both no-ops, so it follows the same None convention as the Bias backward. The conformance harness now checks the invariant the compiled path actually relies on -- the flat Tensor[] slot layout, via the framework's own helpers -- rather than approximating it. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Checks the assumption the whole approach rests on, before the fuser is touched: Dynamo traces a pipeline-level autograd.Function as the autograd_function_apply higher-order op, so the forward and the backward can walk different op groupings -- which is what OperationFuser does, and what an op with per-op autograd could not express. The pipeline runs Bias -> GELU -> Bias through the real registered custom ops, creates context objects inside the forward and reads them in the backward, and carries an FP8 tensor from one op to the next inside the traced region. Forward and backward match eager under fullgraph=True. Also makes the activation's saved tensor a static choice. It used to be returned as None when dequantize and contiguous were both no-ops, but the fake cannot see strides and so cannot predict that; the resulting metadata mismatch surfaced as an inductor assertion on the sentinel's rank. The op now keeps the input only when cache_quantized_input is set, and the backward rebuilds its input from the operation's input otherwise. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The checks were written per operation, so each one restated the same three questions in its own shape. They are now driven by a single _OP_CASES list: adding an operation means adding one entry, and every operation gets the same fake-conformance, matches-eager and fullgraph checks, each run with and without an FP8 output. resolve_fwd_args takes the same arguments on every operation to make that possible; Bias accepts next_op_input_quantizer and ignores it, as its op_forward already did. Compile cases reset Dynamo between runs: the compiled helpers are closures over one operation, so the parametrized runs otherwise walk into the recompilation limit and fall back to eager, which fullgraph=True rejects. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The impls and fakes lived outside the class they belong to -- module-level functions for Bias, closures from a factory for the activations -- and each operation then repeated the same op_forward/op_backward plumbing. They are now classmethods on the operation, and BasicOperation does the rest: __init_subclass__ registers the custom ops for any operation that declares fwd_args_type/bwd_args_type, and op_forward/op_backward are written once in the base. Declaring an operation means two dataclasses, four compute halves, resolve_fwd_args and resolve_bwd_args -- no registration or plumbing. Classmethods rather than static functions because the binding is load bearing: the whole activation family shares one implementation and dispatches to its per-class kernel through cls, so it needs no factory. Two hooks cover what genuinely differs: saved_for_backward, for an operation whose backward needs its input but whose forward produces no distinct tensor for it, and resolve_bwd_args, since each backward container names its own fields. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Linear puts compile_unsupported_reason on its forward args, because the module is the compile boundary and the check sits next to the choice between the custom op and eager. In ops/ the boundary is the fuser group, not the operation, so the check belongs elsewhere: - a pipeline compiles as a whole, so the decision is one per group, aggregated over its operations; - the args only exist after resolve_fwd_args, while most reasons are known earlier and more cheaply from the operation itself; - recipe-level limits belong to whoever reads the recipe, which is the fuser. So the hook goes on BasicOperation. The default refuses an operation without the compute halves, and any quantizer torch.compile cannot specialize on -- delayed scaling holds live scale/amax tensors, so baking its quantizer into the graph would silently freeze stale scales. __init_subclass__ now also checks that the argument containers are dataclasses, which is what the framework actually requires of them: the op schema is built from their fields. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
register_op_halves had grown as a copy of register_custom_op's body minus the autograd wiring -- about two thirds of it was the same code. Both now build on _register_two_tier_pair, which owns everything they genuinely share: schemas from the argument containers, the base kernels, the wrapper ops that flatten QuantizedTensor subclass inputs, and the passthrough registrations. What is left in each entry point is only what they differ on. Renamed to register_custom_op_without_autograd. 'Halves' said nothing about why the function exists; the one thing a reader needs is that, unlike register_custom_op, it leaves autograd to the caller. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
register_custom_op_without_autograd named a thing by what it lacked, which had it backwards: after the two entry points were factored onto a shared registration, the forward/backward pair is the primitive and autograd is what the other one adds on top. So the pair takes the plain name and the wired variant becomes register_custom_op_with_autograd. The file follows the same order, primitive first, and the pair's docstring no longer defines itself by negation. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Wires the fuser's compiled path end to end, so a group whose operations declare their compute halves runs through their custom ops under torch.compile(fullgraph=True). The pipeline-level autograd.Function is traced as a higher-order op, which is what lets its forward and backward walk different op groupings later on. Four things blocked tracing, all of them side effects reaching outside the higher-order op's scope: - OperationContext objects are created in the forward, but the backward is a separate subgraph, so writing to them there mutates an enclosing scope; the backward copies them into its own scope instead; - requires_grad_ on an output, which AOTAutograd's functionalization drops anyway -- autograd marks the outputs of an apply() itself; - _do_not_clear on inputs and outputs; - warnings.warn from the gate, which is not traceable, so the reason is now reported from the eager path only. These are gated on being traced rather than on using the custom ops. Under fullgraph there is no leaving the graph, so an unsupported operation does not 'fall back' -- the pipeline is traced either way and only the choice of implementation changes, which means the tracing constraints hold on both paths. Sequential builds its module groups outside the forward pass, since that constructs nn.Modules. Tested with a test-only operation, so the fuser's path does not depend on which real operations happen to be converted. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Owner
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
First step towards
torch.compile(fullgraph=True)support fortransformer_engine.pytorch.ops.Sequential.The plan is to keep the pipeline-level
_OperationFuserAutogradFunctionand let Dynamo trace it as ahigher-order op, with each fusible operation calling its own custom op inside. That keeps the forward
and backward fusion layouts independent (so the backward-only fusions survive), keeps
OperationContextinside the traced scope, and bounds op registration to one entry per op class.This PR does the per-operation half of that: make operations declarable as custom ops, with
tests, plus a proof of concept for the pipeline-level HOP. The fuser is untouched, so nothing in the
eager path changes yet.
Type of change
Changes
register_op_halves(dynamo/custom_op.py): registers an operation's forward and backward as twoindependent two-tier custom ops and returns callables for both, without registering autograd. The
caller wires the halves into its own
autograd.Function, so the forward and backward passes can begrouped differently -- which is what
OperationFuserdoes. Reuses the existingschema/adapter/
TensorSpecmachinery;register_custom_opis untouched.BasicOperationgained the plumbing, so an operation only declares its twoargument containers and implements four compute classmethods (
forward_compute,forward_fake,backward_compute,backward_fake) plusresolve_fwd_args/resolve_bwd_args.__init_subclass__registers the custom ops;op_forward/op_backwardare written once in the base.Biasand the activation family converted. The activations share one implementation and dispatchto their per-class kernel through
cls, so all ten subclasses get their own registered op without afactory. They are also the first operations here that hand back a quantized tensor: with a
next-operation input quantizer the kernel writes FP8 directly.
compile_unsupported_reasononBasicOperation: refuses an operation without the compute halves,and any quantizer torch.compile cannot specialize on. It sits on the operation rather than on the args
(where
Linearkeeps it) because inops/the compile boundary is the fuser group, not the operation.test_ops_custom_ops.pyis driven by a single_OP_CASESlist -- adding an operation meansadding one entry -- and checks fake-vs-real conformance, numerics against eager, and
fullgraph=True,each with and without an FP8 output.
test_ops_hop_poc.pyis the proof of concept.Proof of concept for the higher-order op
test_ops_hop_poc.pychecks the assumption the whole approach rests on, before the fuser is touched.A pipeline-level
autograd.FunctionrunningBias -> GELU -> Biasthrough the registered custom ops:autograd_function_apply, so both its forward and its backward end up in the graph;autograd could not express;
op schema;
Forward and backward match eager under
fullgraph=True.Notes on two contracts that the implementation had to respect
A custom op may not return one of its own inputs. A backward that passes its gradient through therefore
returns
Nonefor that slot and the caller substitutes; cloning would cost a full-size copy on thecommon path. The same trick does not work for saved tensors: the fake cannot see strides, so it cannot
predict whether
contiguous()will be a no-op, and the resulting metadata mismatch surfaces as aninductor assertion. Those rules have to be static, which is why the activations keep their input only
when
cache_quantized_inputis set.An FP8 output crosses the boundary as its inner buffers (
_data,_scale_inv) and is rebuilt on the farside from the fake's
TensorSpec; the subclass itself never crosses the schema.Testing
test_fusible_ops.py,test_ops_custom_ops.pyandtest_ops_hop_poc.py: 1574 passed, 1163 skipped.Lint: 10.00/10. Run on an RTX Ada workstation, so the FP8 coverage is current scaling only.
Not covered yet, and gated:
BasicLinearand the fused operation classes, grouped operations,userbuffers, delayed scaling, FP8 block scaling.
Checklist: