From c6544d0aae630cb0bdf52bfe5915c3193434d8ad Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 17:36:59 +0200 Subject: [PATCH 01/50] [PyTorch] [torch.compile] torch.compile support for Linear Register the Linear forward/backward as torch.library custom ops on top of the TensorSpec mechanism (#3153), so Linear traces under fullgraph compile with FP8/MXFP8/NVFP4 recipes. - transformer_engine/pytorch/dynamo/custom_op.py: custom-op registration framework (arg bundles, fake impls, autograd wiring) - module/linear.py: split forward into compute + ctx save, fake forward/backward - tests/pytorch/test_torch_compile.py: coverage for the compiled path Signed-off-by: Pawel Gadzinski --- .../distributed/run_layer_with_overlap.py | 28 + tests/pytorch/distributed/run_numerics.py | 51 +- .../distributed/test_comm_gemm_overlap.py | 40 + tests/pytorch/test_torch_compile.py | 384 ++++- transformer_engine/pytorch/dynamo/__init__.py | 2 + .../pytorch/dynamo/custom_op.py | 1518 +++++++++++++++++ .../pytorch/dynamo/quantizer_opaque.py | 29 +- .../pytorch/dynamo/tensor_spec.py | 5 + transformer_engine/pytorch/module/linear.py | 504 +++++- .../pytorch/tensor/_quantization_helpers.py | 37 + .../pytorch/tensor/float8_tensor.py | 6 + .../tensor/storage/float8_tensor_storage.py | 11 +- transformer_engine/pytorch/utils.py | 14 + 13 files changed, 2568 insertions(+), 61 deletions(-) create mode 100644 transformer_engine/pytorch/dynamo/custom_op.py diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 46795415e5..e65824ce85 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -200,6 +200,19 @@ def _parse_args(argv=None, namespace=None): parser.add_argument( "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA Graphs." ) + parser.add_argument( + "--compile", + action="store_true", + default=False, + help="Wrap each layer in torch.compile (tests Userbuffers on the compiled path).", + ) + parser.add_argument( + "--compile-mode", + type=str, + default="default", + choices=["default", "reduce-overhead"], + help="torch.compile mode used when --compile is set.", + ) parser.add_argument( "--ub-cfg", type=str, default=None, help="Optional TP config yaml file input." ) @@ -485,6 +498,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0) dist_print("Copied parameters from test model to reference model...", debug=True) + if opts.compile and opts.use_cuda_graphs: + raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.") + # Fp8 recipe setup fp8_format = Format.HYBRID fp8_recipe = None @@ -535,6 +551,18 @@ def run_fwd_bwd(model, x): loss.backward() return out + if opts.compile: + for i, layer in enumerate(test_model.layers): + # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle + # op arg whose hash chokes on non-nested SymInt (see run_numerics). + test_model.layers[i] = torch.compile( + layer, fullgraph=True, mode=opts.compile_mode, dynamic=False + ) + dist_print( + f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...", + debug=True, + ) + torch_rng_state = torch.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state(torch.device(f"cuda:{LOCAL_RANK}")) if opts.use_cuda_graphs: diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index fe02f990b4..1590979ba3 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -310,22 +310,45 @@ def _copy_params(model_distributed, model_single): def _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed, **kwargs + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=False, + compile_mode="default", + **kwargs, ): _alloc_main_grad(model_single_node, model_distributed) # for fuse_wgrad_accumulation=True input_single_node.requires_grad_() input_distributed.requires_grad_() + forward_single_node = model_single_node + forward_distributed = model_distributed + if use_compile: + # Each parametrized case compiles the same module.forward code object with + # a different shape/recipe; with dynamic=False those guards accumulate and + # eventually trip Dynamo's recompile_limit. Reset so every case starts from + # a clean compile cache (mirrors the single-GPU torch.compile tests). + torch._dynamo.reset() + # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle + # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static + # shapes (recompile per shape) until the bundle handles symbolic shapes. + forward_single_node = torch.compile( + model_single_node, fullgraph=True, mode=compile_mode, dynamic=False + ) + forward_distributed = torch.compile( + model_distributed, fullgraph=True, mode=compile_mode, dynamic=False + ) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), ): - output_single_node = model_single_node(input_single_node, **kwargs) + output_single_node = forward_single_node(input_single_node, **kwargs) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), amax_reduction_group=NCCL_WORLD, ): - output_distributed = model_distributed(input_distributed, **kwargs) + output_distributed = forward_distributed(input_distributed, **kwargs) return output_single_node, output_distributed @@ -641,12 +664,20 @@ def test_quantized_all_gather(): # Linear # ############################################ @run_distributed_test() -def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): +def _test_linear( + parallel_mode=None, + sequence_parallel=False, + use_compile=False, + compile_mode="default", + **kwargs, +): """Test the linear layer with specified parallel mode and sequence parallelization. Args: parallel_mode (str): 'row' or 'column' parallelism. sequence_parallel (bool): Enable sequence parallelism if True. + use_compile (bool): Wrap the modules in ``torch.compile`` before running. + compile_mode (str): ``torch.compile`` mode ("default" or "reduce-overhead"). kwargs (dict): Additional arguments for the linear layer. """ # Set parameter data type @@ -696,7 +727,12 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): # Apply models output_single_node, output_distributed = _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=use_compile, + compile_mode=compile_mode, ) if "return_bias" in kwargs: @@ -740,6 +776,8 @@ def test_linear(): {"params_dtype": torch.float16 if QUANTIZATION != "nvfp4" else torch.bfloat16}, {"delay_wgrad_compute": True}, {"save_original_input": True}, + {"use_compile": True}, + {"use_compile": True, "compile_mode": "reduce-overhead"}, ] for kwargs in kwargs_list: @@ -747,6 +785,9 @@ def test_linear(): continue if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: continue + # debug instrumentation forces the eager fallback, so compile is a no-op there. + if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel, **kwargs) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6b1ad870e9..6521ed7dbc 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -111,6 +111,8 @@ def _run_layer_with_overlap( quantization, num_layers=1, use_cublasmp=False, + compile=False, + compile_mode="default", ): test_path = TEST_ROOT / "run_layer_with_overlap.py" test_cmd = LAUNCH_CMD + [ @@ -129,6 +131,10 @@ def _run_layer_with_overlap( if overlap_rs_dgrad: test_cmd.append("--overlap-rs-dgrad") + if compile: + test_cmd.append("--compile") + test_cmd.append(f"--compile-mode={compile_mode}") + if fp8: if quantization in ("fp8_delayed_scaling", "fp8_current_scaling") and not fp8_available: pytest.skip(reason_for_no_fp8) @@ -281,6 +287,40 @@ def test_layers_with_overlap_bf16( ) +@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize( + "linear_parallel_mode,overlap_rs_dgrad", + [ + ("row", False), + ("column", False), + ("column", True), + ], + ids=[ + "ROW-PARALLEL", + "COL-PARALLEL - BULK DGRAD/WGRAD", + "COL-PARALLEL - DGRAD+RS", + ], +) +def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode): + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16). + + Userbuffers is expected to stay on Linear's compiled custom-op path (the + collective lives inside the opaque op), so this checks that torch.compile + + Userbuffers stays numerically correct against the eager, non-overlap reference. + ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on + top of the Userbuffers collectives. + """ + _run_layer_with_overlap( + te.Linear.__name__, + linear_parallel_mode, + overlap_rs_dgrad, + False, + None, + compile=True, + compile_mode=compile_mode, + ) + + @pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f61e7b4111..68b1174474 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -7,6 +7,7 @@ import pytest import torch +from torch._dynamo.utils import counters from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: @@ -26,9 +27,9 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec @@ -37,9 +38,9 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8Quantizer, Float8BlockQuantizer, MXFP8Quantizer, - NVFP4Quantizer, ) from utils import recipe_id from transformer_engine.pytorch.attention.dot_product_attention.backends import ( @@ -93,6 +94,75 @@ def nvfp4_4over6(): _all_recipes.append(nvfp4_row_scaled()) +# torch.compile modes exercised by the te.Linear tests: the default backend and +# "reduce-overhead" (CUDA-graph trees), to ensure the custom-op path is +# CUDA-graph capturable. +_compile_modes = ["default", "reduce-overhead"] + + +def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: + """ + Force TE's lazily-created global scratch to be allocated before capture. + """ + out = fn(inp) + if backward: + out.sum().backward() + + +@contextlib.contextmanager +def _assert_no_cudagraph_skips(enabled: bool): + """Assert ``torch.compile(mode="reduce-overhead")`` actually captured CUDA + graphs for every graph instead of silently running it eagerly. + + Inductor bumps ``counters["inductor"]["cudagraph_skips"]`` whenever it + declines to capture a cudagraph (input mutation, CPU scalars, cudagraph-unsafe + ops, ...) and falls back to eager for that graph. ``fullgraph=True`` only rules + out *dynamo* graph breaks, not these *inductor*-level skips, so this guards that + the reduce-overhead path didn't degrade to eager. No-op when ``enabled`` is + False (e.g. the default backend, where cudagraphs don't apply). + """ + before = counters["inductor"]["cudagraph_skips"] + yield + if enabled: + skipped = counters["inductor"]["cudagraph_skips"] - before + assert skipped == 0, ( + f"reduce-overhead fell back to eager: {skipped} cudagraph skip(s); " + "see the 'skipping cudagraphs due to ...' log for the reason" + ) + + +# bf16 output tolerance: eager and compiled run the same kernels, so they should +# agree closely; the slack only absorbs reduction-order / cuda-graph differences. +_EAGER_ATOL, _EAGER_RTOL = 1e-2, 1.6e-2 + + +def _assert_close_eager_compiled(fn, compiled, model, base): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the + forward output and the input / weight gradients match. + + Guards the compiled custom-op path against silently diverging from eager + execution -- a wrong-but-same-shape result would slip past shape / grad + presence checks alone. + """ + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = fn(inp_eager) + out_eager.sum().backward() + ref_out = out_eager.detach().clone() + ref_wgrad = model.weight.grad.detach().clone() + ref_igrad = inp_eager.grad.detach().clone() + + inp_compiled = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + # Clone before a later cuda-graph replay overwrites the static output buffer. + out_compiled = compiled(inp_compiled).clone() + out_compiled.sum().backward() + + torch.testing.assert_close(out_compiled, ref_out, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(model.weight.grad, ref_wgrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + # --------------------------------------------------------------------------- # ToyQuantizer – opaque value-type quantizer for torch.compile # (requires torch opaque object support, not available in older PyTorch) @@ -715,17 +785,32 @@ def _hw_available(quantizer): # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ - pytest.param(_mxfp8, id="mxfp8"), - pytest.param(_blockwise, id="float8_blockwise"), - pytest.param(_current_scaling, id="float8_current_scaling"), - pytest.param(_nvfp4, id="nvfp4"), + pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), + pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), + pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param( + _nvfp4, + {"with_rht": False}, + id="nvfp4", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), ] -@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) -def test_quantizer_value_object(factory): +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory, other_kwargs): """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" - a = factory() + a, b = factory(), factory() + # Same config -> equal, same hash, interchangeable as a dict/set key. + assert a is not b + assert a == b + assert hash(a) == hash(b) + assert {a: "x"}[b] == "x" + # Different config -> not equal. + assert a != factory(**other_kwargs) # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. repr_str, globals_ = a.__fx_repr__() @@ -799,8 +884,8 @@ def _qdq_fake(x, q): not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", ) -@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) -def test_quantizer_value_object_fullgraph(factory): +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory, other_kwargs): """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. A custom op quantizes+dequantizes with the (opaque value) quantizer; the @@ -1085,9 +1170,13 @@ def test_tensor_spec_matches_primitives(factory, shape): # Metadata matches the quantizer's. assert spec.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16) - # inner_names + create_inner_tensors match inner_tensor_specs. + # inner_names follows the storage's canonical __tensor_flatten__ order (the + # order the real op flattens its outputs to), while create_inner_tensors + # matches the inner_tensor_specs geometry (a name->shape/dtype mapping). specs = q.inner_tensor_specs(shape) - names = tuple(specs) + direct = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(direct.__tensor_flatten__()[0]) + assert set(names) == set(specs) assert spec.inner_names() == names inner_tensors = spec.create_inner_tensors() assert len(inner_tensors) == len(names) @@ -1097,7 +1186,6 @@ def test_tensor_spec_matches_primitives(factory, shape): assert inner.dtype == exp_dtype # The assembled tensor matches one built directly from the primitives. - direct = _build_from_primitives(q, shape, torch.bfloat16) assert _signature(spec.create_tensor(), names) == _signature(direct, names) @@ -1164,3 +1252,271 @@ def test_to_tensor_spec_quantized(factory, shape): assert _signature(spec.create_tensor(), spec.inner_names()) == _signature( tensor, spec.inner_names() ) + + +# --------------------------------------------------------------------------- +# te.Linear +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize( + "fp8_recipe", + [None, *_all_recipes], + ids=lambda r: "bf16" if r is None else type(r).__name__, +) +def test_te_linear_compiles(fp8_recipe, compile_mode): + """ + torch.compile(fullgraph=True) of ``te.Linear`` under every built-in + recipe (plus the bf16-only baseline with no autocast), for both the default + backend and ``mode="reduce-overhead"`` (CUDA-graph trees). + """ + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + + dtype = torch.bfloat16 + device = "cuda" + + # FP8 GEMMs require leading dimensions divisible by 16. + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + if fp8_recipe is None: + return model(inp) + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + # ``reduce-overhead`` warms up on the first call(s) and replays a captured + # CUDA graph afterwards, so iterate a few times to actually exercise replay. + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_quantized_fp8_weight(compile_mode): + """torch.compile should handle Linear weights initialized as FP8 tensors, + for both the default backend and ``mode="reduce-overhead"``. + + Exercises the two-tier op + ``register_torch_dispatch`` flattening of a + ``Float8Tensor`` weight *input* in + :mod:`transformer_engine.pytorch.dynamo`. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + with te.quantized_model_init(enabled=True, recipe=fp8_recipe): + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + assert isinstance(model.weight, te.Float8Tensor) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_fp8_output(compile_mode): + """torch.compile of ``te.Linear(..., fp8_output=True)`` without gradient: + forward returns a :class:`Float8Tensor`. Covers the default backend and + ``mode="reduce-overhead"``. + + Exercises the output-rewrap path in + :mod:`transformer_engine.pytorch.dynamo`: when an output quantizer is + active, the op returns the flat inner data tensors and the framework + rewraps them into a ``Float8Tensor`` via ``__tensor_unflatten__``. A + differentiable FP8 output is unsupported under compile (``Linear.forward`` + falls back to eager), so this test covers the supported case: an FP8 output + that does not require grad (inference / ``torch.no_grad``). + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, fp8_output=True) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + with torch.no_grad(): + _cudagraph_warmup(fn, torch.randn(32, 64, dtype=dtype, device=device), backward=False) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + inp = torch.randn(32, 64, dtype=dtype, device=device) + with torch.no_grad(): + out_eager = fn(inp) + out = compiled(inp) + assert isinstance( + out, te.Float8Tensor + ), f"expected Float8Tensor output, got {type(out).__name__}" + assert out.shape == (32, 32) + assert ( + out._quantizer is not None + ), "FP8 output lost its quantizer on the torch.compile path" + # The rewrap rebuilt a fully-functional Float8Tensor: dequantizing it + # outside the compiled region exercises scale + data + dtype wiring. + deq = out.dequantize() + assert deq.shape == (32, 32) + assert deq.dtype == dtype + # Compiled FP8 output must match the eager FP8 output value-wise. + torch.testing.assert_close( + deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_is_first_microbatch(compile_mode): + """torch.compile of ``te.Linear`` across a multi-step microbatch schedule that + drives FP8 weight caching via ``is_first_microbatch``, for the default backend + and ``mode="reduce-overhead"`` (CUDA-graph trees). + + ``is_first_microbatch=True`` quantizes and caches the FP8 weight; subsequent + ``False`` steps must reuse the cached FP8 weight instead of re-quantizing. This + exercises that cache path under compile and checks it stays numerically aligned + with eager. ``is_first_microbatch`` is a Python bool, so each distinct value is + its own dynamo guard/graph. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + # First microbatch caches the FP8 weight, the rest reuse the cache. + schedule = [True, False, False] + is_first = schedule[0] # rebound each step; closed over by ``fn``. + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=is_first) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for is_first in schedule: + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_linear_dynamic_shapes(): + """torch.compile(dynamic=True) of ``te.Linear`` with varying batch sizes. + + Verifies that the compiled graph handles symbolic (dynamic) leading + dimensions without graph breaks or recompilations after the initial trace. + Key correctness property: a graph compiled for batch=16 must produce + numerically correct results for batch=32 without triggering a recompile. + + This exercises two fixes for dynamic shapes: + 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle + (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). + 2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output + + weight + SP config instead of relying on the stored ``inp_shape``. + 3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same + tensor sources when it is None (compiled mode). + + FP8 + dynamic=True is tracked separately (requires resolving + ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo). + """ + dtype = torch.bfloat16 + device = "cuda" + in_features, out_features = 64, 32 + model = te.Linear(in_features, out_features, params_dtype=dtype, device=device) + + def fn(inp): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + batch_sizes = [16, 32, 48] + + for i, batch in enumerate(batch_sizes): + inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) + # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes. + torch._dynamo.mark_dynamic(inp, 0) + out = compiled(inp) + assert out.shape == (batch, out_features), f"wrong output shape for batch={batch}" + out.sum().backward() + assert inp.grad is not None, f"no input gradient for batch={batch}" + assert inp.grad.shape == inp.shape, f"wrong grad shape for batch={batch}" + + # Verify numerics against eager on each distinct batch size. + inp_eager = inp.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = model(inp_eager) + out_eager.sum().backward() + torch.testing.assert_close( + out.detach(), + out_eager.detach(), + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"forward mismatch at batch={batch}", + ) + torch.testing.assert_close( + inp.grad, + inp_eager.grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"dgrad mismatch at batch={batch}", + ) + + if i == 0: + # After the first (tracing) call, record the recompile counter + # baseline -- subsequent batch sizes must not trigger recompiles. + recompile_count_baseline = counters["stats"].get("recompile_reasons", 0) + + recompile_count_after = counters["stats"].get("recompile_reasons", 0) + assert recompile_count_after == recompile_count_baseline, ( + "Unexpected recompilation(s) across different batch sizes: " + f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" + ) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 4d5c76e9ce..3598e54daa 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,10 +6,12 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec +from .custom_op import register_custom_op __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", "TensorSpec", "to_tensor_spec", + "register_custom_op", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py new file mode 100644 index 0000000000..2a4f64d9a8 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -0,0 +1,1518 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile custom-op framework for Transformer Engine. + +Turns a TE module's eager forward/backward into ``torch.library`` custom ops so +``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph +break into the eager ``autograd.Function``. ``register_custom_op`` is the entry +point (its docstring documents the per-callable contract); ``module/linear.py`` +is the first user. Internal framework API -- exported from +``transformer_engine.pytorch.dynamo``, not re-exported at the top level. + +A TE op's forward/backward is written as a plain impl over a single *args +dataclass* (``fwd_arg_type`` / ``backward_arg_type``, e.g. ``LinearFwdArgs``): its +fields are a mix of tensors, quantized tensors, quantizers, process groups, +scalars and other Python values. The forward impl returns a tuple (user outputs + +saved-for-backward tensors + ctx metadata); the backward impl returns one gradient +per differentiable input. + +A ``torch.library`` custom op is narrower: it takes a flat list of schema slots +-- tensors / ``Tensor[]`` plus, via torch's opaque-object support, value-opaque +and reference-opaque objects -- and returns a flat ``Tensor[]``. + +Bridging the two takes three parts (below): per-field *adapters* map the args +dataclass onto the op's input slots; *fake impls* on data-free specs give the +output geometry and reassemble the op's flat return; and a *two-tier op* lets a +quantized-tensor subclass be an op input. + +Field <-> slot mapping. This mapping turns each field of the args dataclass into +the op's flat input slots, in a way that suits the field's type. A field's type +annotation selects the one ``_Adapter`` that handles it; that adapter declares the +slot(s) the field needs, packs the field's value into them on the way into the op, +and unpacks it back on the way out. The kinds -- and how each represents its field +as op inputs: + + * ``_TensorAdapter`` -- a plain ``Tensor`` / ``Optional[Tensor]``: one tensor + slot. + * ``_TensorOrQuantizedAdapter`` -- a field that may be a plain tensor, a bare + quantized storage, or ``None``: three slots (the tensor, its flat inner + buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. + * ``_QuantizerAdapter`` -- a quantizer, baked into the graph as a value-opaque + constant. + * ``_ReferenceOpaqueAdapter`` -- a ProcessGroup, carried as a live opaque graph + input. + * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, + sizes, nested collections of them), gathered into one ``OpaqueValueBundle`` + slot. + * ``_UnsupportedAdapter`` -- fallback for a field no adapter can encode; allowed + only when its value is trivial (``None`` / all-``None``) at call time. + +What runs where. Each op registers a data-free fake (``register_fake``) so it +traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward +call through it: + + * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see + ``tensor_spec.py``) to get the outputs' geometry in pure Python; + * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat + ``Tensor[]`` payload; + * rebuilds the structured user outputs from that payload, sliced and reassembled + per the fake's output descriptors (``_spec_reassemble``; + ``_value_to_flat_tensors`` is the pack-side inverse). + +Autograd, registered on the op, drives backward: + + * ``setup_context`` (run when the forward is taped) re-runs ``fwd_fake_impl`` for + the saved-tensor descriptors and a ``ctx_attrs`` dict, reassembles the saved + tensors from the op's flat output, then calls the user ``setup_context`` to + fill the backward args from forward state + ``ctx_attrs`` (e.g. saved-tensor + aliases) and return the tensors to persist; + * on ``backward()`` the backward args container's optional ``setup_saved_tensors`` + hook restores those saved tensors, then the *backward op* runs the real + ``backward_impl`` and returns the flat grads (``bwd_fake_impl`` is its + data-free fake). + +Two-tier op (``base`` + ``wrapper``), so a ``QuantizedTensor`` subclass can be an +op *input*. The ``_base`` op carries the real schema + autograd; a custom op +can't take a tensor-subclass input directly, so the ```` wrapper intercepts +those via ``register_torch_dispatch`` and flattens each into the base op's slots +(``_flatten_subclass_into_slots``) before forwarding. An empty subclass list makes +the wrapper a pass-through (plain / bf16 calls go straight through). +""" + +from __future__ import annotations +import dataclasses +import types as _types +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + get_args, + get_origin, + get_type_hints, +) + +import torch + +from torch._prims_common import make_contiguous_strides_for + +from .quantizer_opaque import warn_compile_unsupported +from .tensor_spec import TensorSpec, to_tensor_spec +from ..quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, + _quantized_tensor_passthrough_ops, + prepare_for_saving, +) + +_TE_OP_NAMESPACE = "transformer_engine_compile" + + +# ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a +# 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for +# ``register_autograd`` to attach a ``grad_fn`` to the outputs. +# +# Once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable +# ``Tensor?[]`` return schema will let ``None`` pass through directly and this +# sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed. +_NONE_SENTINEL_DTYPE = torch.uint8 + + +def _encode_none(t: Optional[torch.Tensor]) -> torch.Tensor: + """Replace ``None`` with a 0-element uint8 sentinel tensor.""" + if t is None: + return torch.empty(0, dtype=_NONE_SENTINEL_DTYPE) + return t + + +def _decode_none(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Inverse of :func:`_encode_none`.""" + if t is None: + return None + if t.numel() == 0 and t.dtype == _NONE_SENTINEL_DTYPE: + return None + return t + + +# --------------------------------------------------------------------------- # +# OpaqueValueBundle: bundle of simple / value-opaque Python values +# --------------------------------------------------------------------------- # + + +class OpaqueValueBundle: + """Opaque value-type bundle of simple Python values. + + Wraps a ``{name: value}`` dict so many small non-Tensor args pass through a + single custom-op input; registered as a torch.compile *value* opaque type + (Dynamo specializes the graph on its contents). Allowed values: primitives + in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, classes, + any registered value-opaque type (e.g. TE quantizers), plus nested tuples / + lists / dicts thereof (so a bundle can carry a ``__tensor_flatten__`` + context verbatim -- including its ``cls`` entry). + """ + + PRIMITIVE_TYPES: Tuple[type, ...] = ( + type(None), + bool, + int, + float, + str, + torch.dtype, + torch.device, + torch.Size, + ) + + @classmethod + def is_simple_value(cls, value: Any) -> bool: + """Whether ``value`` may be stored inside an instance (recursive).""" + if isinstance(value, cls.PRIMITIVE_TYPES): + return True + if isinstance(value, Enum): + return True + if isinstance(value, type): + return True + if _is_opaque_value_type(type(value)): + return True + if isinstance(value, dict): + return all( + isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items() + ) + if isinstance(value, (list, tuple)): + return all(cls.is_simple_value(v) for v in value) + return False + + @classmethod + def _to_hashable(cls, value: Any) -> Any: + if isinstance(value, dict): + return tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items())) + if isinstance(value, (list, tuple, torch.Size)): + return tuple(cls._to_hashable(v) for v in value) + return value + + @classmethod + def _fmt_simple(cls, value: Any) -> str: + """Repr for a value, evaluable in a context with ``torch`` globals.""" + if isinstance(value, torch.dtype): + return f"__import__('torch').{str(value).split('.')[-1]}" + if isinstance(value, torch.device): + return f"__import__('torch').device({str(value)!r})" + if isinstance(value, torch.Size): + return f"__import__('torch').Size({list(value)!r})" + # Enum before primitives: IntEnum is also ``int`` but must render as + # ``EnumName.MEMBER`` (the Enum class is added to globals by ``_collect``). + if isinstance(value, Enum): + return f"{type(value).__name__}.{value.name}" + # Class objects (e.g. the flatten context's ``cls``) render by name; the + # class itself is added to globals by ``_collect``. + if isinstance(value, type): + return value.__name__ + if isinstance(value, dict): + body = ", ".join(f"{k!r}: {cls._fmt_simple(v)}" for k, v in value.items()) + return f"{{{body}}}" + if isinstance(value, list): + return "[" + ", ".join(cls._fmt_simple(v) for v in value) + "]" + if isinstance(value, tuple): + body = ", ".join(cls._fmt_simple(v) for v in value) + return f"({body},)" if len(value) == 1 else f"({body})" + if _is_opaque_value_type(type(value)): + return value.__fx_repr__()[0] + return repr(value) + + def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: + data = dict(data) if data else {} + for k, v in data.items(): + if not OpaqueValueBundle.is_simple_value(v): + raise TypeError( + f"OpaqueValueBundle field '{k}' has unsupported type " + f"{type(v).__name__}; only simple primitives, Enum, " + "torch.Size, registered value-opaque types and nested " + "tuples / lists / dicts thereof are allowed." + ) + self._data: Dict[str, Any] = data + self._frozen: Tuple[Tuple[str, Any], ...] = tuple( + (k, OpaqueValueBundle._to_hashable(v)) for k, v in sorted(data.items()) + ) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __getattr__(self, name: str) -> Any: + try: + return self._data[name] + except KeyError as e: + raise AttributeError(name) from e + + def get(self, key: str, default: Any = None) -> Any: + """Return ``self._data.get(key, default)``.""" + return self._data.get(key, default) + + def as_dict(self) -> Dict[str, Any]: + """Return a shallow copy of the stored mapping.""" + return dict(self._data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, OpaqueValueBundle): + return NotImplemented + return self._frozen == other._frozen + + def __hash__(self) -> int: + return hash(self._frozen) + + def __fx_repr__(self) -> Tuple[str, Dict[str, Any]]: + items = ", ".join( + f"{k!r}: {OpaqueValueBundle._fmt_simple(v)}" for k, v in self._data.items() + ) + globals_: Dict[str, Any] = {"OpaqueValueBundle": OpaqueValueBundle} + + def _collect(value: Any) -> None: + if isinstance(value, dict): + for v in value.values(): + _collect(v) + return + if isinstance(value, (list, tuple)): + for v in value: + _collect(v) + return + if isinstance(value, Enum): + globals_[type(value).__name__] = type(value) + return + if isinstance(value, type): + globals_[value.__name__] = value + return + if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES): + return + if _is_opaque_value_type(type(value)): + _, extra = value.__fx_repr__() + globals_.update(extra) + + for v in self._data.values(): + _collect(v) + return (f"OpaqueValueBundle({{{items}}})", globals_) + + +try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + get_opaque_type_name, + is_opaque_value_type as _is_opaque_value_type, + is_opaque_reference_type as _is_opaque_reference_type, + register_opaque_type, + ) + + register_opaque_type(OpaqueValueBundle, typ="value") + _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) +except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object + warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") + _is_opaque_value_type = None + _is_opaque_reference_type = None + _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None + + +def _pg_pickle_stub(*args: Any) -> None: # pragma: no cover + raise RuntimeError("ProcessGroup cannot be unpickled — cache-key use only") + + +def _ensure_distributed_opaque_types() -> None: + """Register ``torch.distributed.ProcessGroup`` as a *reference* opaque type. + + A process group is live distributed state: unlike a value-opaque quantizer + (which Dynamo bakes into the graph as a constant), it must be carried through + the custom op as a graph *input*. PyTorch supports this via + ``register_opaque_type(ProcessGroup, typ="reference")`` but only auto-runs it + when ``torch.distributed.tensor`` (DTensor) is imported; TE may not import + that, so trigger the same idempotent registration here. Best-effort: on + builds without the opaque-object / distributed APIs this is a no-op and the + process-group field simply falls back to eager under torch.compile. + + Also registers a ``copyreg`` reducer that lets ``FxGraphCachePickler`` hash + graphs containing a ``ProcessGroup`` input without crashing. Without this, + inductor logs "Failed to pickle cache key" warnings and bypasses the FX + graph disk cache for every distributed compiled call. The reducer encodes + the group as (world_size, rank, backend) — enough to distinguish configs — + and raises on reconstruct since deserialization is never needed for hashing. + """ + if _is_opaque_reference_type is None: + return + try: # pylint: disable=import-outside-toplevel + from torch.distributed.device_mesh import _register_distributed_opaque_types + + _register_distributed_opaque_types() + except Exception: # pylint: disable=broad-exception-caught + pass + + # Workaround for PyTorch issue: FxGraphCachePickler handles FakeScriptObject + # but not the real ProcessGroup that appears in example_inputs at inductor + # compile time. Register a copyreg reducer so the pickler can hash the key. + try: # pylint: disable=import-outside-toplevel + import copyreg + import torch.distributed as dist + from torch._C._distributed_c10d import ProcessGroup + + if ProcessGroup not in copyreg.dispatch_table: + + def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] + try: + return _pg_pickle_stub, ( + dist.get_world_size(pg), + dist.get_rank(pg), + dist.get_backend(pg), + ) + except Exception: # pylint: disable=broad-exception-caught + return _pg_pickle_stub, (id(pg),) + + copyreg.pickle(ProcessGroup, _pg_reduce) + except Exception: # pylint: disable=broad-exception-caught + pass + + +_ensure_distributed_opaque_types() + + +# --------------------------------------------------------------------------- # +# Storage flatten / unflatten (value-opaque quantizer; no ProcessGroup) +# --------------------------------------------------------------------------- # + + +def _storage_flatten( + value: Any, extra_meta: Optional[Dict[str, Any]] = None +) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: + """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``. + + The flatten context (embedding the value-opaque quantizer) plus inner names + and -- for a wrapper subclass -- the outer geometry are stashed in the bundle + so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. + ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` + hash key stays consistent) -- used to tag the tensor-or-quantized slot ``__kind__``. + """ + inner_names, ctx = value.__tensor_flatten__() + meta = dict(ctx) + meta["_inner_names"] = list(inner_names) + if isinstance(value, torch.Tensor): + meta["_outer_shape"] = torch.Size(value.shape) + if extra_meta: + meta.update(extra_meta) + tensors = [getattr(value, name) for name in inner_names] + return OpaqueValueBundle(meta), tensors + + +def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any: + """Inverse of :func:`_storage_flatten`.""" + meta_dict = meta.as_dict() if isinstance(meta, OpaqueValueBundle) else dict(meta) + inner_names = meta_dict["_inner_names"] + inner = dict(zip(inner_names, tensors)) + outer_shape = meta_dict.get("_outer_shape") + stride = make_contiguous_strides_for(tuple(outer_shape)) if outer_shape is not None else None + return QuantizedTensorStorage.__tensor_unflatten__(inner, meta_dict, outer_shape, stride) + + +# --------------------------------------------------------------------------- # +# Field adapters: dataclass field <-> flat torch.library slot(s) +# --------------------------------------------------------------------------- # + + +def _is_union(annot: Any) -> bool: + """True for both ``typing.Union[...]`` / ``Optional[...]`` and PEP 604 ``X | Y``. + + ``get_origin`` returns ``typing.Union`` for the former but ``types.UnionType`` + for the latter, so the two syntaxes must be checked separately. + """ + origin = get_origin(annot) + return origin is Union or origin is _types.UnionType + + +def _strip_optional(annot: Any) -> Tuple[Any, bool]: + """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``.""" + if _is_union(annot): + args = get_args(annot) + if type(None) in args: + non_none = [a for a in args if a is not type(None)] + if len(non_none) == 1: + return non_none[0], True + return annot, False + + +class _Adapter: + """Maps one (or, for the aggregating adapter, several) dataclass field(s) + to/from a contiguous run of custom-op schema *slots*. + + A custom op only takes flat, simply-typed arguments, but a TE op takes a + single ``@dataclass`` of mixed fields. Each adapter knows how to translate + its kind of field both ways. ``try_build`` and ``schema_slots`` run once at + registration (to build the op's schema); ``to_slots`` and ``from_slots`` run + on each call and must agree on the slot layout that ``schema_slots`` declares. + """ + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_Adapter"]: + """Decide whether this adapter type handles the field ``name`` given its + type annotation ``annot``; return a configured adapter if so, else + ``None`` so the next candidate is tried. + + Called once per field at registration, in :data:`_FIELD_ADAPTERS` + priority order. + """ + raise NotImplementedError + + def schema_slots(self) -> List[Tuple[str, str]]: + """Declare the schema slots this field occupies, each as a + ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``). + + Concatenated across all adapters to form the op's schema string. + """ + raise NotImplementedError + + def to_slots(self, owner: Any) -> Dict[str, Any]: + """Read this field from the dataclass ``owner`` and produce the concrete + value for each of its schema slots, as a ``{slot_name: value}`` dict. + + Composite values are flattened to fit the (tensor-only) slots: e.g. a + quantized tensor is split into its plain inner buffers plus a metadata + bundle. Inverse of :meth:`from_slots`. + """ + raise NotImplementedError + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + """Read this field's slots back from the op arguments ``args`` and write + the reconstructed field value into ``kwargs`` (rebuilding any flattened + composite). The filled ``kwargs`` are then used to rebuild the original + dataclass for the eager implementation. Inverse of :meth:`to_slots`. + """ + raise NotImplementedError + + def grad_slot(self) -> Optional[int]: + """Index (within this adapter's :meth:`schema_slots`) of the slot that + carries a gradient, or ``None`` if the field is not differentiable. + + Used to map ``input_tensors_for_grad`` names onto backward grad-output + positions. Non-tensor adapters (quantizers, metadata) return ``None``. + """ + return None + + +class _TensorOrQuantizedKind(Enum): + """What a tensor-or-quantized slot group carries, tagged in its ``__meta``.""" + + NONE = "none" + TENSOR = "tensor" + STORAGE = "storage" + + +class _TensorOrQuantizedAdapter(_Adapter): + """``Tensor | QuantizedTensorStorage | None`` (also subclass tensor) field. + + Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass + tensor passes through, ``None`` for bare storage), ``__tensors`` + (``Tensor[]`` flat inner tensors when flattened), ``__meta`` + (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). A ``None`` + field is tagged ``_TensorOrQuantizedKind.NONE`` with the other two slots empty. + """ + + KIND_KEY = "__kind__" + + def __init__(self, name: str) -> None: + self.name = name + + def slot_name(self) -> str: + """Primary slot name for a plain / subclass tensor.""" + return self.name + + def slot_tensors(self) -> str: + """Flat inner-tensor slot name.""" + return self.name + "__tensors" + + def slot_meta(self) -> str: + """Flatten-metadata slot name.""" + return self.name + "__meta" + + def schema_slots(self) -> List[Tuple[str, str]]: + return [ + (self.slot_name(), "Tensor?"), + (self.slot_tensors(), "Tensor[]"), + (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + ] + + # Canonical "plain tensor or quantized tensor" field annotation (the + # ``TensorOrQuantized`` alias in module code). Matched by exact member set, + # so a bare quantized annotation or an accidental extra union member is + # rejected rather than silently taken as a tensor-or-quantized field. + _MEMBERS = frozenset({torch.Tensor, QuantizedTensorStorage}) + + @classmethod + def _is_tensor_storage_union(cls, annot: Any) -> bool: + if not _is_union(annot): + return False + members = frozenset(a for a in get_args(annot) if a is not type(None)) + return members == cls._MEMBERS + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedAdapter"]: + if cls._is_tensor_storage_union(annot): + return cls(name) + return None + + def to_slots(self, owner: Any) -> Dict[str, Any]: + value = getattr(owner, self.name) + if value is None: + return { + self.slot_name(): None, + self.slot_tensors(): [], + self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE}), + } + if isinstance(value, torch.Tensor): + # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the + # ``Tensor?`` slot; subclass flattening (if any) is done by the + # wrapper op's ``register_torch_dispatch`` rule. + return { + self.slot_name(): value, + self.slot_tensors(): [], + self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR}), + } + if isinstance(value, QuantizedTensorStorage): + meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) + return { + self.slot_name(): None, + self.slot_tensors(): tensors, + self.slot_meta(): meta, + } + raise TypeError( + f"field {self.name!r} expected None, torch.Tensor, or " + f"QuantizedTensorStorage, got {type(value).__name__}" + ) + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + meta = args[self.slot_meta()] + kind = meta.get(self.KIND_KEY) + if kind == _TensorOrQuantizedKind.NONE: + kwargs[self.name] = None + elif kind == _TensorOrQuantizedKind.TENSOR: + kwargs[self.name] = args[self.slot_name()] + else: + kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()]) + + def grad_slot(self) -> Optional[int]: + # Gradient flows to the plain / subclass tensor slot (``slot_name``, + # the first of the three). + return 0 + + +class _TensorAdapter(_Adapter): + """``Tensor`` / ``Optional[Tensor]`` -> single ``Tensor`` / ``Tensor?`` slot.""" + + def __init__(self, name: str, is_optional: bool) -> None: + self.name = name + self.type_str = "Tensor?" if is_optional else "Tensor" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_TensorAdapter"]: + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Tensor: + return cls(name, is_optional) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, self.type_str)] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.name: getattr(owner, self.name)} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.name] + + def grad_slot(self) -> Optional[int]: + return 0 + + +class _QuantizerAdapter(_Adapter): + """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. + + Each quantizer gets its own dedicated slot. The field is annotated with the + base ``Quantizer`` (not itself a registered opaque type), so the simple + bundle would not claim it. + """ + + KEY = "q" + + def __init__(self, name: str) -> None: + self.name = name + + def slot(self) -> str: + """Opaque quantizer metadata slot name.""" + return self.name + "__q" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerAdapter"]: + stripped, _ = _strip_optional(annot) + if isinstance(stripped, type) and issubclass(stripped, Quantizer): + return cls(name) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.slot(): OpaqueValueBundle({self.KEY: getattr(owner, self.name)})} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.slot()][self.KEY] + + +class _ReferenceOpaqueAdapter(_Adapter): + """``ProcessGroup`` (or any reference-opaque type) -> one own opaque slot. + + A reference-opaque object is live, stateful black-box data (e.g. a + ``torch.distributed.ProcessGroup``): it cannot be specialized on or baked + into the graph as a constant the way a value-opaque quantizer is. torch.compile + instead carries it through as a graph *input*, so it passes straight through + its own schema slot (no ``OpaqueValueBundle`` wrapper). The field is annotated + with a concrete type registered via ``register_opaque_type(..., typ="reference")``. + + On the fake / setup-context path the slot holds a ``FakeScriptObject`` (or + ``None``); it is assigned to the field verbatim, so the fake impl must never + read the object's contents. + """ + + def __init__(self, name: str, type_name: str, is_optional: bool) -> None: + self.name = name + self.type_str = f"{type_name}?" if is_optional else type_name + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueAdapter"]: + if _is_opaque_reference_type is None: + return None + stripped, is_optional = _strip_optional(annot) + if not isinstance(stripped, type): + return None + if _is_opaque_reference_type(stripped): + return cls(name, get_opaque_type_name(stripped), is_optional) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, self.type_str)] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.name: getattr(owner, self.name)} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.name] + + +class _SimpleBundleAdapter(_Adapter): + """Aggregates every simple-typed field into a single OpaqueValueBundle. + + Unlike the per-field adapters, at most one of these exists per op (none if the + dataclass has no simple-typed fields): it owns the single shared + ``_simple_meta`` slot, and ``_get_adapters`` builds it once from all + simple-typed field names collected across the dataclass. + """ + + SLOT = "_simple_meta" + + def __init__(self, names: List[str]) -> None: + self.names = list(names) + + @classmethod + def matches_field(cls, annot: Any) -> bool: + """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" + annot, _ = _strip_optional(annot) + if annot in OpaqueValueBundle.PRIMITIVE_TYPES: + return True + if isinstance(annot, type) and issubclass(annot, Enum): + return True + if ( + isinstance(annot, type) + and _is_opaque_value_type is not None + and _is_opaque_value_type(annot) + ): + return True + if get_origin(annot) in (tuple, list): + inner = [a for a in get_args(annot) if a is not Ellipsis] + return bool(inner) and all(cls.matches_field(a) for a in inner) + return False + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.SLOT: OpaqueValueBundle({n: getattr(owner, n) for n in self.names})} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + if self.SLOT not in args: + return + meta = args[self.SLOT] + for n in self.names: + kwargs[n] = meta[n] + + +class _UnsupportedAdapter(_Adapter): + """Fallback for fields whose type no other adapter can encode. + + Such a field cannot cross the op boundary, so it emits no slot and is + tolerated only when its runtime value carries nothing: ``to_slots`` accepts + ``None`` / an all-``None`` sequence (e.g. an unset ``Optional[Any]`` field, + or an empty list, on the compiled path) and ``from_slots`` restores it as + ``None``. A non-trivial value means the config is genuinely unsupported + under torch.compile, and ``to_slots`` raises. + + The check must run at call time (not in ``_get_adapters``): the annotation + alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so + only the runtime value can decide. + """ + + def __init__(self, name: str, owner_cls_name: str) -> None: + self.name = name + self.owner_cls_name = owner_cls_name + + @staticmethod + def _is_trivial(value: Any) -> bool: + if value is None: + return True + if isinstance(value, (list, tuple)): + return all(v is None for v in value) + return False + + def schema_slots(self) -> List[Tuple[str, str]]: + return [] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + value = getattr(owner, self.name, None) + if not self._is_trivial(value): + raise TypeError( + f"{self.owner_cls_name} field {self.name!r} has a type not " + "supported by torch.compile (not Tensor, simple, Quantizer, or a " + "reference-opaque type such as ProcessGroup) and carries a " + "non-trivial value; add a matching adapter in dynamo.py to handle it." + ) + return {} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = None + + +# Adapters, in priority order, owning ``try_build`` for a single field. +# These adapters are mutually exclusive on annotations (a plain ``torch.Tensor`` +# matches only ``_TensorAdapter``; the ``TensorOrQuantized`` union only +# ``_TensorOrQuantizedAdapter``; etc.), so the order is just iteration, not a +# priority ranking -- no annotation can be claimed by more than one. +_FIELD_ADAPTERS: Tuple[type, ...] = ( + _TensorOrQuantizedAdapter, + _TensorAdapter, + _ReferenceOpaqueAdapter, + _QuantizerAdapter, +) + + +def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: + """Return ``[(field_name, resolved_type), ...]`` for a dataclass.""" + if not dataclasses.is_dataclass(cls): + raise TypeError(f"{cls.__name__} must be a @dataclass to be a TE op arg container.") + try: + hints = get_type_hints(cls) + except Exception: # pylint: disable=broad-exception-caught + hints = {} + return [(f.name, hints.get(f.name, f.type)) for f in dataclasses.fields(cls)] + + +def _get_adapters(cls: type) -> List[_Adapter]: + """Build the adapter list for a dataclass from its field annotations.""" + if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None: + raise RuntimeError( + f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle " + "is not registered as a torch._library value-opaque type (PyTorch build " + "without opaque-object support)." + ) + adapters: List[_Adapter] = [] + simple_names: List[str] = [] + for name, annot in _resolved_field_annotations(cls): + built: Optional[_Adapter] = None + for adapter_cls in _FIELD_ADAPTERS: + built = adapter_cls.try_build(name, annot) + if built is not None: + break + if built is not None: + adapters.append(built) + elif _SimpleBundleAdapter.matches_field(annot): + simple_names.append(name) + else: + adapters.append(_UnsupportedAdapter(name, cls.__name__)) + if simple_names: + adapters.append(_SimpleBundleAdapter(simple_names)) + return adapters + + +def _tensor_field_names(adapters: List[_Adapter]) -> List[str]: + """Names of fields carrying tensors (for building the spec view).""" + return [b.name for b in adapters if isinstance(b, (_TensorAdapter, _TensorOrQuantizedAdapter))] + + +def _build_schema(adapters: List[_Adapter]) -> Tuple[str, List[str]]: + """Return ``(schema_arg_str, slot_names)`` for a adapter list.""" + spec = [slot for b in adapters for slot in b.schema_slots()] + names = [name for name, _ in spec] + schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")" + return schema_str, names + + +def _args_to_slots(obj: Any, adapters: List[_Adapter]) -> Dict[str, Any]: + """Build the op's flat ``{slot_name: value}`` argument dict from an args + dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every adapter's + packed slot(s). Inverse of :func:`_args_from_slots`. + """ + out: Dict[str, Any] = {} + for adapter in adapters: + out.update(adapter.to_slots(obj)) + return out + + +def _args_from_slots(cls: type, args: Dict[str, Any], adapters: List[_Adapter]) -> Any: + """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the + op's flat slot ``args`` dict, by letting every adapter restore its field(s). + Inverse of :func:`_args_to_slots`. + """ + kwargs: Dict[str, Any] = {} + for adapter in adapters: + adapter.from_slots(args, kwargs) + obj = cls.__new__(cls) + for k, v in kwargs.items(): + object.__setattr__(obj, k, v) + return obj + + +def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any: + """Copy of dataclass ``obj`` with each tensor field replaced by a :class:`TensorSpec`. + + Only tensor fields have a ``TensorSpec`` equivalent, so quantizer / scalar + fields are simply carried over unchanged; the fake impl works purely on + geometry. Built with :func:`dataclasses.replace` (the only such construction + Dynamo can trace). + """ + overrides: Dict[str, Any] = {} + for name in tensor_field_names: + value = getattr(obj, name, None) + if value is not None and not isinstance(value, TensorSpec): + overrides[name] = to_tensor_spec(value) + if not overrides: + return obj + return dataclasses.replace(obj, **overrides) + + +# --------------------------------------------------------------------------- # +# Op outputs <-> flat ``Tensor[]`` payload: this is how an op returns / saves +# quantized tensors (and wrapper subclasses). Outputs are flattened to their +# inner buffers on the way out and rebuilt via ``__tensor_unflatten__`` on the +# way back; on the fake side a TensorSpec supplies the geometry. +# --------------------------------------------------------------------------- # + + +def _spec_slot_count(spec: Optional[TensorSpec]) -> int: + """Flat ``Tensor[]`` slots the value for ``spec`` occupies.""" + if spec is None: + return 1 + return len(spec.inner_names()) + + +def _spec_reassemble( + spec: Optional[TensorSpec], + chunk: List[Optional[torch.Tensor]], +) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]: + """Rebuild the value described by ``spec`` from its flat tensors ``chunk``. + + ``spec is None`` -> ``None`` (op-boundary sentinel for an absent output); + otherwise delegates to :meth:`TensorSpec.assemble`, which returns a plain + tensor as-is or reassembles a quantized tensor from its inner buffers. + """ + if spec is None: + return None + return spec.assemble(chunk) + + +def _value_to_flat_tensors( + value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorSpec]], +) -> List[torch.Tensor]: + """Return the flat ``Tensor[]`` slots that represent one op output ``value``. + + Inverse of :func:`_spec_reassemble`; the slot count matches + :func:`_spec_slot_count`. + """ + if value is None: + return [_encode_none(None)] + if isinstance(value, TensorSpec): + return [_encode_none(t) for t in value.create_inner_tensors()] + if hasattr(value, "__tensor_flatten__"): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + if isinstance(value, torch.Tensor): + return [_encode_none(value)] + raise TypeError( + f"unsupported value type {type(value).__name__}; expected None / " + "torch.Tensor / tensor subclass / bare storage / TensorSpec." + ) + + +# Trailing slots in every fwd-impl return: ``tensors_to_save, ctx_attrs``. +# User-output count is ``len(result) - this``. +_FWD_TRAILING_SLOTS = 2 + + +def _check_fwd_result(result: Any) -> None: + """Validate a fwd-impl return against the + ``(*user_outputs, tensors_to_save, ctx_attrs)`` contract, with a clear + message for op authors (user-output *types* are checked later, by + :func:`_value_to_flat_tensors`). + + Only called on the fake path (:func:`_split_fwd_fake_result`), which runs at + trace/compile time -- so this is a compile-time check with no per-call cost. + The real impl must return the same shape as the fake, so validating the fake + covers both. + """ + if not isinstance(result, tuple) or len(result) < _FWD_TRAILING_SLOTS: + raise TypeError( + f"fwd impl must return a tuple of >= {_FWD_TRAILING_SLOTS} elements " + "(*user_outputs, tensors_to_save, ctx_attrs); " + f"got {type(result).__name__}" + ) + tensors_to_save, ctx_attrs = result[-2], result[-1] + if tensors_to_save is not None and not isinstance(tensors_to_save, (list, tuple)): + raise TypeError("fwd impl 'tensors_to_save' slot must be a list/tuple or None") + if ctx_attrs is not None and not isinstance(ctx_attrs, dict): + raise TypeError("fwd impl 'ctx_attrs' slot must be a dict or None") + + +def _format_fwd_result(result: Any) -> List[torch.Tensor]: + """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. + + User outputs first, then saved-for-backward tensors in declaration order. + """ + num_outputs = len(result) - _FWD_TRAILING_SLOTS + flat: List[torch.Tensor] = [] + for value in result[:num_outputs]: + flat.extend(_value_to_flat_tensors(value)) + saved = result[num_outputs] + if saved is not None: + for value in saved: + flat.extend(_value_to_flat_tensors(value)) + return flat + + +def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: + """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. + + Each grad occupies exactly one slot (validated against ``num_grad_inputs``); + a :class:`TensorSpec` grad is materialized into a single tensor. + """ + grads = list(grads) + if len(grads) != num_grad_inputs: + raise RuntimeError( + f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads " + f"(one per input_tensors_for_grad entry), got {len(grads)}" + ) + out: List[torch.Tensor] = [] + for g in grads: + if isinstance(g, TensorSpec): + out.append(_encode_none(g.create_tensor())) + else: + out.append(_encode_none(g)) + return out + + +def _split_fwd_fake_result( + result: Tuple[Any, ...], +) -> Tuple[List[Any], List[Any], Dict[str, Any]]: + """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" + _check_fwd_result(result) + num_outputs = len(result) - _FWD_TRAILING_SLOTS + saved = result[num_outputs] + ctx_attrs = result[num_outputs + 1] + user_fakes = list(result[:num_outputs]) + saved_fakes = list(saved) if saved is not None else [] + ctx_attrs = dict(ctx_attrs) if ctx_attrs else {} + return user_fakes, saved_fakes, ctx_attrs + + +# --------------------------------------------------------------------------- # +# Op registration +# --------------------------------------------------------------------------- # + + +def _resolve_grad_targets( + fwd_adapters: List[_Adapter], + input_tensors_for_grad: List[str], +) -> Tuple[int, List[int]]: + """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. + + ``fwd_adapters`` already encode the arg dataclass's fields (they are built + from it), so the type itself is not needed here. + + Returns ``(slot_count, grad_targets)``: the total number of input schema + slots and, for each requested input name, the schema-slot index its gradient + maps to. + """ + name_to_slot: Dict[str, int] = {} + slot_offset = 0 + for adapter in fwd_adapters: + slots = adapter.schema_slots() + grad_slot = adapter.grad_slot() + if grad_slot is not None: + name_to_slot[adapter.name] = slot_offset + grad_slot + slot_offset += len(slots) + + non_differentiable = [n for n in input_tensors_for_grad if n not in name_to_slot] + if non_differentiable: + raise ValueError( + f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" + ) + grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] + return slot_offset, grad_targets + + +def _register_kernel( + *, + op_name: str, + schema_str: str, + arg_type: type, + arg_names: List[str], + adapters: List[_Adapter], + tensor_field_names: List[str], + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + format_result: Callable[[Any], List[torch.Tensor]], +) -> Any: + """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the + ``fake_impl`` (spec), returning the ``CustomOpDef``. + + The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel + runs the spec fake impl on the :func:`_spec_view`. Both go through + ``format_result``. + """ + + def _impl(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _args_from_slots(arg_type, kwargs, adapters) + return format_result(impl(obj)) + + def _fake(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _args_from_slots(arg_type, kwargs, adapters) + spec_obj = _spec_view(obj, tensor_field_names) + return format_result(fake_impl(spec_obj)) + + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str + ) + op.register_fake(_fake) + return op + + +def _register_autograd_for_op( + *, + fwd_op: Any, + bwd_op: Any, + fwd_arg_type: type, + fwd_arg_names: List[str], + fwd_adapters: List[_Adapter], + fwd_tensor_field_names: List[str], + bwd_arg_names: List[str], + bwd_adapters: List[_Adapter], + slot_count: int, + grad_targets: List[int], + setup_context_user: Callable[..., Any], + backward_obj_type: type, + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> None: + """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. + + ``setup_context`` re-runs the spec fwd fake impl to recover output / saved + templates, reassembles each flat output chunk, and hands the saved tuple + + ``ctx_attrs`` to the module's ``setup_context``. + """ + + def _setup_context(ctx, inputs, output): + ctx._te_fwd_tensor_list_lengths = { + i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + } + kwargs = dict(zip(fwd_arg_names, inputs)) + fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_adapters) + spec_obj = _spec_view(fwd_obj, fwd_tensor_field_names) + + user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + + cursor = 0 + user_outputs: List[Any] = [] + for spec in user_fakes: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + user_outputs.append(_spec_reassemble(spec, chunk)) + + saved_list: List[Any] = [] + for spec in saved_fakes: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + saved_list.append(_spec_reassemble(spec, chunk)) + + bwd_obj = backward_obj_type() + tensors_to_save_from_setup = setup_context_user( + bwd_obj, + fwd_obj, + user_outputs[0] if len(user_outputs) == 1 else tuple(user_outputs), + ctx_attrs, + tuple(saved_list), + ) + tensors_to_save, tensor_objects = prepare_for_saving( + *(tensors_to_save_from_setup or ()) + ) + ctx.tensor_objects = tensor_objects + ctx.save_for_backward(*tensors_to_save) + ctx.bwd_obj = bwd_obj + + def _autograd_backward(ctx, *grad_outputs): + bwd_obj = ctx.bwd_obj + if hasattr(bwd_obj, "setup_saved_tensors"): + bwd_obj.setup_saved_tensors(ctx) + ctx.tensor_objects = None + per_output_grads = grad_outputs[0] + bwd_obj.grad_output = _decode_none(per_output_grads[0]) + kwargs = _args_to_slots(bwd_obj, bwd_adapters) + bwd_args_flat = [kwargs[name] for name in bwd_arg_names] + grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + # One grad per input schema slot: default None, but a ``Tensor[]`` slot + # (always recorded in ``_te_fwd_tensor_list_lengths``) needs a + # list-shaped no-grad of matching length. + out: List[Any] = [None] * slot_count + tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {}) + for pos, length in tensor_list_lengths.items(): + out[pos] = [None] * length + for pos, g in zip(grad_targets, grads): + out[pos] = g + return tuple(out) + + fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) + + +def _collect_tensor_or_quantized_slot_offsets(adapters: List[_Adapter]) -> List[int]: + """Start index of each ``_TensorOrQuantizedAdapter`` group in the flat args.""" + offsets: List[int] = [] + pos = 0 + for adapter in adapters: + if isinstance(adapter, _TensorOrQuantizedAdapter): + offsets.append(pos) + pos += len(adapter.schema_slots()) + return offsets + + +def _flatten_subclass_into_slots( + new_args: List[Any], slot_offsets: List[int], subclass: type +) -> None: + """Rewrite each tensor-or-quantized-adapter group whose ``Tensor?`` slot holds an + instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). + """ + for offset in slot_offsets: + val = new_args[offset] + if not isinstance(val, subclass): + continue + meta, tensors = _storage_flatten( + val, {_TensorOrQuantizedAdapter.KIND_KEY: _TensorOrQuantizedKind.STORAGE} + ) + new_args[offset] = None + new_args[offset + 1] = tensors + new_args[offset + 2] = meta + + +def _register_wrapper_op( + *, + wrapper_op_name: str, + schema_str: str, + base_op: Any, + adapters: Optional[List[_Adapter]] = None, +) -> Any: + """Define the wrapper op via ``torch.library.custom_op``: forward to the base + op, first flattening any ``QuantizedTensor`` subclass input into the base op's + slots. + + A ``torch.library`` op cannot take a tensor subclass directly, so each such + input is unpacked into its tensor-or-quantized slots + (``_flatten_subclass_into_slots``) before forwarding to the base op -- see the + two-tier op note in the module docstring. The subclasses to flatten are the + live ``QuantizedTensor`` wrappers (e.g. ``Float8Tensor``), obtained from the + registry via ``_all_quantized_tensor_subclasses()``. With no adapters the + wrapper is a plain pass-through. Returns the ``CustomOpDef``. + """ + subclass_list = _all_quantized_tensor_subclasses() + input_flatten_enabled = bool(subclass_list) and adapters is not None + slot_offsets = _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] + + def _forward(*flat: Any) -> List[torch.Tensor]: + if not input_flatten_enabled: + return base_op(*flat) + new_args = list(flat) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, slot_offsets, sub) + return base_op(*new_args) + + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + ) + op.register_fake(_forward) + return op + + +def _all_quantized_tensor_subclasses() -> List[type]: + """Return every imported ``QuantizedTensor`` wrapper subclass.""" + import transformer_engine.pytorch.tensor # noqa: F401 pylint: disable=import-outside-toplevel,unused-import + + found: List[type] = [] + stack = list(QuantizedTensor.__subclasses__()) + while stack: + cls = stack.pop() + if cls not in found: + found.append(cls) + stack.extend(cls.__subclasses__()) + return found + + +def register_custom_op( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Optional[Callable[..., Any]]: + """Register a TE module's forward + backward as torch custom ops. + + Always two-tier: an base ``_base`` op carries the real schema / + autograd, and an wrapper ```` op forwards to it, flattening any + quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an + empty subclass list simply makes the wrapper op a pass-through, so a pure + plain-tensor / bf16 call goes straight through). + + Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for + ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches + through the wrapper op and returns the user-facing outputs. + + Arg containers. ``fwd_arg_type`` and ``backward_arg_type`` are ``@dataclass``es + whose *field annotations* define the op schema: each field maps to one or more + flat schema slots (tensor fields cross the boundary as tensors, quantizers ride + as value-opaque objects, simple values are bundled -- see the ``_Adapter`` + classes). The caller builds a ``fwd_arg_type`` instance and passes it to the + returned ``forward_fn``. + + How the backward container is populated. ``setup_context`` fills the + ``backward_arg_type`` instance's non-tensor fields (quantizers, config) from + forward state and returns the tensors to persist; the framework saves them + (``ctx.save_for_backward``). Before ``backward_impl`` runs, the framework + restores those tensors into the container's *tensor* fields by calling its + optional ``setup_saved_tensors(self, ctx)`` hook (invoked only if defined), + and sets ``grad_output`` directly. So ``backward_impl`` receives a + fully-populated ``backward_arg_type``. + + Callable contracts: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` -- the + real forward. ``user_outputs``: op outputs (tensor / quantized / ``None``); + ``tensors_to_save``: list/tuple (or ``None``) of tensors for backward; + ``ctx_attrs``: dict (or ``None``) of plain metadata for ``setup_context``. + The trailing two slots are fixed (``_FWD_TRAILING_SLOTS``); everything + before them is a user output. + * ``fwd_fake_impl(fwd_args)`` -- data-free traceable twin of ``fwd_impl``: + same return shape, but tensor outputs are :class:`TensorSpec`. Must match + ``fwd_impl``'s shape (checked at compile time by ``_check_fwd_result``). + * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved) + -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the + tensors to persist across the boundary. + * ``backward_impl(bwd_args) -> grads`` -- exactly one grad per + ``input_tensors_for_grad`` entry, in that order (``None`` for a + non-differentiable input). + * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``backward_impl`` returning + :class:`TensorSpec` grads. + * ``backward_arg_type.setup_saved_tensors(ctx)`` -- optional hook on the backward + container (see above); skipped if absent. + + ``input_tensors_for_grad`` lists the ``fwd_arg_type`` fields that receive + gradients (this fixes the backward grad order). ``backward_arg_type`` is both + the schema source and the type instantiated (``backward_arg_type()``) to hold + the backward args, so it must be constructible with no arguments. + + Registration touches experimental ``torch.library`` / opaque-object APIs + that may be missing on older PyTorch. If it fails, this warns once and + returns ``None`` instead of raising, so callers can fall back to eager under + ``torch.compile`` (a graph break) rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + input_tensors_for_grad=input_tensors_for_grad, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + setup_context=setup_context, + backward_arg_type=backward_arg_type, + backward_impl=backward_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_fake_impl=bwd_fake_impl, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warn_compile_unsupported( + f"could not register the custom op '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Callable[..., Any]: + """Body of :func:`register_custom_op`; see it for semantics.""" + # Existence check at the API boundary: every ``input_tensors_for_grad`` name + # must be an actual field of ``fwd_arg_type`` (differentiability -- whether + # that field can carry a gradient -- is checked later in + # :func:`_resolve_grad_targets`). + fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} + missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] + if missing: + raise ValueError( + f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}" + ) + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + subclass_list = _all_quantized_tensor_subclasses() + + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(backward_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + + num_grad_inputs = len(input_tensors_for_grad) + slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) + + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_kernel( + op_name=base_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + adapters=fwd_adapters, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=base_bwd_name, + schema_str=bwd_schema, + arg_type=backward_arg_type, + arg_names=bwd_arg_names, + adapters=bwd_adapters, + tensor_field_names=bwd_tensor_field_names, + impl=backward_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + adapters=fwd_adapters, + ) + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + ) + + autograd_common = { + "fwd_arg_type": fwd_arg_type, + "fwd_arg_names": fwd_arg_names, + "fwd_adapters": fwd_adapters, + "fwd_tensor_field_names": fwd_tensor_field_names, + "bwd_arg_names": bwd_arg_names, + "bwd_adapters": bwd_adapters, + "slot_count": slot_count, + "grad_targets": grad_targets, + "setup_context_user": setup_context, + "backward_obj_type": backward_arg_type, + "fwd_fake_impl": fwd_fake_impl, + } + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) + _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) + + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return base_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return base_bwd_op(*new_args) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) + _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) + _quantized_tensor_passthrough_ops.add(base_fwd_op.default) + _quantized_tensor_passthrough_ops.add(base_bwd_op.default) + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) + user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, fwd_adapters) + flat_in = [kwargs[name] for name in fwd_arg_names] + result = wrapper_fwd_op(*flat_in) + + cursor = 0 + outputs: List[Any] = [] + for spec in user_fakes: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in result[cursor : cursor + n]] + cursor += n + outputs.append(_spec_reassemble(spec, chunk)) + + if len(outputs) == 1: + return outputs[0] + return tuple(outputs) + + return forward_fn diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index a342dd1e6c..96770e5e4e 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -6,11 +6,34 @@ from __future__ import annotations import enum +import warnings from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType +_warned_compile_unsupported = False + + +def warn_compile_unsupported(reason: str) -> None: + """Warn once per process that TE's torch.compile path is off. + + The registrations below all work or all fail together, so one message is + enough. + """ + global _warned_compile_unsupported # pylint: disable=global-statement + if _warned_compile_unsupported: + return + _warned_compile_unsupported = True + warnings.warn( + "Transformer Engine torch.compile support is disabled: " + f"{reason}. Modules will fall back to eager execution under " + "torch.compile, i.e. a graph break, which is incompatible with " + "fullgraph=True. Use a newer PyTorch build.", + stacklevel=3, + ) + + # Qualnames of the registered quantizer classes. The set holds strings rather # than the classes themselves so that ``is_value_opaque_quantizer`` can be # called inside a ``torch.compile``'d function without a graph break: Dynamo @@ -117,18 +140,20 @@ def register_value_opaque_quantizer(cls: type) -> None: register_opaque_type, is_opaque_value_type, ) - except (ImportError, AttributeError): + except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. + warn_compile_unsupported(f"this PyTorch build has no opaque-object API ({e})") return try: if not is_opaque_value_type(cls): register_opaque_type(cls, typ="value") - except (RuntimeError, TypeError): + except (RuntimeError, TypeError) as e: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. + warn_compile_unsupported(f"could not register {cls.__name__} as an opaque type ({e})") return _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py index 4cfe225952..8c156766e5 100644 --- a/transformer_engine/pytorch/dynamo/tensor_spec.py +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -151,6 +151,11 @@ def to_tensor_spec(tensor: Any) -> TensorSpec: Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` / ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via ``_dtype`` rather than ``.dtype``. + + Not for re-describing a ``TensorSpec``: a spec holds its quantizer as + ``quantizer``, not ``_quantizer``, so it would come back unquantized. Fake + impls already receive specs from ``_spec_view`` -- copy those with + ``dataclasses.replace``. """ requires_grad = bool(getattr(tensor, "requires_grad", False)) dtype = getattr(tensor, "dtype", None) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6..219263773d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,10 +3,12 @@ # See LICENSE for license information. """Linear API""" -from dataclasses import dataclass + +from dataclasses import dataclass, replace as dataclass_replace from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op +import math import warnings import weakref @@ -44,10 +46,10 @@ divide, init_method_constant, needs_quantized_gemm, - assert_dim_for_fp8_exec, nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, + warn_compile_eager_fallback, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -70,9 +72,10 @@ from ..cpp_extensions import ( general_gemm, ) +from ..cpp_extensions.gemm import get_cublas_workspace from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type -from ..jit import no_torch_dynamo from ..graph import is_graph_capturing +from ..jit import no_torch_dynamo from ..quantized_tensor import ( QuantizedTensor, QuantizedTensorStorage, @@ -80,6 +83,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorSpec, register_custom_op, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -95,6 +99,9 @@ __all__ = ["Linear"] +# Fields with this union may hold a *bare* ``QuantizedTensorStorage`` (the +# internal-quantizer optimization; see its docstring), not just a plain / +# subclass tensor -- hence the union rather than a plain ``torch.Tensor``. TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] @@ -108,7 +115,18 @@ class LinearFwdArgs: bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- - weight_workspace: Optional[torch.Tensor] + # Same union as ``weight`` so a cached quantized workspace is flattened to its + # inner tensors on the way into the op (symmetric with ``new_weight_workspace`` + # on the way out); a plain ``Tensor?`` slot can't carry a quantized subclass + # across the torch.compile custom-op boundary. + weight_workspace: Optional[TensorOrQuantized] + + # Workspace pinning (torch.compile). The process-global, lru_cached cuBLAS + # workspace is fetched in the traced forward and threaded in as an op input so + # it is allocated at trace time rather than lazily inside the op. The op body + # never reads it (general_gemm fetches the same global by address). None on + # eager / non-compiled paths. + cublas_workspace: Optional[torch.Tensor] # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool @@ -142,7 +160,9 @@ class LinearFwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] - tp_group: Optional[Any] + # ProcessGroup is a *reference*-opaque type: carried through the torch.compile + # custom op as a graph input (never baked into the graph as a constant). + tp_group: Optional[dist_group_type] tp_size: int tensor_parallel: bool sequence_parallel: bool @@ -170,6 +190,39 @@ class LinearFwdArgs: cpu_offloading: bool is_grad_enabled: bool + def compile_unsupported_reason(self) -> Optional[str]: + """Reason this config can't use the torch.compile custom-op path (else None).""" + if self.debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if self.fsdp_group is not None and self.is_grad_enabled: + return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" + if ( + self.fp8_output + and self.is_grad_enabled + and (self.input_requires_grad or self.weight_requires_grad) + ): + return "differentiable fp8_output=True" + if self.cpu_offloading: + return "CPU activation offloading" + if self.wgrad_store is not None: + # Non-None only when delayed wgrad compute is on (see Linear.forward). + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + for quantizer in ( + self.input_quantizer, + self.weight_quantizer, + self.output_quantizer, + self.grad_input_quantizer, + self.grad_weight_quantizer, + self.grad_output_quantizer, + ): + # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe + # quantizers are not value-opaque and can't cross the custom-op boundary. + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + @dataclass(slots=True) class LinearBwdArgs: @@ -207,7 +260,8 @@ class LinearBwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] = None - tp_group: Optional[Any] = None + # Reference-opaque ProcessGroup (graph input), see LinearFwdArgs.tp_group. + tp_group: Optional[dist_group_type] = None tp_size: int = 1 tensor_parallel: bool = False sequence_parallel: bool = False @@ -267,13 +321,13 @@ def _check_fp8_reduce_and_update(): def _linear_forward_impl( args: LinearFwdArgs, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], None, Optional[Dict]]: +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: """Forward implementation for the linear layer. - Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, None, + Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight workspace (returned alongside ``out`` so the caller can refresh its - cache). The last three are ``None`` when gradients are disabled. + cache). The last two are ``None`` when gradients are disabled. """ weight = args.weight @@ -338,12 +392,11 @@ def _linear_forward_impl( if ub_name is not None: nvtx_label = f"{nvtx_label}.{ub_name}" - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - assert inp.shape[-1] == in_features, "GEMM not possible" + out_features = weight.shape[0] # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) + backward_needs_input = is_grad_enabled and args.weight_requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -372,8 +425,6 @@ def _linear_forward_impl( inputmat = inp # Input tensor to save for backward (maybe sharded) inputmat_total = None # Input tensor to pass to GEMM (gathered) own_quantized_input = False - if fp8: - assert_dim_for_fp8_exec(inputmat, weight) if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor @@ -464,7 +515,7 @@ def _linear_forward_impl( # No need to set the quantizer states if weight is already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug): - columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 + columnwise_usage = is_grad_enabled and args.input_requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -647,12 +698,24 @@ def _linear_forward_impl( if is_dist_weight: wt_save = None - # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` - # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. - # Needed for torch.compile to work correctly. + # Dedup save slots that alias forward inputs or other op returns; + # ``_linear_setup_ctx`` rebuilds the refs. A custom op may not return a + # tensor aliasing an input or another return, and the cached FP8 weight + # is the same object as ``new_weight_workspace`` (cache miss) or + # ``weight_workspace`` (cache hit). + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - "weight" if wt_save is weight else None, + wt_alias, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -668,13 +731,223 @@ def _linear_forward_impl( "saved_tensor_aliases": saved_tensor_aliases, } - return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _linear_forward_impl_fake( + args: LinearFwdArgs, +) -> Tuple[TensorSpec, Optional[TensorSpec], Optional[Tuple[Any, ...]], Optional[Dict]]: + """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile, + returning ``TensorSpec`` descriptors for the outputs and saved tensors instead + of allocating real data.""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time Linear forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.weight + inp = args.inp + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + fp8 = args.fp8 + debug = args.debug + fp8_or_debug = fp8 or debug + is_grad_enabled = args.is_grad_enabled + activation_dtype = args.activation_dtype + save_original_input = args.save_original_input + if args.backward_override == "high_precision": + save_original_input = True + elif args.backward_override == "dequantized": + save_original_input = False + + out_features, _ = weight.shape + backward_needs_input = is_grad_enabled and args.weight_requires_grad + + own_quantized_input = False + inputmat_is_storage = False + inputmat_aliases_inp = False + if fp8_or_debug: + if inp.is_quantized: + # Primary-quantized input reused as-is. + inputmat_is_storage = True + inputmat_aliases_inp = True + else: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + input_quantizer.set_usage( + rowwise=True, + columnwise=( + backward_needs_input + and not save_original_input + and args.backward_override is None + ), + ) + own_quantized_input = True + inputmat_is_storage = True + else: + inputmat_aliases_inp = inp.dtype == activation_dtype + + if save_original_input: + inputmat_aliases_inp = True + inputmat_is_storage = False + + # ------------------------------------------------------ + # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ``new_weight_workspace`` is a fresh fake storage only on the + # cache-miss + ``cache_weight`` path, else ``None``. + # ------------------------------------------------------ + new_weight_workspace = None + weightmat = None + weightmat_is_storage = False + weightmat_aliases_weight = False + if fp8_or_debug: + if weight_quantizer is not None and (not weight.is_quantized or debug): + columnwise_usage = is_grad_enabled and args.input_requires_grad and not args.is_fsdp2 + if args.backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() + and not in_fp8_activation_recompute_phase() + ) + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif weight.is_quantized: + weight_quantizer = weight.quantizer + + if weight.is_quantized: + # Primary-quantized weight: the impl reuses it as ``weightmat``. + weightmat = weight + weightmat_is_storage = True + weightmat_aliases_weight = True + else: + weightmat_is_storage = True + workspace = args.weight_workspace + if workspace is not None: + # Copy, so the ``update_usage`` below stays off the input spec. + weightmat = dataclass_replace(workspace) + else: + weightmat = TensorSpec( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=weight_quantizer, + device=weight.device, + ) + if args.cache_weight: + # Persistent cache entries are wrappers, not bare storages. + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_weight_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) + else: + weightmat_aliases_weight = weight.dtype == activation_dtype + weightmat = TensorSpec( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). + # ------------------------------------------------------ + out_leading = inp.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + out_leading = out_leading * args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + out_leading = out_leading // args.tp_size + out = TensorSpec( + shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), + dtype=activation_dtype, + quantizer=output_quantizer, + requires_grad=is_grad_enabled + and (args.input_requires_grad or args.weight_requires_grad or args.bias_requires_grad), + device=inp.device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout + # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing. + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if is_grad_enabled: + # Slot 0 -- ``saved_inputmat``. + inputmat_alias = None + saved_inputmat = None + if backward_needs_input: + if inputmat_aliases_inp: + inputmat_alias = "inp" + elif inputmat_is_storage: + saved_inputmat = TensorSpec( + shape=tuple(inp.shape), + dtype=activation_dtype, + quantizer=input_quantizer, + device=inp.device, + ) + # Mirror ``_linear_forward_impl``'s post-quantization + # ``inputmat.update_usage(...)`` so the saved input's buffer layout + # matches -- driven by the same conditions as the real impl. + if own_quantized_input and not save_original_input: + if args.backward_override is not None: + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + elif ( + args.backward_input_needs_gather + and weight_quantizer is not None + and weight_quantizer.supports_only_rowwise_all_gather() + ): + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + saved_inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + else: + saved_inputmat = TensorSpec( + shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device + ) + + # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached + # FP8 weight is shared with ``new_weight_workspace`` (a return, on a cache + # miss) or the ``weight_workspace`` input (on a cache hit), so it is + # reconstructed in ``_linear_setup_ctx`` rather than saved twice. + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered weight in backward. + elif weightmat_is_storage and new_weight_workspace is not None: + wt_alias = "new_workspace" + elif weightmat_is_storage and args.weight_workspace is not None: + wt_alias = "weight_workspace" + elif weightmat_is_storage: + wt_save = weightmat + else: + wt_save = TensorSpec( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + # Slot 2 -- ``saved_weight`` (always aliased to ``weight``). + # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent). + saved_tensor_aliases = ( + inputmat_alias, + wt_alias, + "weight", + "bias" if bias is not None else None, + ) + tensors_to_save_from_forward = (saved_inputmat, wt_save, None, None) + ctx_attrs = { + "fsdp_shapes": [], + "saved_tensor_aliases": saved_tensor_aliases, + } + + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs def _linear_setup_ctx( bwd_args: LinearBwdArgs, fwd_args: LinearFwdArgs, - out: torch.Tensor, + fwd_outputs: Tuple[Any, ...], ctx_attrs: Dict, tensors_to_save_from_forward: Tuple[Any, ...], ) -> Tuple[Any, ...]: @@ -687,7 +960,8 @@ def _linear_setup_ctx( for FSDP2 re-quantization) without having to mutate the structured metadata returned by ``prepare_for_saving``. """ - del out # No-op; kept for symmetry with the compile-time helper signature. + # ``fwd_outputs`` is ``(out, new_weight_workspace)``; only the latter is used, + # to rebuild the deduped weight save slot. inp = fwd_args.inp weight = fwd_args.weight @@ -710,7 +984,10 @@ def _linear_setup_ctx( bwd_args.use_bias = bias is not None bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad - bwd_args.inp_shape = inp.shape + # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) + # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle. + # The backward reconstructs inp_shape from grad_output + weight + SP config. + bwd_args.inp_shape = None # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype @@ -779,6 +1056,10 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight + elif wt_save_alias == "new_workspace": + wt_save = fwd_outputs[1] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace if saved_weight_alias == "weight": saved_weight = weight if bias_alias == "bias": @@ -858,6 +1139,18 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") + # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). + if bwd_args.inp_shape is None: + in_features = saved_weight.shape[-1] + go_leading = grad_output.shape[0] + if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: + inp_leading = go_leading // bwd_args.tp_size + elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: + inp_leading = go_leading * bwd_args.tp_size + else: + inp_leading = go_leading + bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) + # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1387,6 +1680,92 @@ def wgrad_gemm( ) +def _linear_backward_impl_fake( + args: LinearBwdArgs, +) -> Tuple[Optional[TensorSpec], Optional[TensorSpec], Optional[TensorSpec]]: + """Allocation-free fake of :func:`_linear_backward` on ``TensorSpec``. + + The saved-tensor fields of ``args`` carry + :class:`~transformer_engine.pytorch.dynamo.TensorSpec` instances. Returns + ``(wgrad, dgrad, grad_bias)`` specs describing the nature of the gradients, + mirroring the real backward's return contract without allocating storage. + + Tensor-/sequence-parallel gather/scatter happens inside the eager backward + custom op and is opaque to ``torch.compile``: ``dgrad`` always carries the + rank-local input shape and ``wgrad`` the local weight shape, so no extra + shape modeling is needed here. + """ + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake Linear backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.saved_weight + out_dtype = args.activation_dtype + out_features, in_features = weight.shape + + # Mirror ``_linear_backward``: ``set_usage`` on ``grad_input_quantizer`` + # influences ``dgrad``'s buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + # dgrad has the logical input shape and may be quantized for the next op. + # Derive shape from grad_output + weight + SP config instead of args.inp_shape: + # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is + # not hashable in OpaqueValueBundle), so we reconstruct it here. + go_leading = args.grad_output.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + dgrad_leading = go_leading // args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + dgrad_leading = go_leading * args.tp_size + else: + dgrad_leading = go_leading + dgrad = TensorSpec( + shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), + dtype=out_dtype, + quantizer=args.grad_input_quantizer, + device=args.grad_output.device, + ) + + wgrad = None + if args.requires_wgrad and not args.fuse_wgrad_accumulation: + # wgrad has the weight's shape; quantized iff an fp8 wgrad output is + # requested (mirrors ``quantization_params=grad_weight_quantizer``), + # otherwise high precision. Under fuse_wgrad_accumulation the grad is + # written into ``main_grad`` in place and no wgrad tensor is returned. + wgrad = TensorSpec( + shape=(out_features, in_features), + dtype=out_dtype, + quantizer=args.grad_weight_quantizer, + device=weight.device, + ) + + grad_bias = None + if args.use_bias and args.requires_wgrad: + grad_bias = TensorSpec( + shape=(out_features,), dtype=out_dtype, device=args.grad_output.device + ) + + return wgrad, dgrad, grad_bias + + +# Custom op used under ``torch.compile``. +_linear_op = register_custom_op( + op_name="linear", + input_tensors_for_grad=["weight", "inp", "bias"], + fwd_arg_type=LinearFwdArgs, + fwd_impl=_linear_forward_impl, + fwd_fake_impl=_linear_forward_impl_fake, + setup_context=_linear_setup_ctx, + backward_arg_type=LinearBwdArgs, + backward_impl=_linear_backward, + bwd_fake_impl=_linear_backward_impl_fake, +) + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1418,7 +1797,6 @@ def forward( out, new_weight_workspace, tensors_to_save_from_forward, - _, ctx_attrs, ) = _linear_forward_impl(fwd_args) if ctx is not None: @@ -1426,7 +1804,7 @@ def forward( tensors_to_save_from_setup = _linear_setup_ctx( bwd_args, fwd_args, - out, + (out, new_weight_workspace), ctx_attrs, tensors_to_save_from_forward, ) @@ -1471,6 +1849,20 @@ def backward( return result +@no_torch_dynamo() +def _linear_eager( + weight_tensor: torch.Tensor, + inp: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LinearFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run ``_Linear`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _Linear.apply(weight_tensor, inp, bias, fwd_args) + return _Linear.forward(None, weight_tensor, inp, bias, fwd_args) + + class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` @@ -1869,7 +2261,6 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -1948,12 +2339,7 @@ def forward( weight_quantizer, weight_tensor ) - if is_grad_enabled: - linear_fn = _Linear.apply - autograd_ctx = [] - else: - linear_fn = _Linear.forward - autograd_ctx = [None] + use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( @@ -1993,16 +2379,50 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad + torch._check( + inp.shape[-1] == weight_tensor.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + if self.fp8: + torch._check( + math.prod(inp.shape[:-1]) % 8 == 0, + lambda: ( + "FP8 execution requires the product of all input dimensions except" + " the last to be divisible by 8" + ), + ) + torch._check( + inp.shape[-1] % 16 == 0, + lambda: "FP8 execution requires the input last dimension to be divisible by 16", + ) + torch._check( + weight_tensor.shape[0] % 16 == 0, + lambda: "FP8 execution requires out_features to be divisible by 16", + ) + torch._check( + weight_tensor.shape[1] % 16 == 0, + lambda: "FP8 execution requires in_features to be divisible by 16", + ) + linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + + # Pin the lazily-cached cuBLAS workspace as an op input so it is + # materialized at trace time (external to the cudagraph pool) rather + # than inside the op during capture. See LinearFwdArgs for details. + cublas_workspace = None + if use_compiled_op: + cublas_workspace = get_cublas_workspace(inp.device.index, False, False) + fwd_args = LinearFwdArgs( # tensors weight=weight_tensor, inp=inp, bias=linear_bias_tensor, weight_workspace=weight_workspace, + cublas_workspace=cublas_workspace, # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, @@ -2057,13 +2477,19 @@ def forward( cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, ) - out, new_weight_workspace = linear_fn( - *autograd_ctx, - weight_tensor, - inp, - linear_bias_tensor, - fwd_args, - ) + + if use_compiled_op: + fallback_reason = fwd_args.compile_unsupported_reason() + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + use_compiled_op = False + + if use_compiled_op: + out, new_weight_workspace = _linear_op(fwd_args) + else: + out, new_weight_workspace = _linear_eager( + weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled + ) if new_weight_workspace is not None and cache_name is not None: if isinstance(new_weight_workspace, torch.Tensor): diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 10672bbcfb..6161faaddc 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -138,6 +138,43 @@ def _stride_from_shape(shape: list[int]): return list(reversed(rstride)) +def tensor_can_be_materialized(t) -> bool: + """Whether ``t`` holds concrete data that ``.item()`` / ``.tolist()`` can read + without side effects. + + A ``__repr__`` must never mutate tracing state. On a fake / meta / functional + tensor (torch.compile / export tracing) ``.item()`` does *not* raise -- it + silently allocates an *unbacked* SymInt/SymFloat into the active ShapeEnv, + which later crashes inductor with ``PendingUnbackedSymbolNotFound``. (torch's + AOTAutograd repr's the fake quantized tensor while logging graph metadata, so + a scalar-materializing ``__repr__`` leaks an unbacked symbol during compile.) + So detect those tensors and fall back to a metadata-only repr instead. + """ + if not isinstance(t, torch.Tensor): + return False + if getattr(t, "is_meta", False): + return False + try: + from torch._subclasses.fake_tensor import ( # pylint: disable=import-outside-toplevel + FakeTensor, + ) + + if isinstance(t, FakeTensor): + return False + except Exception: # pylint: disable=broad-except + pass + try: + from torch._subclasses.functional_tensor import ( # pylint: disable=import-outside-toplevel + FunctionalTensor, + ) + + if isinstance(t, FunctionalTensor): + return False + except Exception: # pylint: disable=broad-except + pass + return True + + def safe_quantized_repr(obj, cls_name, extras=None, error=None): """Metadata-only repr fallback for quantized tensors whose data cannot be materialized for any reason. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 5c31022123..2f89326750 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -23,6 +23,7 @@ _IdentityFunc, _resolve_view_shape, safe_quantized_repr, + tensor_can_be_materialized, ) from ..constants import dist_group_type, DType @@ -473,6 +474,11 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): amax_reduction_group: Optional[dist_group_type] = None def __repr__(self, *, tensor_contents=None): + # A fake/meta/functional scale_inv cannot be materialized without leaking + # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); + # fall back to a metadata-only repr under tracing. + if not tensor_can_be_materialized(self._scale_inv): + return safe_quantized_repr(self, "Float8Tensor") try: return ( "Float8Tensor(" diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 16bba391c4..1c3ce68c6d 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -11,7 +11,11 @@ import transformer_engine_torch as tex from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer -from .._quantization_helpers import _resolve_view_shape, safe_quantized_repr +from .._quantization_helpers import ( + _resolve_view_shape, + safe_quantized_repr, + tensor_can_be_materialized, +) from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -253,6 +257,11 @@ def view(self, shape: torch.Size): ) def __repr__(self): + # A fake/meta/functional scale_inv cannot be materialized without leaking + # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); + # fall back to a metadata-only repr under tracing. + if not tensor_can_be_materialized(self._scale_inv): + return safe_quantized_repr(self, "Float8TensorStorage") try: return ( "Float8TensorStorage(" diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 39162f8311..bd9231c460 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,6 +26,20 @@ ] +def warn_compile_eager_fallback(reason: str) -> None: + """Warn that a TE module is running eagerly under ``torch.compile``. + + Emitted when ``reason`` is unsupported on the module's compiled custom-op + path. Python's default warning filter dedups identical messages, so each + distinct ``reason`` is surfaced once. + """ + warnings.warn( + f"Falling back to eager execution under torch.compile: {reason} is " + "unsupported on the compiled path (graph-breaks under fullgraph=True).", + stacklevel=2, + ) + + @functools.lru_cache(maxsize=None) def get_cached_ones_tensor( num_elements: int, From dfa79c2c0428c6de92b64ed2b6028b27f3885981 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:09:02 +0000 Subject: [PATCH 02/50] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/dynamo/custom_op.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 2a4f64d9a8..ed4e0e1c8c 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -183,9 +183,7 @@ def is_simple_value(cls, value: Any) -> bool: if _is_opaque_value_type(type(value)): return True if isinstance(value, dict): - return all( - isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items() - ) + return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()) if isinstance(value, (list, tuple)): return all(cls.is_simple_value(v) for v in value) return False @@ -309,7 +307,9 @@ def _collect(value: Any) -> None: register_opaque_type(OpaqueValueBundle, typ="value") _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) -except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object +except ( + Exception +) as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") _is_opaque_value_type = None _is_opaque_reference_type = None @@ -1166,9 +1166,7 @@ def _setup_context(ctx, inputs, output): ctx_attrs, tuple(saved_list), ) - tensors_to_save, tensor_objects = prepare_for_saving( - *(tensors_to_save_from_setup or ()) - ) + tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ())) ctx.tensor_objects = tensor_objects ctx.save_for_backward(*tensors_to_save) ctx.bwd_obj = bwd_obj @@ -1247,7 +1245,9 @@ def _register_wrapper_op( """ subclass_list = _all_quantized_tensor_subclasses() input_flatten_enabled = bool(subclass_list) and adapters is not None - slot_offsets = _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] + slot_offsets = ( + _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] + ) def _forward(*flat: Any) -> List[torch.Tensor]: if not input_flatten_enabled: @@ -1389,9 +1389,7 @@ def _register_custom_op_impl( fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] if missing: - raise ValueError( - f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}" - ) + raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") wrapper_fwd_name = op_name wrapper_bwd_name = f"{op_name}_backward" From 116d4774ee0e25a244cfa712af3a54354648603d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 18:21:53 +0200 Subject: [PATCH 03/50] [PyTorch] Keep the broad-except pylint disable on the anchored line black wrapped the 122-char except clause, moving Exception onto its own line while the disable comment stayed on the closing paren, so pylint's W0718 no longer saw it. Shorten the line instead. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index ed4e0e1c8c..b529deb75a 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -307,9 +307,8 @@ def _collect(value: Any) -> None: register_opaque_type(OpaqueValueBundle, typ="value") _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) -except ( - Exception -) as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object +# Older torch without opaque_object support. +except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") _is_opaque_value_type = None _is_opaque_reference_type = None From dc4ff77e2ed37c0a58b37ee1499ec7358beeed6a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 23:00:12 +0200 Subject: [PATCH 04/50] [PyTorch] [torch.compile] Style pass on the Linear custom-op path Naming consistency and de-duplication in the torch.compile custom-op framework and its Linear user. No functional change. Naming: - unify the register_custom_op API on fwd_*/bwd_* (backward_arg_type, backward_impl, backward_obj_type -> bwd_arg_type, bwd_impl) - _register_kernel -> _register_base_op, pairing with _register_wrapper_op - _format_*_result / _split_fwd_fake_result -> _pack_*_result / _unpack_fwd_fake_result - _value_to_flat_tensors / _spec_reassemble -> _flatten_value / _unflatten_value, matching _storage_flatten / _storage_unflatten - adapter slots: tensor_slot / inner_slot / meta_slot, META_SLOT, QUANTIZER_KEY - _linear_backward -> _linear_backward_impl and *_fake twins, so the real and fake implementations pair up by name - ctx attrs: drop the lone _te_ prefix, and use ctx.backward_objects as the eager path already does - move warn_compile_unsupported to utils as warn_compile_disabled, next to warn_compile_eager_fallback, so the two "unsupported" meanings are distinguishable - move the TensorOrQuantized alias next to the adapter that matches it De-duplication: - _unflatten_values() replaces three copies of the cursor/reassemble loop - _make_slot_forwarder() / _make_dispatch_rule() replace three copies of the subclass-flattening forward path - _sp_out_leading() / _sp_inp_leading() replace three copies of the sequence-parallel leading-dim arithmetic (two of them inverses) - check_gemm_dims() moves the fp8 dimension checks to utils - drop the duplicate backward_needs_input assignment in the forward impl Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_hybrid_quantization.py | 2 +- tests/pytorch/test_torch_compile.py | 4 +- transformer_engine/pytorch/dynamo/__init__.py | 3 +- .../pytorch/dynamo/custom_op.py | 345 ++++++++++-------- .../pytorch/dynamo/quantizer_opaque.py | 28 +- transformer_engine/pytorch/module/linear.py | 122 +++---- transformer_engine/pytorch/utils.py | 50 +++ 7 files changed, 296 insertions(+), 258 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 74ec0a05ec..a7edf87b13 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -495,7 +495,7 @@ def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self): ``gather_along_first_dim`` cannot operate on a columnwise-only NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather`` must return True in this case so ``_linear_forward_impl`` / - ``_linear_backward`` preserve rowwise data (which NVFP4 can + ``_linear_backward_impl`` preserve rowwise data (which NVFP4 can dequantize) instead. """ hq = HybridQuantizer( diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 68b1174474..a1935f2d93 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1459,9 +1459,9 @@ def test_te_linear_dynamic_shapes(): This exercises two fixes for dynamic shapes: 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). - 2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output + + 2. ``_linear_backward_fake`` derives dgrad shape from grad_output + weight + SP config instead of relying on the stored ``inp_shape``. - 3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same + 3. ``_linear_backward_impl`` reconstructs ``inp_shape`` on-the-fly from the same tensor sources when it is None (compiled mode). FP8 + dynamic=True is tracked separately (requires resolving diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..e42eb8f9f6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op +from .custom_op import register_custom_op, TensorOrQuantized __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "TensorOrQuantized", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b529deb75a..bde9cca351 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -12,7 +12,7 @@ ``transformer_engine.pytorch.dynamo``, not re-exported at the top level. A TE op's forward/backward is written as a plain impl over a single *args -dataclass* (``fwd_arg_type`` / ``backward_arg_type``, e.g. ``LinearFwdArgs``): its +dataclass* (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``): its fields are a mix of tensors, quantized tensors, quantizers, process groups, scalars and other Python values. The forward impl returns a tuple (user outputs + saved-for-backward tensors + ctx metadata); the backward impl returns one gradient @@ -59,8 +59,8 @@ * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat ``Tensor[]`` payload; * rebuilds the structured user outputs from that payload, sliced and reassembled - per the fake's output descriptors (``_spec_reassemble``; - ``_value_to_flat_tensors`` is the pack-side inverse). + per the fake's output descriptors (``_unflatten_value``; + ``_flatten_value`` is the pack-side inverse). Autograd, registered on the op, drives backward: @@ -71,7 +71,7 @@ aliases) and return the tensors to persist; * on ``backward()`` the backward args container's optional ``setup_saved_tensors`` hook restores those saved tensors, then the *backward op* runs the real - ``backward_impl`` and returns the flat grads (``bwd_fake_impl`` is its + ``bwd_impl`` and returns the flat grads (``bwd_fake_impl`` is its data-free fake). Two-tier op (``base`` + ``wrapper``), so a ``QuantizedTensor`` subclass can be an @@ -84,7 +84,7 @@ from __future__ import annotations import dataclasses -import types as _types +import types as _types # aliased: torch_dispatch rules take a ``types`` param from enum import Enum from typing import ( Any, @@ -104,7 +104,6 @@ from torch._prims_common import make_contiguous_strides_for -from .quantizer_opaque import warn_compile_unsupported from .tensor_spec import TensorSpec, to_tensor_spec from ..quantized_tensor import ( QuantizedTensor, @@ -113,9 +112,15 @@ _quantized_tensor_passthrough_ops, prepare_for_saving, ) +from ..utils import warn_compile_disabled _TE_OP_NAMESPACE = "transformer_engine_compile" +# Annotation for an op arg field that may hold a plain tensor, a quantized +# tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer +# optimization). Matched exactly by ``_TensorOrQuantizedAdapter``. +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a # 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for @@ -309,7 +314,7 @@ def _collect(value: Any) -> None: _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) # Older torch without opaque_object support. except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") + warn_compile_disabled(f"could not register OpaqueValueBundle as an opaque type ({e})") _is_opaque_value_type = None _is_opaque_reference_type = None _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None @@ -519,30 +524,29 @@ class _TensorOrQuantizedAdapter(_Adapter): def __init__(self, name: str) -> None: self.name = name - def slot_name(self) -> str: + def tensor_slot(self) -> str: """Primary slot name for a plain / subclass tensor.""" return self.name - def slot_tensors(self) -> str: + def inner_slot(self) -> str: """Flat inner-tensor slot name.""" return self.name + "__tensors" - def slot_meta(self) -> str: + def meta_slot(self) -> str: """Flatten-metadata slot name.""" return self.name + "__meta" def schema_slots(self) -> List[Tuple[str, str]]: return [ - (self.slot_name(), "Tensor?"), - (self.slot_tensors(), "Tensor[]"), - (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + (self.tensor_slot(), "Tensor?"), + (self.inner_slot(), "Tensor[]"), + (self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), ] - # Canonical "plain tensor or quantized tensor" field annotation (the - # ``TensorOrQuantized`` alias in module code). Matched by exact member set, - # so a bare quantized annotation or an accidental extra union member is - # rejected rather than silently taken as a tensor-or-quantized field. - _MEMBERS = frozenset({torch.Tensor, QuantizedTensorStorage}) + # Matched by exact member set, so a bare quantized annotation or an + # accidental extra union member is rejected rather than silently taken as a + # tensor-or-quantized field. + _MEMBERS = frozenset(get_args(TensorOrQuantized)) @classmethod def _is_tensor_storage_union(cls, annot: Any) -> bool: @@ -561,25 +565,25 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: value = getattr(owner, self.name) if value is None: return { - self.slot_name(): None, - self.slot_tensors(): [], - self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE}), + self.tensor_slot(): None, + self.inner_slot(): [], + self.meta_slot(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE}), } if isinstance(value, torch.Tensor): # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the # ``Tensor?`` slot; subclass flattening (if any) is done by the # wrapper op's ``register_torch_dispatch`` rule. return { - self.slot_name(): value, - self.slot_tensors(): [], - self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR}), + self.tensor_slot(): value, + self.inner_slot(): [], + self.meta_slot(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR}), } if isinstance(value, QuantizedTensorStorage): meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) return { - self.slot_name(): None, - self.slot_tensors(): tensors, - self.slot_meta(): meta, + self.tensor_slot(): None, + self.inner_slot(): tensors, + self.meta_slot(): meta, } raise TypeError( f"field {self.name!r} expected None, torch.Tensor, or " @@ -587,14 +591,14 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: ) def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - meta = args[self.slot_meta()] + meta = args[self.meta_slot()] kind = meta.get(self.KIND_KEY) if kind == _TensorOrQuantizedKind.NONE: kwargs[self.name] = None elif kind == _TensorOrQuantizedKind.TENSOR: - kwargs[self.name] = args[self.slot_name()] + kwargs[self.name] = args[self.tensor_slot()] else: - kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()]) + kwargs[self.name] = _storage_unflatten(meta, args[self.inner_slot()]) def grad_slot(self) -> Optional[int]: # Gradient flows to the plain / subclass tensor slot (``slot_name``, @@ -637,12 +641,12 @@ class _QuantizerAdapter(_Adapter): bundle would not claim it. """ - KEY = "q" + QUANTIZER_KEY = "q" def __init__(self, name: str) -> None: self.name = name - def slot(self) -> str: + def meta_slot(self) -> str: """Opaque quantizer metadata slot name.""" return self.name + "__q" @@ -654,13 +658,15 @@ def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerAdapter"]: return None def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + return [(self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] def to_slots(self, owner: Any) -> Dict[str, Any]: - return {self.slot(): OpaqueValueBundle({self.KEY: getattr(owner, self.name)})} + return { + self.meta_slot(): OpaqueValueBundle({self.QUANTIZER_KEY: getattr(owner, self.name)}) + } def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - kwargs[self.name] = args[self.slot()][self.KEY] + kwargs[self.name] = args[self.meta_slot()][self.QUANTIZER_KEY] class _ReferenceOpaqueAdapter(_Adapter): @@ -712,7 +718,7 @@ class _SimpleBundleAdapter(_Adapter): simple-typed field names collected across the dataclass. """ - SLOT = "_simple_meta" + META_SLOT = "_simple_meta" def __init__(self, names: List[str]) -> None: self.names = list(names) @@ -737,15 +743,15 @@ def matches_field(cls, annot: Any) -> bool: return False def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + return [(self.META_SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] def to_slots(self, owner: Any) -> Dict[str, Any]: - return {self.SLOT: OpaqueValueBundle({n: getattr(owner, n) for n in self.names})} + return {self.META_SLOT: OpaqueValueBundle({n: getattr(owner, n) for n in self.names})} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - if self.SLOT not in args: + if self.META_SLOT not in args: return - meta = args[self.SLOT] + meta = args[self.META_SLOT] for n in self.names: kwargs[n] = meta[n] @@ -787,7 +793,7 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: f"{self.owner_cls_name} field {self.name!r} has a type not " "supported by torch.compile (not Tensor, simple, Quantizer, or a " "reference-opaque type such as ProcessGroup) and carries a " - "non-trivial value; add a matching adapter in dynamo.py to handle it." + "non-trivial value; add a matching adapter in custom_op.py to handle it." ) return {} @@ -852,7 +858,7 @@ def _tensor_field_names(adapters: List[_Adapter]) -> List[str]: def _build_schema(adapters: List[_Adapter]) -> Tuple[str, List[str]]: - """Return ``(schema_arg_str, slot_names)`` for a adapter list.""" + """Return ``(schema_arg_str, slot_names)`` for an adapter list.""" spec = [slot for b in adapters for slot in b.schema_slots()] names = [name for name, _ in spec] schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")" @@ -917,7 +923,7 @@ def _spec_slot_count(spec: Optional[TensorSpec]) -> int: return len(spec.inner_names()) -def _spec_reassemble( +def _unflatten_value( spec: Optional[TensorSpec], chunk: List[Optional[torch.Tensor]], ) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]: @@ -932,12 +938,31 @@ def _spec_reassemble( return spec.assemble(chunk) -def _value_to_flat_tensors( +def _unflatten_values( + specs: Sequence[Optional[TensorSpec]], + flat: Sequence[Optional[torch.Tensor]], + cursor: int = 0, +) -> Tuple[List[Any], int]: + """Rebuild one group of values from an op's flat return, starting at ``cursor``. + + Returns the values and the new cursor, so consecutive groups (user outputs, + then saved tensors) can walk the same payload. + """ + values: List[Any] = [] + for spec in specs: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in flat[cursor : cursor + n]] + cursor += n + values.append(_unflatten_value(spec, chunk)) + return values, cursor + + +def _flatten_value( value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorSpec]], ) -> List[torch.Tensor]: """Return the flat ``Tensor[]`` slots that represent one op output ``value``. - Inverse of :func:`_spec_reassemble`; the slot count matches + Inverse of :func:`_unflatten_value`; the slot count matches :func:`_spec_slot_count`. """ if value is None: @@ -964,9 +989,9 @@ def _check_fwd_result(result: Any) -> None: """Validate a fwd-impl return against the ``(*user_outputs, tensors_to_save, ctx_attrs)`` contract, with a clear message for op authors (user-output *types* are checked later, by - :func:`_value_to_flat_tensors`). + :func:`_flatten_value`). - Only called on the fake path (:func:`_split_fwd_fake_result`), which runs at + Only called on the fake path (:func:`_unpack_fwd_fake_result`), which runs at trace/compile time -- so this is a compile-time check with no per-call cost. The real impl must return the same shape as the fake, so validating the fake covers both. @@ -984,7 +1009,7 @@ def _check_fwd_result(result: Any) -> None: raise TypeError("fwd impl 'ctx_attrs' slot must be a dict or None") -def _format_fwd_result(result: Any) -> List[torch.Tensor]: +def _pack_fwd_result(result: Any) -> List[torch.Tensor]: """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. User outputs first, then saved-for-backward tensors in declaration order. @@ -992,15 +1017,15 @@ def _format_fwd_result(result: Any) -> List[torch.Tensor]: num_outputs = len(result) - _FWD_TRAILING_SLOTS flat: List[torch.Tensor] = [] for value in result[:num_outputs]: - flat.extend(_value_to_flat_tensors(value)) + flat.extend(_flatten_value(value)) saved = result[num_outputs] if saved is not None: for value in saved: - flat.extend(_value_to_flat_tensors(value)) + flat.extend(_flatten_value(value)) return flat -def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: +def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. Each grad occupies exactly one slot (validated against ``num_grad_inputs``); @@ -1009,7 +1034,7 @@ def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> Li grads = list(grads) if len(grads) != num_grad_inputs: raise RuntimeError( - f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads " + f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads " f"(one per input_tensors_for_grad entry), got {len(grads)}" ) out: List[torch.Tensor] = [] @@ -1021,7 +1046,7 @@ def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> Li return out -def _split_fwd_fake_result( +def _unpack_fwd_fake_result( result: Tuple[Any, ...], ) -> Tuple[List[Any], List[Any], Dict[str, Any]]: """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" @@ -1071,7 +1096,7 @@ def _resolve_grad_targets( return slot_offset, grad_targets -def _register_kernel( +def _register_base_op( *, op_name: str, schema_str: str, @@ -1081,26 +1106,26 @@ def _register_kernel( tensor_field_names: List[str], impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any], - format_result: Callable[[Any], List[torch.Tensor]], + pack_result: Callable[[Any], List[torch.Tensor]], ) -> Any: """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the ``fake_impl`` (spec), returning the ``CustomOpDef``. The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel runs the spec fake impl on the :func:`_spec_view`. Both go through - ``format_result``. + ``pack_result``. """ def _impl(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) obj = _args_from_slots(arg_type, kwargs, adapters) - return format_result(impl(obj)) + return pack_result(impl(obj)) def _fake(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) obj = _args_from_slots(arg_type, kwargs, adapters) spec_obj = _spec_view(obj, tensor_field_names) - return format_result(fake_impl(spec_obj)) + return pack_result(fake_impl(spec_obj)) op = torch.library.custom_op( f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str @@ -1122,7 +1147,7 @@ def _register_autograd_for_op( slot_count: int, grad_targets: List[int], setup_context_user: Callable[..., Any], - backward_obj_type: type, + bwd_arg_type: type, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> None: """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. @@ -1133,31 +1158,19 @@ def _register_autograd_for_op( """ def _setup_context(ctx, inputs, output): - ctx._te_fwd_tensor_list_lengths = { + ctx.fwd_tensor_list_lengths = { i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) } kwargs = dict(zip(fwd_arg_names, inputs)) fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_adapters) spec_obj = _spec_view(fwd_obj, fwd_tensor_field_names) - user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) - - cursor = 0 - user_outputs: List[Any] = [] - for spec in user_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in output[cursor : cursor + n]] - cursor += n - user_outputs.append(_spec_reassemble(spec, chunk)) + user_fakes, saved_fakes, ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) - saved_list: List[Any] = [] - for spec in saved_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in output[cursor : cursor + n]] - cursor += n - saved_list.append(_spec_reassemble(spec, chunk)) + user_outputs, cursor = _unflatten_values(user_fakes, output) + saved_list, _ = _unflatten_values(saved_fakes, output, cursor) - bwd_obj = backward_obj_type() + bwd_obj = bwd_arg_type() tensors_to_save_from_setup = setup_context_user( bwd_obj, fwd_obj, @@ -1168,24 +1181,23 @@ def _setup_context(ctx, inputs, output): tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ())) ctx.tensor_objects = tensor_objects ctx.save_for_backward(*tensors_to_save) - ctx.bwd_obj = bwd_obj + ctx.backward_objects = bwd_obj def _autograd_backward(ctx, *grad_outputs): - bwd_obj = ctx.bwd_obj + bwd_obj = ctx.backward_objects if hasattr(bwd_obj, "setup_saved_tensors"): bwd_obj.setup_saved_tensors(ctx) ctx.tensor_objects = None - per_output_grads = grad_outputs[0] - bwd_obj.grad_output = _decode_none(per_output_grads[0]) + flat_grads = grad_outputs[0] + bwd_obj.grad_output = _decode_none(flat_grads[0]) kwargs = _args_to_slots(bwd_obj, bwd_adapters) bwd_args_flat = [kwargs[name] for name in bwd_arg_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] # One grad per input schema slot: default None, but a ``Tensor[]`` slot - # (always recorded in ``_te_fwd_tensor_list_lengths``) needs a + # (always recorded in ``fwd_tensor_list_lengths``) needs a # list-shaped no-grad of matching length. out: List[Any] = [None] * slot_count - tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {}) - for pos, length in tensor_list_lengths.items(): + for pos, length in ctx.fwd_tensor_list_lengths.items(): out[pos] = [None] * length for pos, g in zip(grad_targets, grads): out[pos] = g @@ -1194,7 +1206,7 @@ def _autograd_backward(ctx, *grad_outputs): fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) -def _collect_tensor_or_quantized_slot_offsets(adapters: List[_Adapter]) -> List[int]: +def _tensor_or_quantized_offsets(adapters: List[_Adapter]) -> List[int]: """Start index of each ``_TensorOrQuantizedAdapter`` group in the flat args.""" offsets: List[int] = [] pos = 0 @@ -1223,44 +1235,64 @@ def _flatten_subclass_into_slots( new_args[offset + 2] = meta +def _make_slot_forwarder( + base_op: Any, slot_offsets: Sequence[int], subclasses: Sequence[type] +) -> Callable[[Sequence[Any]], List[torch.Tensor]]: + """Return ``call(args)`` forwarding to ``base_op``, first flattening any + ``subclasses`` instance sitting in the tensor-or-quantized slot groups at + ``slot_offsets``. + + A ``torch.library`` op cannot take a tensor subclass directly, so the wrapper + op body and its ``register_torch_dispatch`` rules all funnel through this one + path -- see the two-tier op note in the module docstring. With no slots or no + subclasses to flatten it is a plain pass-through. + """ + enabled = bool(slot_offsets) and bool(subclasses) + + def call(args: Sequence[Any]) -> List[torch.Tensor]: + if not enabled: + return base_op(*args) + new_args = list(args) + for sub in subclasses: + _flatten_subclass_into_slots(new_args, slot_offsets, sub) + return base_op(*new_args) + + return call + + +def _make_dispatch_rule( + forward: Callable[[Sequence[Any]], List[torch.Tensor]] +) -> Callable[..., Any]: + """Adapt a slot forwarder to the ``register_torch_dispatch`` signature.""" + + def _rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + return forward(args) + + return _rule + + def _register_wrapper_op( *, wrapper_op_name: str, schema_str: str, base_op: Any, - adapters: Optional[List[_Adapter]] = None, + slot_offsets: Sequence[int] = (), + subclasses: Sequence[type] = (), ) -> Any: """Define the wrapper op via ``torch.library.custom_op``: forward to the base - op, first flattening any ``QuantizedTensor`` subclass input into the base op's - slots. - - A ``torch.library`` op cannot take a tensor subclass directly, so each such - input is unpacked into its tensor-or-quantized slots - (``_flatten_subclass_into_slots``) before forwarding to the base op -- see the - two-tier op note in the module docstring. The subclasses to flatten are the - live ``QuantizedTensor`` wrappers (e.g. ``Float8Tensor``), obtained from the - registry via ``_all_quantized_tensor_subclasses()``. With no adapters the - wrapper is a plain pass-through. Returns the ``CustomOpDef``. + op through :func:`_make_slot_forwarder`. Returns the ``CustomOpDef``. """ - subclass_list = _all_quantized_tensor_subclasses() - input_flatten_enabled = bool(subclass_list) and adapters is not None - slot_offsets = ( - _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] - ) + forward = _make_slot_forwarder(base_op, slot_offsets, subclasses) def _forward(*flat: Any) -> List[torch.Tensor]: - if not input_flatten_enabled: - return base_op(*flat) - new_args = list(flat) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, slot_offsets, sub) - return base_op(*new_args) + return forward(flat) - op = torch.library.custom_op( + op_def = torch.library.custom_op( f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str ) - op.register_fake(_forward) - return op + op_def.register_fake(_forward) + return op_def def _all_quantized_tensor_subclasses() -> List[type]: @@ -1284,15 +1316,15 @@ def register_custom_op( fwd_arg_type: type, fwd_impl: Callable[[Any], Any], setup_context: Callable[..., Any], - backward_arg_type: type, - backward_impl: Callable[[Any], Any], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Optional[Callable[..., Any]]: """Register a TE module's forward + backward as torch custom ops. - Always two-tier: an base ``_base`` op carries the real schema / - autograd, and an wrapper ```` op forwards to it, flattening any + Always two-tier: a base ``_base`` op carries the real schema / + autograd, and a wrapper ```` op forwards to it, flattening any quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an empty subclass list simply makes the wrapper op a pass-through, so a pure plain-tensor / bf16 call goes straight through). @@ -1301,7 +1333,7 @@ def register_custom_op( ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches through the wrapper op and returns the user-facing outputs. - Arg containers. ``fwd_arg_type`` and ``backward_arg_type`` are ``@dataclass``es + Arg containers. ``fwd_arg_type`` and ``bwd_arg_type`` are ``@dataclass``es whose *field annotations* define the op schema: each field maps to one or more flat schema slots (tensor fields cross the boundary as tensors, quantizers ride as value-opaque objects, simple values are bundled -- see the ``_Adapter`` @@ -1309,13 +1341,13 @@ def register_custom_op( returned ``forward_fn``. How the backward container is populated. ``setup_context`` fills the - ``backward_arg_type`` instance's non-tensor fields (quantizers, config) from + ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from forward state and returns the tensors to persist; the framework saves them - (``ctx.save_for_backward``). Before ``backward_impl`` runs, the framework + (``ctx.save_for_backward``). Before ``bwd_impl`` runs, the framework restores those tensors into the container's *tensor* fields by calling its optional ``setup_saved_tensors(self, ctx)`` hook (invoked only if defined), - and sets ``grad_output`` directly. So ``backward_impl`` receives a - fully-populated ``backward_arg_type``. + and sets ``grad_output`` directly. So ``bwd_impl`` receives a + fully-populated ``bwd_arg_type``. Callable contracts: @@ -1331,17 +1363,17 @@ def register_custom_op( * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved) -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the tensors to persist across the boundary. - * ``backward_impl(bwd_args) -> grads`` -- exactly one grad per + * ``bwd_impl(bwd_args) -> grads`` -- exactly one grad per ``input_tensors_for_grad`` entry, in that order (``None`` for a non-differentiable input). - * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``backward_impl`` returning + * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning :class:`TensorSpec` grads. - * ``backward_arg_type.setup_saved_tensors(ctx)`` -- optional hook on the backward + * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook on the backward container (see above); skipped if absent. ``input_tensors_for_grad`` lists the ``fwd_arg_type`` fields that receive - gradients (this fixes the backward grad order). ``backward_arg_type`` is both - the schema source and the type instantiated (``backward_arg_type()``) to hold + gradients (this fixes the backward grad order). ``bwd_arg_type`` is both + the schema source and the type instantiated (``bwd_arg_type()``) to hold the backward args, so it must be constructible with no arguments. Registration touches experimental ``torch.library`` / opaque-object APIs @@ -1356,13 +1388,13 @@ def register_custom_op( fwd_arg_type=fwd_arg_type, fwd_impl=fwd_impl, setup_context=setup_context, - backward_arg_type=backward_arg_type, - backward_impl=backward_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, fwd_fake_impl=fwd_fake_impl, bwd_fake_impl=bwd_fake_impl, ) except (ImportError, AttributeError, RuntimeError, TypeError) as e: - warn_compile_unsupported( + warn_compile_disabled( f"could not register the custom op '{op_name}' ({type(e).__name__}: {e})" ) return None @@ -1375,8 +1407,8 @@ def _register_custom_op_impl( fwd_arg_type: type, fwd_impl: Callable[[Any], Any], setup_context: Callable[..., Any], - backward_arg_type: type, - backward_impl: Callable[[Any], Any], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: @@ -1397,7 +1429,7 @@ def _register_custom_op_impl( subclass_list = _all_quantized_tensor_subclasses() fwd_adapters = _get_adapters(fwd_arg_type) - bwd_adapters = _get_adapters(backward_arg_type) + bwd_adapters = _get_adapters(bwd_arg_type) fwd_tensor_field_names = _tensor_field_names(fwd_adapters) bwd_tensor_field_names = _tensor_field_names(bwd_adapters) @@ -1412,7 +1444,7 @@ def _register_custom_op_impl( base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - base_fwd_def = _register_kernel( + base_fwd_def = _register_base_op( op_name=base_fwd_name, schema_str=fwd_schema, arg_type=fwd_arg_type, @@ -1421,29 +1453,35 @@ def _register_custom_op_impl( tensor_field_names=fwd_tensor_field_names, impl=fwd_impl, fake_impl=fwd_fake_impl, - format_result=_format_fwd_result, + pack_result=_pack_fwd_result, ) - _register_kernel( + _register_base_op( op_name=base_bwd_name, schema_str=bwd_schema, - arg_type=backward_arg_type, + arg_type=bwd_arg_type, arg_names=bwd_arg_names, adapters=bwd_adapters, tensor_field_names=bwd_tensor_field_names, - impl=backward_impl, + impl=bwd_impl, fake_impl=bwd_fake_impl, - format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname), ) base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + fwd_slot_offsets = _tensor_or_quantized_offsets(fwd_adapters) + bwd_slot_offsets = _tensor_or_quantized_offsets(bwd_adapters) + wrapper_fwd_def = _register_wrapper_op( wrapper_op_name=wrapper_fwd_name, schema_str=fwd_schema, base_op=base_fwd_op, - adapters=fwd_adapters, + slot_offsets=fwd_slot_offsets, + subclasses=subclass_list, ) + # Pass-through: a subclass input reaches the base op through the dispatch + # rule below, never through the wrapper body. wrapper_bwd_def = _register_wrapper_op( wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op ) @@ -1458,7 +1496,7 @@ def _register_custom_op_impl( "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, - "backward_obj_type": backward_arg_type, + "bwd_arg_type": bwd_arg_type, "fwd_fake_impl": fwd_fake_impl, } wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) @@ -1467,22 +1505,12 @@ def _register_custom_op_impl( _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) - bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) - - def _fwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) - return base_fwd_op(*new_args) - - def _bwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) - return base_bwd_op(*new_args) + _fwd_rule = _make_dispatch_rule( + _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list) + ) + _bwd_rule = _make_dispatch_rule( + _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list) + ) for sub in subclass_list: wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) @@ -1495,19 +1523,12 @@ def _bwd_rule(mode, func, types, args, kwargs): def forward_fn(fwd_args): spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) - user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + user_fakes, _saved_fakes, _ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) kwargs = _args_to_slots(fwd_args, fwd_adapters) flat_in = [kwargs[name] for name in fwd_arg_names] result = wrapper_fwd_op(*flat_in) - cursor = 0 - outputs: List[Any] = [] - for spec in user_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in result[cursor : cursor + n]] - cursor += n - outputs.append(_spec_reassemble(spec, chunk)) - + outputs, _ = _unflatten_values(user_fakes, result) if len(outputs) == 1: return outputs[0] return tuple(outputs) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 96770e5e4e..ae8849097c 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -6,32 +6,10 @@ from __future__ import annotations import enum -import warnings from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType - - -_warned_compile_unsupported = False - - -def warn_compile_unsupported(reason: str) -> None: - """Warn once per process that TE's torch.compile path is off. - - The registrations below all work or all fail together, so one message is - enough. - """ - global _warned_compile_unsupported # pylint: disable=global-statement - if _warned_compile_unsupported: - return - _warned_compile_unsupported = True - warnings.warn( - "Transformer Engine torch.compile support is disabled: " - f"{reason}. Modules will fall back to eager execution under " - "torch.compile, i.e. a graph break, which is incompatible with " - "fullgraph=True. Use a newer PyTorch build.", - stacklevel=3, - ) +from ..utils import warn_compile_disabled # Qualnames of the registered quantizer classes. The set holds strings rather @@ -143,7 +121,7 @@ def register_value_opaque_quantizer(cls: type) -> None: except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. - warn_compile_unsupported(f"this PyTorch build has no opaque-object API ({e})") + warn_compile_disabled(f"this PyTorch build has no opaque-object API ({e})") return try: @@ -153,7 +131,7 @@ def register_value_opaque_quantizer(cls: type) -> None: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. - warn_compile_unsupported(f"could not register {cls.__name__} as an opaque type ({e})") + warn_compile_disabled(f"could not register {cls.__name__} as an opaque type ({e})") return _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 219263773d..868aced1f9 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -8,7 +8,6 @@ from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op -import math import warnings import weakref @@ -50,6 +49,7 @@ nvtx_range_push, get_nvtx_range_context, warn_compile_eager_fallback, + check_gemm_dims, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -83,7 +83,12 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorSpec, register_custom_op, is_value_opaque_quantizer +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op, + is_value_opaque_quantizer, +) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -99,12 +104,6 @@ __all__ = ["Linear"] -# Fields with this union may hold a *bare* ``QuantizedTensorStorage`` (the -# internal-quantizer optimization; see its docstring), not just a plain / -# subclass tensor -- hence the union rather than a plain ``torch.Tensor``. -TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] - - @dataclass(slots=True) class LinearFwdArgs: """Single-argument bag for the forward path of :class:`_Linear`.""" @@ -294,7 +293,7 @@ class LinearBwdArgs: cpu_offloading: bool = False owns_input: bool = False - # --- Per-backward scratch state (populated inside _linear_backward) --- + # --- Per-backward scratch state (populated inside _linear_backward_impl) --- ub_obj_gradout: Optional[Any] = None def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: @@ -319,6 +318,35 @@ def _check_fp8_reduce_and_update(): return result +def _sp_out_leading(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Leading (sequence) dim of the output, given the input's. + + Under sequence parallelism a column-parallel layer gathers that dim and a + row-parallel one scatters it; without SP it passes through. + """ + if not args.sequence_parallel: + return leading + if args.parallel_mode == "column": + return leading * args.tp_size + if args.parallel_mode == "row": + return leading // args.tp_size + return leading + + +def _sp_inp_leading(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Inverse of :func:`_sp_out_leading`: input's leading dim from the output's. + + Used by backward, which reconstructs the input geometry from ``grad_output``. + """ + if not args.sequence_parallel: + return leading + if args.parallel_mode == "column": + return leading // args.tp_size + if args.parallel_mode == "row": + return leading * args.tp_size + return leading + + def _linear_forward_impl( args: LinearFwdArgs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: @@ -354,7 +382,7 @@ def _linear_forward_impl( debug = args.debug backward_override = args.backward_override is_fsdp2 = args.is_fsdp2 - backward_needs_input = is_grad_enabled and weight.requires_grad + backward_needs_input = is_grad_enabled and args.weight_requires_grad if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": @@ -396,7 +424,6 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and args.weight_requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -708,7 +735,7 @@ def _linear_forward_impl( elif wt_save is weight: wt_alias = "weight" elif new_weight_workspace is not None and wt_save is new_weight_workspace: - wt_alias = "new_workspace" + wt_alias = "new_weight_workspace" elif args.weight_workspace is not None and wt_save is args.weight_workspace: wt_alias = "weight_workspace" else: @@ -734,7 +761,7 @@ def _linear_forward_impl( return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs -def _linear_forward_impl_fake( +def _linear_forward_fake( args: LinearFwdArgs, ) -> Tuple[TensorSpec, Optional[TensorSpec], Optional[Tuple[Any, ...]], Optional[Dict]]: """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile, @@ -853,11 +880,7 @@ def _linear_forward_impl_fake( # ------------------------------------------------------ # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). # ------------------------------------------------------ - out_leading = inp.shape[0] - if args.parallel_mode == "column" and args.sequence_parallel: - out_leading = out_leading * args.tp_size - elif args.parallel_mode == "row" and args.sequence_parallel: - out_leading = out_leading // args.tp_size + out_leading = _sp_out_leading(inp.shape[0], args) out = TensorSpec( shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), dtype=activation_dtype, @@ -917,7 +940,7 @@ def _linear_forward_impl_fake( elif args.is_fsdp2: pass # FSDP2 re-quantizes from the gathered weight in backward. elif weightmat_is_storage and new_weight_workspace is not None: - wt_alias = "new_workspace" + wt_alias = "new_weight_workspace" elif weightmat_is_storage and args.weight_workspace is not None: wt_alias = "weight_workspace" elif weightmat_is_storage: @@ -1056,7 +1079,7 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight - elif wt_save_alias == "new_workspace": + elif wt_save_alias == "new_weight_workspace": wt_save = fwd_outputs[1] elif wt_save_alias == "weight_workspace": wt_save = fwd_args.weight_workspace @@ -1067,7 +1090,7 @@ def _linear_setup_ctx( return (saved_inputmat, wt_save, saved_weight, saved_bias) -def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: +def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: """Backward implementation for the linear layer. Caller must have populated ``args.grad_output`` and run @@ -1142,13 +1165,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). if bwd_args.inp_shape is None: in_features = saved_weight.shape[-1] - go_leading = grad_output.shape[0] - if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: - inp_leading = go_leading // bwd_args.tp_size - elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: - inp_leading = go_leading * bwd_args.tp_size - else: - inp_leading = go_leading + inp_leading = _sp_inp_leading(grad_output.shape[0], bwd_args) bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) # Configure Userbuffers communication (comm+GEMM overlap) @@ -1680,10 +1697,10 @@ def wgrad_gemm( ) -def _linear_backward_impl_fake( +def _linear_backward_fake( args: LinearBwdArgs, ) -> Tuple[Optional[TensorSpec], Optional[TensorSpec], Optional[TensorSpec]]: - """Allocation-free fake of :func:`_linear_backward` on ``TensorSpec``. + """Allocation-free fake of :func:`_linear_backward_impl` on ``TensorSpec``. The saved-tensor fields of ``args`` carry :class:`~transformer_engine.pytorch.dynamo.TensorSpec` instances. Returns @@ -1705,7 +1722,7 @@ def _linear_backward_impl_fake( out_dtype = args.activation_dtype out_features, in_features = weight.shape - # Mirror ``_linear_backward``: ``set_usage`` on ``grad_input_quantizer`` + # Mirror ``_linear_backward_impl``: ``set_usage`` on ``grad_input_quantizer`` # influences ``dgrad``'s buffer layout. if args.grad_input_quantizer is not None: args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) @@ -1716,13 +1733,7 @@ def _linear_backward_impl_fake( # Derive shape from grad_output + weight + SP config instead of args.inp_shape: # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is # not hashable in OpaqueValueBundle), so we reconstruct it here. - go_leading = args.grad_output.shape[0] - if args.parallel_mode == "column" and args.sequence_parallel: - dgrad_leading = go_leading // args.tp_size - elif args.parallel_mode == "row" and args.sequence_parallel: - dgrad_leading = go_leading * args.tp_size - else: - dgrad_leading = go_leading + dgrad_leading = _sp_inp_leading(args.grad_output.shape[0], args) dgrad = TensorSpec( shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), dtype=out_dtype, @@ -1758,11 +1769,11 @@ def _linear_backward_impl_fake( input_tensors_for_grad=["weight", "inp", "bias"], fwd_arg_type=LinearFwdArgs, fwd_impl=_linear_forward_impl, - fwd_fake_impl=_linear_forward_impl_fake, + fwd_fake_impl=_linear_forward_fake, setup_context=_linear_setup_ctx, - backward_arg_type=LinearBwdArgs, - backward_impl=_linear_backward, - bwd_fake_impl=_linear_backward_impl_fake, + bwd_arg_type=LinearBwdArgs, + bwd_impl=_linear_backward_impl, + bwd_fake_impl=_linear_backward_fake, ) @@ -1836,7 +1847,7 @@ def backward( nvtx_label = "transformer_engine._Linear.backward" if bwd_args.ub_name is not None: nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" - result = _linear_backward(bwd_args) + (None,) # fwd_args grad slot + result = _linear_backward_impl(bwd_args) + (None,) # fwd_args grad slot reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, # main_grad closure) so they don't outlive backward via ctx under retain_graph. @@ -2379,30 +2390,7 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad - torch._check( - inp.shape[-1] == weight_tensor.shape[-1], - lambda: "GEMM not possible: input last dim must equal in_features", - ) - if self.fp8: - torch._check( - math.prod(inp.shape[:-1]) % 8 == 0, - lambda: ( - "FP8 execution requires the product of all input dimensions except" - " the last to be divisible by 8" - ), - ) - torch._check( - inp.shape[-1] % 16 == 0, - lambda: "FP8 execution requires the input last dimension to be divisible by 16", - ) - torch._check( - weight_tensor.shape[0] % 16 == 0, - lambda: "FP8 execution requires out_features to be divisible by 16", - ) - torch._check( - weight_tensor.shape[1] % 16 == 0, - lambda: "FP8 execution requires in_features to be divisible by 16", - ) + check_gemm_dims(inp, weight_tensor, self.fp8) linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index bd9231c460..47a7db9f7e 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,6 +26,30 @@ ] +_warned_compile_disabled = False + + +def warn_compile_disabled(reason: str) -> None: + """Warn once per process that TE's torch.compile custom-op path is off. + + Registration of the torch.compile machinery either works or fails as a + whole, so one message is enough. Distinct from + :func:`warn_compile_eager_fallback`, which reports a single *configuration* + falling back while the path itself is available. + """ + global _warned_compile_disabled # pylint: disable=global-statement + if _warned_compile_disabled: + return + _warned_compile_disabled = True + warnings.warn( + "Transformer Engine torch.compile support is disabled: " + f"{reason}. Modules will fall back to eager execution under " + "torch.compile, i.e. a graph break, which is incompatible with " + "fullgraph=True. Use a newer PyTorch build.", + stacklevel=3, + ) + + def warn_compile_eager_fallback(reason: str) -> None: """Warn that a TE module is running eagerly under ``torch.compile``. @@ -638,6 +662,32 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: ) +def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: + """Validate the dims of a TN GEMM pair (``y = x @ w^T``) for ``inp``/``weight``. + + The torch.compile-friendly counterpart of :func:`assert_dim_for_fp8_exec`: + uses ``torch._check`` so under dynamic shapes the constraints become guards + instead of being silently baked into the trace. + """ + # pylint: disable=protected-access + torch._check( + inp.shape[-1] == weight.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + if not fp8: + return + for ok, requirement in ( + ( + math.prod(inp.shape[:-1]) % 8 == 0, + "the product of all input dims except the last to be divisible by 8", + ), + (inp.shape[-1] % 16 == 0, "the input last dim to be divisible by 16"), + (weight.shape[0] % 16 == 0, "out_features to be divisible by 16"), + (weight.shape[1] % 16 == 0, "in_features to be divisible by 16"), + ): + torch._check(ok, lambda r=requirement: f"FP8 execution requires {r}") + + def is_bf16_compatible() -> bool: """Replaces torch.cuda.is_bf16_compatible() with an explicit check on device compute capability to enforce sm_80 or higher. From 250bd71f08a00f8e65486b76fdf4761721e86012 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:01:26 +0000 Subject: [PATCH 05/50] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/dynamo/custom_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index bde9cca351..d22c97bb25 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1261,7 +1261,7 @@ def call(args: Sequence[Any]) -> List[torch.Tensor]: def _make_dispatch_rule( - forward: Callable[[Sequence[Any]], List[torch.Tensor]] + forward: Callable[[Sequence[Any]], List[torch.Tensor]], ) -> Callable[..., Any]: """Adapt a slot forwarder to the ``register_torch_dispatch`` signature.""" From 39d10f8c57ffe0366c417ebd1dbc98976a335342 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 16:19:23 +0200 Subject: [PATCH 06/50] Address review: fix recompile assert, empty-batch sentinel collision, dim checks and cleanups - check_gemm_dims: restore assert_dim_for_fp8_exec semantics (per-tensor leading%8 / last%16, out_features%8 not %16); rich error messages with dims on the eager path, constant torch._check messages under compile (Dynamo forbids tensor closures in _check message lambdas). - test_te_linear_dynamic_shapes: the recompile assertion compared a nonexistent counter (always 0==0); use stats/unique_graphs and absorb the one-time lazy is_fsdp2 hasattr-guard recompile with a warmup. - custom_op: None-sentinel dtype uint8 -> complex32; a genuinely empty FP8 uint8 buffer (batch=0) decoded as None and broke compilation. - OpaqueValueBundle: type-tag _to_hashable (list/tuple/Size no longer compare equal), guard __getattr__ against copy/pickle recursion on underscored probes, render non-finite floats evaluably in __fx_repr__. - Linear.forward: fetch the cuBLAS workspace only after the eager-fallback decision; explicit torch._dynamo.graph_break(msg=...) so fullgraph=True errors carry the fallback reason instead of breaking on warnings.warn. - warn_compile_disabled: move the 'use a newer PyTorch build' advice to the version-related call sites only. - Comment/docstring/typography/pylint-disable cleanups in custom_op; test cosmetics (use_compile arg name, argparse-time validation of --compile/--use-cuda-graphs, merged NVINSPECT skips, docstring fixes); export get_cublas_workspace from cpp_extensions. Signed-off-by: Pawel Gadzinski --- .../distributed/run_layer_with_overlap.py | 6 +- tests/pytorch/distributed/run_numerics.py | 9 +-- .../distributed/test_comm_gemm_overlap.py | 6 +- tests/pytorch/test_torch_compile.py | 33 ++++++---- .../pytorch/cpp_extensions/gemm.py | 1 + .../pytorch/dynamo/custom_op.py | 64 +++++++++++++------ .../pytorch/dynamo/quantizer_opaque.py | 4 +- transformer_engine/pytorch/module/linear.py | 26 ++++---- transformer_engine/pytorch/utils.py | 35 ++++++---- 9 files changed, 111 insertions(+), 73 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index e65824ce85..e99c99c739 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -298,6 +298,9 @@ def _parse_args(argv=None, namespace=None): ) args = parser.parse_args(argv, namespace) + if args.compile and args.use_cuda_graphs: + parser.error("--compile and --use-cuda-graphs are mutually exclusive.") + if args.use_cuda_graphs and args.layer_type in [te.MultiheadAttention, te.TransformerLayer]: warnings.warn(f"{args.layer_type.__name__} does not support CUDA Graphs!") args.use_cuda_graphs = False @@ -498,9 +501,6 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0) dist_print("Copied parameters from test model to reference model...", debug=True) - if opts.compile and opts.use_cuda_graphs: - raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.") - # Fp8 recipe setup fp8_format = Format.HYBRID fp8_recipe = None diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 1590979ba3..319c46d75d 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -783,10 +783,11 @@ def test_linear(): for kwargs in kwargs_list: if kwargs.get("save_original_input", False) and QUANTIZATION == "fp8": continue - if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: - continue - # debug instrumentation forces the eager fallback, so compile is a no-op there. - if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED: + # use_compile: debug instrumentation forces the eager fallback, so + # compile is a no-op there. + if NVTE_TEST_NVINSPECT_ENABLED and ( + kwargs.get("delay_wgrad_compute", False) or kwargs.get("use_compile", False) + ): continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6521ed7dbc..358b948d61 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -111,7 +111,7 @@ def _run_layer_with_overlap( quantization, num_layers=1, use_cublasmp=False, - compile=False, + use_compile=False, compile_mode="default", ): test_path = TEST_ROOT / "run_layer_with_overlap.py" @@ -131,7 +131,7 @@ def _run_layer_with_overlap( if overlap_rs_dgrad: test_cmd.append("--overlap-rs-dgrad") - if compile: + if use_compile: test_cmd.append("--compile") test_cmd.append(f"--compile-mode={compile_mode}") @@ -316,7 +316,7 @@ def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, com overlap_rs_dgrad, False, None, - compile=True, + use_compile=True, compile_mode=compile_mode, ) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index a1935f2d93..91e931e7dd 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -101,9 +101,7 @@ def nvfp4_4over6(): def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: - """ - Force TE's lazily-created global scratch to be allocated before capture. - """ + """Force TE's lazily-created global scratch to be allocated before capture.""" out = fn(inp) if backward: out.sum().backward() @@ -1449,14 +1447,14 @@ def fn(inp): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") def test_te_linear_dynamic_shapes(): - """torch.compile(dynamic=True) of ``te.Linear`` with varying batch sizes. + """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension. Verifies that the compiled graph handles symbolic (dynamic) leading dimensions without graph breaks or recompilations after the initial trace. Key correctness property: a graph compiled for batch=16 must produce numerically correct results for batch=32 without triggering a recompile. - This exercises two fixes for dynamic shapes: + This exercises three fixes for dynamic shapes: 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). 2. ``_linear_backward_fake`` derives dgrad shape from grad_output + @@ -1480,7 +1478,19 @@ def fn(inp): batch_sizes = [16, 32, 48] - for i, batch in enumerate(batch_sizes): + # Warm up with two calls at the first batch size: the first call traces; the + # second absorbs the one-time recompile caused by module attributes lazily + # created during call one (the cached ``is_fsdp2``), which flip a ``hasattr`` + # guard and are unrelated to dynamic shapes. Baseline the graph count after + # that -- any further recompile across batch sizes is a real failure. + for _ in range(2): + warm = torch.randn(batch_sizes[0], in_features, dtype=dtype, device=device) + torch._dynamo.mark_dynamic(warm, 0) + compiled(warm.requires_grad_(True)).sum().backward() + model.zero_grad(set_to_none=True) + unique_graphs_baseline = counters["stats"]["unique_graphs"] + + for batch in batch_sizes: inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes. torch._dynamo.mark_dynamic(inp, 0) @@ -1510,13 +1520,8 @@ def fn(inp): msg=f"dgrad mismatch at batch={batch}", ) - if i == 0: - # After the first (tracing) call, record the recompile counter - # baseline -- subsequent batch sizes must not trigger recompiles. - recompile_count_baseline = counters["stats"].get("recompile_reasons", 0) - - recompile_count_after = counters["stats"].get("recompile_reasons", 0) - assert recompile_count_after == recompile_count_baseline, ( + unique_graphs_after = counters["stats"]["unique_graphs"] + assert unique_graphs_after == unique_graphs_baseline, ( "Unexpected recompilation(s) across different batch sizes: " - f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" + f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3d97b7269..03f31e28b5 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -31,6 +31,7 @@ "general_gemm", "general_grouped_gemm", "general_grouped_gemm_for_grouped_tensor", + "get_cublas_workspace", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index d22c97bb25..d33006c74c 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -84,6 +84,7 @@ from __future__ import annotations import dataclasses +import math import types as _types # aliased: torch_dispatch rules take a ``types`` param from enum import Enum from typing import ( @@ -123,17 +124,21 @@ # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a -# 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for -# ``register_autograd`` to attach a ``grad_fn`` to the outputs. +# 0-element sentinel tensor: a non-nullable ``Tensor[]`` schema is required for +# ``register_autograd`` to attach a ``grad_fn`` to the outputs. The sentinel is +# recognized by (numel == 0, dtype), so its dtype must be one no real payload +# tensor can have -- complex is never a TE payload (quantized data / scales are +# uint8 / fp32, outputs are float), while e.g. a uint8 sentinel would collide +# with a genuinely empty FP8 data buffer (empty batch). # # Once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable # ``Tensor?[]`` return schema will let ``None`` pass through directly and this # sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed. -_NONE_SENTINEL_DTYPE = torch.uint8 +_NONE_SENTINEL_DTYPE = torch.complex32 def _encode_none(t: Optional[torch.Tensor]) -> torch.Tensor: - """Replace ``None`` with a 0-element uint8 sentinel tensor.""" + """Replace ``None`` with a 0-element sentinel tensor.""" if t is None: return torch.empty(0, dtype=_NONE_SENTINEL_DTYPE) return t @@ -195,11 +200,13 @@ def is_simple_value(cls, value: Any) -> bool: @classmethod def _to_hashable(cls, value: Any) -> Any: + # Tag with the concrete type so e.g. [1] / (1,) / Size([1]) or True / 1 + # stay distinct under __eq__ / __hash__ (graph guards compare bundles). if isinstance(value, dict): - return tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items())) - if isinstance(value, (list, tuple, torch.Size)): - return tuple(cls._to_hashable(v) for v in value) - return value + return ("dict", tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items()))) + if isinstance(value, (list, tuple)): # incl. torch.Size + return (type(value).__name__, tuple(cls._to_hashable(v) for v in value)) + return (type(value).__name__, value) @classmethod def _fmt_simple(cls, value: Any) -> str: @@ -228,6 +235,9 @@ def _fmt_simple(cls, value: Any) -> str: return f"({body},)" if len(value) == 1 else f"({body})" if _is_opaque_value_type(type(value)): return value.__fx_repr__()[0] + # repr(float('inf')) is 'inf', which is not an evaluable literal. + if isinstance(value, float) and not math.isfinite(value): + return f"float({str(value)!r})" return repr(value) def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: @@ -249,6 +259,10 @@ def __getitem__(self, key: str) -> Any: return self._data[key] def __getattr__(self, name: str) -> Any: + # Underscored names raise cleanly: copy/pickle probe dunders on a clone + # created without __init__, where reading ``self._data`` would recurse. + if name.startswith("_"): + raise AttributeError(name) try: return self._data[name] except KeyError as e: @@ -303,7 +317,7 @@ def _collect(value: Any) -> None: try: - from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + from torch._library.opaque_object import ( get_opaque_type_name, is_opaque_value_type as _is_opaque_value_type, is_opaque_reference_type as _is_opaque_reference_type, @@ -314,14 +328,16 @@ def _collect(value: Any) -> None: _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) # Older torch without opaque_object support. except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - warn_compile_disabled(f"could not register OpaqueValueBundle as an opaque type ({e})") + warn_compile_disabled( + f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build" + ) _is_opaque_value_type = None _is_opaque_reference_type = None _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None def _pg_pickle_stub(*args: Any) -> None: # pragma: no cover - raise RuntimeError("ProcessGroup cannot be unpickled — cache-key use only") + raise RuntimeError("ProcessGroup cannot be unpickled -- cache-key use only") def _ensure_distributed_opaque_types() -> None: @@ -337,16 +353,18 @@ def _ensure_distributed_opaque_types() -> None: process-group field simply falls back to eager under torch.compile. Also registers a ``copyreg`` reducer that lets ``FxGraphCachePickler`` hash - graphs containing a ``ProcessGroup`` input without crashing. Without this, + graphs containing a ``ProcessGroup`` input without crashing. Without this, inductor logs "Failed to pickle cache key" warnings and bypasses the FX - graph disk cache for every distributed compiled call. The reducer encodes - the group as (world_size, rank, backend) — enough to distinguish configs — + graph disk cache for every distributed compiled call. The reducer encodes + the group as (world_size, rank, backend) -- enough to distinguish configs -- and raises on reconstruct since deserialization is never needed for hashing. """ if _is_opaque_reference_type is None: return - try: # pylint: disable=import-outside-toplevel - from torch.distributed.device_mesh import _register_distributed_opaque_types + try: + from torch.distributed.device_mesh import ( # pylint: disable=import-outside-toplevel + _register_distributed_opaque_types, + ) _register_distributed_opaque_types() except Exception: # pylint: disable=broad-exception-caught @@ -354,12 +372,15 @@ def _ensure_distributed_opaque_types() -> None: # Workaround for PyTorch issue: FxGraphCachePickler handles FakeScriptObject # but not the real ProcessGroup that appears in example_inputs at inductor - # compile time. Register a copyreg reducer so the pickler can hash the key. - try: # pylint: disable=import-outside-toplevel + # compile time. Register a copyreg reducer so the pickler can hash the key. + try: + # pylint: disable=import-outside-toplevel import copyreg import torch.distributed as dist from torch._C._distributed_c10d import ProcessGroup + # pylint: enable=import-outside-toplevel + if ProcessGroup not in copyreg.dispatch_table: def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] @@ -460,8 +481,9 @@ def try_build(cls, name: str, annot: Any) -> Optional["_Adapter"]: type annotation ``annot``; return a configured adapter if so, else ``None`` so the next candidate is tried. - Called once per field at registration, in :data:`_FIELD_ADAPTERS` - priority order. + Called once per field at registration, iterating :data:`_FIELD_ADAPTERS` + (adapters are mutually exclusive on annotations, so the order is not a + ranking). """ raise NotImplementedError @@ -601,7 +623,7 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = _storage_unflatten(meta, args[self.inner_slot()]) def grad_slot(self) -> Optional[int]: - # Gradient flows to the plain / subclass tensor slot (``slot_name``, + # Gradient flows to the plain / subclass tensor slot (``tensor_slot()``, # the first of the three). return 0 diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index ae8849097c..856dee2022 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -121,7 +121,9 @@ def register_value_opaque_quantizer(cls: type) -> None: except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. - warn_compile_disabled(f"this PyTorch build has no opaque-object API ({e})") + warn_compile_disabled( + f"this PyTorch build has no opaque-object API ({e}); use a newer build" + ) return try: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 868aced1f9..f6ee48dd6e 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -71,8 +71,8 @@ ) from ..cpp_extensions import ( general_gemm, + get_cublas_workspace, ) -from ..cpp_extensions.gemm import get_cublas_workspace from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..graph import is_graph_capturing from ..jit import no_torch_dynamo @@ -120,11 +120,9 @@ class LinearFwdArgs: # across the torch.compile custom-op boundary. weight_workspace: Optional[TensorOrQuantized] - # Workspace pinning (torch.compile). The process-global, lru_cached cuBLAS - # workspace is fetched in the traced forward and threaded in as an op input so - # it is allocated at trace time rather than lazily inside the op. The op body - # never reads it (general_gemm fetches the same global by address). None on - # eager / non-compiled paths. + # Process-global cuBLAS workspace, fetched at trace time so it isn't first + # allocated during CUDA-graph capture. The op never reads it (general_gemm + # fetches the same global); None outside the compiled path. cublas_workspace: Optional[torch.Tensor] # --- requires_grad flags (cached so backward does not re-query) --- @@ -2397,20 +2395,13 @@ def forward( ) wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None - # Pin the lazily-cached cuBLAS workspace as an op input so it is - # materialized at trace time (external to the cudagraph pool) rather - # than inside the op during capture. See LinearFwdArgs for details. - cublas_workspace = None - if use_compiled_op: - cublas_workspace = get_cublas_workspace(inp.device.index, False, False) - fwd_args = LinearFwdArgs( # tensors weight=weight_tensor, inp=inp, bias=linear_bias_tensor, weight_workspace=weight_workspace, - cublas_workspace=cublas_workspace, + cublas_workspace=None, # set below, only on the compiled path # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, @@ -2469,10 +2460,17 @@ def forward( if use_compiled_op: fallback_reason = fwd_args.compile_unsupported_reason() if fallback_reason is not None: + # Explicit break so fullgraph=True errors show the reason + # (warnings.warn below would break the graph inscrutably). + torch._dynamo.graph_break( + msg=f"te.Linear falling back to eager: {fallback_reason}" + ) warn_compile_eager_fallback(fallback_reason) use_compiled_op = False if use_compiled_op: + # See LinearFwdArgs.cublas_workspace. + fwd_args.cublas_workspace = get_cublas_workspace(inp.device.index, False, False) out, new_weight_workspace = _linear_op(fwd_args) else: out, new_weight_workspace = _linear_eager( diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 47a7db9f7e..712895204c 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -45,7 +45,7 @@ def warn_compile_disabled(reason: str) -> None: "Transformer Engine torch.compile support is disabled: " f"{reason}. Modules will fall back to eager execution under " "torch.compile, i.e. a graph break, which is incompatible with " - "fullgraph=True. Use a newer PyTorch build.", + "fullgraph=True.", stacklevel=3, ) @@ -666,9 +666,20 @@ def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: """Validate the dims of a TN GEMM pair (``y = x @ w^T``) for ``inp``/``weight``. The torch.compile-friendly counterpart of :func:`assert_dim_for_fp8_exec`: - uses ``torch._check`` so under dynamic shapes the constraints become guards - instead of being silently baked into the trace. + under compile it uses ``torch._check`` so with dynamic shapes the + constraints become guards instead of being silently baked into the trace. + Dynamo forbids tensor closures in ``torch._check`` message lambdas, so the + compiled-path messages omit the dims; eager keeps the full messages. """ + if not torch.compiler.is_compiling(): + if inp.shape[-1] != weight.shape[-1]: + raise ValueError( + "GEMM not possible: input last dim must equal in_features, but got " + f"input dims={list(inp.shape)} and weight dims={list(weight.shape)}" + ) + if fp8: + assert_dim_for_fp8_exec(inp, weight) + return # pylint: disable=protected-access torch._check( inp.shape[-1] == weight.shape[-1], @@ -676,16 +687,14 @@ def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: ) if not fp8: return - for ok, requirement in ( - ( - math.prod(inp.shape[:-1]) % 8 == 0, - "the product of all input dims except the last to be divisible by 8", - ), - (inp.shape[-1] % 16 == 0, "the input last dim to be divisible by 16"), - (weight.shape[0] % 16 == 0, "out_features to be divisible by 16"), - (weight.shape[1] % 16 == 0, "in_features to be divisible by 16"), - ): - torch._check(ok, lambda r=requirement: f"FP8 execution requires {r}") + for tensor, name in ((inp, "input"), (weight, "weight")): + torch._check( + math.prod(tensor.shape[:-1]) % 8 == 0 and tensor.shape[-1] % 16 == 0, + lambda n=name: ( + f"FP8 execution requires the {n}'s product of all dimensions except the" + " last to be divisible by 8 and its last dimension to be divisible by 16" + ), + ) def is_bf16_compatible() -> bool: From 2c373508dac994407ea6ee7c8b1bf1343a074b3c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 23:07:12 +0200 Subject: [PATCH 07/50] Address review: simplify check_gemm_dims, trim test comments, restore eager dim asserts - check_gemm_dims is now a compile-only torch._check guard emitter, called from the compiled-op branch; eager dim validation returns to the op impl (assert + assert_dim_for_fp8_exec, as on main) so eager pays no overhead and keeps full error messages with dims. - Trim verbose test docstrings/comments (te.Linear section, warmup helper, cudagraph-skip helper); describe the dynamic-shape scope (leading dims) instead of the fix history. - Drop the stale 'FP8 with symbolic shapes unsupported' comments: FP8 with a mark_dynamic batch works on current nightly (verified: one graph reused across batch sizes, numerics match eager). Signed-off-by: Pawel Gadzinski --- .../distributed/run_layer_with_overlap.py | 3 +- tests/pytorch/distributed/run_numerics.py | 9 +- .../distributed/test_comm_gemm_overlap.py | 10 +- tests/pytorch/test_torch_compile.py | 104 +++++------------- transformer_engine/pytorch/module/linear.py | 10 +- transformer_engine/pytorch/utils.py | 22 +--- 6 files changed, 44 insertions(+), 114 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index e99c99c739..fb5579d583 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -553,8 +553,7 @@ def run_fwd_bwd(model, x): if opts.compile: for i, layer in enumerate(test_model.layers): - # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle - # op arg whose hash chokes on non-nested SymInt (see run_numerics). + # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py. test_model.layers[i] = torch.compile( layer, fullgraph=True, mode=opts.compile_mode, dynamic=False ) diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 319c46d75d..d9dccc6f0a 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -324,14 +324,9 @@ def _apply_models( forward_single_node = model_single_node forward_distributed = model_distributed if use_compile: - # Each parametrized case compiles the same module.forward code object with - # a different shape/recipe; with dynamic=False those guards accumulate and - # eventually trip Dynamo's recompile_limit. Reset so every case starts from - # a clean compile cache (mirrors the single-GPU torch.compile tests). + # Reset the compile cache so parametrized cases don't trip recompile_limit. torch._dynamo.reset() - # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle - # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static - # shapes (recompile per shape) until the bundle handles symbolic shapes. + # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py. forward_single_node = torch.compile( model_single_node, fullgraph=True, mode=compile_mode, dynamic=False ) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 358b948d61..25331ba85a 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -302,14 +302,8 @@ def test_layers_with_overlap_bf16( ], ) def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode): - """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16). - - Userbuffers is expected to stay on Linear's compiled custom-op path (the - collective lives inside the opaque op), so this checks that torch.compile + - Userbuffers stays numerically correct against the eager, non-overlap reference. - ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on - top of the Userbuffers collectives. - """ + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16), + checked numerically against the eager, non-overlap reference.""" _run_layer_with_overlap( te.Linear.__name__, linear_parallel_mode, diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 91e931e7dd..3336f44f22 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -94,14 +94,13 @@ def nvfp4_4over6(): _all_recipes.append(nvfp4_row_scaled()) -# torch.compile modes exercised by the te.Linear tests: the default backend and -# "reduce-overhead" (CUDA-graph trees), to ensure the custom-op path is -# CUDA-graph capturable. +# Modes exercised by the te.Linear tests; "reduce-overhead" = CUDA-graph trees. _compile_modes = ["default", "reduce-overhead"] def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: - """Force TE's lazily-created global scratch to be allocated before capture.""" + """One eager iteration so lazily-initialized TE state (fp8 meta, workspaces) + is allocated before any CUDA-graph capture.""" out = fn(inp) if backward: out.sum().backward() @@ -109,16 +108,9 @@ def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: @contextlib.contextmanager def _assert_no_cudagraph_skips(enabled: bool): - """Assert ``torch.compile(mode="reduce-overhead")`` actually captured CUDA - graphs for every graph instead of silently running it eagerly. - - Inductor bumps ``counters["inductor"]["cudagraph_skips"]`` whenever it - declines to capture a cudagraph (input mutation, CPU scalars, cudagraph-unsafe - ops, ...) and falls back to eager for that graph. ``fullgraph=True`` only rules - out *dynamo* graph breaks, not these *inductor*-level skips, so this guards that - the reduce-overhead path didn't degrade to eager. No-op when ``enabled`` is - False (e.g. the default backend, where cudagraphs don't apply). - """ + """Assert reduce-overhead really captured CUDA graphs: inductor may skip + capture and silently fall back to eager, which ``fullgraph=True`` does not + catch. No-op when ``enabled`` is False.""" before = counters["inductor"]["cudagraph_skips"] yield if enabled: @@ -129,19 +121,13 @@ def _assert_no_cudagraph_skips(enabled: bool): ) -# bf16 output tolerance: eager and compiled run the same kernels, so they should -# agree closely; the slack only absorbs reduction-order / cuda-graph differences. +# Eager and compiled run the same kernels; slack only for reduction-order noise. _EAGER_ATOL, _EAGER_RTOL = 1e-2, 1.6e-2 def _assert_close_eager_compiled(fn, compiled, model, base): """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the - forward output and the input / weight gradients match. - - Guards the compiled custom-op path against silently diverging from eager - execution -- a wrong-but-same-shape result would slip past shape / grad - presence checks alone. - """ + forward output and the input / weight gradients match.""" inp_eager = base.detach().clone().requires_grad_(True) model.zero_grad(set_to_none=True) out_eager = fn(inp_eager) @@ -1295,8 +1281,7 @@ def fn(inp): model.zero_grad(set_to_none=True) compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) - # ``reduce-overhead`` warms up on the first call(s) and replays a captured - # CUDA graph afterwards, so iterate a few times to actually exercise replay. + # Iterate a few times so reduce-overhead actually replays a captured graph. n_iters = 3 if compile_mode == "reduce-overhead" else 1 with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): for _ in range(n_iters): @@ -1308,13 +1293,8 @@ def fn(inp): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("compile_mode", _compile_modes) def test_te_linear_compile_with_quantized_fp8_weight(compile_mode): - """torch.compile should handle Linear weights initialized as FP8 tensors, - for both the default backend and ``mode="reduce-overhead"``. - - Exercises the two-tier op + ``register_torch_dispatch`` flattening of a - ``Float8Tensor`` weight *input* in - :mod:`transformer_engine.pytorch.dynamo`. - """ + """torch.compile of Linear with the weight initialized as an FP8 tensor + (exercises the wrapper op's ``register_torch_dispatch`` input flattening).""" dtype = torch.bfloat16 device = "cuda" fp8_recipe = recipe.Float8CurrentScaling() @@ -1349,18 +1329,10 @@ def fn(inp): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("compile_mode", _compile_modes) def test_te_linear_compile_with_fp8_output(compile_mode): - """torch.compile of ``te.Linear(..., fp8_output=True)`` without gradient: - forward returns a :class:`Float8Tensor`. Covers the default backend and - ``mode="reduce-overhead"``. - - Exercises the output-rewrap path in - :mod:`transformer_engine.pytorch.dynamo`: when an output quantizer is - active, the op returns the flat inner data tensors and the framework - rewraps them into a ``Float8Tensor`` via ``__tensor_unflatten__``. A - differentiable FP8 output is unsupported under compile (``Linear.forward`` - falls back to eager), so this test covers the supported case: an FP8 output - that does not require grad (inference / ``torch.no_grad``). - """ + """torch.compile of ``te.Linear(..., fp8_output=True)`` under no_grad: + forward must return a working :class:`Float8Tensor` (exercises the output + rewrap path). The differentiable case falls back to eager, so it is not + covered here.""" dtype = torch.bfloat16 device = "cuda" fp8_recipe = recipe.Float8CurrentScaling() @@ -1391,12 +1363,9 @@ def fn(inp): assert ( out._quantizer is not None ), "FP8 output lost its quantizer on the torch.compile path" - # The rewrap rebuilt a fully-functional Float8Tensor: dequantizing it - # outside the compiled region exercises scale + data + dtype wiring. deq = out.dequantize() assert deq.shape == (32, 32) assert deq.dtype == dtype - # Compiled FP8 output must match the eager FP8 output value-wise. torch.testing.assert_close( deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL ) @@ -1406,16 +1375,9 @@ def fn(inp): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("compile_mode", _compile_modes) def test_te_linear_compile_is_first_microbatch(compile_mode): - """torch.compile of ``te.Linear`` across a multi-step microbatch schedule that - drives FP8 weight caching via ``is_first_microbatch``, for the default backend - and ``mode="reduce-overhead"`` (CUDA-graph trees). - - ``is_first_microbatch=True`` quantizes and caches the FP8 weight; subsequent - ``False`` steps must reuse the cached FP8 weight instead of re-quantizing. This - exercises that cache path under compile and checks it stays numerically aligned - with eager. ``is_first_microbatch`` is a Python bool, so each distinct value is - its own dynamo guard/graph. - """ + """torch.compile of ``te.Linear`` across a microbatch schedule: + ``is_first_microbatch=True`` caches the FP8 weight, later steps must reuse + it and stay numerically aligned with eager.""" dtype = torch.bfloat16 device = "cuda" fp8_recipe = recipe.Float8CurrentScaling() @@ -1447,23 +1409,12 @@ def fn(inp): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") def test_te_linear_dynamic_shapes(): - """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension. - - Verifies that the compiled graph handles symbolic (dynamic) leading - dimensions without graph breaks or recompilations after the initial trace. - Key correctness property: a graph compiled for batch=16 must produce - numerically correct results for batch=32 without triggering a recompile. - - This exercises three fixes for dynamic shapes: - 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle - (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). - 2. ``_linear_backward_fake`` derives dgrad shape from grad_output + - weight + SP config instead of relying on the stored ``inp_shape``. - 3. ``_linear_backward_impl`` reconstructs ``inp_shape`` on-the-fly from the same - tensor sources when it is None (compiled mode). - - FP8 + dynamic=True is tracked separately (requires resolving - ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo). + """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension: + one graph must serve all batch sizes -- no recompiles -- and match eager + numerically. + + Only the leading (batch/sequence) dims may be dynamic; the last dim is + fixed by the weight's ``in_features``. """ dtype = torch.bfloat16 device = "cuda" @@ -1478,11 +1429,8 @@ def fn(inp): batch_sizes = [16, 32, 48] - # Warm up with two calls at the first batch size: the first call traces; the - # second absorbs the one-time recompile caused by module attributes lazily - # created during call one (the cached ``is_fsdp2``), which flip a ``hasattr`` - # guard and are unrelated to dynamic shapes. Baseline the graph count after - # that -- any further recompile across batch sizes is a real failure. + # Two warmup calls: the second absorbs the one-time recompile from module + # attributes lazily created during call one (e.g. the cached ``is_fsdp2``). for _ in range(2): warm = torch.randn(batch_sizes[0], in_features, dtype=dtype, device=device) torch._dynamo.mark_dynamic(warm, 0) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index f6ee48dd6e..0244a06cd7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -45,6 +45,7 @@ divide, init_method_constant, needs_quantized_gemm, + assert_dim_for_fp8_exec, nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, @@ -418,7 +419,9 @@ def _linear_forward_impl( if ub_name is not None: nvtx_label = f"{nvtx_label}.{ub_name}" - out_features = weight.shape[0] + # Make sure input dimensions are compatible + out_features, in_features = weight.shape + assert inp.shape[-1] == in_features, "GEMM not possible" # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) @@ -450,6 +453,8 @@ def _linear_forward_impl( inputmat = inp # Input tensor to save for backward (maybe sharded) inputmat_total = None # Input tensor to pass to GEMM (gathered) own_quantized_input = False + if fp8: + assert_dim_for_fp8_exec(inputmat, weight) if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor @@ -2388,8 +2393,6 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad - check_gemm_dims(inp, weight_tensor, self.fp8) - linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) @@ -2469,6 +2472,7 @@ def forward( use_compiled_op = False if use_compiled_op: + check_gemm_dims(inp, weight_tensor, self.fp8) # See LinearFwdArgs.cublas_workspace. fwd_args.cublas_workspace = get_cublas_workspace(inp.device.index, False, False) out, new_weight_workspace = _linear_op(fwd_args) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 712895204c..e5479a0fa5 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -663,23 +663,13 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: - """Validate the dims of a TN GEMM pair (``y = x @ w^T``) for ``inp``/``weight``. - - The torch.compile-friendly counterpart of :func:`assert_dim_for_fp8_exec`: - under compile it uses ``torch._check`` so with dynamic shapes the - constraints become guards instead of being silently baked into the trace. - Dynamo forbids tensor closures in ``torch._check`` message lambdas, so the - compiled-path messages omit the dims; eager keeps the full messages. + """Emit the TN GEMM (``y = x @ w^T``) dim constraints as ``torch._check`` + guards at trace time, so dynamic shapes stay constrained instead of being + silently baked into the graph. Call only on the torch.compile path; eager + validation (with the offending dims in the message) lives in the op impl + (:func:`assert_dim_for_fp8_exec`). Messages here are constant strings -- + Dynamo forbids tensor closures in ``torch._check`` message lambdas. """ - if not torch.compiler.is_compiling(): - if inp.shape[-1] != weight.shape[-1]: - raise ValueError( - "GEMM not possible: input last dim must equal in_features, but got " - f"input dims={list(inp.shape)} and weight dims={list(weight.shape)}" - ) - if fp8: - assert_dim_for_fp8_exec(inp, weight) - return # pylint: disable=protected-access torch._check( inp.shape[-1] == weight.shape[-1], From eb3f50fdac151930b0930e8e9d5d7fabf4f0e8cc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 23:11:14 +0200 Subject: [PATCH 08/50] Shorten check_gemm_dims docstring Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/utils.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index e5479a0fa5..dd74a9c401 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -664,11 +664,8 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: """Emit the TN GEMM (``y = x @ w^T``) dim constraints as ``torch._check`` - guards at trace time, so dynamic shapes stay constrained instead of being - silently baked into the graph. Call only on the torch.compile path; eager - validation (with the offending dims in the message) lives in the op impl - (:func:`assert_dim_for_fp8_exec`). Messages here are constant strings -- - Dynamo forbids tensor closures in ``torch._check`` message lambdas. + guards at trace time. torch.compile path only; eager validation lives in + the op impl. Messages are constant: Dynamo forbids tensor closures here. """ # pylint: disable=protected-access torch._check( From 9695013475ae07bd603152a458b015bb14b1c342 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 23:16:50 +0200 Subject: [PATCH 09/50] Drop tensor_can_be_materialized: inline an exact-class check in the two float8 reprs Signed-off-by: Pawel Gadzinski --- .../pytorch/tensor/_quantization_helpers.py | 37 ------------------- .../pytorch/tensor/float8_tensor.py | 9 ++--- .../tensor/storage/float8_tensor_storage.py | 14 +++---- 3 files changed, 9 insertions(+), 51 deletions(-) diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 6161faaddc..10672bbcfb 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -138,43 +138,6 @@ def _stride_from_shape(shape: list[int]): return list(reversed(rstride)) -def tensor_can_be_materialized(t) -> bool: - """Whether ``t`` holds concrete data that ``.item()`` / ``.tolist()`` can read - without side effects. - - A ``__repr__`` must never mutate tracing state. On a fake / meta / functional - tensor (torch.compile / export tracing) ``.item()`` does *not* raise -- it - silently allocates an *unbacked* SymInt/SymFloat into the active ShapeEnv, - which later crashes inductor with ``PendingUnbackedSymbolNotFound``. (torch's - AOTAutograd repr's the fake quantized tensor while logging graph metadata, so - a scalar-materializing ``__repr__`` leaks an unbacked symbol during compile.) - So detect those tensors and fall back to a metadata-only repr instead. - """ - if not isinstance(t, torch.Tensor): - return False - if getattr(t, "is_meta", False): - return False - try: - from torch._subclasses.fake_tensor import ( # pylint: disable=import-outside-toplevel - FakeTensor, - ) - - if isinstance(t, FakeTensor): - return False - except Exception: # pylint: disable=broad-except - pass - try: - from torch._subclasses.functional_tensor import ( # pylint: disable=import-outside-toplevel - FunctionalTensor, - ) - - if isinstance(t, FunctionalTensor): - return False - except Exception: # pylint: disable=broad-except - pass - return True - - def safe_quantized_repr(obj, cls_name, extras=None, error=None): """Metadata-only repr fallback for quantized tensors whose data cannot be materialized for any reason. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 2f89326750..2799ec3edf 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -23,7 +23,6 @@ _IdentityFunc, _resolve_view_shape, safe_quantized_repr, - tensor_can_be_materialized, ) from ..constants import dist_group_type, DType @@ -474,10 +473,10 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): amax_reduction_group: Optional[dist_group_type] = None def __repr__(self, *, tensor_contents=None): - # A fake/meta/functional scale_inv cannot be materialized without leaking - # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); - # fall back to a metadata-only repr under tracing. - if not tensor_can_be_materialized(self._scale_inv): + # Data-free repr for a fake/meta scale_inv (exact class check: fake / + # functional tensors subclass Tensor); materializing it under tracing + # would leak an unbacked symbol into the ShapeEnv. + if self._scale_inv.__class__ is not torch.Tensor or self._scale_inv.is_meta: return safe_quantized_repr(self, "Float8Tensor") try: return ( diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 1c3ce68c6d..89b6c2cfa9 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -11,11 +11,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer -from .._quantization_helpers import ( - _resolve_view_shape, - safe_quantized_repr, - tensor_can_be_materialized, -) +from .._quantization_helpers import _resolve_view_shape, safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -257,10 +253,10 @@ def view(self, shape: torch.Size): ) def __repr__(self): - # A fake/meta/functional scale_inv cannot be materialized without leaking - # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); - # fall back to a metadata-only repr under tracing. - if not tensor_can_be_materialized(self._scale_inv): + # Data-free repr for a fake/meta scale_inv (exact class check: fake / + # functional tensors subclass Tensor); materializing it under tracing + # would leak an unbacked symbol into the ShapeEnv. + if self._scale_inv.__class__ is not torch.Tensor or self._scale_inv.is_meta: return safe_quantized_repr(self, "Float8TensorStorage") try: return ( From c64a0b0ea7c64853e8c4ad388a4968eea3833f7e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 23:19:48 +0200 Subject: [PATCH 10/50] Trim paraphrase comments in the Linear fake impls Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 37 +++++---------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 0244a06cd7..49ea0ff032 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -826,8 +826,6 @@ def _linear_forward_fake( # ------------------------------------------------------ # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. - # ``new_weight_workspace`` is a fresh fake storage only on the - # cache-miss + ``cache_weight`` path, else ``None``. # ------------------------------------------------------ new_weight_workspace = None weightmat = None @@ -913,9 +911,7 @@ def _linear_forward_fake( quantizer=input_quantizer, device=inp.device, ) - # Mirror ``_linear_forward_impl``'s post-quantization - # ``inputmat.update_usage(...)`` so the saved input's buffer layout - # matches -- driven by the same conditions as the real impl. + # Mirror the impl's post-quantization ``update_usage``. if own_quantized_input and not save_original_input: if args.backward_override is not None: saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) @@ -932,10 +928,8 @@ def _linear_forward_fake( shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device ) - # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached - # FP8 weight is shared with ``new_weight_workspace`` (a return, on a cache - # miss) or the ``weight_workspace`` input (on a cache hit), so it is - # reconstructed in ``_linear_setup_ctx`` rather than saved twice. + # Slot 1 -- ``wt_save``, with the impl's alias dedup (rebuilt in + # ``_linear_setup_ctx`` instead of being saved twice). wt_alias = None wt_save = None if weightmat_aliases_weight: @@ -1705,15 +1699,8 @@ def _linear_backward_fake( ) -> Tuple[Optional[TensorSpec], Optional[TensorSpec], Optional[TensorSpec]]: """Allocation-free fake of :func:`_linear_backward_impl` on ``TensorSpec``. - The saved-tensor fields of ``args`` carry - :class:`~transformer_engine.pytorch.dynamo.TensorSpec` instances. Returns - ``(wgrad, dgrad, grad_bias)`` specs describing the nature of the gradients, - mirroring the real backward's return contract without allocating storage. - - Tensor-/sequence-parallel gather/scatter happens inside the eager backward - custom op and is opaque to ``torch.compile``: ``dgrad`` always carries the - rank-local input shape and ``wgrad`` the local weight shape, so no extra - shape modeling is needed here. + Returns ``(wgrad, dgrad, grad_bias)`` specs. TP/SP gather/scatter happens + inside the eager op, so the specs carry rank-local shapes. """ if args.fsdp_group is not None: raise NotImplementedError( @@ -1725,17 +1712,14 @@ def _linear_backward_fake( out_dtype = args.activation_dtype out_features, in_features = weight.shape - # Mirror ``_linear_backward_impl``: ``set_usage`` on ``grad_input_quantizer`` - # influences ``dgrad``'s buffer layout. + # Mirrors the impl; affects dgrad's buffer layout. if args.grad_input_quantizer is not None: args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) dgrad = None if args.requires_dgrad: - # dgrad has the logical input shape and may be quantized for the next op. - # Derive shape from grad_output + weight + SP config instead of args.inp_shape: - # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is - # not hashable in OpaqueValueBundle), so we reconstruct it here. + # Input shape rederived from grad_output + SP config (inp_shape is not + # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). dgrad_leading = _sp_inp_leading(args.grad_output.shape[0], args) dgrad = TensorSpec( shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), @@ -1745,11 +1729,8 @@ def _linear_backward_fake( ) wgrad = None + # Under fuse_wgrad_accumulation the grad goes into main_grad in place. if args.requires_wgrad and not args.fuse_wgrad_accumulation: - # wgrad has the weight's shape; quantized iff an fp8 wgrad output is - # requested (mirrors ``quantization_params=grad_weight_quantizer``), - # otherwise high precision. Under fuse_wgrad_accumulation the grad is - # written into ``main_grad`` in place and no wgrad tensor is returned. wgrad = TensorSpec( shape=(out_features, in_features), dtype=out_dtype, From b9bf6934ed4288dbe729e2c49febb6c0438a7654 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 10 Aug 2026 23:25:26 +0200 Subject: [PATCH 11/50] Rename SP leading-dim helpers for direction clarity, trim two comments _sp_out_leading/_sp_inp_leading -> _out_leading_from_inp/_inp_leading_from_out; shorten the weight_workspace field comment; drop the to_tensor_spec caveat paragraph. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/tensor_spec.py | 5 ---- transformer_engine/pytorch/module/linear.py | 27 +++++++------------ 2 files changed, 9 insertions(+), 23 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py index 8c156766e5..4cfe225952 100644 --- a/transformer_engine/pytorch/dynamo/tensor_spec.py +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -151,11 +151,6 @@ def to_tensor_spec(tensor: Any) -> TensorSpec: Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` / ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via ``_dtype`` rather than ``.dtype``. - - Not for re-describing a ``TensorSpec``: a spec holds its quantizer as - ``quantizer``, not ``_quantizer``, so it would come back unquantized. Fake - impls already receive specs from ``_spec_view`` -- copy those with - ``dataclasses.replace``. """ requires_grad = bool(getattr(tensor, "requires_grad", False)) dtype = getattr(tensor, "dtype", None) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 49ea0ff032..34772fcd6a 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -115,10 +115,7 @@ class LinearFwdArgs: bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- - # Same union as ``weight`` so a cached quantized workspace is flattened to its - # inner tensors on the way into the op (symmetric with ``new_weight_workspace`` - # on the way out); a plain ``Tensor?`` slot can't carry a quantized subclass - # across the torch.compile custom-op boundary. + # TensorOrQuantized so a cached quantized workspace can cross the op boundary. weight_workspace: Optional[TensorOrQuantized] # Process-global cuBLAS workspace, fetched at trace time so it isn't first @@ -317,12 +314,9 @@ def _check_fp8_reduce_and_update(): return result -def _sp_out_leading(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: - """Leading (sequence) dim of the output, given the input's. - - Under sequence parallelism a column-parallel layer gathers that dim and a - row-parallel one scatters it; without SP it passes through. - """ +def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Output's leading (sequence) dim from the input's: sequence parallelism + gathers it (column-parallel) or scatters it (row-parallel).""" if not args.sequence_parallel: return leading if args.parallel_mode == "column": @@ -332,11 +326,8 @@ def _sp_out_leading(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> return leading -def _sp_inp_leading(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: - """Inverse of :func:`_sp_out_leading`: input's leading dim from the output's. - - Used by backward, which reconstructs the input geometry from ``grad_output``. - """ +def _inp_leading_from_out(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Inverse of :func:`_out_leading_from_inp`.""" if not args.sequence_parallel: return leading if args.parallel_mode == "column": @@ -881,7 +872,7 @@ def _linear_forward_fake( # ------------------------------------------------------ # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). # ------------------------------------------------------ - out_leading = _sp_out_leading(inp.shape[0], args) + out_leading = _out_leading_from_inp(inp.shape[0], args) out = TensorSpec( shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), dtype=activation_dtype, @@ -1162,7 +1153,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). if bwd_args.inp_shape is None: in_features = saved_weight.shape[-1] - inp_leading = _sp_inp_leading(grad_output.shape[0], bwd_args) + inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args) bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) # Configure Userbuffers communication (comm+GEMM overlap) @@ -1720,7 +1711,7 @@ def _linear_backward_fake( if args.requires_dgrad: # Input shape rederived from grad_output + SP config (inp_shape is not # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). - dgrad_leading = _sp_inp_leading(args.grad_output.shape[0], args) + dgrad_leading = _inp_leading_from_out(args.grad_output.shape[0], args) dgrad = TensorSpec( shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), dtype=out_dtype, From aeb34811f6b26d311262540c53fb4ecebc7fdf41 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 00:52:53 +0200 Subject: [PATCH 12/50] Tighten custom_op module docstring intro Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index d33006c74c..d16f63739c 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -4,12 +4,9 @@ """torch.compile custom-op framework for Transformer Engine. -Turns a TE module's eager forward/backward into ``torch.library`` custom ops so -``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph -break into the eager ``autograd.Function``. ``register_custom_op`` is the entry -point (its docstring documents the per-callable contract); ``module/linear.py`` -is the first user. Internal framework API -- exported from -``transformer_engine.pytorch.dynamo``, not re-exported at the top level. +Registers TE modules' eager forward/backward as ``torch.library`` custom ops so +``torch.compile(fullgraph=True)`` traces them as single graph nodes. +``register_custom_op`` is the entry point; ``module/linear.py`` is the first user. A TE op's forward/backward is written as a plain impl over a single *args dataclass* (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``): its From 53b811bcdca6162d83c93c1037b3d25a5b7147be Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 00:57:09 +0200 Subject: [PATCH 13/50] Merge and simplify custom_op docstring paragraphs 2-3 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index d16f63739c..8311500d42 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -8,16 +8,11 @@ ``torch.compile(fullgraph=True)`` traces them as single graph nodes. ``register_custom_op`` is the entry point; ``module/linear.py`` is the first user. -A TE op's forward/backward is written as a plain impl over a single *args -dataclass* (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``): its -fields are a mix of tensors, quantized tensors, quantizers, process groups, -scalars and other Python values. The forward impl returns a tuple (user outputs + -saved-for-backward tensors + ctx metadata); the backward impl returns one gradient -per differentiable input. - -A ``torch.library`` custom op is narrower: it takes a flat list of schema slots --- tensors / ``Tensor[]`` plus, via torch's opaque-object support, value-opaque -and reference-opaque objects -- and returns a flat ``Tensor[]``. +A TE forward/backward implementation takes one dataclass argument +(``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix +tensors, quantized tensors, quantizers, process groups and plain Python values, +while a ``torch.library`` custom op only accepts flat schema slots (tensors plus +opaque objects) and returns a flat ``Tensor[]``. Bridging the two takes three parts (below): per-field *adapters* map the args dataclass onto the op's input slots; *fake impls* on data-free specs give the From 948f94e08cb9eaf2f4c2d2647b0343d4ded7a75d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 00:58:08 +0200 Subject: [PATCH 14/50] Keep the impl-vs-op contrast as two paragraphs Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 8311500d42..13606f54fb 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -10,9 +10,10 @@ A TE forward/backward implementation takes one dataclass argument (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix -tensors, quantized tensors, quantizers, process groups and plain Python values, -while a ``torch.library`` custom op only accepts flat schema slots (tensors plus -opaque objects) and returns a flat ``Tensor[]``. +tensors, quantized tensors, quantizers, process groups and plain Python values. + +A ``torch.library`` custom op is narrower: it only accepts flat schema slots +(tensors plus opaque objects) and returns a flat ``Tensor[]``. Bridging the two takes three parts (below): per-field *adapters* map the args dataclass onto the op's input slots; *fake impls* on data-free specs give the From 872387facad9934dafc747ef312071ed7dc57be2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 01:01:01 +0200 Subject: [PATCH 15/50] Drop reference to a PyTorch PR that will not land Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 13606f54fb..213d2f5d2d 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -123,10 +123,6 @@ # tensor can have -- complex is never a TE payload (quantized data / scales are # uint8 / fp32, outputs are float), while e.g. a uint8 sentinel would collide # with a genuinely empty FP8 data buffer (empty batch). -# -# Once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable -# ``Tensor?[]`` return schema will let ``None`` pass through directly and this -# sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed. _NONE_SENTINEL_DTYPE = torch.complex32 From dcc3c9acb0160c303de5e16cbe9a86727cf88b99 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 01:12:10 +0200 Subject: [PATCH 16/50] Shorten _ensure_distributed_opaque_types docstring Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 30 +++++++------------ 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 213d2f5d2d..a3bd37aada 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -330,23 +330,14 @@ def _pg_pickle_stub(*args: Any) -> None: # pragma: no cover def _ensure_distributed_opaque_types() -> None: - """Register ``torch.distributed.ProcessGroup`` as a *reference* opaque type. - - A process group is live distributed state: unlike a value-opaque quantizer - (which Dynamo bakes into the graph as a constant), it must be carried through - the custom op as a graph *input*. PyTorch supports this via - ``register_opaque_type(ProcessGroup, typ="reference")`` but only auto-runs it - when ``torch.distributed.tensor`` (DTensor) is imported; TE may not import - that, so trigger the same idempotent registration here. Best-effort: on - builds without the opaque-object / distributed APIs this is a no-op and the - process-group field simply falls back to eager under torch.compile. - - Also registers a ``copyreg`` reducer that lets ``FxGraphCachePickler`` hash - graphs containing a ``ProcessGroup`` input without crashing. Without this, - inductor logs "Failed to pickle cache key" warnings and bypasses the FX - graph disk cache for every distributed compiled call. The reducer encodes - the group as (world_size, rank, backend) -- enough to distinguish configs -- - and raises on reconstruct since deserialization is never needed for hashing. + """Register ``ProcessGroup`` as a *reference* opaque type: live state carried + as a graph input, not baked in as a constant. PyTorch runs this registration + only when DTensor is imported, so trigger it here too; best-effort no-op on + builds without the APIs (the field then falls back to eager). + + Also register a ``copyreg`` reducer so ``FxGraphCachePickler`` can hash a + graph with a ``ProcessGroup`` input -- without it, inductor bypasses the FX + disk cache for every distributed compiled call. Hash-only, never unpickled. """ if _is_opaque_reference_type is None: return @@ -359,9 +350,8 @@ def _ensure_distributed_opaque_types() -> None: except Exception: # pylint: disable=broad-exception-caught pass - # Workaround for PyTorch issue: FxGraphCachePickler handles FakeScriptObject - # but not the real ProcessGroup that appears in example_inputs at inductor - # compile time. Register a copyreg reducer so the pickler can hash the key. + # Workaround for a PyTorch gap: its cache-key pickler handles + # FakeScriptObject but not the real ProcessGroup in example_inputs. try: # pylint: disable=import-outside-toplevel import copyreg From d87eda6c822ffe2494d98eaa3014bfdcd393c055 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 01:21:03 +0200 Subject: [PATCH 17/50] Fall back to eager cleanly when ProcessGroup opaque registration is unavailable PG_REFERENCE_OPAQUE is computed once at import (Dynamo-friendly constant); compile_unsupported_reason reports a tp_group it cannot carry instead of the misleading _UnsupportedAdapter TypeError at trace time. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 ++- .../pytorch/dynamo/custom_op.py | 19 +++++++++++++++++++ transformer_engine/pytorch/module/linear.py | 4 ++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index e42eb8f9f6..95d97b7e08 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op, TensorOrQuantized +from .custom_op import register_custom_op, TensorOrQuantized, PG_REFERENCE_OPAQUE __all__ = [ "register_value_opaque_quantizer", @@ -15,4 +15,5 @@ "to_tensor_spec", "register_custom_op", "TensorOrQuantized", + "PG_REFERENCE_OPAQUE", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index a3bd37aada..78f7859971 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -380,6 +380,25 @@ def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] _ensure_distributed_opaque_types() +def _compute_pg_reference_opaque() -> bool: + if _is_opaque_reference_type is None: + return False + try: + from torch._C._distributed_c10d import ( # pylint: disable=import-outside-toplevel + ProcessGroup, + ) + + return bool(_is_opaque_reference_type(ProcessGroup)) + except Exception: # pylint: disable=broad-exception-caught + return False + + +# Whether ProcessGroup ended up registered as a reference-opaque type, i.e. +# whether a process group can cross the op boundary as a live graph input. +# Process-global and fixed at import, so a plain constant (Dynamo-friendly). +PG_REFERENCE_OPAQUE: bool = _compute_pg_reference_opaque() + + # --------------------------------------------------------------------------- # # Storage flatten / unflatten (value-opaque quantizer; no ProcessGroup) # --------------------------------------------------------------------------- # diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 34772fcd6a..034fff0b12 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -89,6 +89,7 @@ TensorOrQuantized, register_custom_op, is_value_opaque_quantizer, + PG_REFERENCE_OPAQUE, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -204,6 +205,9 @@ def compile_unsupported_reason(self) -> Optional[str]: return "delayed wgrad compute (wgrad_store)" if self.fuse_wgrad_accumulation: return "fuse_wgrad_accumulation (main_grad)" + if self.tp_group is not None and not PG_REFERENCE_OPAQUE: + # ProcessGroup's reference-opaque registration failed at import. + return "a tp_group not registered as a torch.compile reference-opaque type" for quantizer in ( self.input_quantizer, self.weight_quantizer, From 1950585f94325b5bb25b81821495d3f96135e38c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 01:23:36 +0200 Subject: [PATCH 18/50] Fix leftover 'priority order' wording at _FIELD_ADAPTERS Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 78f7859971..a8a6675135 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -821,11 +821,8 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = None -# Adapters, in priority order, owning ``try_build`` for a single field. -# These adapters are mutually exclusive on annotations (a plain ``torch.Tensor`` -# matches only ``_TensorAdapter``; the ``TensorOrQuantized`` union only -# ``_TensorOrQuantizedAdapter``; etc.), so the order is just iteration, not a -# priority ranking -- no annotation can be claimed by more than one. +# Adapter candidates for a single field, tried via ``try_build``. Mutually +# exclusive on annotations, so the order is not a ranking. _FIELD_ADAPTERS: Tuple[type, ...] = ( _TensorOrQuantizedAdapter, _TensorAdapter, From 9c4fadfe61293dc60afd5aed27b43a86f6a56e1e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 01:29:10 +0200 Subject: [PATCH 19/50] Restructure register_custom_op docstring: caller contract first, drop module-docstring duplication Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 45 +++++++------------ 1 file changed, 17 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index a8a6675135..91c2679322 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1340,31 +1340,17 @@ def register_custom_op( ) -> Optional[Callable[..., Any]]: """Register a TE module's forward + backward as torch custom ops. - Always two-tier: a base ``_base`` op carries the real schema / - autograd, and a wrapper ```` op forwards to it, flattening any - quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an - empty subclass list simply makes the wrapper op a pass-through, so a pure - plain-tensor / bf16 call goes straight through). - Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches - through the wrapper op and returns the user-facing outputs. - - Arg containers. ``fwd_arg_type`` and ``bwd_arg_type`` are ``@dataclass``es - whose *field annotations* define the op schema: each field maps to one or more - flat schema slots (tensor fields cross the boundary as tensors, quantizers ride - as value-opaque objects, simple values are bundled -- see the ``_Adapter`` - classes). The caller builds a ``fwd_arg_type`` instance and passes it to the - returned ``forward_fn``. + through the op and returns the user-facing outputs. - How the backward container is populated. ``setup_context`` fills the - ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from - forward state and returns the tensors to persist; the framework saves them - (``ctx.save_for_backward``). Before ``bwd_impl`` runs, the framework - restores those tensors into the container's *tensor* fields by calling its - optional ``setup_saved_tensors(self, ctx)`` hook (invoked only if defined), - and sets ``grad_output`` directly. So ``bwd_impl`` receives a - fully-populated ``bwd_arg_type``. + ``fwd_arg_type`` and ``bwd_arg_type`` are ``@dataclass``es whose *field + annotations* define the op schema (see the module docstring for the + field <-> slot mapping). The caller builds a ``fwd_arg_type`` instance and + passes it to ``forward_fn``. ``input_tensors_for_grad`` lists the + ``fwd_arg_type`` fields that receive gradients and fixes the backward grad + order. ``bwd_arg_type`` is also instantiated by the framework + (``bwd_arg_type()``), so it must be constructible with no arguments. Callable contracts: @@ -1385,13 +1371,16 @@ def register_custom_op( non-differentiable input). * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning :class:`TensorSpec` grads. - * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook on the backward - container (see above); skipped if absent. + * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if + absent. - ``input_tensors_for_grad`` lists the ``fwd_arg_type`` fields that receive - gradients (this fixes the backward grad order). ``bwd_arg_type`` is both - the schema source and the type instantiated (``bwd_arg_type()``) to hold - the backward args, so it must be constructible with no arguments. + How the backward container is populated: ``setup_context`` fills the + ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from + forward state and returns the tensors to persist; the framework saves them + via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework + restores them into the container's tensor fields through the + ``setup_saved_tensors`` hook and sets ``grad_output`` directly, so + ``bwd_impl`` receives a fully-populated ``bwd_arg_type``. Registration touches experimental ``torch.library`` / opaque-object APIs that may be missing on older PyTorch. If it fails, this warns once and From a8e4fb75e6bff892dfd3d9ba473d69becfb6b097 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 13:30:43 +0200 Subject: [PATCH 20/50] Fix two fake/impl divergences found in multi-agent review - Backward fake now returns grad_bias whenever bias is used on the FP8 backward path (grad_output_preprocess computes bgrad independent of requires_wgrad); previously a frozen weight silently dropped the bias gradient under torch.compile. - Forward fake now mirrors quantize_weight's workspace invalidation: a cached workspace missing buffers for the quantizer's current usage is dropped and a fresh new_weight_workspace is declared, instead of always assuming a cache hit (previously crashed with an output size/stride mismatch when a rowwise-only cache met a training step). Both verified against eager on RTX Ada. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 23 +++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 034fff0b12..581dc45501 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -341,6 +341,17 @@ def _inp_leading_from_out(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs return leading +def _fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer]) -> bool: + """Spec-level mirror of ``_is_weight_workspace_valid``: the cached workspace + must already hold every inner buffer the quantizer's current usage needs.""" + if quantizer is None: + return True + required = TensorSpec( + shape=workspace.shape, dtype=workspace.dtype, quantizer=quantizer, device=workspace.device + ).inner_names() + return set(required) <= set(workspace.inner_names()) + + def _linear_forward_impl( args: LinearFwdArgs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: @@ -823,6 +834,7 @@ def _linear_forward_fake( # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. # ------------------------------------------------------ new_weight_workspace = None + workspace = None # args.weight_workspace after validation weightmat = None weightmat_is_storage = False weightmat_aliases_weight = False @@ -848,6 +860,9 @@ def _linear_forward_fake( else: weightmat_is_storage = True workspace = args.weight_workspace + if workspace is not None and not _fake_workspace_valid(workspace, weight_quantizer): + # quantize_weight drops a stale workspace and builds a new one. + workspace = None if workspace is not None: # Copy, so the ``update_usage`` below stays off the input spec. weightmat = dataclass_replace(workspace) @@ -933,7 +948,7 @@ def _linear_forward_fake( pass # FSDP2 re-quantizes from the gathered weight in backward. elif weightmat_is_storage and new_weight_workspace is not None: wt_alias = "new_weight_workspace" - elif weightmat_is_storage and args.weight_workspace is not None: + elif weightmat_is_storage and workspace is not None: wt_alias = "weight_workspace" elif weightmat_is_storage: wt_save = weightmat @@ -1734,7 +1749,11 @@ def _linear_backward_fake( ) grad_bias = None - if args.use_bias and args.requires_wgrad: + # FP8 backward computes bgrad in grad_output_preprocess whenever bias is + # used; in high precision it is fused into the wgrad GEMM, so it only + # exists when wgrad runs. + fp8_bwd = args.fp8 and args.backward_override is None + if args.use_bias and (args.requires_wgrad or fp8_bwd): grad_bias = TensorSpec( shape=(out_features,), dtype=out_dtype, device=args.grad_output.device ) From f47a6370c42a37e08c1a5b7aa0766c14ad339fd7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 14:35:37 +0200 Subject: [PATCH 21/50] Drop the global copyreg ProcessGroup reducer copyreg.pickle is process-wide: with the reducer installed, torch.save of any object graph reaching a ProcessGroup silently succeeded and the checkpoint failed only at torch.load (the reconstruct stub raises). Restore the loud failure at save time; the cost is that inductor bypasses the FX disk cache for compiled distributed graphs (with its own warning) until the cache-key pickler handles real opaque objects upstream. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 37 ++----------------- 1 file changed, 4 insertions(+), 33 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 91c2679322..fa8eef6032 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -325,19 +325,16 @@ def _collect(value: Any) -> None: _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None -def _pg_pickle_stub(*args: Any) -> None: # pragma: no cover - raise RuntimeError("ProcessGroup cannot be unpickled -- cache-key use only") - - def _ensure_distributed_opaque_types() -> None: """Register ``ProcessGroup`` as a *reference* opaque type: live state carried as a graph input, not baked in as a constant. PyTorch runs this registration only when DTensor is imported, so trigger it here too; best-effort no-op on builds without the APIs (the field then falls back to eager). - Also register a ``copyreg`` reducer so ``FxGraphCachePickler`` can hash a - graph with a ``ProcessGroup`` input -- without it, inductor bypasses the FX - disk cache for every distributed compiled call. Hash-only, never unpickled. + Note: PyTorch's cache-key pickler (``FxGraphCachePickler``) cannot hash the + real ``ProcessGroup`` in ``example_inputs`` yet, so inductor bypasses the FX + disk cache for compiled distributed calls (it logs a warning). Fixing that + belongs upstream; keeping TE free of process-wide pickle overrides. """ if _is_opaque_reference_type is None: return @@ -350,32 +347,6 @@ def _ensure_distributed_opaque_types() -> None: except Exception: # pylint: disable=broad-exception-caught pass - # Workaround for a PyTorch gap: its cache-key pickler handles - # FakeScriptObject but not the real ProcessGroup in example_inputs. - try: - # pylint: disable=import-outside-toplevel - import copyreg - import torch.distributed as dist - from torch._C._distributed_c10d import ProcessGroup - - # pylint: enable=import-outside-toplevel - - if ProcessGroup not in copyreg.dispatch_table: - - def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] - try: - return _pg_pickle_stub, ( - dist.get_world_size(pg), - dist.get_rank(pg), - dist.get_backend(pg), - ) - except Exception: # pylint: disable=broad-exception-caught - return _pg_pickle_stub, (id(pg),) - - copyreg.pickle(ProcessGroup, _pg_reduce) - except Exception: # pylint: disable=broad-exception-caught - pass - _ensure_distributed_opaque_types() From 23b093d55a23ed13134814d4cdc132e543b88c9d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 14:44:16 +0200 Subject: [PATCH 22/50] Skip use_compile numerics cases for DelayedScaling DelayedScaling quantizers are not value-opaque, so the compiled path falls back to eager, which errors under fullgraph=True. Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_numerics.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index d9dccc6f0a..f7af97a962 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -784,6 +784,8 @@ def test_linear(): kwargs.get("delay_wgrad_compute", False) or kwargs.get("use_compile", False) ): continue + if kwargs.get("use_compile", False) and QUANTIZATION == "fp8": + continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel, **kwargs) From e5a8ba80e69caf5ccb10d1040c2917b287613ee1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 14:52:12 +0200 Subject: [PATCH 23/50] Fix backward fake for UB reduce-scatter dgrad and extend the compiled UB test to FP8 Under ub_overlap_rs_dgrad the impl returns the plain high-precision reduce-scatter output as dgrad (the grad_input_quantizer only feeds the communication buffer), while the fake declared a quantized dgrad spec -- an op output-contract mismatch. test_linear_with_overlap_compile now also runs fp8_current_scaling and mxfp8 for the column-parallel cases (bulk and DGRAD+RS); FP8 row-parallel stays skipped (forced differentiable fp8_output is unsupported under compile) and delayed scaling is excluded like elsewhere. Signed-off-by: Pawel Gadzinski --- .../distributed/test_comm_gemm_overlap.py | 19 +++++++++++++++---- transformer_engine/pytorch/module/linear.py | 5 ++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 25331ba85a..12d3747c3a 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -288,6 +288,11 @@ def test_layers_with_overlap_bf16( @pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize( + "quantization", + [None, "fp8_current_scaling", "mxfp8"], + ids=["bf16", "fp8_current_scaling", "mxfp8"], +) @pytest.mark.parametrize( "linear_parallel_mode,overlap_rs_dgrad", [ @@ -301,15 +306,21 @@ def test_layers_with_overlap_bf16( "COL-PARALLEL - DGRAD+RS", ], ) -def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode): - """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16), +def test_linear_with_overlap_compile( + linear_parallel_mode, overlap_rs_dgrad, quantization, compile_mode +): + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile, checked numerically against the eager, non-overlap reference.""" + if quantization is not None and linear_parallel_mode == "row": + pytest.skip( + "FP8 row-parallel UB forces differentiable fp8_output, unsupported under compile." + ) _run_layer_with_overlap( te.Linear.__name__, linear_parallel_mode, overlap_rs_dgrad, - False, - None, + quantization is not None, + quantization, use_compile=True, compile_mode=compile_mode, ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 581dc45501..7ea4b5285f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1731,10 +1731,13 @@ def _linear_backward_fake( # Input shape rederived from grad_output + SP config (inp_shape is not # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). dgrad_leading = _inp_leading_from_out(args.grad_output.shape[0], args) + # Under UB reduce-scatter overlap the returned dgrad is the plain + # reduce-scatter output; the quantizer only feeds the comm buffer. + dgrad_quantizer = None if args.ub_overlap_rs_dgrad else args.grad_input_quantizer dgrad = TensorSpec( shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), dtype=out_dtype, - quantizer=args.grad_input_quantizer, + quantizer=dgrad_quantizer, device=args.grad_output.device, ) From 2ff3301d3e363a60515118cdad22663121b3bbfd Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:02:25 +0200 Subject: [PATCH 24/50] Fall back to eager for DistributedWeight (GTP) under torch.compile The compiled path handles neither the external weight subclass at the op boundary nor the materialize/refresh logic in the fakes; gate it in compile_unsupported_reason like the other unsupported configs. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7ea4b5285f..73d409d284 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -190,6 +190,8 @@ def compile_unsupported_reason(self) -> Optional[str]: """Reason this config can't use the torch.compile custom-op path (else None).""" if self.debug: return "debug instrumentation (nvidia-dlfw-inspect)" + if is_distributed_weight(self.weight): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" if self.fsdp_group is not None and self.is_grad_enabled: return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" if ( From def25166fb9d56fb882b672a4a75910f7963eaae Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:15:43 +0200 Subject: [PATCH 25/50] Widen two eager-fallback conditions - fsdp_group: fall back regardless of grad mode; the adapter rejects the field for any non-trivial value, so inference with manual TE FSDP could reach the op and fail there instead. - fp8_output: also fall back when only the bias requires grad; the backward tangent for the quantized output was mis-guessed by AOTAutograd (RuntimeError: Expected a Float8Tensor tangent but got a plain Tensor). Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 73d409d284..3a8cb003cf 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -192,12 +192,12 @@ def compile_unsupported_reason(self) -> Optional[str]: return "debug instrumentation (nvidia-dlfw-inspect)" if is_distributed_weight(self.weight): return "a DistributedWeight (custom weight parallelism, e.g. GTP)" - if self.fsdp_group is not None and self.is_grad_enabled: + if self.fsdp_group is not None: return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" if ( self.fp8_output and self.is_grad_enabled - and (self.input_requires_grad or self.weight_requires_grad) + and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad) ): return "differentiable fp8_output=True" if self.cpu_offloading: From 1d1ba9baa4627fe10ac4430e7cebfbc51a827a1a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:21:44 +0200 Subject: [PATCH 26/50] Defer the compile-disabled warning from import time to first compile use Registration failures now only record the reason; importing TE on a build without opaque-object support stays silent. The warning is emitted from Linear.forward when a compiled call finds the op unregistered; under fullgraph=True the resulting error names warn_if_compile_disabled. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 6 ++-- .../pytorch/dynamo/quantizer_opaque.py | 6 ++-- transformer_engine/pytorch/module/linear.py | 3 ++ transformer_engine/pytorch/utils.py | 35 ++++++++++++------- 4 files changed, 31 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index fa8eef6032..631caf8828 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -106,7 +106,7 @@ _quantized_tensor_passthrough_ops, prepare_for_saving, ) -from ..utils import warn_compile_disabled +from ..utils import record_compile_disabled _TE_OP_NAMESPACE = "transformer_engine_compile" @@ -317,7 +317,7 @@ def _collect(value: Any) -> None: _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) # Older torch without opaque_object support. except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - warn_compile_disabled( + record_compile_disabled( f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build" ) _is_opaque_value_type = None @@ -1371,7 +1371,7 @@ def register_custom_op( bwd_fake_impl=bwd_fake_impl, ) except (ImportError, AttributeError, RuntimeError, TypeError) as e: - warn_compile_disabled( + record_compile_disabled( f"could not register the custom op '{op_name}' ({type(e).__name__}: {e})" ) return None diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 856dee2022..dc689258d1 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -9,7 +9,7 @@ from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType -from ..utils import warn_compile_disabled +from ..utils import record_compile_disabled # Qualnames of the registered quantizer classes. The set holds strings rather @@ -121,7 +121,7 @@ def register_value_opaque_quantizer(cls: type) -> None: except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. - warn_compile_disabled( + record_compile_disabled( f"this PyTorch build has no opaque-object API ({e}); use a newer build" ) return @@ -133,7 +133,7 @@ def register_value_opaque_quantizer(cls: type) -> None: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. - warn_compile_disabled(f"could not register {cls.__name__} as an opaque type ({e})") + record_compile_disabled(f"could not register {cls.__name__} as an opaque type ({e})") return _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3a8cb003cf..64f34b9833 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -50,6 +50,7 @@ nvtx_range_push, get_nvtx_range_context, warn_compile_eager_fallback, + warn_if_compile_disabled, check_gemm_dims, ) from ..distributed import ( @@ -2354,6 +2355,8 @@ def forward( ) use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None + if _linear_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index dd74a9c401..f499da1dfe 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,27 +26,36 @@ ] -_warned_compile_disabled = False +_compile_disabled_reason: Optional[str] = None +_compile_disabled_warned = False -def warn_compile_disabled(reason: str) -> None: - """Warn once per process that TE's torch.compile custom-op path is off. - - Registration of the torch.compile machinery either works or fails as a - whole, so one message is enough. Distinct from +def record_compile_disabled(reason: str) -> None: + """Record why TE's torch.compile custom-op path is off; the warning is + emitted only when a compiled TE module actually runs (see + :func:`warn_if_compile_disabled`), so a plain import stays silent. + The first recorded reason wins. Distinct from :func:`warn_compile_eager_fallback`, which reports a single *configuration* falling back while the path itself is available. """ - global _warned_compile_disabled # pylint: disable=global-statement - if _warned_compile_disabled: + global _compile_disabled_reason # pylint: disable=global-statement + if _compile_disabled_reason is None: + _compile_disabled_reason = reason + + +@torch._dynamo.disable # graph-breaks cleanly so the warning actually fires +def warn_if_compile_disabled() -> None: + """Warn once, at the first compile attempt, that the path is off.""" + global _compile_disabled_warned # pylint: disable=global-statement + if _compile_disabled_warned: return - _warned_compile_disabled = True + _compile_disabled_warned = True warnings.warn( "Transformer Engine torch.compile support is disabled: " - f"{reason}. Modules will fall back to eager execution under " - "torch.compile, i.e. a graph break, which is incompatible with " - "fullgraph=True.", - stacklevel=3, + f"{_compile_disabled_reason or 'custom-op registration unavailable'}. " + "Modules will fall back to eager execution under torch.compile, i.e. " + "a graph break, which is incompatible with fullgraph=True.", + stacklevel=2, ) From 1822afcc38dde423a98d9628686e67fe0d7cf56c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:24:35 +0200 Subject: [PATCH 27/50] Log the compile-disabled reason at registration time (INFO, TransformerEngine logger) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index f499da1dfe..c6e0c63ebf 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -5,6 +5,7 @@ """Utility functions for Transformer Engine modules""" from __future__ import annotations import functools +import logging import math import os import warnings @@ -41,6 +42,9 @@ def record_compile_disabled(reason: str) -> None: global _compile_disabled_reason # pylint: disable=global-statement if _compile_disabled_reason is None: _compile_disabled_reason = reason + logging.getLogger("TransformerEngine").info( + "torch.compile custom-op path disabled: %s", reason + ) @torch._dynamo.disable # graph-breaks cleanly so the warning actually fires From dc5cecd6976c603965f9f248c491c677cb001e16 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:53:22 +0200 Subject: [PATCH 28/50] Release ctx.backward_objects after the compiled-op backward Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 631caf8828..de4123de1b 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1181,6 +1181,7 @@ def _autograd_backward(ctx, *grad_outputs): kwargs = _args_to_slots(bwd_obj, bwd_adapters) bwd_args_flat = [kwargs[name] for name in bwd_arg_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + ctx.backward_objects = None # One grad per input schema slot: default None, but a ``Tensor[]`` slot # (always recorded in ``fwd_tensor_list_lengths``) needs a # list-shaped no-grad of matching length. From 6586d5926f24c8cde4bcfb7b0ae579a027620467 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 15:58:52 +0200 Subject: [PATCH 29/50] Mirror the save_original_input runtime flip in the forward fake The impl disables save_original_input when the input quantizer cannot reconstruct the wgrad operand from the original input (e.g. NVFP4 with stochastic rounding); the fake kept it on and declared the saved-input slot as an alias while the impl saved a quantized storage. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 64f34b9833..4b71b9a80d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -804,6 +804,16 @@ def _linear_forward_fake( out_features, _ = weight.shape backward_needs_input = is_grad_enabled and args.weight_requires_grad + if ( + args.backward_override is None + and save_original_input + and backward_needs_input + and input_quantizer is not None + ): + if not isinstance( + input_quantizer, Float8Quantizer + ) and not can_reconstruct_wgrad_input_from_original(input_quantizer): + save_original_input = False own_quantized_input = False inputmat_is_storage = False From da11b9a2e8fd8628bfb834bfedf5bb958a8d964a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 16:07:07 +0200 Subject: [PATCH 30/50] Fall back to eager for quantized input tensors under torch.compile The inp field crosses the op boundary as a plain Tensor slot, so a quantized activation (e.g. the fp8_output of a previous layer) breaks fake propagation even under no_grad; gate it until the boundary supports quantized inputs. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 4b71b9a80d..eea0f3430d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -193,6 +193,8 @@ def compile_unsupported_reason(self) -> Optional[str]: return "debug instrumentation (nvidia-dlfw-inspect)" if is_distributed_weight(self.weight): return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)): + return "a quantized input tensor" if self.fsdp_group is not None: return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" if ( From 95d5271640e2b06fcbab0f4bf77bd4affc93b54b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 16:30:44 +0200 Subject: [PATCH 31/50] Test hardening: exact eager-vs-compiled comparison, counters guard, real microbatch cache checks - Tolerances tightened to exact: all compute runs inside the op and the loss grad is ones, so eager and compiled are bit-identical (measured on both compile modes, bf16 and fp8). - torch._dynamo counters reads degrade with a warning instead of failing when the private API changes. - is_first_microbatch test: eager reference on a separate module (shared cache made it unable to catch corruption or rebuilds), structural asserts that the compiled step creates the cache and later steps reuse the same object, eager priming of FP8 state before tracing (in-graph quantizer creation breaks recompiles; upstream Dynamo bug). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 89 +++++++++++++++++++++++------ 1 file changed, 72 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 3336f44f22..eab5f5114a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,10 +4,15 @@ import abc import contextlib +import warnings import pytest import torch -from torch._dynamo.utils import counters + +try: + from torch._dynamo.utils import counters +except ImportError: # pragma: no cover + counters = None from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: @@ -106,23 +111,33 @@ def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: out.sum().backward() +def _dynamo_counter(group: str, key: str): + """Read a torch._dynamo counter; None (with a warning) if the private + counters API is gone, so CI degrades instead of failing.""" + try: + return counters[group][key] + except Exception: # pylint: disable=broad-except + warnings.warn(f"torch._dynamo.utils.counters[{group!r}][{key!r}] unavailable") + return None + + @contextlib.contextmanager def _assert_no_cudagraph_skips(enabled: bool): """Assert reduce-overhead really captured CUDA graphs: inductor may skip capture and silently fall back to eager, which ``fullgraph=True`` does not catch. No-op when ``enabled`` is False.""" - before = counters["inductor"]["cudagraph_skips"] + before = _dynamo_counter("inductor", "cudagraph_skips") yield - if enabled: - skipped = counters["inductor"]["cudagraph_skips"] - before + if enabled and before is not None: + skipped = _dynamo_counter("inductor", "cudagraph_skips") - before assert skipped == 0, ( f"reduce-overhead fell back to eager: {skipped} cudagraph skip(s); " "see the 'skipping cudagraphs due to ...' log for the reason" ) -# Eager and compiled run the same kernels; slack only for reduction-order noise. -_EAGER_ATOL, _EAGER_RTOL = 1e-2, 1.6e-2 +# All compute runs inside the op and the loss grad is ones, so bit-exact. +_EAGER_ATOL, _EAGER_RTOL = 0.0, 0.0 def _assert_close_eager_compiled(fn, compiled, model, base): @@ -1377,20 +1392,34 @@ def fn(inp): def test_te_linear_compile_is_first_microbatch(compile_mode): """torch.compile of ``te.Linear`` across a microbatch schedule: ``is_first_microbatch=True`` caches the FP8 weight, later steps must reuse - it and stay numerically aligned with eager.""" + it and stay numerically aligned with eager. The eager reference runs on a + separate module so it cannot mask a corrupted or rebuilt cache.""" dtype = torch.bfloat16 device = "cuda" fp8_recipe = recipe.Float8CurrentScaling() model = te.Linear(64, 32, params_dtype=dtype, device=device) + ref_model = te.Linear(64, 32, params_dtype=dtype, device=device) + with torch.no_grad(): + ref_model.weight.copy_(model.weight) + ref_model.bias.copy_(model.bias) - # First microbatch caches the FP8 weight, the rest reuse the cache. schedule = [True, False, False] - is_first = schedule[0] # rebound each step; closed over by ``fn``. + is_first = schedule[0] # rebound each step; closed over by the fns. def fn(inp): with te.autocast(recipe=fp8_recipe): return model(inp, is_first_microbatch=is_first) + def ref_fn(inp): + with te.autocast(recipe=fp8_recipe): + return ref_model(inp, is_first_microbatch=is_first) + + # Eager priming: FP8 state must exist before tracing (creating quantizers + # in-graph breaks later recompiles; upstream Dynamo bug). + is_first = None + fn(torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True)) + is_first = schedule[0] + torch._dynamo.reset() if compile_mode == "reduce-overhead": _cudagraph_warmup( @@ -1401,10 +1430,33 @@ def fn(inp): model.zero_grad(set_to_none=True) compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + cached_workspace = None with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): - for is_first in schedule: + for step, is_first in enumerate(schedule): base = torch.randn(32, 64, dtype=dtype, device=device) - _assert_close_eager_compiled(fn, compiled, model, base) + + inp_ref = base.detach().clone().requires_grad_(True) + ref_model.zero_grad(set_to_none=True) + out_ref = ref_fn(inp_ref) + out_ref.sum().backward() + + inp = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out = compiled(inp).clone() + out.sum().backward() + + torch.testing.assert_close(out, out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close( + model.weight.grad, ref_model.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + workspace = model._fp8_workspaces.get("weight") + assert workspace is not None, f"no cached FP8 weight after step {step}" + if step == 0: + cached_workspace = workspace + else: + assert workspace is cached_workspace, f"cache rebuilt at step {step}" @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @@ -1436,7 +1488,9 @@ def fn(inp): torch._dynamo.mark_dynamic(warm, 0) compiled(warm.requires_grad_(True)).sum().backward() model.zero_grad(set_to_none=True) - unique_graphs_baseline = counters["stats"]["unique_graphs"] + unique_graphs_baseline = _dynamo_counter("stats", "unique_graphs") + if not unique_graphs_baseline: + warnings.warn("unique_graphs counter is stale; skipping the recompile check") for batch in batch_sizes: inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) @@ -1468,8 +1522,9 @@ def fn(inp): msg=f"dgrad mismatch at batch={batch}", ) - unique_graphs_after = counters["stats"]["unique_graphs"] - assert unique_graphs_after == unique_graphs_baseline, ( - "Unexpected recompilation(s) across different batch sizes: " - f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" - ) + if unique_graphs_baseline: + unique_graphs_after = _dynamo_counter("stats", "unique_graphs") + assert unique_graphs_after == unique_graphs_baseline, ( + "Unexpected recompilation(s) across different batch sizes: " + f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" + ) From 0e9e1bedd9f39a19a6a87d5bdf14ed70dc379285 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 16:36:20 +0200 Subject: [PATCH 32/50] Distributed tests: compare input gradients; exercise cudagraph replay under reduce-overhead run_numerics now checks dgrad (gathered per parallel mode) in every linear case; run_layer_with_overlap warms the compiled model up under reduce-overhead so the measured run replays captured graphs, and asserts inductor recorded no cudagraph skips. Signed-off-by: Pawel Gadzinski --- .../pytorch/distributed/run_layer_with_overlap.py | 14 ++++++++++++++ tests/pytorch/distributed/run_numerics.py | 11 +++++++++++ 2 files changed, 25 insertions(+) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index fb5579d583..058a78a756 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -18,6 +18,11 @@ import torch import torch.distributed as dist +try: + from torch._dynamo.utils import counters as dynamo_counters +except ImportError: # pragma: no cover + dynamo_counters = None + import transformer_engine.pytorch as te from transformer_engine.common.recipe import ( DelayedScaling, @@ -572,7 +577,16 @@ def run_fwd_bwd(model, x): if not opts.benchmark: del test_graph else: + if opts.compile and opts.compile_mode == "reduce-overhead": + # Warm up so the measured run below replays captured CUDA graphs. + for _ in range(2): + run_fwd_bwd(test_model, test_x) + test_model.zero_grad(set_to_none=True) + test_x.grad = None test_out = run_fwd_bwd(test_model, test_x) + if opts.compile and opts.compile_mode == "reduce-overhead" and dynamo_counters is not None: + skips = dynamo_counters["inductor"]["cudagraph_skips"] + assert skips == 0, f"reduce-overhead fell back to eager: {skips} cudagraph skip(s)" test_grads = [test_out, test_x.grad] names = ["output", "input.grad"] for test_name, test_param in test_model.named_parameters(): diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index f7af97a962..6160f0373f 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -292,6 +292,15 @@ def _check_gradients(model_distributed, model_single, main_grad_check=False): assert not bool(numerics_failed.item()) +def _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel): + grad_d = input_distributed.grad + if parallel_mode == "row": + grad_d = _gather(grad_d, dim=1) + elif sequence_parallel: + grad_d = _gather(grad_d, dim=0) + _check_outputs(input_single_node.grad, grad_d) + + def _copy_params(model_distributed, model_single): for dist_param, single_param in zip(model_distributed.parameters(), model_single.parameters()): with torch.no_grad(): @@ -759,6 +768,8 @@ def _test_linear( main_grad_check=("fuse_wgrad_accumulation" in kwargs), ) + _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel) + def test_linear(): """Run linear layer tests with various configurations.""" From e6b847e525313f48bb0ffb0693932b7738155926 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 16:36:20 +0200 Subject: [PATCH 33/50] Address review comments: re-drop the value-equality boilerplate, remove a redundant skip The a==b/hash/dict-key block and its other_kwargs parametrization were already removed once (6f66c3e2) as covered by the __fx_repr__ round-trip; a rebase resurrected them. The fp8_available skip in test_te_linear_compiles is dead: _all_recipes is availability-gated at construction. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 28 ++++++++-------------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index eab5f5114a..bf54e7dc19 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -782,14 +782,12 @@ def _hw_available(quantizer): return fp8_available # Float8CurrentScalingQuantizer -# (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ - pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), - pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), - pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param(_mxfp8, id="mxfp8"), + pytest.param(_blockwise, id="float8_blockwise"), + pytest.param(_current_scaling, id="float8_current_scaling"), pytest.param( _nvfp4, - {"with_rht": False}, id="nvfp4", marks=pytest.mark.skipif( not torch.cuda.is_available(), @@ -799,17 +797,10 @@ def _hw_available(quantizer): ] -@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) -def test_quantizer_value_object(factory, other_kwargs): +@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory): """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" - a, b = factory(), factory() - # Same config -> equal, same hash, interchangeable as a dict/set key. - assert a is not b - assert a == b - assert hash(a) == hash(b) - assert {a: "x"}[b] == "x" - # Different config -> not equal. - assert a != factory(**other_kwargs) + a = factory() # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. repr_str, globals_ = a.__fx_repr__() @@ -883,8 +874,8 @@ def _qdq_fake(x, q): not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", ) -@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) -def test_quantizer_value_object_fullgraph(factory, other_kwargs): +@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory): """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. A custom op quantizes+dequantizes with the (opaque value) quantizer; the @@ -1271,9 +1262,6 @@ def test_te_linear_compiles(fp8_recipe, compile_mode): recipe (plus the bf16-only baseline with no autocast), for both the default backend and ``mode="reduce-overhead"`` (CUDA-graph trees). """ - if fp8_recipe is not None and not fp8_available: - pytest.skip(reason_for_no_fp8) - dtype = torch.bfloat16 device = "cuda" From 1a00946027e8c98a8a91bca73199659fd2c87224 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 11 Aug 2026 20:46:31 +0200 Subject: [PATCH 34/50] Fix reduce-overhead UB warmup: mark step boundaries and drop grads between iterations The warmup iterations kept warmup gradients alive in the cudagraph pool, tripping cudagraph_trees' check_memory_pool on the next capture (Detected N tensor(s) in the cudagraph pool not tracked as outputs). Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_layer_with_overlap.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 058a78a756..adbce681c0 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -580,9 +580,11 @@ def run_fwd_bwd(model, x): if opts.compile and opts.compile_mode == "reduce-overhead": # Warm up so the measured run below replays captured CUDA graphs. for _ in range(2): + torch.compiler.cudagraph_mark_step_begin() run_fwd_bwd(test_model, test_x) - test_model.zero_grad(set_to_none=True) - test_x.grad = None + test_model.zero_grad(set_to_none=True) + test_x.grad = None + torch.compiler.cudagraph_mark_step_begin() test_out = run_fwd_bwd(test_model, test_x) if opts.compile and opts.compile_mode == "reduce-overhead" and dynamo_counters is not None: skips = dynamo_counters["inductor"]["cudagraph_skips"] From c0955825710acdce119b946de7052283ca3b0069 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 12 Aug 2026 13:44:03 +0200 Subject: [PATCH 35/50] Carry ProcessGroup through the op boundary by c10d registry name Replace the reference-opaque ProcessGroup graph input with the pattern traceable functional collectives use: the adapter ships pg.group_name (a plain string in the value bundle) and re-resolves the live group from the c10d registry inside the op, in the same process -- from_slots(to_slots(pg)) is the identical object by construction. This removes the opaque PG from example_inputs entirely, so inductor's FX disk cache works for compiled distributed graphs again (verified: second process gets fxgraph_cache_hit=2, no pickle bypass) without any upstream change. The reference-opaque registration machinery and the PG_REFERENCE_OPAQUE fallback gate are no longer needed. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 +- .../pytorch/dynamo/custom_op.py | 111 ++++++------------ transformer_engine/pytorch/module/linear.py | 9 +- 3 files changed, 39 insertions(+), 84 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 95d97b7e08..e42eb8f9f6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op, TensorOrQuantized, PG_REFERENCE_OPAQUE +from .custom_op import register_custom_op, TensorOrQuantized __all__ = [ "register_value_opaque_quantizer", @@ -15,5 +15,4 @@ "to_tensor_spec", "register_custom_op", "TensorOrQuantized", - "PG_REFERENCE_OPAQUE", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index de4123de1b..178892c74d 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -34,8 +34,8 @@ buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. * ``_QuantizerAdapter`` -- a quantizer, baked into the graph as a value-opaque constant. - * ``_ReferenceOpaqueAdapter`` -- a ProcessGroup, carried as a live opaque graph - input. + * ``_ProcessGroupAdapter`` -- a ProcessGroup, carried as its c10d registry + name and re-resolved inside the op. * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, sizes, nested collections of them), gathered into one ``OpaqueValueBundle`` slot. @@ -309,7 +309,6 @@ def _collect(value: Any) -> None: from torch._library.opaque_object import ( get_opaque_type_name, is_opaque_value_type as _is_opaque_value_type, - is_opaque_reference_type as _is_opaque_reference_type, register_opaque_type, ) @@ -321,53 +320,14 @@ def _collect(value: Any) -> None: f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build" ) _is_opaque_value_type = None - _is_opaque_reference_type = None _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None - -def _ensure_distributed_opaque_types() -> None: - """Register ``ProcessGroup`` as a *reference* opaque type: live state carried - as a graph input, not baked in as a constant. PyTorch runs this registration - only when DTensor is imported, so trigger it here too; best-effort no-op on - builds without the APIs (the field then falls back to eager). - - Note: PyTorch's cache-key pickler (``FxGraphCachePickler``) cannot hash the - real ``ProcessGroup`` in ``example_inputs`` yet, so inductor bypasses the FX - disk cache for compiled distributed calls (it logs a warning). Fixing that - belongs upstream; keeping TE free of process-wide pickle overrides. - """ - if _is_opaque_reference_type is None: - return - try: - from torch.distributed.device_mesh import ( # pylint: disable=import-outside-toplevel - _register_distributed_opaque_types, - ) - - _register_distributed_opaque_types() - except Exception: # pylint: disable=broad-exception-caught - pass - - -_ensure_distributed_opaque_types() - - -def _compute_pg_reference_opaque() -> bool: - if _is_opaque_reference_type is None: - return False - try: - from torch._C._distributed_c10d import ( # pylint: disable=import-outside-toplevel - ProcessGroup, - ) - - return bool(_is_opaque_reference_type(ProcessGroup)) - except Exception: # pylint: disable=broad-exception-caught - return False - - -# Whether ProcessGroup ended up registered as a reference-opaque type, i.e. -# whether a process group can cross the op boundary as a live graph input. -# Process-global and fixed at import, so a plain constant (Dynamo-friendly). -PG_REFERENCE_OPAQUE: bool = _compute_pg_reference_opaque() +try: + from torch._C._distributed_c10d import ProcessGroup as _PROCESS_GROUP_TYPE + from torch._C._distributed_c10d import _resolve_process_group +except ImportError: # pragma: no cover + _PROCESS_GROUP_TYPE = None + _resolve_process_group = None # --------------------------------------------------------------------------- # @@ -660,44 +620,45 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.meta_slot()][self.QUANTIZER_KEY] -class _ReferenceOpaqueAdapter(_Adapter): - """``ProcessGroup`` (or any reference-opaque type) -> one own opaque slot. - - A reference-opaque object is live, stateful black-box data (e.g. a - ``torch.distributed.ProcessGroup``): it cannot be specialized on or baked - into the graph as a constant the way a value-opaque quantizer is. torch.compile - instead carries it through as a graph *input*, so it passes straight through - its own schema slot (no ``OpaqueValueBundle`` wrapper). The field is annotated - with a concrete type registered via ``register_opaque_type(..., typ="reference")``. +class _ProcessGroupAdapter(_Adapter): + """``ProcessGroup`` -> its c10d registry name in one ``OpaqueValueBundle`` slot. - On the fake / setup-context path the slot holds a ``FakeScriptObject`` (or - ``None``); it is assigned to the field verbatim, so the fake impl must never - read the object's contents. + Mirrors traceable functional collectives: the graph carries the group's + *name* (a plain string, so guards and the FX cache key are trivial) and the + live group is re-resolved from the registry inside the op, in the same + process -- ``from_slots(to_slots(pg))`` returns the very group the caller + passed. Groups created outside the c10d registry fail the resolve loudly. """ - def __init__(self, name: str, type_name: str, is_optional: bool) -> None: + NAME_KEY = "group_name" + + def __init__(self, name: str) -> None: self.name = name - self.type_str = f"{type_name}?" if is_optional else type_name + + def meta_slot(self) -> str: + """Group-name slot name.""" + return self.name + "__pg" @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueAdapter"]: - if _is_opaque_reference_type is None: + def try_build(cls, name: str, annot: Any) -> Optional["_ProcessGroupAdapter"]: + if _PROCESS_GROUP_TYPE is None: return None - stripped, is_optional = _strip_optional(annot) - if not isinstance(stripped, type): - return None - if _is_opaque_reference_type(stripped): - return cls(name, get_opaque_type_name(stripped), is_optional) + stripped, _ = _strip_optional(annot) + if stripped is _PROCESS_GROUP_TYPE: + return cls(name) return None def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.name, self.type_str)] + return [(self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] def to_slots(self, owner: Any) -> Dict[str, Any]: - return {self.name: getattr(owner, self.name)} + pg = getattr(owner, self.name) + name = None if pg is None else pg.group_name + return {self.meta_slot(): OpaqueValueBundle({self.NAME_KEY: name})} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - kwargs[self.name] = args[self.name] + name = args[self.meta_slot()][self.NAME_KEY] + kwargs[self.name] = None if name is None else _resolve_process_group(name) class _SimpleBundleAdapter(_Adapter): @@ -782,8 +743,8 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: if not self._is_trivial(value): raise TypeError( f"{self.owner_cls_name} field {self.name!r} has a type not " - "supported by torch.compile (not Tensor, simple, Quantizer, or a " - "reference-opaque type such as ProcessGroup) and carries a " + "supported by torch.compile (not Tensor, simple, Quantizer, or " + "ProcessGroup) and carries a " "non-trivial value; add a matching adapter in custom_op.py to handle it." ) return {} @@ -797,7 +758,7 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: _FIELD_ADAPTERS: Tuple[type, ...] = ( _TensorOrQuantizedAdapter, _TensorAdapter, - _ReferenceOpaqueAdapter, + _ProcessGroupAdapter, _QuantizerAdapter, ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index eea0f3430d..d5759321be 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -90,7 +90,6 @@ TensorOrQuantized, register_custom_op, is_value_opaque_quantizer, - PG_REFERENCE_OPAQUE, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -157,8 +156,7 @@ class LinearFwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] - # ProcessGroup is a *reference*-opaque type: carried through the torch.compile - # custom op as a graph input (never baked into the graph as a constant). + # Crosses the op boundary as its c10d registry name, re-resolved in the op. tp_group: Optional[dist_group_type] tp_size: int tensor_parallel: bool @@ -210,9 +208,6 @@ def compile_unsupported_reason(self) -> Optional[str]: return "delayed wgrad compute (wgrad_store)" if self.fuse_wgrad_accumulation: return "fuse_wgrad_accumulation (main_grad)" - if self.tp_group is not None and not PG_REFERENCE_OPAQUE: - # ProcessGroup's reference-opaque registration failed at import. - return "a tp_group not registered as a torch.compile reference-opaque type" for quantizer in ( self.input_quantizer, self.weight_quantizer, @@ -264,7 +259,7 @@ class LinearBwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] = None - # Reference-opaque ProcessGroup (graph input), see LinearFwdArgs.tp_group. + # See LinearFwdArgs.tp_group. tp_group: Optional[dist_group_type] = None tp_size: int = 1 tensor_parallel: bool = False From 0dcaadd44c24944e3e0ff5914aee7418261ca8ca Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 12 Aug 2026 14:00:17 +0200 Subject: [PATCH 36/50] Add train/eval mode-switch test as xfail Blocked by an upstream Dynamo bug: FP8 state created inside the first compiled call comes back as FakeScriptObject/None in the graph outputs, so any later recompile crashes. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 44 +++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index bf54e7dc19..8946b5c8f9 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1447,6 +1447,50 @@ def ref_fn(inp): assert workspace is cached_workspace, f"cache rebuilt at step {step}" +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.xfail(reason="waiting for a PyTorch fix", strict=False) +def test_te_linear_compile_train_eval_switch(): + """train -> eval -> train on the same compiled ``te.Linear``, vs eager.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=True) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + def train_step(): + inp = torch.randn(16, 64, dtype=dtype, device=device, requires_grad=True) + out = compiled(inp) + out.sum().backward() + inp_ref = inp.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_ref = fn(inp_ref) + out_ref.sum().backward() + torch.testing.assert_close( + out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + model.zero_grad(set_to_none=True) + + train_step() + + model.eval() + x = torch.randn(16, 64, dtype=dtype, device=device) + with torch.no_grad(): + out_eval = compiled(x) + ref_eval = fn(x) + torch.testing.assert_close(out_eval, ref_eval, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + model.train() + train_step() + + @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") def test_te_linear_dynamic_shapes(): """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension: From 577b307bf57933d7f1dbe9dacd29a1e629dc03e3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 12 Aug 2026 17:16:25 +0200 Subject: [PATCH 37/50] Remove dynamic=False from distributed compile runners Each parametrized case resets dynamo and compiles once with a single shape, so automatic dynamic shapes never trigger; verified on 4xGB200 (run_numerics 126/126, comm-GEMM overlap compile suite green). Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_layer_with_overlap.py | 5 +---- tests/pytorch/distributed/run_numerics.py | 9 ++------- 2 files changed, 3 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index adbce681c0..1209b34c36 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -558,10 +558,7 @@ def run_fwd_bwd(model, x): if opts.compile: for i, layer in enumerate(test_model.layers): - # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py. - test_model.layers[i] = torch.compile( - layer, fullgraph=True, mode=opts.compile_mode, dynamic=False - ) + test_model.layers[i] = torch.compile(layer, fullgraph=True, mode=opts.compile_mode) dist_print( f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...", debug=True, diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 6160f0373f..7d331ccb55 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -335,13 +335,8 @@ def _apply_models( if use_compile: # Reset the compile cache so parametrized cases don't trip recompile_limit. torch._dynamo.reset() - # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py. - forward_single_node = torch.compile( - model_single_node, fullgraph=True, mode=compile_mode, dynamic=False - ) - forward_distributed = torch.compile( - model_distributed, fullgraph=True, mode=compile_mode, dynamic=False - ) + forward_single_node = torch.compile(model_single_node, fullgraph=True, mode=compile_mode) + forward_distributed = torch.compile(model_distributed, fullgraph=True, mode=compile_mode) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), From ff8bf74270c335a6867bb1edaab1ef11231e924e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 12 Aug 2026 17:16:25 +0200 Subject: [PATCH 38/50] Eagerly pre-allocate cuBLAS workspaces in Linear.reset_parameters The previous fetch in forward never executed for real: Dynamo ignores the lru_cache wrapper and traces the wrapped function, so the workspace was first allocated by the op impl at runtime - under reduce-overhead on capture-first torch builds that lands in the CUDA-graph pool and trips 'cudagraph pool not tracked as outputs'. Allocate both variants (plain and UB) in reset_parameters, which always runs eagerly, drop the dead cublas_workspace bundle field, and fail fast if a workspace would first be allocated during stream capture. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/cpp_extensions/gemm.py | 4 ++++ transformer_engine/pytorch/module/linear.py | 17 +++++++++-------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 03f31e28b5..6d013ec547 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -50,6 +50,10 @@ def get_cublas_workspace_size_bytes() -> None: def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor: """Returns workspace for cublas GEMM.""" assert not (ub and grouped_gemm), "UB is unsupported for grouped GEMM." + assert not torch.cuda.is_current_stream_capturing(), ( + "cuBLAS workspace would be first allocated during CUDA-graph capture and would live in" + " the graph's memory pool; run a warmup iteration outside the graph first." + ) if ub: return torch.empty( diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index d5759321be..8e1ac1bbfb 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -119,11 +119,6 @@ class LinearFwdArgs: # TensorOrQuantized so a cached quantized workspace can cross the op boundary. weight_workspace: Optional[TensorOrQuantized] - # Process-global cuBLAS workspace, fetched at trace time so it isn't first - # allocated during CUDA-graph capture. The op never reads it (general_gemm - # fetches the same global); None outside the compiled path. - cublas_workspace: Optional[torch.Tensor] - # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool weight_requires_grad: bool @@ -2283,6 +2278,15 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) + # Allocate the process-global cuBLAS workspaces eagerly: under + # torch.compile the first GEMM can run inside CUDA-graph capture, + # and a workspace first allocated there would live in the graph pool. + device = getattr(self, self.weight_names[0]).device + if device.type == "cuda": + get_cublas_workspace(device.index, False, False) + if self.ub_name is not None: + get_cublas_workspace(device.index, True, False) + def forward( self, inp: torch.Tensor, @@ -2414,7 +2418,6 @@ def forward( inp=inp, bias=linear_bias_tensor, weight_workspace=weight_workspace, - cublas_workspace=None, # set below, only on the compiled path # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, @@ -2483,8 +2486,6 @@ def forward( if use_compiled_op: check_gemm_dims(inp, weight_tensor, self.fp8) - # See LinearFwdArgs.cublas_workspace. - fwd_args.cublas_workspace = get_cublas_workspace(inp.device.index, False, False) out, new_weight_workspace = _linear_op(fwd_args) else: out, new_weight_workspace = _linear_eager( From 31b385e3d569de1c2d99a42922cf888539adba9c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 12:33:06 +0200 Subject: [PATCH 39/50] [PyTorch] [torch.compile] Address trivial review comments - drop unused Float8Quantizer import; import all quantizers from their tensor modules - import the local tests utils.py by explicit sys.path so a cutedsl top-level utils package cannot shadow it - cache OpaqueValueBundle hash at construction - guard is_simple_value when the opaque-object API is unavailable - inline the one-line _unflatten_value helper - point the --compile/--use-cuda-graphs error at --compile-mode reduce-overhead Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../distributed/run_layer_with_overlap.py | 6 +++- tests/pytorch/test_torch_compile.py | 13 +++++++-- .../pytorch/dynamo/custom_op.py | 28 ++++++------------- 3 files changed, 23 insertions(+), 24 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 1209b34c36..0b9b0aa13c 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -304,7 +304,11 @@ def _parse_args(argv=None, namespace=None): args = parser.parse_args(argv, namespace) if args.compile and args.use_cuda_graphs: - parser.error("--compile and --use-cuda-graphs are mutually exclusive.") + parser.error( + "--compile and --use-cuda-graphs are mutually exclusive; to test" + " torch.compile with CUDA graphs use --compile --compile-mode" + " reduce-overhead." + ) if args.use_cuda_graphs and args.layer_type in [te.MultiheadAttention, te.TransformerLayer]: warnings.warn(f"{args.layer_type.__name__} does not support CUDA Graphs!") diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8946b5c8f9..ae748b00e2 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,6 +4,8 @@ import abc import contextlib +import os +import sys import warnings import pytest @@ -35,6 +37,8 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec @@ -43,11 +47,14 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, - Float8Quantizer, - Float8BlockQuantizer, - MXFP8Quantizer, ) + +# Import from the local utils.py by explicit path: importing cutedsl makes a +# top-level ``utils`` package visible that would shadow it. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from utils import recipe_id + +sys.path.pop(0) from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 178892c74d..c5d2a4c2fd 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -52,7 +52,7 @@ * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat ``Tensor[]`` payload; * rebuilds the structured user outputs from that payload, sliced and reassembled - per the fake's output descriptors (``_unflatten_value``; + per the fake's output descriptors (``_unflatten_values``; ``_flatten_value`` is the pack-side inverse). Autograd, registered on the op, drives backward: @@ -179,7 +179,7 @@ def is_simple_value(cls, value: Any) -> bool: return True if isinstance(value, type): return True - if _is_opaque_value_type(type(value)): + if _is_opaque_value_type is not None and _is_opaque_value_type(type(value)): return True if isinstance(value, dict): return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()) @@ -243,6 +243,8 @@ def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: self._frozen: Tuple[Tuple[str, Any], ...] = tuple( (k, OpaqueValueBundle._to_hashable(v)) for k, v in sorted(data.items()) ) + # Precomputed: Dynamo guards hash bundles on every compiled call. + self._hash: int = hash(self._frozen) def __getitem__(self, key: str) -> Any: return self._data[key] @@ -271,7 +273,7 @@ def __eq__(self, other: object) -> bool: return self._frozen == other._frozen def __hash__(self) -> int: - return hash(self._frozen) + return self._hash def __fx_repr__(self) -> Tuple[str, Dict[str, Any]]: items = ", ".join( @@ -872,21 +874,6 @@ def _spec_slot_count(spec: Optional[TensorSpec]) -> int: return len(spec.inner_names()) -def _unflatten_value( - spec: Optional[TensorSpec], - chunk: List[Optional[torch.Tensor]], -) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]: - """Rebuild the value described by ``spec`` from its flat tensors ``chunk``. - - ``spec is None`` -> ``None`` (op-boundary sentinel for an absent output); - otherwise delegates to :meth:`TensorSpec.assemble`, which returns a plain - tensor as-is or reassembles a quantized tensor from its inner buffers. - """ - if spec is None: - return None - return spec.assemble(chunk) - - def _unflatten_values( specs: Sequence[Optional[TensorSpec]], flat: Sequence[Optional[torch.Tensor]], @@ -902,7 +889,8 @@ def _unflatten_values( n = _spec_slot_count(spec) chunk = [_decode_none(t) for t in flat[cursor : cursor + n]] cursor += n - values.append(_unflatten_value(spec, chunk)) + # ``spec is None`` is the op-boundary sentinel for an absent output. + values.append(spec.assemble(chunk) if spec is not None else None) return values, cursor @@ -911,7 +899,7 @@ def _flatten_value( ) -> List[torch.Tensor]: """Return the flat ``Tensor[]`` slots that represent one op output ``value``. - Inverse of :func:`_unflatten_value`; the slot count matches + Inverse of :func:`_unflatten_values`; the slot count matches :func:`_spec_slot_count`. """ if value is None: From 3bdce4ec291b2b2e3416e7462460fa692ca777b3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 15:35:50 +0200 Subject: [PATCH 40/50] Reference the upstream PyTorch fixes in the train/eval xfail pytorch/pytorch#187041; #187057 fixes the cold-compile path (merged), #193190 fixes the FX-graph-cache-hit path (in review). Both verified against this test on nightly 2.15.0.dev20260815. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index ae748b00e2..3790ba0957 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1456,7 +1456,14 @@ def ref_fn(inp): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.xfail(reason="waiting for a PyTorch fix", strict=False) +@pytest.mark.xfail( + reason=( + "value-opaque module state comes back as None on recompile" + " (pytorch/pytorch#187041; fixed by #187057 (cold compile, merged)" + " + #193190 (FX-graph-cache hit, in review))" + ), + strict=False, +) def test_te_linear_compile_train_eval_switch(): """train -> eval -> train on the same compiled ``te.Linear``, vs eager.""" dtype = torch.bfloat16 From 47fee780b8d062461c0e5c0fd00783bbca8e7255 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 15:35:50 +0200 Subject: [PATCH 41/50] Fold quantizers into the shared simple-value bundle The dedicated _QuantizerAdapter only existed because the bundle matches fields by annotation and quantizer fields are annotated with the abstract Quantizer base, which is not a registered opaque type itself. Match the base class in _SimpleBundleAdapter instead and drop the adapter; the per-quantizer schema slots carried no gradients. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 47 +++---------------- 1 file changed, 6 insertions(+), 41 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index c5d2a4c2fd..b8dcfdaf5e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -32,13 +32,11 @@ * ``_TensorOrQuantizedAdapter`` -- a field that may be a plain tensor, a bare quantized storage, or ``None``: three slots (the tensor, its flat inner buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. - * ``_QuantizerAdapter`` -- a quantizer, baked into the graph as a value-opaque - constant. * ``_ProcessGroupAdapter`` -- a ProcessGroup, carried as its c10d registry name and re-resolved inside the op. * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, - sizes, nested collections of them), gathered into one ``OpaqueValueBundle`` - slot. + sizes, quantizers -- value-opaque constants baked into the graph -- and + nested collections of them), gathered into one ``OpaqueValueBundle`` slot. * ``_UnsupportedAdapter`` -- fallback for a field no adapter can encode; allowed only when its value is trivial (``None`` / all-``None``) at call time. @@ -586,42 +584,6 @@ def grad_slot(self) -> Optional[int]: return 0 -class _QuantizerAdapter(_Adapter): - """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. - - Each quantizer gets its own dedicated slot. The field is annotated with the - base ``Quantizer`` (not itself a registered opaque type), so the simple - bundle would not claim it. - """ - - QUANTIZER_KEY = "q" - - def __init__(self, name: str) -> None: - self.name = name - - def meta_slot(self) -> str: - """Opaque quantizer metadata slot name.""" - return self.name + "__q" - - @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerAdapter"]: - stripped, _ = _strip_optional(annot) - if isinstance(stripped, type) and issubclass(stripped, Quantizer): - return cls(name) - return None - - def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - - def to_slots(self, owner: Any) -> Dict[str, Any]: - return { - self.meta_slot(): OpaqueValueBundle({self.QUANTIZER_KEY: getattr(owner, self.name)}) - } - - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - kwargs[self.name] = args[self.meta_slot()][self.QUANTIZER_KEY] - - class _ProcessGroupAdapter(_Adapter): """``ProcessGroup`` -> its c10d registry name in one ``OpaqueValueBundle`` slot. @@ -685,6 +647,10 @@ def matches_field(cls, annot: Any) -> bool: return True if isinstance(annot, type) and issubclass(annot, Enum): return True + # Quantizers are value-opaque constants; the abstract ``Quantizer`` + # annotation itself is not a registered opaque type, so match by base. + if isinstance(annot, type) and issubclass(annot, Quantizer): + return True if ( isinstance(annot, type) and _is_opaque_value_type is not None @@ -761,7 +727,6 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: _TensorOrQuantizedAdapter, _TensorAdapter, _ProcessGroupAdapter, - _QuantizerAdapter, ) From 6815201194d52fee7e8253c5eec7d8a006a4ab6f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 16:08:03 +0200 Subject: [PATCH 42/50] Make the eager-fallback warnings actually fire under torch.compile Dynamo silently drops warnings.warn in traced code, and the graph break inside the forward's try/finally makes it skip the whole frame and re-run it with is_compiling() == False, so neither warn_compile_eager_fallback nor warn_if_compile_disabled ever emitted. Emit them at trace time via torch._dynamo.comptime instead (once per compilation) and warn before the explicit graph break, which ends the trace. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 7 ++-- transformer_engine/pytorch/utils.py | 41 ++++++++++++++++----- 2 files changed, 36 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 8e1ac1bbfb..2330a8c314 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -2476,12 +2476,13 @@ def forward( if use_compiled_op: fallback_reason = fwd_args.compile_unsupported_reason() if fallback_reason is not None: - # Explicit break so fullgraph=True errors show the reason - # (warnings.warn below would break the graph inscrutably). + # Warn first: the break below makes Dynamo skip this frame, + # so anything after it is never traced. Explicit break so + # fullgraph=True errors show the reason. + warn_compile_eager_fallback(fallback_reason) torch._dynamo.graph_break( msg=f"te.Linear falling back to eager: {fallback_reason}" ) - warn_compile_eager_fallback(fallback_reason) use_compiled_op = False if use_compiled_op: diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index c6e0c63ebf..4fb9f0e2b2 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -30,6 +30,25 @@ _compile_disabled_reason: Optional[str] = None _compile_disabled_warned = False +try: + from torch._dynamo.comptime import comptime as _comptime +except ImportError: # pragma: no cover + _comptime = None + + +def _trace_time_warn(msg: str) -> None: + """Emit a warning from code being traced by Dynamo. + + Dynamo silently drops a traced ``warnings.warn``, and TE forwards wrap + everything in try/finally, so a graph break there makes Dynamo skip the + whole frame and re-run it with ``is_compiling() == False`` -- a runtime + warning branch is never reached. ``comptime`` runs for real inside the + compiler instead, so the warning fires once per compilation. The message + is read back via ``get_local`` (a traced closure would capture a + VariableTracker, not the value). + """ + _comptime(lambda ctx: warnings.warn(ctx.get_local("msg").as_python_constant())) + def record_compile_disabled(reason: str) -> None: """Record why TE's torch.compile custom-op path is off; the warning is @@ -47,34 +66,38 @@ def record_compile_disabled(reason: str) -> None: ) -@torch._dynamo.disable # graph-breaks cleanly so the warning actually fires def warn_if_compile_disabled() -> None: """Warn once, at the first compile attempt, that the path is off.""" global _compile_disabled_warned # pylint: disable=global-statement if _compile_disabled_warned: return _compile_disabled_warned = True - warnings.warn( + msg = ( "Transformer Engine torch.compile support is disabled: " f"{_compile_disabled_reason or 'custom-op registration unavailable'}. " "Modules will fall back to eager execution under torch.compile, i.e. " - "a graph break, which is incompatible with fullgraph=True.", - stacklevel=2, + "a graph break, which is incompatible with fullgraph=True." ) + if torch.compiler.is_compiling() and _comptime is not None: + _trace_time_warn(msg) + else: + warnings.warn(msg, stacklevel=2) def warn_compile_eager_fallback(reason: str) -> None: """Warn that a TE module is running eagerly under ``torch.compile``. Emitted when ``reason`` is unsupported on the module's compiled custom-op - path. Python's default warning filter dedups identical messages, so each - distinct ``reason`` is surfaced once. + path -- once per compilation (see :func:`_trace_time_warn`). """ - warnings.warn( + msg = ( f"Falling back to eager execution under torch.compile: {reason} is " - "unsupported on the compiled path (graph-breaks under fullgraph=True).", - stacklevel=2, + "unsupported on the compiled path (graph-breaks under fullgraph=True)." ) + if torch.compiler.is_compiling() and _comptime is not None: + _trace_time_warn(msg) + else: + warnings.warn(msg, stacklevel=2) @functools.lru_cache(maxsize=None) From db8c13cc7395515c2670a95f39ccbbdd7ad356e0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 16:08:03 +0200 Subject: [PATCH 43/50] Test the eager-fallback paths of the compiled Linear op Parametrized test over the single-GPU-constructible reasons rejected by compile_unsupported_reason (differentiable fp8_output, wgrad fusion / delay, quantized input): fallback warning fires, numerics match eager, fullgraph=True fails with the explicit reason. Delayed scaling is a hard error under fullgraph (check_recipe_support) and is tested separately. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 119 ++++++++++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 3790ba0957..cc384db137 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -5,6 +5,7 @@ import abc import contextlib import os +import re import sys import warnings @@ -1381,6 +1382,124 @@ def fn(inp): ) +# Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a +# single-GPU unit test can construct. Distributed-only reasons (fsdp_group, +# DistributedWeight) and CPU offloading need machinery this file doesn't have; +# delayed scaling is a hard error (check_recipe_support), tested separately. +# Modes: "bwd" = fwd+bwd vs eager; "fwd_grad" = grad-enabled forward only +# (differentiable fp8_output backward hits a PyTorch limitation: the Float8 +# output crossing the graph-break boundary gets a plain-tensor tangent); +# "no_grad" = forward under no_grad. +_FALLBACK_CASES = [ + "fp8_output_differentiable", + "fuse_wgrad_accumulation", + "delayed_wgrad", + "quantized_input", +] + + +def _fallback_case(case, dtype, device): + """Build ``(model, fn, mode, post_backward, reason)`` for one case.""" + model_kwargs = {} + if case == "fuse_wgrad_accumulation": + model_kwargs["fuse_wgrad_accumulation"] = True + elif case == "delayed_wgrad": + model_kwargs["delay_wgrad_compute"] = True + model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs) + + if case == "fp8_output_differentiable": + fp8_recipe = recipe.Float8CurrentScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, fp8_output=True).dequantize() + + return model, fn, "fwd_grad", None, "differentiable fp8_output=True" + if case == "fuse_wgrad_accumulation": + model.weight.main_grad = torch.zeros_like(model.weight, dtype=torch.float32) + return model, model, "bwd", None, "fuse_wgrad_accumulation" + if case == "delayed_wgrad": + return model, model, "bwd", model.backward_dw, "delayed wgrad compute" + if case == "quantized_input": + fp8_recipe = recipe.Float8CurrentScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + return model, fn, "no_grad", None, "a quantized input tensor" + raise ValueError(case) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("case", _FALLBACK_CASES) +def test_te_linear_compile_eager_fallback(case): + """Configs unsupported on the compiled custom-op path must fall back to + eager under ``torch.compile`` -- warning + numerics identical to eager -- + and graph-break with the explicit reason under ``fullgraph=True``.""" + dtype, device = torch.bfloat16, "cuda" + torch.manual_seed(0) + model_ref, fn_ref, mode, post_bwd_ref, _ = _fallback_case(case, dtype, device) + torch.manual_seed(0) + _model, fn, _, post_bwd, reason = _fallback_case(case, dtype, device) + + def make_inp(): + torch.manual_seed(1) + x = torch.randn(32, 64, dtype=dtype, device=device) + if case == "quantized_input": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, device=device + ) + return quantizer(x) + return x.requires_grad_(mode != "no_grad") + + torch._dynamo.reset() + compiled = torch.compile(fn) + grad_ctx = torch.no_grad() if mode == "no_grad" else contextlib.nullcontext() + + inp_ref, inp = make_inp(), make_inp() + with grad_ctx: + out_ref = fn_ref(inp_ref) + with pytest.warns(UserWarning, match="Falling back to eager execution under torch.compile"): + out = compiled(inp) + if mode == "bwd": + out_ref.sum().backward() + if post_bwd_ref is not None: + post_bwd_ref() + out.sum().backward() + if post_bwd is not None: + post_bwd() + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + torch._dynamo.reset() + compiled_fg = torch.compile(fn, fullgraph=True) + with pytest.raises(Exception, match=re.escape(reason)): + with grad_ctx: + compiled_fg(make_inp()) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_linear_compile_delayed_scaling_raises(): + """Delayed scaling is rejected under torch.compile with a hard error + (``check_recipe_support`` in ``te.autocast.__enter__``), not a fallback. + Without fullgraph the raising frame is skipped and re-run eagerly (where + the guard passes), so only ``fullgraph=True`` surfaces the error.""" + dtype, device = torch.bfloat16, "cuda" + model = te.Linear(64, 32, params_dtype=dtype, device=device) + fp8_recipe = recipe.DelayedScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + inp = torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True) + with pytest.raises(Exception, match="DelayedScaling is not supported under torch.compile"): + torch.compile(fn, fullgraph=True)(inp) + + @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("compile_mode", _compile_modes) From 43a078384e3429af62262f5da7ea94cc2f5bd5ff Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 16:14:19 +0200 Subject: [PATCH 44/50] Fold the compile-aware branching into the warn helper Both warn functions duplicated the is_compiling()/comptime dispatch; move it into _compile_safe_warn so callers just pass the message. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/utils.py | 31 +++++++++++++---------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 4fb9f0e2b2..c4c6024f47 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -36,18 +36,22 @@ _comptime = None -def _trace_time_warn(msg: str) -> None: - """Emit a warning from code being traced by Dynamo. +def _compile_safe_warn(msg: str) -> None: + """``warnings.warn`` that also works from code being traced by Dynamo. Dynamo silently drops a traced ``warnings.warn``, and TE forwards wrap everything in try/finally, so a graph break there makes Dynamo skip the whole frame and re-run it with ``is_compiling() == False`` -- a runtime - warning branch is never reached. ``comptime`` runs for real inside the - compiler instead, so the warning fires once per compilation. The message - is read back via ``get_local`` (a traced closure would capture a - VariableTracker, not the value). + warning branch is never reached. Under compilation ``comptime`` runs for + real inside the compiler instead, so the warning fires once per + compilation; the message is read back via ``get_local`` (a traced closure + would capture a VariableTracker, not the value). In eager this is a plain + ``warnings.warn``. """ - _comptime(lambda ctx: warnings.warn(ctx.get_local("msg").as_python_constant())) + if torch.compiler.is_compiling() and _comptime is not None: + _comptime(lambda ctx: warnings.warn(ctx.get_local("msg").as_python_constant())) + else: + warnings.warn(msg, stacklevel=3) def record_compile_disabled(reason: str) -> None: @@ -78,26 +82,19 @@ def warn_if_compile_disabled() -> None: "Modules will fall back to eager execution under torch.compile, i.e. " "a graph break, which is incompatible with fullgraph=True." ) - if torch.compiler.is_compiling() and _comptime is not None: - _trace_time_warn(msg) - else: - warnings.warn(msg, stacklevel=2) + _compile_safe_warn(msg) def warn_compile_eager_fallback(reason: str) -> None: """Warn that a TE module is running eagerly under ``torch.compile``. Emitted when ``reason`` is unsupported on the module's compiled custom-op - path -- once per compilation (see :func:`_trace_time_warn`). + path -- once per compilation (see :func:`_compile_safe_warn`). """ - msg = ( + _compile_safe_warn( f"Falling back to eager execution under torch.compile: {reason} is " "unsupported on the compiled path (graph-breaks under fullgraph=True)." ) - if torch.compiler.is_compiling() and _comptime is not None: - _trace_time_warn(msg) - else: - warnings.warn(msg, stacklevel=2) @functools.lru_cache(maxsize=None) From 0dc980a141af6d8cc482ea671be0d47262f3e990 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 17:00:15 +0200 Subject: [PATCH 45/50] Simplify adapter selection and drop the bundle's __getattr__ Replace the _FIELD_ADAPTERS registry + per-class try_build with a single _build_field_adapter factory dispatching on the annotation. Remove OpaqueValueBundle.__getattr__: nothing uses attribute access (consumers go through __getitem__/get/as_dict), and without it default copy/deepcopy/pickle work with no special-casing -- which matters since Dynamo's guard machinery deepcopies value-opaque objects. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 79 +++++-------------- 1 file changed, 19 insertions(+), 60 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b8dcfdaf5e..459c7ec336 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -247,16 +247,6 @@ def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: def __getitem__(self, key: str) -> Any: return self._data[key] - def __getattr__(self, name: str) -> Any: - # Underscored names raise cleanly: copy/pickle probe dunders on a clone - # created without __init__, where reading ``self._data`` would recurse. - if name.startswith("_"): - raise AttributeError(name) - try: - return self._data[name] - except KeyError as e: - raise AttributeError(name) from e - def get(self, key: str, default: Any = None) -> Any: """Return ``self._data.get(key, default)``.""" return self._data.get(key, default) @@ -399,23 +389,12 @@ class _Adapter: A custom op only takes flat, simply-typed arguments, but a TE op takes a single ``@dataclass`` of mixed fields. Each adapter knows how to translate - its kind of field both ways. ``try_build`` and ``schema_slots`` run once at - registration (to build the op's schema); ``to_slots`` and ``from_slots`` run - on each call and must agree on the slot layout that ``schema_slots`` declares. + its kind of field both ways; :func:`_build_field_adapter` picks the kind + from the field's annotation. ``schema_slots`` runs once at registration + (to build the op's schema); ``to_slots`` and ``from_slots`` run on each + call and must agree on the slot layout that ``schema_slots`` declares. """ - @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_Adapter"]: - """Decide whether this adapter type handles the field ``name`` given its - type annotation ``annot``; return a configured adapter if so, else - ``None`` so the next candidate is tried. - - Called once per field at registration, iterating :data:`_FIELD_ADAPTERS` - (adapters are mutually exclusive on annotations, so the order is not a - ranking). - """ - raise NotImplementedError - def schema_slots(self) -> List[Tuple[str, str]]: """Declare the schema slots this field occupies, each as a ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``). @@ -500,18 +479,13 @@ def schema_slots(self) -> List[Tuple[str, str]]: _MEMBERS = frozenset(get_args(TensorOrQuantized)) @classmethod - def _is_tensor_storage_union(cls, annot: Any) -> bool: + def is_tensor_storage_union(cls, annot: Any) -> bool: + """Whether ``annot`` is exactly the tensor-or-quantized union.""" if not _is_union(annot): return False members = frozenset(a for a in get_args(annot) if a is not type(None)) return members == cls._MEMBERS - @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedAdapter"]: - if cls._is_tensor_storage_union(annot): - return cls(name) - return None - def to_slots(self, owner: Any) -> Dict[str, Any]: value = getattr(owner, self.name) if value is None: @@ -564,13 +538,6 @@ def __init__(self, name: str, is_optional: bool) -> None: self.name = name self.type_str = "Tensor?" if is_optional else "Tensor" - @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_TensorAdapter"]: - stripped, is_optional = _strip_optional(annot) - if stripped is torch.Tensor: - return cls(name, is_optional) - return None - def schema_slots(self) -> List[Tuple[str, str]]: return [(self.name, self.type_str)] @@ -603,15 +570,6 @@ def meta_slot(self) -> str: """Group-name slot name.""" return self.name + "__pg" - @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_ProcessGroupAdapter"]: - if _PROCESS_GROUP_TYPE is None: - return None - stripped, _ = _strip_optional(annot) - if stripped is _PROCESS_GROUP_TYPE: - return cls(name) - return None - def schema_slots(self) -> List[Tuple[str, str]]: return [(self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] @@ -721,13 +679,18 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = None -# Adapter candidates for a single field, tried via ``try_build``. Mutually -# exclusive on annotations, so the order is not a ranking. -_FIELD_ADAPTERS: Tuple[type, ...] = ( - _TensorOrQuantizedAdapter, - _TensorAdapter, - _ProcessGroupAdapter, -) +def _build_field_adapter(name: str, annot: Any) -> Optional[_Adapter]: + """Pick the per-field adapter for one dataclass field from its annotation + (``None`` -> not a per-field kind; the caller falls back to the simple + bundle / unsupported).""" + if _TensorOrQuantizedAdapter.is_tensor_storage_union(annot): + return _TensorOrQuantizedAdapter(name) + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Tensor: + return _TensorAdapter(name, is_optional) + if _PROCESS_GROUP_TYPE is not None and stripped is _PROCESS_GROUP_TYPE: + return _ProcessGroupAdapter(name) + return None def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: @@ -752,11 +715,7 @@ def _get_adapters(cls: type) -> List[_Adapter]: adapters: List[_Adapter] = [] simple_names: List[str] = [] for name, annot in _resolved_field_annotations(cls): - built: Optional[_Adapter] = None - for adapter_cls in _FIELD_ADAPTERS: - built = adapter_cls.try_build(name, annot) - if built is not None: - break + built = _build_field_adapter(name, annot) if built is not None: adapters.append(built) elif _SimpleBundleAdapter.matches_field(annot): From 355a2cdd2bf55c60e7f88c74427f5ce84750ba14 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 17:11:45 +0200 Subject: [PATCH 46/50] Trim OpaqueValueBundle to its actual consumers Drop get(): the __kind__ tag is set at every construction site, so plain indexing (loud KeyError) is the right access. _storage_unflatten's only caller always passes a bundle, so drop the dead dict(meta) branch. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 459c7ec336..ea77ac34ad 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -247,10 +247,6 @@ def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: def __getitem__(self, key: str) -> Any: return self._data[key] - def get(self, key: str, default: Any = None) -> Any: - """Return ``self._data.get(key, default)``.""" - return self._data.get(key, default) - def as_dict(self) -> Dict[str, Any]: """Return a shallow copy of the stored mapping.""" return dict(self._data) @@ -347,9 +343,9 @@ def _storage_flatten( return OpaqueValueBundle(meta), tensors -def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any: +def _storage_unflatten(meta: "OpaqueValueBundle", tensors: List[torch.Tensor]) -> Any: """Inverse of :func:`_storage_flatten`.""" - meta_dict = meta.as_dict() if isinstance(meta, OpaqueValueBundle) else dict(meta) + meta_dict = meta.as_dict() inner_names = meta_dict["_inner_names"] inner = dict(zip(inner_names, tensors)) outer_shape = meta_dict.get("_outer_shape") @@ -517,7 +513,7 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: meta = args[self.meta_slot()] - kind = meta.get(self.KIND_KEY) + kind = meta[self.KIND_KEY] if kind == _TensorOrQuantizedKind.NONE: kwargs[self.name] = None elif kind == _TensorOrQuantizedKind.TENSOR: From d3f684927bf32c36c3fbe4d18b6f943010f7e4ae Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 17:22:38 +0200 Subject: [PATCH 47/50] Fold ProcessGroup fields into the shared simple-value bundle A live group can't cross as a value, so the bundle stores its c10d registry name and the op re-resolves it (same scheme the dedicated adapter used). Drops _ProcessGroupAdapter and the tp_group__pg schema slot; per-field adapters are now only the two tensor kinds. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 70 ++++++++----------- 1 file changed, 30 insertions(+), 40 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index ea77ac34ad..bba61afd49 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -32,11 +32,11 @@ * ``_TensorOrQuantizedAdapter`` -- a field that may be a plain tensor, a bare quantized storage, or ``None``: three slots (the tensor, its flat inner buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. - * ``_ProcessGroupAdapter`` -- a ProcessGroup, carried as its c10d registry - name and re-resolved inside the op. * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, sizes, quantizers -- value-opaque constants baked into the graph -- and nested collections of them), gathered into one ``OpaqueValueBundle`` slot. + A ProcessGroup field rides here too, as its c10d registry name, + re-resolved inside the op. * ``_UnsupportedAdapter`` -- fallback for a field no adapter can encode; allowed only when its value is trivial (``None`` / all-``None``) at call time. @@ -547,38 +547,6 @@ def grad_slot(self) -> Optional[int]: return 0 -class _ProcessGroupAdapter(_Adapter): - """``ProcessGroup`` -> its c10d registry name in one ``OpaqueValueBundle`` slot. - - Mirrors traceable functional collectives: the graph carries the group's - *name* (a plain string, so guards and the FX cache key are trivial) and the - live group is re-resolved from the registry inside the op, in the same - process -- ``from_slots(to_slots(pg))`` returns the very group the caller - passed. Groups created outside the c10d registry fail the resolve loudly. - """ - - NAME_KEY = "group_name" - - def __init__(self, name: str) -> None: - self.name = name - - def meta_slot(self) -> str: - """Group-name slot name.""" - return self.name + "__pg" - - def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - - def to_slots(self, owner: Any) -> Dict[str, Any]: - pg = getattr(owner, self.name) - name = None if pg is None else pg.group_name - return {self.meta_slot(): OpaqueValueBundle({self.NAME_KEY: name})} - - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - name = args[self.meta_slot()][self.NAME_KEY] - kwargs[self.name] = None if name is None else _resolve_process_group(name) - - class _SimpleBundleAdapter(_Adapter): """Aggregates every simple-typed field into a single OpaqueValueBundle. @@ -586,12 +554,19 @@ class _SimpleBundleAdapter(_Adapter): dataclass has no simple-typed fields): it owns the single shared ``_simple_meta`` slot, and ``_get_adapters`` builds it once from all simple-typed field names collected across the dataclass. + + ``pg_names`` marks the fields carrying a ProcessGroup: a live group can't + cross as a value, so -- mirroring traceable functional collectives -- the + bundle stores its c10d registry *name* and the op re-resolves the very + group the caller passed, in the same process. Groups created outside the + c10d registry fail the resolve loudly. """ META_SLOT = "_simple_meta" - def __init__(self, names: List[str]) -> None: + def __init__(self, names: List[str], pg_names: Sequence[str] = ()) -> None: self.names = list(names) + self.pg_names = frozenset(pg_names) @classmethod def matches_field(cls, annot: Any) -> bool: @@ -620,14 +595,19 @@ def schema_slots(self) -> List[Tuple[str, str]]: return [(self.META_SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] def to_slots(self, owner: Any) -> Dict[str, Any]: - return {self.META_SLOT: OpaqueValueBundle({n: getattr(owner, n) for n in self.names})} + data: Dict[str, Any] = {} + for n in self.names: + v = getattr(owner, n) + data[n] = v.group_name if n in self.pg_names and v is not None else v + return {self.META_SLOT: OpaqueValueBundle(data)} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: if self.META_SLOT not in args: return meta = args[self.META_SLOT] for n in self.names: - kwargs[n] = meta[n] + v = meta[n] + kwargs[n] = _resolve_process_group(v) if n in self.pg_names and v is not None else v class _UnsupportedAdapter(_Adapter): @@ -684,11 +664,17 @@ def _build_field_adapter(name: str, annot: Any) -> Optional[_Adapter]: stripped, is_optional = _strip_optional(annot) if stripped is torch.Tensor: return _TensorAdapter(name, is_optional) - if _PROCESS_GROUP_TYPE is not None and stripped is _PROCESS_GROUP_TYPE: - return _ProcessGroupAdapter(name) return None +def _is_process_group_annot(annot: Any) -> bool: + """Whether the field annotation is (Optional) ProcessGroup.""" + if _PROCESS_GROUP_TYPE is None: + return False + stripped, _ = _strip_optional(annot) + return stripped is _PROCESS_GROUP_TYPE + + def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: """Return ``[(field_name, resolved_type), ...]`` for a dataclass.""" if not dataclasses.is_dataclass(cls): @@ -710,16 +696,20 @@ def _get_adapters(cls: type) -> List[_Adapter]: ) adapters: List[_Adapter] = [] simple_names: List[str] = [] + pg_names: List[str] = [] for name, annot in _resolved_field_annotations(cls): built = _build_field_adapter(name, annot) if built is not None: adapters.append(built) + elif _is_process_group_annot(annot): + simple_names.append(name) + pg_names.append(name) elif _SimpleBundleAdapter.matches_field(annot): simple_names.append(name) else: adapters.append(_UnsupportedAdapter(name, cls.__name__)) if simple_names: - adapters.append(_SimpleBundleAdapter(simple_names)) + adapters.append(_SimpleBundleAdapter(simple_names, pg_names)) return adapters From 727def7c9151ef5c5f70b97c890370bcdee0411f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 19 Aug 2026 17:39:47 +0200 Subject: [PATCH 48/50] Reset FP8 global state between torch.compile tests Tests merged from main leave pending delayed-scaling amax reductions in FP8GlobalStateManager; a later autocast __exit__ then calls raw tex bindings, graph-breaking the fullgraph=True Linear tests. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index d2faccee76..8c136dd797 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -35,7 +35,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx from transformer_engine.pytorch.module.base import TransformerEngineBaseModule -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer @@ -68,6 +68,15 @@ nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True) +@pytest.fixture(autouse=True) +def _reset_fp8_global_state(): + """Pending FP8 global state (e.g. delayed-scaling amax reductions) must not + leak between tests: a leftover buffer makes a later autocast __exit__ call + raw tex bindings, which graph-breaks fullgraph=True tests.""" + yield + FP8GlobalStateManager.reset() + + def nvfp4_row_scaled(): nvfp4_recipe = recipe.NVFP4BlockScaling( disable_rht=True, From 97825f3fd50cf0816fa46b45565b20c60dfec7f4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 12:06:51 +0200 Subject: [PATCH 49/50] Rewrite the custom-op arg boundary as a parsed plan Parse the args dataclass's annotations once, at registration, into an immutable _ArgPlan: per-field _FieldPlan records (_FieldKind + schema slots) plus the derived layout -- schema string, slot order, gradient placement, tensor-or-quantized offsets -- with duplicate-slot-name validation. pack/unpack interpret the plan on each call. Replaces the adapter classes and the four layout helpers that each re-walked them; the op schema and Linear semantics are unchanged. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 753 ++++++++---------- 1 file changed, 312 insertions(+), 441 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index bba61afd49..1d4bb670cd 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -15,29 +15,28 @@ A ``torch.library`` custom op is narrower: it only accepts flat schema slots (tensors plus opaque objects) and returns a flat ``Tensor[]``. -Bridging the two takes three parts (below): per-field *adapters* map the args -dataclass onto the op's input slots; *fake impls* on data-free specs give the -output geometry and reassemble the op's flat return; and a *two-tier op* lets a -quantized-tensor subclass be an op input. - -Field <-> slot mapping. This mapping turns each field of the args dataclass into -the op's flat input slots, in a way that suits the field's type. A field's type -annotation selects the one ``_Adapter`` that handles it; that adapter declares the -slot(s) the field needs, packs the field's value into them on the way into the op, -and unpacks it back on the way out. The kinds -- and how each represents its field -as op inputs: - - * ``_TensorAdapter`` -- a plain ``Tensor`` / ``Optional[Tensor]``: one tensor - slot. - * ``_TensorOrQuantizedAdapter`` -- a field that may be a plain tensor, a bare +Bridging the two takes three parts (below): a parsed per-op *arg plan* maps the +args dataclass onto the op's input slots; *fake impls* on data-free specs give +the output geometry and reassemble the op's flat return; and a *two-tier op* +lets a quantized-tensor subclass be an op input. + +Field <-> slot mapping. ``_parse_arg_type`` parses the dataclass's field +annotations once, at registration, into an immutable ``_ArgPlan``: per field a +``_FieldPlan`` (its ``_FieldKind`` plus the schema slots it occupies), and the +derived layout -- schema string, slot order, gradient placement and +tensor-or-quantized offsets. ``_ArgPlan.pack`` / ``unpack`` interpret the plan +on each call. The kinds -- and how each represents its field as op inputs: + + * ``TENSOR`` -- a plain ``Tensor`` / ``Optional[Tensor]``: one tensor slot. + * ``TENSOR_OR_QUANTIZED`` -- a field that may be a plain tensor, a bare quantized storage, or ``None``: three slots (the tensor, its flat inner buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. - * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, - sizes, quantizers -- value-opaque constants baked into the graph -- and - nested collections of them), gathered into one ``OpaqueValueBundle`` slot. - A ProcessGroup field rides here too, as its c10d registry name, - re-resolved inside the op. - * ``_UnsupportedAdapter`` -- fallback for a field no adapter can encode; allowed + * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes, + quantizers -- value-opaque constants baked into the graph -- and nested + collections of them), gathered into one shared ``OpaqueValueBundle`` slot. + * ``PROCESS_GROUP`` -- rides in the shared bundle too, as its c10d registry + name, re-resolved inside the op. + * ``UNSUPPORTED`` -- a field no kind can encode; emits no slot and is allowed only when its value is trivial (``None`` / all-``None``) at call time. What runs where. Each op registers a data-free fake (``register_fake``) so it @@ -354,7 +353,10 @@ def _storage_unflatten(meta: "OpaqueValueBundle", tensors: List[torch.Tensor]) - # --------------------------------------------------------------------------- # -# Field adapters: dataclass field <-> flat torch.library slot(s) +# Arg plans: dataclass annotations are parsed once, at registration, into an +# immutable per-op plan (schema string, slot order, gradient placement, +# tensor-or-quantized offsets); packing / unpacking interpret that plan on +# each call. # --------------------------------------------------------------------------- # @@ -379,52 +381,14 @@ def _strip_optional(annot: Any) -> Tuple[Any, bool]: return annot, False -class _Adapter: - """Maps one (or, for the aggregating adapter, several) dataclass field(s) - to/from a contiguous run of custom-op schema *slots*. +class _FieldKind(Enum): + """How one dataclass field crosses the custom-op boundary.""" - A custom op only takes flat, simply-typed arguments, but a TE op takes a - single ``@dataclass`` of mixed fields. Each adapter knows how to translate - its kind of field both ways; :func:`_build_field_adapter` picks the kind - from the field's annotation. ``schema_slots`` runs once at registration - (to build the op's schema); ``to_slots`` and ``from_slots`` run on each - call and must agree on the slot layout that ``schema_slots`` declares. - """ - - def schema_slots(self) -> List[Tuple[str, str]]: - """Declare the schema slots this field occupies, each as a - ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``). - - Concatenated across all adapters to form the op's schema string. - """ - raise NotImplementedError - - def to_slots(self, owner: Any) -> Dict[str, Any]: - """Read this field from the dataclass ``owner`` and produce the concrete - value for each of its schema slots, as a ``{slot_name: value}`` dict. - - Composite values are flattened to fit the (tensor-only) slots: e.g. a - quantized tensor is split into its plain inner buffers plus a metadata - bundle. Inverse of :meth:`from_slots`. - """ - raise NotImplementedError - - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - """Read this field's slots back from the op arguments ``args`` and write - the reconstructed field value into ``kwargs`` (rebuilding any flattened - composite). The filled ``kwargs`` are then used to rebuild the original - dataclass for the eager implementation. Inverse of :meth:`to_slots`. - """ - raise NotImplementedError - - def grad_slot(self) -> Optional[int]: - """Index (within this adapter's :meth:`schema_slots`) of the slot that - carries a gradient, or ``None`` if the field is not differentiable. - - Used to map ``input_tensors_for_grad`` names onto backward grad-output - positions. Non-tensor adapters (quantizers, metadata) return ``None``. - """ - return None + TENSOR = "tensor" # one ``Tensor`` / ``Tensor?`` slot + TENSOR_OR_QUANTIZED = "tensor_or_quantized" # 3 slots: tensor / inner / meta + PROCESS_GROUP = "process_group" # c10d group name inside the shared bundle + SIMPLE = "simple" # value carried verbatim inside the shared bundle + UNSUPPORTED = "unsupported" # no slots; only a trivial value may cross class _TensorOrQuantizedKind(Enum): @@ -435,244 +399,238 @@ class _TensorOrQuantizedKind(Enum): STORAGE = "storage" -class _TensorOrQuantizedAdapter(_Adapter): - """``Tensor | QuantizedTensorStorage | None`` (also subclass tensor) field. - - Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass - tensor passes through, ``None`` for bare storage), ``__tensors`` - (``Tensor[]`` flat inner tensors when flattened), ``__meta`` - (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). A ``None`` - field is tagged ``_TensorOrQuantizedKind.NONE`` with the other two slots empty. - """ - - KIND_KEY = "__kind__" - - def __init__(self, name: str) -> None: - self.name = name - - def tensor_slot(self) -> str: - """Primary slot name for a plain / subclass tensor.""" - return self.name +_TQ_KIND_KEY = "__kind__" +_SIMPLE_META_SLOT = "_simple_meta" - def inner_slot(self) -> str: - """Flat inner-tensor slot name.""" - return self.name + "__tensors" +# Matched by exact member set, so a bare quantized annotation or an accidental +# extra union member is rejected rather than silently taken as tensor-or-quantized. +_TQ_MEMBERS = frozenset(get_args(TensorOrQuantized)) - def meta_slot(self) -> str: - """Flatten-metadata slot name.""" - return self.name + "__meta" - def schema_slots(self) -> List[Tuple[str, str]]: - return [ - (self.tensor_slot(), "Tensor?"), - (self.inner_slot(), "Tensor[]"), - (self.meta_slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), - ] - - # Matched by exact member set, so a bare quantized annotation or an - # accidental extra union member is rejected rather than silently taken as a - # tensor-or-quantized field. - _MEMBERS = frozenset(get_args(TensorOrQuantized)) - - @classmethod - def is_tensor_storage_union(cls, annot: Any) -> bool: - """Whether ``annot`` is exactly the tensor-or-quantized union.""" - if not _is_union(annot): - return False - members = frozenset(a for a in get_args(annot) if a is not type(None)) - return members == cls._MEMBERS - - def to_slots(self, owner: Any) -> Dict[str, Any]: - value = getattr(owner, self.name) - if value is None: - return { - self.tensor_slot(): None, - self.inner_slot(): [], - self.meta_slot(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE}), - } - if isinstance(value, torch.Tensor): - # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the - # ``Tensor?`` slot; subclass flattening (if any) is done by the - # wrapper op's ``register_torch_dispatch`` rule. - return { - self.tensor_slot(): value, - self.inner_slot(): [], - self.meta_slot(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR}), - } - if isinstance(value, QuantizedTensorStorage): - meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) - return { - self.tensor_slot(): None, - self.inner_slot(): tensors, - self.meta_slot(): meta, - } - raise TypeError( - f"field {self.name!r} expected None, torch.Tensor, or " - f"QuantizedTensorStorage, got {type(value).__name__}" - ) +@dataclasses.dataclass(frozen=True) +class _SlotSpec: + """One schema slot: its name and torch.library type string.""" - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - meta = args[self.meta_slot()] - kind = meta[self.KIND_KEY] - if kind == _TensorOrQuantizedKind.NONE: - kwargs[self.name] = None - elif kind == _TensorOrQuantizedKind.TENSOR: - kwargs[self.name] = args[self.tensor_slot()] - else: - kwargs[self.name] = _storage_unflatten(meta, args[self.inner_slot()]) + name: str + type_str: str - def grad_slot(self) -> Optional[int]: - # Gradient flows to the plain / subclass tensor slot (``tensor_slot()``, - # the first of the three). - return 0 +@dataclasses.dataclass(frozen=True) +class _FieldPlan: + """Parsed record for one dataclass field. -class _TensorAdapter(_Adapter): - """``Tensor`` / ``Optional[Tensor]`` -> single ``Tensor`` / ``Tensor?`` slot.""" + ``slots`` are the schema slots the field occupies (empty for the kinds that + ride in the shared simple bundle, or cross nothing); ``grad_slot`` is the + index within ``slots`` of the slot carrying the field's gradient, or ``None``. + """ - def __init__(self, name: str, is_optional: bool) -> None: - self.name = name - self.type_str = "Tensor?" if is_optional else "Tensor" + name: str + kind: _FieldKind + slots: Tuple[_SlotSpec, ...] + grad_slot: Optional[int] - def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.name, self.type_str)] - def to_slots(self, owner: Any) -> Dict[str, Any]: - return {self.name: getattr(owner, self.name)} +def _is_tensor_storage_union(annot: Any) -> bool: + """Whether ``annot`` is exactly the tensor-or-quantized union.""" + if not _is_union(annot): + return False + members = frozenset(a for a in get_args(annot) if a is not type(None)) + return members == _TQ_MEMBERS - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - kwargs[self.name] = args[self.name] - def grad_slot(self) -> Optional[int]: - return 0 +def _is_process_group_annot(annot: Any) -> bool: + """Whether the field annotation is (Optional) ProcessGroup.""" + if _PROCESS_GROUP_TYPE is None: + return False + stripped, _ = _strip_optional(annot) + return stripped is _PROCESS_GROUP_TYPE -class _SimpleBundleAdapter(_Adapter): - """Aggregates every simple-typed field into a single OpaqueValueBundle. +def _is_simple_annot(annot: Any) -> bool: + """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" + annot, _ = _strip_optional(annot) + if annot in OpaqueValueBundle.PRIMITIVE_TYPES: + return True + if isinstance(annot, type) and issubclass(annot, Enum): + return True + # Quantizers are value-opaque constants; the abstract ``Quantizer`` + # annotation itself is not a registered opaque type, so match by base. + if isinstance(annot, type) and issubclass(annot, Quantizer): + return True + if ( + isinstance(annot, type) + and _is_opaque_value_type is not None + and _is_opaque_value_type(annot) + ): + return True + if get_origin(annot) in (tuple, list): + inner = [a for a in get_args(annot) if a is not Ellipsis] + return bool(inner) and all(_is_simple_annot(a) for a in inner) + return False + + +def _parse_field(name: str, annot: Any) -> _FieldPlan: + """Parse one field's annotation into its :class:`_FieldPlan`.""" + if _is_tensor_storage_union(annot): + slots = ( + _SlotSpec(name, "Tensor?"), + _SlotSpec(name + "__tensors", "Tensor[]"), + _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + ) + return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots, grad_slot=0) + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Tensor: + slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor") + return _FieldPlan(name, _FieldKind.TENSOR, (slot,), grad_slot=0) + if _is_process_group_annot(annot): + return _FieldPlan(name, _FieldKind.PROCESS_GROUP, (), grad_slot=None) + if _is_simple_annot(annot): + return _FieldPlan(name, _FieldKind.SIMPLE, (), grad_slot=None) + return _FieldPlan(name, _FieldKind.UNSUPPORTED, (), grad_slot=None) - Unlike the per-field adapters, at most one of these exists per op (none if the - dataclass has no simple-typed fields): it owns the single shared - ``_simple_meta`` slot, and ``_get_adapters`` builds it once from all - simple-typed field names collected across the dataclass. - ``pg_names`` marks the fields carrying a ProcessGroup: a live group can't - cross as a value, so -- mirroring traceable functional collectives -- the - bundle stores its c10d registry *name* and the op re-resolves the very - group the caller passed, in the same process. Groups created outside the - c10d registry fail the resolve loudly. - """ +def _is_trivial(value: Any) -> bool: + """Whether an unsupported field's runtime value carries nothing.""" + if value is None: + return True + if isinstance(value, (list, tuple)): + return all(v is None for v in value) + return False - META_SLOT = "_simple_meta" - def __init__(self, names: List[str], pg_names: Sequence[str] = ()) -> None: - self.names = list(names) - self.pg_names = frozenset(pg_names) +def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, Any]) -> None: + """Fill a tensor-or-quantized field's three slots from its runtime value.""" + tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots) + if value is None: + slots[tensor_slot] = None + slots[inner_slot] = [] + slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE}) + elif isinstance(value, torch.Tensor): + # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the + # ``Tensor?`` slot; subclass flattening (if any) is done by the + # wrapper op's ``register_torch_dispatch`` rule. + slots[tensor_slot] = value + slots[inner_slot] = [] + slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.TENSOR}) + elif isinstance(value, QuantizedTensorStorage): + meta, tensors = _storage_flatten(value, {_TQ_KIND_KEY: _TensorOrQuantizedKind.STORAGE}) + slots[tensor_slot] = None + slots[inner_slot] = tensors + slots[meta_slot] = meta + else: + raise TypeError( + f"field {field.name!r} expected None, torch.Tensor, or " + f"QuantizedTensorStorage, got {type(value).__name__}" + ) - @classmethod - def matches_field(cls, annot: Any) -> bool: - """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" - annot, _ = _strip_optional(annot) - if annot in OpaqueValueBundle.PRIMITIVE_TYPES: - return True - if isinstance(annot, type) and issubclass(annot, Enum): - return True - # Quantizers are value-opaque constants; the abstract ``Quantizer`` - # annotation itself is not a registered opaque type, so match by base. - if isinstance(annot, type) and issubclass(annot, Quantizer): - return True - if ( - isinstance(annot, type) - and _is_opaque_value_type is not None - and _is_opaque_value_type(annot) - ): - return True - if get_origin(annot) in (tuple, list): - inner = [a for a in get_args(annot) if a is not Ellipsis] - return bool(inner) and all(cls.matches_field(a) for a in inner) - return False - def schema_slots(self) -> List[Tuple[str, str]]: - return [(self.META_SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - - def to_slots(self, owner: Any) -> Dict[str, Any]: - data: Dict[str, Any] = {} - for n in self.names: - v = getattr(owner, n) - data[n] = v.group_name if n in self.pg_names and v is not None else v - return {self.META_SLOT: OpaqueValueBundle(data)} - - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - if self.META_SLOT not in args: - return - meta = args[self.META_SLOT] - for n in self.names: - v = meta[n] - kwargs[n] = _resolve_process_group(v) if n in self.pg_names and v is not None else v - - -class _UnsupportedAdapter(_Adapter): - """Fallback for fields whose type no other adapter can encode. - - Such a field cannot cross the op boundary, so it emits no slot and is - tolerated only when its runtime value carries nothing: ``to_slots`` accepts - ``None`` / an all-``None`` sequence (e.g. an unset ``Optional[Any]`` field, - or an empty list, on the compiled path) and ``from_slots`` restores it as - ``None``. A non-trivial value means the config is genuinely unsupported - under torch.compile, and ``to_slots`` raises. - - The check must run at call time (not in ``_get_adapters``): the annotation - alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so - only the runtime value can decide. +def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any: + """Inverse of :func:`_pack_tensor_or_quantized`.""" + tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots) + meta = slots[meta_slot] + kind = meta[_TQ_KIND_KEY] + if kind == _TensorOrQuantizedKind.NONE: + return None + if kind == _TensorOrQuantizedKind.TENSOR: + return slots[tensor_slot] + return _storage_unflatten(meta, slots[inner_slot]) + + +@dataclasses.dataclass(frozen=True) +class _ArgPlan: + """The parsed plan for one args dataclass: the single source of truth for + the schema string, slot order, packing / unpacking, gradient placement and + the tensor-or-quantized slot offsets used for subclass flattening. + + Built once per registration by :func:`_parse_arg_type`; :meth:`pack` and + :meth:`unpack` interpret it on each call. ProcessGroup fields ride in the + shared bundle as their c10d registry *name* -- mirroring traceable + functional collectives -- and the op re-resolves the very group the caller + passed; groups created outside the c10d registry fail the resolve loudly. """ - def __init__(self, name: str, owner_cls_name: str) -> None: - self.name = name - self.owner_cls_name = owner_cls_name - - @staticmethod - def _is_trivial(value: Any) -> bool: - if value is None: - return True - if isinstance(value, (list, tuple)): - return all(v is None for v in value) - return False - - def schema_slots(self) -> List[Tuple[str, str]]: - return [] - - def to_slots(self, owner: Any) -> Dict[str, Any]: - value = getattr(owner, self.name, None) - if not self._is_trivial(value): - raise TypeError( - f"{self.owner_cls_name} field {self.name!r} has a type not " - "supported by torch.compile (not Tensor, simple, Quantizer, or " - "ProcessGroup) and carries a " - "non-trivial value; add a matching adapter in custom_op.py to handle it." + arg_type: type + fields: Tuple[_FieldPlan, ...] + slot_names: Tuple[str, ...] + schema_str: str + simple_slot: Optional[str] + tensor_field_names: Tuple[str, ...] + tq_offsets: Tuple[int, ...] + grad_slot_index: Dict[str, int] + + @property + def slot_count(self) -> int: + """Total number of input schema slots.""" + return len(self.slot_names) + + def resolve_grad_targets(self, input_tensors_for_grad: Sequence[str]) -> List[int]: + """Absolute schema-slot index receiving each requested field's gradient.""" + non_differentiable = [n for n in input_tensors_for_grad if n not in self.grad_slot_index] + if non_differentiable: + raise ValueError( + f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" ) - return {} + return [self.grad_slot_index[n] for n in input_tensors_for_grad] - def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: - kwargs[self.name] = None + def pack(self, obj: Any) -> Dict[str, Any]: + """Flatten an ``arg_type`` instance into the op's ``{slot: value}`` dict. - -def _build_field_adapter(name: str, annot: Any) -> Optional[_Adapter]: - """Pick the per-field adapter for one dataclass field from its annotation - (``None`` -> not a per-field kind; the caller falls back to the simple - bundle / unsupported).""" - if _TensorOrQuantizedAdapter.is_tensor_storage_union(annot): - return _TensorOrQuantizedAdapter(name) - stripped, is_optional = _strip_optional(annot) - if stripped is torch.Tensor: - return _TensorAdapter(name, is_optional) - return None - - -def _is_process_group_annot(annot: Any) -> bool: - """Whether the field annotation is (Optional) ProcessGroup.""" - if _PROCESS_GROUP_TYPE is None: - return False - stripped, _ = _strip_optional(annot) - return stripped is _PROCESS_GROUP_TYPE + Inverse of :meth:`unpack`. + """ + slots: Dict[str, Any] = {} + simple: Dict[str, Any] = {} + for field in self.fields: + value = getattr(obj, field.name, None) + match field.kind: + case _FieldKind.TENSOR: + slots[field.slots[0].name] = value + case _FieldKind.TENSOR_OR_QUANTIZED: + _pack_tensor_or_quantized(field, value, slots) + case _FieldKind.PROCESS_GROUP: + simple[field.name] = None if value is None else value.group_name + case _FieldKind.SIMPLE: + simple[field.name] = value + case _FieldKind.UNSUPPORTED: + # Annotation alone (e.g. Optional[Any]) can't decide; only + # the runtime value can, so the check runs at pack time. + if not _is_trivial(value): + raise TypeError( + f"{self.arg_type.__name__} field {field.name!r} has a type not " + "supported by torch.compile (not Tensor, simple, Quantizer, or " + "ProcessGroup) and carries a " + "non-trivial value; add a matching field kind in custom_op.py " + "to handle it." + ) + if self.simple_slot is not None: + slots[self.simple_slot] = OpaqueValueBundle(simple) + return slots + + def unpack(self, slots: Dict[str, Any]) -> Any: + """Rebuild a fresh ``arg_type`` instance from the op's flat slot dict. + + Inverse of :meth:`pack`. + """ + kwargs: Dict[str, Any] = {} + bundle = slots.get(self.simple_slot) if self.simple_slot is not None else None + for field in self.fields: + match field.kind: + case _FieldKind.TENSOR: + kwargs[field.name] = slots[field.slots[0].name] + case _FieldKind.TENSOR_OR_QUANTIZED: + kwargs[field.name] = _unpack_tensor_or_quantized(field, slots) + case _FieldKind.PROCESS_GROUP: + if bundle is not None: + name = bundle[field.name] + kwargs[field.name] = None if name is None else _resolve_process_group(name) + case _FieldKind.SIMPLE: + if bundle is not None: + kwargs[field.name] = bundle[field.name] + case _FieldKind.UNSUPPORTED: + kwargs[field.name] = None + obj = self.arg_type.__new__(self.arg_type) + for k, v in kwargs.items(): + object.__setattr__(obj, k, v) + return obj def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: @@ -686,69 +644,55 @@ def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: return [(f.name, hints.get(f.name, f.type)) for f in dataclasses.fields(cls)] -def _get_adapters(cls: type) -> List[_Adapter]: - """Build the adapter list for a dataclass from its field annotations.""" +def _parse_arg_type(cls: type) -> _ArgPlan: + """Parse an args ``@dataclass`` into its immutable :class:`_ArgPlan`. + + The layout pass assigns absolute slot positions (the shared simple bundle, + if any field needs it, takes the last slot) and validates the result: + duplicate slot names are rejected here, before any op is registered. + """ if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None: raise RuntimeError( f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle " "is not registered as a torch._library value-opaque type (PyTorch build " "without opaque-object support)." ) - adapters: List[_Adapter] = [] - simple_names: List[str] = [] - pg_names: List[str] = [] - for name, annot in _resolved_field_annotations(cls): - built = _build_field_adapter(name, annot) - if built is not None: - adapters.append(built) - elif _is_process_group_annot(annot): - simple_names.append(name) - pg_names.append(name) - elif _SimpleBundleAdapter.matches_field(annot): - simple_names.append(name) - else: - adapters.append(_UnsupportedAdapter(name, cls.__name__)) - if simple_names: - adapters.append(_SimpleBundleAdapter(simple_names, pg_names)) - return adapters - - -def _tensor_field_names(adapters: List[_Adapter]) -> List[str]: - """Names of fields carrying tensors (for building the spec view).""" - return [b.name for b in adapters if isinstance(b, (_TensorAdapter, _TensorOrQuantizedAdapter))] - - -def _build_schema(adapters: List[_Adapter]) -> Tuple[str, List[str]]: - """Return ``(schema_arg_str, slot_names)`` for an adapter list.""" - spec = [slot for b in adapters for slot in b.schema_slots()] - names = [name for name, _ in spec] - schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")" - return schema_str, names - - -def _args_to_slots(obj: Any, adapters: List[_Adapter]) -> Dict[str, Any]: - """Build the op's flat ``{slot_name: value}`` argument dict from an args - dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every adapter's - packed slot(s). Inverse of :func:`_args_from_slots`. - """ - out: Dict[str, Any] = {} - for adapter in adapters: - out.update(adapter.to_slots(obj)) - return out - - -def _args_from_slots(cls: type, args: Dict[str, Any], adapters: List[_Adapter]) -> Any: - """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the - op's flat slot ``args`` dict, by letting every adapter restore its field(s). - Inverse of :func:`_args_to_slots`. - """ - kwargs: Dict[str, Any] = {} - for adapter in adapters: - adapter.from_slots(args, kwargs) - obj = cls.__new__(cls) - for k, v in kwargs.items(): - object.__setattr__(obj, k, v) - return obj + fields = tuple(_parse_field(name, annot) for name, annot in _resolved_field_annotations(cls)) + + slot_specs: List[_SlotSpec] = [] + tq_offsets: List[int] = [] + tensor_field_names: List[str] = [] + grad_slot_index: Dict[str, int] = {} + for field in fields: + if field.kind is _FieldKind.TENSOR_OR_QUANTIZED: + tq_offsets.append(len(slot_specs)) + if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): + tensor_field_names.append(field.name) + if field.grad_slot is not None: + grad_slot_index[field.name] = len(slot_specs) + field.grad_slot + slot_specs.extend(field.slots) + if any(f.kind in (_FieldKind.SIMPLE, _FieldKind.PROCESS_GROUP) for f in fields): + simple_slot: Optional[str] = _SIMPLE_META_SLOT + slot_specs.append(_SlotSpec(_SIMPLE_META_SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)) + else: + simple_slot = None + + slot_names = tuple(s.name for s in slot_specs) + if len(set(slot_names)) != len(slot_names): + dupes = sorted(n for n in set(slot_names) if slot_names.count(n) > 1) + raise ValueError(f"{cls.__name__}: duplicate schema slot names: {dupes}") + schema_str = "(" + ", ".join(f"{s.type_str} {s.name}" for s in slot_specs) + ")" + + return _ArgPlan( + arg_type=cls, + fields=fields, + slot_names=slot_names, + schema_str=schema_str, + simple_slot=simple_slot, + tensor_field_names=tuple(tensor_field_names), + tq_offsets=tuple(tq_offsets), + grad_slot_index=grad_slot_index, + ) def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any: @@ -912,45 +856,11 @@ def _unpack_fwd_fake_result( # --------------------------------------------------------------------------- # -def _resolve_grad_targets( - fwd_adapters: List[_Adapter], - input_tensors_for_grad: List[str], -) -> Tuple[int, List[int]]: - """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. - - ``fwd_adapters`` already encode the arg dataclass's fields (they are built - from it), so the type itself is not needed here. - - Returns ``(slot_count, grad_targets)``: the total number of input schema - slots and, for each requested input name, the schema-slot index its gradient - maps to. - """ - name_to_slot: Dict[str, int] = {} - slot_offset = 0 - for adapter in fwd_adapters: - slots = adapter.schema_slots() - grad_slot = adapter.grad_slot() - if grad_slot is not None: - name_to_slot[adapter.name] = slot_offset + grad_slot - slot_offset += len(slots) - - non_differentiable = [n for n in input_tensors_for_grad if n not in name_to_slot] - if non_differentiable: - raise ValueError( - f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" - ) - grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] - return slot_offset, grad_targets - - def _register_base_op( *, op_name: str, schema_str: str, - arg_type: type, - arg_names: List[str], - adapters: List[_Adapter], - tensor_field_names: List[str], + plan: _ArgPlan, impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any], pack_result: Callable[[Any], List[torch.Tensor]], @@ -964,14 +874,12 @@ def _register_base_op( """ def _impl(*flat: Any) -> List[torch.Tensor]: - kwargs = dict(zip(arg_names, flat)) - obj = _args_from_slots(arg_type, kwargs, adapters) + obj = plan.unpack(dict(zip(plan.slot_names, flat))) return pack_result(impl(obj)) def _fake(*flat: Any) -> List[torch.Tensor]: - kwargs = dict(zip(arg_names, flat)) - obj = _args_from_slots(arg_type, kwargs, adapters) - spec_obj = _spec_view(obj, tensor_field_names) + obj = plan.unpack(dict(zip(plan.slot_names, flat))) + spec_obj = _spec_view(obj, plan.tensor_field_names) return pack_result(fake_impl(spec_obj)) op = torch.library.custom_op( @@ -985,16 +893,10 @@ def _register_autograd_for_op( *, fwd_op: Any, bwd_op: Any, - fwd_arg_type: type, - fwd_arg_names: List[str], - fwd_adapters: List[_Adapter], - fwd_tensor_field_names: List[str], - bwd_arg_names: List[str], - bwd_adapters: List[_Adapter], - slot_count: int, + fwd_plan: _ArgPlan, + bwd_plan: _ArgPlan, grad_targets: List[int], setup_context_user: Callable[..., Any], - bwd_arg_type: type, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> None: """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. @@ -1008,16 +910,15 @@ def _setup_context(ctx, inputs, output): ctx.fwd_tensor_list_lengths = { i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) } - kwargs = dict(zip(fwd_arg_names, inputs)) - fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_adapters) - spec_obj = _spec_view(fwd_obj, fwd_tensor_field_names) + fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs))) + spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names) user_fakes, saved_fakes, ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) user_outputs, cursor = _unflatten_values(user_fakes, output) saved_list, _ = _unflatten_values(saved_fakes, output, cursor) - bwd_obj = bwd_arg_type() + bwd_obj = bwd_plan.arg_type() tensors_to_save_from_setup = setup_context_user( bwd_obj, fwd_obj, @@ -1037,14 +938,14 @@ def _autograd_backward(ctx, *grad_outputs): ctx.tensor_objects = None flat_grads = grad_outputs[0] bwd_obj.grad_output = _decode_none(flat_grads[0]) - kwargs = _args_to_slots(bwd_obj, bwd_adapters) - bwd_args_flat = [kwargs[name] for name in bwd_arg_names] + kwargs = bwd_plan.pack(bwd_obj) + bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] ctx.backward_objects = None # One grad per input schema slot: default None, but a ``Tensor[]`` slot # (always recorded in ``fwd_tensor_list_lengths``) needs a # list-shaped no-grad of matching length. - out: List[Any] = [None] * slot_count + out: List[Any] = [None] * fwd_plan.slot_count for pos, length in ctx.fwd_tensor_list_lengths.items(): out[pos] = [None] * length for pos, g in zip(grad_targets, grads): @@ -1054,30 +955,17 @@ def _autograd_backward(ctx, *grad_outputs): fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) -def _tensor_or_quantized_offsets(adapters: List[_Adapter]) -> List[int]: - """Start index of each ``_TensorOrQuantizedAdapter`` group in the flat args.""" - offsets: List[int] = [] - pos = 0 - for adapter in adapters: - if isinstance(adapter, _TensorOrQuantizedAdapter): - offsets.append(pos) - pos += len(adapter.schema_slots()) - return offsets - - def _flatten_subclass_into_slots( - new_args: List[Any], slot_offsets: List[int], subclass: type + new_args: List[Any], slot_offsets: Sequence[int], subclass: type ) -> None: - """Rewrite each tensor-or-quantized-adapter group whose ``Tensor?`` slot holds an + """Rewrite each tensor-or-quantized slot group whose ``Tensor?`` slot holds an instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). """ for offset in slot_offsets: val = new_args[offset] if not isinstance(val, subclass): continue - meta, tensors = _storage_flatten( - val, {_TensorOrQuantizedAdapter.KIND_KEY: _TensorOrQuantizedKind.STORAGE} - ) + meta, tensors = _storage_flatten(val, {_TQ_KIND_KEY: _TensorOrQuantizedKind.STORAGE}) new_args[offset] = None new_args[offset + 1] = tensors new_args[offset + 2] = meta @@ -1252,8 +1140,8 @@ def _register_custom_op_impl( """Body of :func:`register_custom_op`; see it for semantics.""" # Existence check at the API boundary: every ``input_tensors_for_grad`` name # must be an actual field of ``fwd_arg_type`` (differentiability -- whether - # that field can carry a gradient -- is checked later in - # :func:`_resolve_grad_targets`). + # that field can carry a gradient -- is checked later, in + # :meth:`_ArgPlan.resolve_grad_targets`). fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] if missing: @@ -1265,29 +1153,21 @@ def _register_custom_op_impl( base_bwd_name = f"{wrapper_bwd_name}_base" subclass_list = _all_quantized_tensor_subclasses() - fwd_adapters = _get_adapters(fwd_arg_type) - bwd_adapters = _get_adapters(bwd_arg_type) - fwd_tensor_field_names = _tensor_field_names(fwd_adapters) - bwd_tensor_field_names = _tensor_field_names(bwd_adapters) - - fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) - bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + fwd_plan = _parse_arg_type(fwd_arg_type) + bwd_plan = _parse_arg_type(bwd_arg_type) num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) + grad_targets = fwd_plan.resolve_grad_targets(input_tensors_for_grad) - fwd_schema = f"{fwd_schema_args} -> Tensor[]" - bwd_schema = f"{bwd_schema_args} -> Tensor[]" + fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]" + bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]" base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" base_fwd_def = _register_base_op( op_name=base_fwd_name, schema_str=fwd_schema, - arg_type=fwd_arg_type, - arg_names=fwd_arg_names, - adapters=fwd_adapters, - tensor_field_names=fwd_tensor_field_names, + plan=fwd_plan, impl=fwd_impl, fake_impl=fwd_fake_impl, pack_result=_pack_fwd_result, @@ -1295,10 +1175,7 @@ def _register_custom_op_impl( _register_base_op( op_name=base_bwd_name, schema_str=bwd_schema, - arg_type=bwd_arg_type, - arg_names=bwd_arg_names, - adapters=bwd_adapters, - tensor_field_names=bwd_tensor_field_names, + plan=bwd_plan, impl=bwd_impl, fake_impl=bwd_fake_impl, pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname), @@ -1307,8 +1184,8 @@ def _register_custom_op_impl( base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) - fwd_slot_offsets = _tensor_or_quantized_offsets(fwd_adapters) - bwd_slot_offsets = _tensor_or_quantized_offsets(bwd_adapters) + fwd_slot_offsets = fwd_plan.tq_offsets + bwd_slot_offsets = bwd_plan.tq_offsets wrapper_fwd_def = _register_wrapper_op( wrapper_op_name=wrapper_fwd_name, @@ -1324,16 +1201,10 @@ def _register_custom_op_impl( ) autograd_common = { - "fwd_arg_type": fwd_arg_type, - "fwd_arg_names": fwd_arg_names, - "fwd_adapters": fwd_adapters, - "fwd_tensor_field_names": fwd_tensor_field_names, - "bwd_arg_names": bwd_arg_names, - "bwd_adapters": bwd_adapters, - "slot_count": slot_count, + "fwd_plan": fwd_plan, + "bwd_plan": bwd_plan, "grad_targets": grad_targets, "setup_context_user": setup_context, - "bwd_arg_type": bwd_arg_type, "fwd_fake_impl": fwd_fake_impl, } wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) @@ -1359,10 +1230,10 @@ def _register_custom_op_impl( _quantized_tensor_passthrough_ops.add(base_bwd_op.default) def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) + spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names) user_fakes, _saved_fakes, _ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) - kwargs = _args_to_slots(fwd_args, fwd_adapters) - flat_in = [kwargs[name] for name in fwd_arg_names] + kwargs = fwd_plan.pack(fwd_args) + flat_in = [kwargs[name] for name in fwd_plan.slot_names] result = wrapper_fwd_op(*flat_in) outputs, _ = _unflatten_values(user_fakes, result) From 97bb4beec14f0a0595ba025842ded218dc0719e3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 12:13:12 +0200 Subject: [PATCH 50/50] Add an output plan and lift the single-grad-output limit Parse the fwd fake-impl result into a per-trace _OutputPlan (logical outputs / saved tensors with their flat Tensor[] ranges) and use it as the single structure behind forward_fn, setup_context and backward. Backward now slices grads per user output from the plan stashed on ctx: a grad_outputs field on the backward args receives the whole tuple, otherwise grad_output receives the first output's grad -- removing the flat_grads[0] single-output assumption. Also reject unions mixing tensor types with other members at registration. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 237 ++++++++++++------ 1 file changed, 156 insertions(+), 81 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1d4bb670cd..1378e65acb 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -16,9 +16,10 @@ (tensors plus opaque objects) and returns a flat ``Tensor[]``. Bridging the two takes three parts (below): a parsed per-op *arg plan* maps the -args dataclass onto the op's input slots; *fake impls* on data-free specs give -the output geometry and reassemble the op's flat return; and a *two-tier op* -lets a quantized-tensor subclass be an op input. +args dataclass onto the op's input slots; a per-trace *output plan*, parsed from +the data-free fake impl's result, maps the logical outputs / saved tensors / +grads onto ranges of the flat return; and a *two-tier op* lets a +quantized-tensor subclass be an op input. Field <-> slot mapping. ``_parse_arg_type`` parses the dataclass's field annotations once, at registration, into an immutable ``_ArgPlan``: per field a @@ -45,24 +46,26 @@ call through it: * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see - ``tensor_spec.py``) to get the outputs' geometry in pure Python; + ``tensor_spec.py``) and parses its result into an ``_OutputPlan`` -- the + outputs' geometry and their ranges in the flat payload, in pure Python; * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat ``Tensor[]`` payload; - * rebuilds the structured user outputs from that payload, sliced and reassembled - per the fake's output descriptors (``_unflatten_values``; - ``_flatten_value`` is the pack-side inverse). + * rebuilds the structured user outputs from that payload per the plan + (``_OutputPlan.user_outputs``; ``_flatten_value`` is the pack-side inverse). Autograd, registered on the op, drives backward: - * ``setup_context`` (run when the forward is taped) re-runs ``fwd_fake_impl`` for - the saved-tensor descriptors and a ``ctx_attrs`` dict, reassembles the saved - tensors from the op's flat output, then calls the user ``setup_context`` to - fill the backward args from forward state + ``ctx_attrs`` (e.g. saved-tensor - aliases) and return the tensors to persist; - * on ``backward()`` the backward args container's optional ``setup_saved_tensors`` - hook restores those saved tensors, then the *backward op* runs the real - ``bwd_impl`` and returns the flat grads (``bwd_fake_impl`` is its - data-free fake). + * ``setup_context`` (run when the forward is taped) re-runs ``fwd_fake_impl``, + parses the ``_OutputPlan``, reassembles the saved tensors from the op's flat + output, then calls the user ``setup_context`` to fill the backward args from + forward state + ``ctx_attrs`` (e.g. saved-tensor aliases) and return the + tensors to persist; the plan is stashed on ``ctx``; + * on ``backward()`` the incoming flat grads are sliced per user output from the + stashed plan (a ``grad_outputs`` field on the backward args receives the + whole tuple; otherwise ``grad_output`` receives the first output's grad), + the container's optional ``setup_saved_tensors`` hook restores the saved + tensors, then the *backward op* runs the real ``bwd_impl`` and returns the + flat grads (``bwd_fake_impl`` is its data-free fake). Two-tier op (``base`` + ``wrapper``), so a ``QuantizedTensor`` subclass can be an op *input*. The ``_base`` op carries the real schema + autograd; a custom op @@ -420,14 +423,12 @@ class _FieldPlan: """Parsed record for one dataclass field. ``slots`` are the schema slots the field occupies (empty for the kinds that - ride in the shared simple bundle, or cross nothing); ``grad_slot`` is the - index within ``slots`` of the slot carrying the field's gradient, or ``None``. + ride in the shared simple bundle, or cross nothing). """ name: str kind: _FieldKind slots: Tuple[_SlotSpec, ...] - grad_slot: Optional[int] def _is_tensor_storage_union(annot: Any) -> bool: @@ -477,16 +478,29 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan: _SlotSpec(name + "__tensors", "Tensor[]"), _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME), ) - return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots, grad_slot=0) + return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots) stripped, is_optional = _strip_optional(annot) if stripped is torch.Tensor: slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor") - return _FieldPlan(name, _FieldKind.TENSOR, (slot,), grad_slot=0) + return _FieldPlan(name, _FieldKind.TENSOR, (slot,)) + # A union mixing tensor types with anything else is a malformed signature + # (e.g. a bare quantized-storage Optional, or Tensor | int): reject it at + # registration instead of silently degrading to an unsupported field. + if _is_union(annot): + members = [a for a in get_args(annot) if a is not type(None)] + if any( + isinstance(m, type) and issubclass(m, (torch.Tensor, QuantizedTensorStorage)) + for m in members + ): + raise TypeError( + f"field {name!r}: union {annot!r} is not a supported tensor " + "signature; use Tensor, Optional[Tensor], or TensorOrQuantized." + ) if _is_process_group_annot(annot): - return _FieldPlan(name, _FieldKind.PROCESS_GROUP, (), grad_slot=None) + return _FieldPlan(name, _FieldKind.PROCESS_GROUP, ()) if _is_simple_annot(annot): - return _FieldPlan(name, _FieldKind.SIMPLE, (), grad_slot=None) - return _FieldPlan(name, _FieldKind.UNSUPPORTED, (), grad_slot=None) + return _FieldPlan(name, _FieldKind.SIMPLE, ()) + return _FieldPlan(name, _FieldKind.UNSUPPORTED, ()) def _is_trivial(value: Any) -> bool: @@ -556,7 +570,6 @@ class _ArgPlan: simple_slot: Optional[str] tensor_field_names: Tuple[str, ...] tq_offsets: Tuple[int, ...] - grad_slot_index: Dict[str, int] @property def slot_count(self) -> int: @@ -564,13 +577,26 @@ def slot_count(self) -> int: return len(self.slot_names) def resolve_grad_targets(self, input_tensors_for_grad: Sequence[str]) -> List[int]: - """Absolute schema-slot index receiving each requested field's gradient.""" - non_differentiable = [n for n in input_tensors_for_grad if n not in self.grad_slot_index] + """Absolute schema-slot index receiving each requested field's gradient. + + Derived from ``fields`` on demand -- called once per registration, so + the plan doesn't cache the mapping. + """ + index: Dict[str, int] = {} + offset = 0 + for field in self.fields: + if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): + # The gradient flows to the group's first slot -- for + # tensor-or-quantized that is the ``Tensor?`` slot, the one + # autograd sees the (subclass) tensor in. + index[field.name] = offset + offset += len(field.slots) + non_differentiable = [n for n in input_tensors_for_grad if n not in index] if non_differentiable: raise ValueError( f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" ) - return [self.grad_slot_index[n] for n in input_tensors_for_grad] + return [index[n] for n in input_tensors_for_grad] def pack(self, obj: Any) -> Dict[str, Any]: """Flatten an ``arg_type`` instance into the op's ``{slot: value}`` dict. @@ -662,14 +688,11 @@ def _parse_arg_type(cls: type) -> _ArgPlan: slot_specs: List[_SlotSpec] = [] tq_offsets: List[int] = [] tensor_field_names: List[str] = [] - grad_slot_index: Dict[str, int] = {} for field in fields: if field.kind is _FieldKind.TENSOR_OR_QUANTIZED: tq_offsets.append(len(slot_specs)) if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): tensor_field_names.append(field.name) - if field.grad_slot is not None: - grad_slot_index[field.name] = len(slot_specs) + field.grad_slot slot_specs.extend(field.slots) if any(f.kind in (_FieldKind.SIMPLE, _FieldKind.PROCESS_GROUP) for f in fields): simple_slot: Optional[str] = _SIMPLE_META_SLOT @@ -691,7 +714,6 @@ def _parse_arg_type(cls: type) -> _ArgPlan: simple_slot=simple_slot, tensor_field_names=tuple(tensor_field_names), tq_offsets=tuple(tq_offsets), - grad_slot_index=grad_slot_index, ) @@ -728,33 +750,13 @@ def _spec_slot_count(spec: Optional[TensorSpec]) -> int: return len(spec.inner_names()) -def _unflatten_values( - specs: Sequence[Optional[TensorSpec]], - flat: Sequence[Optional[torch.Tensor]], - cursor: int = 0, -) -> Tuple[List[Any], int]: - """Rebuild one group of values from an op's flat return, starting at ``cursor``. - - Returns the values and the new cursor, so consecutive groups (user outputs, - then saved tensors) can walk the same payload. - """ - values: List[Any] = [] - for spec in specs: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in flat[cursor : cursor + n]] - cursor += n - # ``spec is None`` is the op-boundary sentinel for an absent output. - values.append(spec.assemble(chunk) if spec is not None else None) - return values, cursor - - def _flatten_value( value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorSpec]], ) -> List[torch.Tensor]: """Return the flat ``Tensor[]`` slots that represent one op output ``value``. - Inverse of :func:`_unflatten_values`; the slot count matches - :func:`_spec_slot_count`. + Pack-side inverse of :meth:`_OutputPlan.user_outputs`; the slot count + matches :func:`_spec_slot_count`. """ if value is None: return [_encode_none(None)] @@ -782,7 +784,7 @@ def _check_fwd_result(result: Any) -> None: message for op authors (user-output *types* are checked later, by :func:`_flatten_value`). - Only called on the fake path (:func:`_unpack_fwd_fake_result`), which runs at + Only called on the fake path (:meth:`_OutputPlan.parse`), which runs at trace/compile time -- so this is a compile-time check with no per-call cost. The real impl must return the same shape as the fake, so validating the fake covers both. @@ -837,18 +839,82 @@ def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List return out -def _unpack_fwd_fake_result( - result: Tuple[Any, ...], -) -> Tuple[List[Any], List[Any], Dict[str, Any]]: - """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" - _check_fwd_result(result) - num_outputs = len(result) - _FWD_TRAILING_SLOTS - saved = result[num_outputs] - ctx_attrs = result[num_outputs + 1] - user_fakes = list(result[:num_outputs]) - saved_fakes = list(saved) if saved is not None else [] - ctx_attrs = dict(ctx_attrs) if ctx_attrs else {} - return user_fakes, saved_fakes, ctx_attrs +@dataclasses.dataclass(frozen=True) +class _OutputPlan: + """Per-trace layout of an op's flat ``Tensor[]`` return. + + Parsed from a fwd fake-impl result: the logical user outputs and the + saved-for-backward tensors, each with its range in the flat payload. The + single source of truth for rebuilding forward outputs and saved tensors and + for slicing backward grad_outputs per user output. Per-trace rather than + per-registration because a quantized output's inner-tensor count is only + known from the fake result. + """ + + user_specs: Tuple[Optional[TensorSpec], ...] + saved_specs: Tuple[Optional[TensorSpec], ...] + ctx_attrs: Dict[str, Any] + user_ranges: Tuple[Tuple[int, int], ...] + saved_start: int + + @classmethod + def parse(cls, result: Tuple[Any, ...]) -> "_OutputPlan": + """Slice a fwd fake-impl return into the plan (validating the contract).""" + _check_fwd_result(result) + num_outputs = len(result) - _FWD_TRAILING_SLOTS + user_specs = tuple(result[:num_outputs]) + saved = result[num_outputs] + ctx_attrs = result[num_outputs + 1] + cursor = 0 + user_ranges: List[Tuple[int, int]] = [] + for spec in user_specs: + n = _spec_slot_count(spec) + user_ranges.append((cursor, cursor + n)) + cursor += n + return cls( + user_specs=user_specs, + saved_specs=tuple(saved) if saved is not None else (), + ctx_attrs=dict(ctx_attrs) if ctx_attrs else {}, + user_ranges=tuple(user_ranges), + saved_start=cursor, + ) + + @staticmethod + def _assemble( + spec: Optional[TensorSpec], flat: Sequence[Optional[torch.Tensor]], start: int, stop: int + ) -> Any: + chunk = [_decode_none(t) for t in flat[start:stop]] + # ``spec is None`` is the op-boundary sentinel for an absent output. + return spec.assemble(chunk) if spec is not None else None + + def user_outputs(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Rebuild the structured user outputs from the op's flat return.""" + return [ + self._assemble(spec, flat, start, stop) + for spec, (start, stop) in zip(self.user_specs, self.user_ranges) + ] + + def saved_tensors(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Rebuild the saved-for-backward tensors from the op's flat return.""" + values: List[Any] = [] + cursor = self.saved_start + for spec in self.saved_specs: + n = _spec_slot_count(spec) + values.append(self._assemble(spec, flat, cursor, cursor + n)) + cursor += n + return values + + def user_grads(self, flat_grads: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Gradient of each user output, sliced from the op's flat grad list. + + A single-slot output yields its tensor grad; a flattened quantized + output yields the tuple of its inner-buffer grads. + """ + grads: List[Any] = [] + for start, stop in self.user_ranges: + chunk = [_decode_none(g) for g in flat_grads[start:stop]] + grads.append(chunk[0] if stop - start == 1 else tuple(chunk)) + return grads # --------------------------------------------------------------------------- # @@ -901,10 +967,12 @@ def _register_autograd_for_op( ) -> None: """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. - ``setup_context`` re-runs the spec fwd fake impl to recover output / saved - templates, reassembles each flat output chunk, and hands the saved tuple + - ``ctx_attrs`` to the module's ``setup_context``. + ``setup_context`` re-runs the spec fwd fake impl to parse the + :class:`_OutputPlan`, reassembles the outputs / saved tensors from it, hands + the saved tuple + ``ctx_attrs`` to the module's ``setup_context`` and stashes + the plan on ``ctx`` so backward can slice its grads per user output. """ + bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields) def _setup_context(ctx, inputs, output): ctx.fwd_tensor_list_lengths = { @@ -913,31 +981,35 @@ def _setup_context(ctx, inputs, output): fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs))) spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names) - user_fakes, saved_fakes, ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) - - user_outputs, cursor = _unflatten_values(user_fakes, output) - saved_list, _ = _unflatten_values(saved_fakes, output, cursor) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + user_outputs = out_plan.user_outputs(output) + saved_list = out_plan.saved_tensors(output) bwd_obj = bwd_plan.arg_type() tensors_to_save_from_setup = setup_context_user( bwd_obj, fwd_obj, user_outputs[0] if len(user_outputs) == 1 else tuple(user_outputs), - ctx_attrs, + out_plan.ctx_attrs, tuple(saved_list), ) tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ())) ctx.tensor_objects = tensor_objects ctx.save_for_backward(*tensors_to_save) ctx.backward_objects = bwd_obj + ctx.output_plan = out_plan def _autograd_backward(ctx, *grad_outputs): bwd_obj = ctx.backward_objects if hasattr(bwd_obj, "setup_saved_tensors"): bwd_obj.setup_saved_tensors(ctx) ctx.tensor_objects = None - flat_grads = grad_outputs[0] - bwd_obj.grad_output = _decode_none(flat_grads[0]) + user_grads = ctx.output_plan.user_grads(grad_outputs[0]) + ctx.output_plan = None + if bwd_takes_grad_tuple: + bwd_obj.grad_outputs = tuple(user_grads) + else: + bwd_obj.grad_output = user_grads[0] kwargs = bwd_plan.pack(bwd_obj) bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] @@ -1098,8 +1170,11 @@ def register_custom_op( forward state and returns the tensors to persist; the framework saves them via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework restores them into the container's tensor fields through the - ``setup_saved_tensors`` hook and sets ``grad_output`` directly, so - ``bwd_impl`` receives a fully-populated ``bwd_arg_type``. + ``setup_saved_tensors`` hook and sets the incoming gradient directly -- + into a ``grad_outputs`` field (tuple, one grad per user output) if + ``bwd_arg_type`` declares one, else into ``grad_output`` (the first user + output's grad) -- so ``bwd_impl`` receives a fully-populated + ``bwd_arg_type``. Registration touches experimental ``torch.library`` / opaque-object APIs that may be missing on older PyTorch. If it fails, this warns once and @@ -1231,12 +1306,12 @@ def _register_custom_op_impl( def forward_fn(fwd_args): spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names) - user_fakes, _saved_fakes, _ctx_attrs = _unpack_fwd_fake_result(fwd_fake_impl(spec_obj)) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) kwargs = fwd_plan.pack(fwd_args) flat_in = [kwargs[name] for name in fwd_plan.slot_names] result = wrapper_fwd_op(*flat_in) - outputs, _ = _unflatten_values(user_fakes, result) + outputs = out_plan.user_outputs(result) if len(outputs) == 1: return outputs[0] return tuple(outputs)