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_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 f61e7b4111..a1935f2d93 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_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). + """ + 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..e42eb8f9f6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,10 +6,13 @@ 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 __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", "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 new file mode 100644 index 0000000000..d22c97bb25 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -0,0 +1,1536 @@ +# 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`` / ``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[]``. + +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 (``_unflatten_value``; + ``_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). + +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 # aliased: torch_dispatch rules take a ``types`` param +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 .tensor_spec import TensorSpec, to_tensor_spec +from ..quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, + _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 +# ``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) +# 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})") + _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 tensor_slot(self) -> str: + """Primary slot name for a plain / subclass tensor.""" + return self.name + + def inner_slot(self) -> str: + """Flat inner-tensor slot name.""" + return self.name + "__tensors" + + 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: + 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.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__}" + ) + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + 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.tensor_slot()] + else: + 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``, + # 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. + """ + + 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 _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. + """ + + META_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.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})} + + 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] + + +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 custom_op.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 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 + + +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 _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]], + 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:`_unflatten_value`; 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:`_flatten_value`). + + 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. + """ + 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 _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. + """ + num_outputs = len(result) - _FWD_TRAILING_SLOTS + flat: List[torch.Tensor] = [] + for value in result[:num_outputs]: + flat.extend(_flatten_value(value)) + saved = result[num_outputs] + if saved is not None: + for value in saved: + flat.extend(_flatten_value(value)) + return flat + + +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``); + 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 bwd_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 _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 + + +# --------------------------------------------------------------------------- # +# 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_base_op( + *, + 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], + 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 + ``pack_result``. + """ + + def _impl(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _args_from_slots(arg_type, kwargs, adapters) + 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 pack_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], + 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``. + + ``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.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 = _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() + 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.backward_objects = bwd_obj + + 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]) + 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 ``fwd_tensor_list_lengths``) needs a + # list-shaped no-grad of matching length. + out: List[Any] = [None] * slot_count + 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 + return tuple(out) + + 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 +) -> 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 _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, + slot_offsets: Sequence[int] = (), + subclasses: Sequence[type] = (), +) -> Any: + """Define the wrapper op via ``torch.library.custom_op``: forward to the base + op through :func:`_make_slot_forwarder`. Returns the ``CustomOpDef``. + """ + forward = _make_slot_forwarder(base_op, slot_offsets, subclasses) + + def _forward(*flat: Any) -> List[torch.Tensor]: + return forward(flat) + + op_def = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + ) + op_def.register_fake(_forward) + return op_def + + +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], + 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: 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``. + + 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``. + + 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. + * ``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 ``bwd_impl`` returning + :class:`TensorSpec` grads. + * ``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). ``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 + 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, + 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_disabled( + 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], + 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]: + """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(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) + + 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_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, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + pack_result=_pack_fwd_result, + ) + _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, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + 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, + 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 + ) + + 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, + "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) + 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_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) + 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 = _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) + + outputs, _ = _unflatten_values(user_fakes, result) + 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..ae8849097c 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -9,6 +9,7 @@ from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType +from ..utils import warn_compile_disabled # Qualnames of the registered quantizer classes. The set holds strings rather @@ -117,18 +118,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_disabled(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_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/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..868aced1f9 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,7 +3,8 @@ # 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 @@ -44,10 +45,11 @@ 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, + check_gemm_dims, ) 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,12 @@ prepare_for_saving, restore_from_func_ctx, ) +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 @@ -95,9 +104,6 @@ __all__ = ["Linear"] -TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] - - @dataclass(slots=True) class LinearFwdArgs: """Single-argument bag for the forward path of :class:`_Linear`.""" @@ -108,7 +114,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 +159,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 +189,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 +259,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 @@ -240,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: @@ -265,15 +318,44 @@ 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], 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 @@ -300,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": @@ -338,9 +420,7 @@ 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) @@ -372,8 +452,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 +542,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 +725,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_weight_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 +758,219 @@ 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_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 = _sp_out_leading(inp.shape[0], args) + 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_weight_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 +983,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 +1007,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 +1079,10 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight + elif wt_save_alias == "new_weight_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": @@ -786,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 @@ -858,6 +1162,12 @@ 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] + 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) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1387,6 +1697,86 @@ def wgrad_gemm( ) +def _linear_backward_fake( + args: LinearBwdArgs, +) -> 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. + """ + 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_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) + + 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. + 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, + 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_fake, + setup_context=_linear_setup_ctx, + bwd_arg_type=LinearBwdArgs, + bwd_impl=_linear_backward_impl, + bwd_fake_impl=_linear_backward_fake, +) + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1418,7 +1808,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 +1815,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, ) @@ -1458,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. @@ -1471,6 +1860,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 +2272,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 +2350,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 +2390,27 @@ 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 ) 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 +2465,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..47a7db9f7e 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,6 +26,44 @@ ] +_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``. + + 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, @@ -624,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.