diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 46795415e5..e65824ce85 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -200,6 +200,19 @@ def _parse_args(argv=None, namespace=None): parser.add_argument( "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA Graphs." ) + parser.add_argument( + "--compile", + action="store_true", + default=False, + help="Wrap each layer in torch.compile (tests Userbuffers on the compiled path).", + ) + parser.add_argument( + "--compile-mode", + type=str, + default="default", + choices=["default", "reduce-overhead"], + help="torch.compile mode used when --compile is set.", + ) parser.add_argument( "--ub-cfg", type=str, default=None, help="Optional TP config yaml file input." ) @@ -485,6 +498,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0) dist_print("Copied parameters from test model to reference model...", debug=True) + if opts.compile and opts.use_cuda_graphs: + raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.") + # Fp8 recipe setup fp8_format = Format.HYBRID fp8_recipe = None @@ -535,6 +551,18 @@ def run_fwd_bwd(model, x): loss.backward() return out + if opts.compile: + for i, layer in enumerate(test_model.layers): + # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle + # op arg whose hash chokes on non-nested SymInt (see run_numerics). + test_model.layers[i] = torch.compile( + layer, fullgraph=True, mode=opts.compile_mode, dynamic=False + ) + dist_print( + f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...", + debug=True, + ) + torch_rng_state = torch.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state(torch.device(f"cuda:{LOCAL_RANK}")) if opts.use_cuda_graphs: diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index fe02f990b4..1590979ba3 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -310,22 +310,45 @@ def _copy_params(model_distributed, model_single): def _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed, **kwargs + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=False, + compile_mode="default", + **kwargs, ): _alloc_main_grad(model_single_node, model_distributed) # for fuse_wgrad_accumulation=True input_single_node.requires_grad_() input_distributed.requires_grad_() + forward_single_node = model_single_node + forward_distributed = model_distributed + if use_compile: + # Each parametrized case compiles the same module.forward code object with + # a different shape/recipe; with dynamic=False those guards accumulate and + # eventually trip Dynamo's recompile_limit. Reset so every case starts from + # a clean compile cache (mirrors the single-GPU torch.compile tests). + torch._dynamo.reset() + # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle + # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static + # shapes (recompile per shape) until the bundle handles symbolic shapes. + forward_single_node = torch.compile( + model_single_node, fullgraph=True, mode=compile_mode, dynamic=False + ) + forward_distributed = torch.compile( + model_distributed, fullgraph=True, mode=compile_mode, dynamic=False + ) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), ): - output_single_node = model_single_node(input_single_node, **kwargs) + output_single_node = forward_single_node(input_single_node, **kwargs) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), amax_reduction_group=NCCL_WORLD, ): - output_distributed = model_distributed(input_distributed, **kwargs) + output_distributed = forward_distributed(input_distributed, **kwargs) return output_single_node, output_distributed @@ -641,12 +664,20 @@ def test_quantized_all_gather(): # Linear # ############################################ @run_distributed_test() -def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): +def _test_linear( + parallel_mode=None, + sequence_parallel=False, + use_compile=False, + compile_mode="default", + **kwargs, +): """Test the linear layer with specified parallel mode and sequence parallelization. Args: parallel_mode (str): 'row' or 'column' parallelism. sequence_parallel (bool): Enable sequence parallelism if True. + use_compile (bool): Wrap the modules in ``torch.compile`` before running. + compile_mode (str): ``torch.compile`` mode ("default" or "reduce-overhead"). kwargs (dict): Additional arguments for the linear layer. """ # Set parameter data type @@ -696,7 +727,12 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): # Apply models output_single_node, output_distributed = _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=use_compile, + compile_mode=compile_mode, ) if "return_bias" in kwargs: @@ -740,6 +776,8 @@ def test_linear(): {"params_dtype": torch.float16 if QUANTIZATION != "nvfp4" else torch.bfloat16}, {"delay_wgrad_compute": True}, {"save_original_input": True}, + {"use_compile": True}, + {"use_compile": True, "compile_mode": "reduce-overhead"}, ] for kwargs in kwargs_list: @@ -747,6 +785,9 @@ def test_linear(): continue if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: continue + # debug instrumentation forces the eager fallback, so compile is a no-op there. + if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel, **kwargs) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6b1ad870e9..6521ed7dbc 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -111,6 +111,8 @@ def _run_layer_with_overlap( quantization, num_layers=1, use_cublasmp=False, + compile=False, + compile_mode="default", ): test_path = TEST_ROOT / "run_layer_with_overlap.py" test_cmd = LAUNCH_CMD + [ @@ -129,6 +131,10 @@ def _run_layer_with_overlap( if overlap_rs_dgrad: test_cmd.append("--overlap-rs-dgrad") + if compile: + test_cmd.append("--compile") + test_cmd.append(f"--compile-mode={compile_mode}") + if fp8: if quantization in ("fp8_delayed_scaling", "fp8_current_scaling") and not fp8_available: pytest.skip(reason_for_no_fp8) @@ -281,6 +287,40 @@ def test_layers_with_overlap_bf16( ) +@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize( + "linear_parallel_mode,overlap_rs_dgrad", + [ + ("row", False), + ("column", False), + ("column", True), + ], + ids=[ + "ROW-PARALLEL", + "COL-PARALLEL - BULK DGRAD/WGRAD", + "COL-PARALLEL - DGRAD+RS", + ], +) +def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode): + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16). + + Userbuffers is expected to stay on Linear's compiled custom-op path (the + collective lives inside the opaque op), so this checks that torch.compile + + Userbuffers stays numerically correct against the eager, non-overlap reference. + ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on + top of the Userbuffers collectives. + """ + _run_layer_with_overlap( + te.Linear.__name__, + linear_parallel_mode, + overlap_rs_dgrad, + False, + None, + compile=True, + compile_mode=compile_mode, + ) + + @pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..87dfd171f6 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,9 +3,12 @@ # See LICENSE for license information. import abc +import contextlib import pytest import torch +from torch._dynamo.utils import counters +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: from torch._opaque_base import OpaqueBaseMeta @@ -24,14 +27,20 @@ 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, _STORAGE_REGISTRY +from transformer_engine.pytorch.dynamo.traceable_utils import make_empty_traceable, _contiguous_stride from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8Quantizer, + Float8BlockQuantizer, + MXFP8Quantizer, ) from utils import recipe_id @@ -80,6 +89,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) @@ -384,3 +462,674 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# Value-opaque quantizers +# --------------------------------------------------------------------------- + + +def _mxfp8(dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=dtype) + + +def _blockwise(force_pow_2_scales=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=force_pow_2_scales, + ) + + +def _current_scaling(amax_epsilon=0.0): + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cpu"), + amax_epsilon=amax_epsilon, + ) + + +def _nvfp4(with_rht=True): + # Default with_rht=True so the quantize round-trip below exercises the + # derived ``rht_matrix`` tensor (the field most likely to be dropped on + # value-key reconstruction). + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=with_rht, + with_post_rht_amax=with_rht, + ) + + +def _hw_available(quantizer): + """Whether this HW can actually run the quantize kernel for *quantizer*.""" + if isinstance(quantizer, MXFP8Quantizer): + return mxfp8_available + if isinstance(quantizer, NVFP4Quantizer): + return nvfp4_available + if isinstance(quantizer, Float8BlockQuantizer): + return fp8_block_scaling_available + return fp8_available # Float8CurrentScalingQuantizer + + +# (factory, kwargs producing a different-but-valid config) +_VALUE_QUANTIZERS = [ + pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), + pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), + pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param( + _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, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory, other_kwargs): + """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" + a, b = factory(), factory() + # Same config -> equal, same hash, interchangeable as a dict/set key. + assert a is not b + assert a == b + assert hash(a) == hash(b) + assert {a: "x"}[b] == "x" + # Different config -> not equal. + assert a != factory(**other_kwargs) + + # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. + repr_str, globals_ = a.__fx_repr__() + rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used + assert rebuilt == a and rebuilt is not a + assert hash(rebuilt) == hash(a) + # The deprecated amax-reduction group is never part of the value. + assert getattr(rebuilt, "amax_reduction_group", None) is None + + # The rebuilt quantizer must also *behave* identically, not just compare + # equal: equality only looks at the value key, so a field the kernel needs + # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would + # slip through the checks above and only blow up at quantize time. Run the + # real quantize kernel on both and require bit-exact results. + if torch.cuda.is_available() and _hw_available(a): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) + + +def test_value_quantizer_rejects_process_group(): + """A value quantizer holding a live ProcessGroup must refuse to be turned + into a value key / FX constant (raise), not silently drop the group.""" + import torch.distributed as dist # pylint: disable=import-outside-toplevel + + created = not dist.is_initialized() + if created: + dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + q.amax_reduction_group = dist.group.WORLD + # Every value-materialization path must reject it (hash, eq, __fx_repr__). + with pytest.raises(TypeError): + hash(q) + with pytest.raises(TypeError): + q.__fx_repr__() + finally: + if created: + dist.destroy_process_group() + + +if _opaque_available: + # A minimal custom op taking a tensor and a value-opaque quantizer that + # quantizes + dequantizes inside it, one per production quantizer class. + # ``test_quantizer_value_object_fullgraph`` drives this under + # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the + # graph -- proving the opaque-type registration took effect (a graph break + # would make ``fullgraph=True`` raise). + _qdq_lib = torch.library.Library("test_te_qdq", "DEF") + _QDQ_OPS = {} + for _qcls in ( + MXFP8Quantizer, + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + NVFP4Quantizer, + ): + _op = f"qdq_{_qcls.__name__}" + _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor") + + @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib) + def _qdq_impl(x, q): + return q(x).dequantize() + + @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib) + def _qdq_fake(x, q): + return torch.empty_like(x) + + _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op) + + +@pytest.mark.skipif( + not _opaque_available, + reason="torch.compile opaque-object support requires PyTorch >= 2.11", +) +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory, other_kwargs): + """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. + + A custom op quantizes+dequantizes with the (opaque value) quantizer; the + compiled result must match eager. ``fullgraph=True`` raises on any graph + break, so this proves the opaque-type registration actually took effect -- + unlike merely passing the quantizer through. + """ + q = factory() + if not (torch.cuda.is_available() and _hw_available(q)): + pytest.skip("format not supported on this HW") + + op = _QDQ_OPS[type(q)] + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def fn(inp): + return op(inp, q) + + ref = fn(x) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# torch.compile-traceable allocation primitives +# --------------------------------------------------------------------------- + + +# (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128) +# / NVFP4 (mult. of 16) constraints. +_PROTO_QUANTIZERS = [ + pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"), + pytest.param(_mxfp8, (64, 128), id="mxfp8"), + pytest.param(_blockwise, (128, 256), id="fp8_blockwise"), + pytest.param( + _nvfp4, + (64, 128), + id="nvfp4", + marks=pytest.mark.skipif( + not nvfp4_available, + reason="NVFP4 is not available", + ), + ), +] + + +def _build_from_primitives(quantizer, shape, dtype, device="cpu"): + """Assemble a quantized tensor straight from the quantizer primitives: + ``alloc_tensors`` (buffers) + ``create_metadata`` (ctx) + the storage's + ``__tensor_unflatten__``.""" + names = tuple(quantizer._describe_buffers(shape)) # pylint: disable=protected-access + ctx = quantizer.create_metadata(shape, dtype=dtype) + buffers = quantizer.alloc_tensors(shape, device=device) + inner = {name: buffers[name] for name in names} + storage_cls = _STORAGE_REGISTRY[ctx["cls"]] + # Row-major (contiguous) outer stride for ``__tensor_unflatten__``; ``meta`` + # device computes it without allocating storage. + outer_stride = torch.empty(tuple(shape), device="meta").stride() + return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), outer_stride) + + +def _signature(tensor, names): + """Comparable shape/dtype fingerprint of a tensor and its inner buffers.""" + sig = {"__shape__": tuple(tensor.shape), "__dtype__": tensor.dtype} + for name in names: + buf = getattr(tensor, name) + sig[name] = (tuple(buf.shape), buf.dtype) + return sig + + +def _skip_if_dequantize_unsupported(q): + """Skip when this HW can't run ``dequantize()`` for the quantizer's format. + + ``dequantize()`` runs the real kernel on CUDA, so each format has its own + availability gate (mirrors the ``is_*_available`` checks in test_numerics). + """ + if isinstance(q, MXFP8Quantizer): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif isinstance(q, NVFP4Quantizer): + if not nvfp4_available: + pytest.skip("NVFP4 is not available") + elif isinstance(q, Float8BlockQuantizer): + if not fp8_block_scaling_available: + pytest.skip("FP8 block scaling is not available") + elif not fp8_available: # Float8 current scaling + pytest.skip(reason_for_no_fp8) + + +# ----- Quantizer primitives ----- + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_primitives_unflatten_compiles(factory, shape): + """create_metadata + alloc_tensors + __tensor_unflatten__ compose and trace + under ``fullgraph=True`` (CPU).""" + q = factory() + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + + def fn(x): + t = _build_from_primitives(q, shape, x.dtype, device=x.device) + # Read every buffer into the result so the alloc + unflatten can't be + # eliminated as dead code -- forces the whole build path into the graph. + acc = x.new_zeros(()) + for name in names: + acc = acc + getattr(t, name).float().sum() + return acc + + x = torch.zeros(*shape, dtype=torch.bfloat16) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + assert out.shape == () + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_alloc_tensors_fake(factory, shape): + """``alloc_tensors`` produces FakeTensors with the described shapes/dtypes.""" + q = factory() + bufs = q._describe_buffers(shape) # pylint: disable=protected-access + with FakeTensorMode(): + alloc = q.alloc_tensors(shape, device="cpu") + assert set(alloc) == set(bufs) + for name, (buf_shape, buf_dtype) in bufs.items(): + assert isinstance(alloc[name], FakeTensor) + assert tuple(alloc[name].shape) == tuple(buf_shape) + assert alloc[name].dtype == buf_dtype + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_storage_flatten_unflatten_roundtrip(factory, shape): + """Storage ``__tensor_flatten__`` / ``__tensor_unflatten__`` round-trips. + + Build a tensor from ``alloc_tensors`` + ``create_metadata``, flatten it, then + unflatten and verify shape/dtype and every inner buffer match before vs after. + """ + q = factory() + _skip_if_dequantize_unsupported(q) + + tensor = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + # Fill buffers with deterministic data (empty() may contain NaNs) so the + # round-trip can be checked by value via dequantize(). + for name in names: + buf = getattr(tensor, name) + buf.copy_(torch.arange(buf.numel(), device=buf.device).reshape(buf.shape)) + before = _signature(tensor, names) + expected = tensor.dequantize() + + flat_names, flat_ctx = tensor.__tensor_flatten__() + assert set(flat_names) == set(names) + inner = {name: getattr(tensor, name) for name in flat_names} + rebuilt = type(tensor).__tensor_unflatten__( + inner, flat_ctx, tuple(tensor.shape), tensor.stride() + ) + + assert isinstance(rebuilt, QuantizedTensor) + assert _signature(rebuilt, flat_names) == before + # The reconstructed tensor dequantizes to the same values. + torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True) + + +# ----- make_empty_traceable ----- + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_make_empty_traceable_matches_primitives(factory, shape): + """make_empty_traceable is equivalent to building from quantizer primitives + directly (alloc_tensors + create_metadata + __tensor_unflatten__).""" + q = factory() + + # Build via make_empty_traceable. + tensor = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu") + + # Build directly from primitives. + direct = _build_from_primitives(q, shape, torch.bfloat16) + + # Both should produce the same buffer layout. + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + assert _signature(tensor, names) == _signature(direct, names) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_make_empty_traceable_eager(factory, shape): + """make_empty_traceable (no fake) yields a real quantized tensor.""" + q = factory() + out = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu") + assert isinstance(out, QuantizedTensor) + assert tuple(out.shape) == tuple(shape) + assert out.dtype == torch.bfloat16 + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + for name in names: + assert not isinstance(getattr(out, name), FakeTensor) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_make_empty_traceable_fake(factory, shape): + """make_empty_traceable under FakeTensorMode yields a fake-backed quantized + tensor with the right shape/dtype and fake inner buffers.""" + q = factory() + with FakeTensorMode(): + out = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu") + assert isinstance(out, QuantizedTensor) + assert tuple(out.shape) == tuple(shape) + assert out.dtype == torch.bfloat16 + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + for name in names: + assert isinstance(getattr(out, name), FakeTensor) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_make_empty_traceable_compiles(factory, shape): + """make_empty_traceable traces under torch.compile(fullgraph=True) (CPU).""" + q = factory() + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + + def fn(x): + t = make_empty_traceable(q, tuple(x.shape), dtype=x.dtype, device=x.device) + acc = x.new_zeros(()) + for name in names: + acc = acc + getattr(t, name).float().sum() + return acc + + x = torch.zeros(*shape, dtype=torch.bfloat16) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + assert out.shape == () + + +def test_make_empty_traceable_plain_tensor(): + """For non-quantized tensors, make_empty_traceable produces a plain tensor.""" + t = make_empty_traceable(None, (2, 3), dtype=torch.float32, device="cpu") + assert not isinstance(t, QuantizedTensor) + assert t.shape == (2, 3) + assert t.dtype == torch.float32 + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_make_empty_traceable_roundtrip(factory, shape): + """A tensor built via make_empty_traceable can be flattened and unflattened.""" + q = factory() + tensor = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu") + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + + # Flatten. + flat_names, flat_ctx = tensor.__tensor_flatten__() + assert set(flat_names) == set(names) + inner = {name: getattr(tensor, name) for name in flat_names} + + # Unflatten. + rebuilt = type(tensor).__tensor_unflatten__( + inner, flat_ctx, tuple(tensor.shape), tensor.stride() + ) + assert isinstance(rebuilt, QuantizedTensor) + assert _signature(rebuilt, flat_names) == _signature(tensor, names) + + +# --------------------------------------------------------------------------- +# te.Linear +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize( + "fp8_recipe", + [None, *_all_recipes], + ids=lambda r: "bf16" if r is None else type(r).__name__, +) +def test_te_linear_compiles(fp8_recipe, compile_mode): + """ + torch.compile(fullgraph=True) of ``te.Linear`` under every built-in + recipe (plus the bf16-only baseline with no autocast), for both the default + backend and ``mode="reduce-overhead"`` (CUDA-graph trees). + """ + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + + dtype = torch.bfloat16 + device = "cuda" + + # FP8 GEMMs require leading dimensions divisible by 16. + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + if fp8_recipe is None: + return model(inp) + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + # ``reduce-overhead`` warms up on the first call(s) and replays a captured + # CUDA graph afterwards, so iterate a few times to actually exercise replay. + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_quantized_fp8_weight(compile_mode): + """torch.compile should handle Linear weights initialized as FP8 tensors, + for both the default backend and ``mode="reduce-overhead"``. + + Exercises the two-tier op + ``register_torch_dispatch`` flattening of a + ``Float8Tensor`` weight *input* in + :mod:`transformer_engine.pytorch.dynamo`. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + with te.quantized_model_init(enabled=True, recipe=fp8_recipe): + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + assert isinstance(model.weight, te.Float8Tensor) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_fp8_output(compile_mode): + """torch.compile of ``te.Linear(..., fp8_output=True)`` without gradient: + forward returns a :class:`Float8Tensor`. Covers the default backend and + ``mode="reduce-overhead"``. + + Exercises the output-rewrap path in + :mod:`transformer_engine.pytorch.dynamo`: when an output quantizer is + active, the op returns the flat inner data tensors and the framework + rewraps them into a ``Float8Tensor`` via ``__tensor_unflatten__``. A + differentiable FP8 output is unsupported under compile (``Linear.forward`` + falls back to eager), so this test covers the supported case: an FP8 output + that does not require grad (inference / ``torch.no_grad``). + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, fp8_output=True) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + with torch.no_grad(): + _cudagraph_warmup(fn, torch.randn(32, 64, dtype=dtype, device=device), backward=False) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + inp = torch.randn(32, 64, dtype=dtype, device=device) + with torch.no_grad(): + out_eager = fn(inp) + out = compiled(inp) + assert isinstance( + out, te.Float8Tensor + ), f"expected Float8Tensor output, got {type(out).__name__}" + assert out.shape == (32, 32) + assert ( + out._quantizer is not None + ), "FP8 output lost its quantizer on the torch.compile path" + # The rewrap rebuilt a fully-functional Float8Tensor: dequantizing it + # outside the compiled region exercises scale + data + dtype wiring. + deq = out.dequantize() + assert deq.shape == (32, 32) + assert deq.dtype == dtype + # Compiled FP8 output must match the eager FP8 output value-wise. + torch.testing.assert_close( + deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_is_first_microbatch(compile_mode): + """torch.compile of ``te.Linear`` across a multi-step microbatch schedule that + drives FP8 weight caching via ``is_first_microbatch``, for the default backend + and ``mode="reduce-overhead"`` (CUDA-graph trees). + + ``is_first_microbatch=True`` quantizes and caches the FP8 weight; subsequent + ``False`` steps must reuse the cached FP8 weight instead of re-quantizing. This + exercises that cache path under compile and checks it stays numerically aligned + with eager. ``is_first_microbatch`` is a Python bool, so each distinct value is + its own dynamo guard/graph. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + # First microbatch caches the FP8 weight, the rest reuse the cache. + schedule = [True, False, False] + is_first = schedule[0] # rebound each step; closed over by ``fn``. + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=is_first) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for is_first in schedule: + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_linear_dynamic_shapes(): + """torch.compile(dynamic=True) of ``te.Linear`` with varying batch sizes. + + Verifies that the compiled graph handles symbolic (dynamic) leading + dimensions without graph breaks or recompilations after the initial trace. + Key correctness property: a graph compiled for batch=16 must produce + numerically correct results for batch=32 without triggering a recompile. + + This exercises two fixes for dynamic shapes: + 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle + (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). + 2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output + + weight + SP config instead of relying on the stored ``inp_shape``. + 3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same + tensor sources when it is None (compiled mode). + + FP8 + dynamic=True is tracked separately (requires resolving + ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo). + """ + dtype = torch.bfloat16 + device = "cuda" + in_features, out_features = 64, 32 + model = te.Linear(in_features, out_features, params_dtype=dtype, device=device) + + def fn(inp): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + batch_sizes = [16, 32, 48] + + for i, batch in enumerate(batch_sizes): + inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) + # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes. + torch._dynamo.mark_dynamic(inp, 0) + out = compiled(inp) + assert out.shape == (batch, out_features), f"wrong output shape for batch={batch}" + out.sum().backward() + assert inp.grad is not None, f"no input gradient for batch={batch}" + assert inp.grad.shape == inp.shape, f"wrong grad shape for batch={batch}" + + # Verify numerics against eager on each distinct batch size. + inp_eager = inp.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = model(inp_eager) + out_eager.sum().backward() + torch.testing.assert_close( + out.detach(), out_eager.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + msg=f"forward mismatch at batch={batch}", + ) + torch.testing.assert_close( + inp.grad, inp_eager.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + msg=f"dgrad mismatch at batch={batch}", + ) + + if i == 0: + # After the first (tracing) call, record the recompile counter + # baseline -- subsequent batch sizes must not trigger recompiles. + recompile_count_baseline = counters["stats"].get("recompile_reasons", 0) + + recompile_count_after = counters["stats"].get("recompile_reasons", 0) + assert recompile_count_after == recompile_count_baseline, ( + f"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 new file mode 100644 index 0000000000..f4864cda5b --- /dev/null +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine.""" + +from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer +from .traceable_utils import make_empty_traceable +from .custom_op import register_custom_op + +__all__ = [ + "register_value_opaque_quantizer", + "is_value_opaque_quantizer", + "make_empty_traceable", + "register_custom_op", +] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py new file mode 100644 index 0000000000..9aca5b866f --- /dev/null +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -0,0 +1,1252 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile custom-op framework for Transformer Engine.""" + +from __future__ import annotations +import dataclasses +import warnings +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Union, + get_args, + get_origin, + get_type_hints, +) + +import torch + +from .traceable_utils import _contiguous_stride, _slot_count, _maybe_reassemble_tensor_subclass +from ..quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, + _STORAGE_REGISTRY, + _quantized_tensor_passthrough_ops, + prepare_for_saving, +) + +_TE_OP_NAMESPACE = "transformer_engine_compile" + + +# ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a +# 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for +# ``register_autograd`` to attach a ``grad_fn`` to the outputs. +# +# TODO: once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable +# ``Tensor?[]`` return schema lets ``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``, 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). + """ + + 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 _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}" + 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, OpaqueValueBundle.PRIMITIVE_TYPES): + return + if _is_opaque_value_type(type(value)): + _, extra = value.__fx_repr__() + globals_.update(extra) + + for v in self._data.values(): + _collect(v) + return (f"OpaqueValueBundle({{{items}}})", globals_) + + +try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + get_opaque_type_name, + is_opaque_value_type as _is_opaque_value_type, + is_opaque_reference_type as _is_opaque_reference_type, + register_opaque_type, + ) + + register_opaque_type(OpaqueValueBundle, typ="value") + _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) +except Exception: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object + _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) -> 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``. + """ + 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) + 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 = _contiguous_stride(tuple(outer_shape)) if outer_shape is not None else None + return QuantizedTensorStorage.__tensor_unflatten__(inner, meta_dict, outer_shape, stride) + + +# --------------------------------------------------------------------------- # +# Field buckets: dataclass field <-> flat torch.library slot(s) +# --------------------------------------------------------------------------- # + + +def _strip_optional(annot: Any) -> Tuple[Any, bool]: + """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``.""" + if get_origin(annot) is Union: + 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 _Bucket: + """Maps one (or, for the aggregating bucket, 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 bucket 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); ``pack`` and ``unpack`` 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["_Bucket"]: + """Decide whether this bucket type handles the field ``name`` given its + type annotation ``annot``; return a configured bucket if so, else + ``None`` so the next candidate is tried. + + Called once per field at registration, in :data:`_FIELD_BUCKETS` + 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 buckets to form the op's schema string. + """ + raise NotImplementedError + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + """Read this field from the dataclass ``owner`` and produce the concrete + value for each of its schema slots, as ``(slot_name, value)`` pairs. + + 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:`unpack`. + """ + raise NotImplementedError + + def unpack(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:`pack`. + """ + raise NotImplementedError + + def grad_slot(self) -> Optional[int]: + """Index (within this bucket'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 buckets (quantizers, metadata) return ``None``. + """ + return None + + +class _UniversalKind(Enum): + """What a universal-tensor slot group carries, tagged in its ``__meta``.""" + + NONE = "none" + TENSOR = "tensor" + STORAGE = "storage" + + +class _UniversalTensorBucket(_Bucket): + """``Tensor | QuantizedTensorStorage`` (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). + """ + + KIND_KEY = "__kind__" + + def __init__(self, name: str) -> None: + self.name = name + + def slot_name(self) -> str: + """Primary slot name for a plain / subclass tensor.""" + return self.name + + def slot_tensors(self) -> str: + """Flat inner-tensor slot name.""" + return self.name + "__tensors" + + def slot_meta(self) -> str: + """Flatten-metadata slot name.""" + return self.name + "__meta" + + def schema_slots(self) -> List[Tuple[str, str]]: + return [ + (self.slot_name(), "Tensor?"), + (self.slot_tensors(), "Tensor[]"), + (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + ] + + @staticmethod + def _is_tensor_storage_union(annot: Any) -> bool: + if get_origin(annot) is not Union: + return False + members = [a for a in get_args(annot) if a is not type(None)] + if torch.Tensor not in members: + return False + return any( + isinstance(m, type) and issubclass(m, QuantizedTensorStorage) for m in members + ) + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]: + if cls._is_tensor_storage_union(annot): + return cls(name) + return None + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + value = getattr(owner, self.name) + if value is None: + return [ + (self.slot_name(), None), + (self.slot_tensors(), []), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.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 + # outer op's ``register_torch_dispatch`` rule. + return [ + (self.slot_name(), value), + (self.slot_tensors(), []), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.TENSOR})), + ] + if isinstance(value, QuantizedTensorStorage): + meta, tensors = _storage_flatten(value) + meta._data[self.KIND_KEY] = _UniversalKind.STORAGE + return [ + (self.slot_name(), None), + (self.slot_tensors(), list(tensors)), + (self.slot_meta(), meta), + ] + raise TypeError( + f"field {self.name!r} expected None, torch.Tensor, or " + f"QuantizedTensorStorage, got {type(value).__name__}" + ) + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + meta = args[self.slot_meta()] + kind = meta.get(self.KIND_KEY) + if kind == _UniversalKind.NONE: + kwargs[self.name] = None + elif kind == _UniversalKind.TENSOR: + kwargs[self.name] = args[self.slot_name()] + else: + kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()]) + + def grad_slot(self) -> Optional[int]: + # Gradient flows to the plain / subclass tensor slot (``slot_name``, + # the first of the three). + return 0 + + +class _TensorBucket(_Bucket): + """``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["_TensorBucket"]: + 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 pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.name, getattr(owner, self.name))] + + def unpack(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 _QuantizerBucket(_Bucket): + """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. + + Each quantizer gets its own dedicated slot. The field is annotated with the + base ``Quantizer`` (not itself a registered opaque type), so the simple + bundle would not claim it. + """ + + KEY = "q" + + def __init__(self, name: str) -> None: + self.name = name + + def slot(self) -> str: + """Opaque quantizer metadata slot name.""" + return self.name + "__q" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: + stripped, _ = _strip_optional(annot) + if not isinstance(stripped, type): + return None + if issubclass(stripped, Quantizer): + return cls(name) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.slot(), OpaqueValueBundle({self.KEY: getattr(owner, self.name)}))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.slot()][self.KEY] + + +class _ReferenceOpaqueBucket(_Bucket): + """``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["_ReferenceOpaqueBucket"]: + 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 pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.name, getattr(owner, self.name))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.name] + + +class _SimpleBundleBucket(_Bucket): + """Aggregates every simple-typed field into a single OpaqueValueBundle.""" + + SLOT = "_simple_meta" + + def __init__(self, names: List[str]) -> None: + self.names = list(names) + + @classmethod + def matches_field(cls, annot: Any) -> bool: + """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" + annot, _ = _strip_optional(annot) + if annot in OpaqueValueBundle.PRIMITIVE_TYPES: + return True + if isinstance(annot, type) and issubclass(annot, Enum): + return True + if ( + isinstance(annot, type) + and _is_opaque_value_type is not None + and _is_opaque_value_type(annot) + ): + return True + if get_origin(annot) in (tuple, list): + inner = [a for a in get_args(annot) if a is not Ellipsis] + return bool(inner) and all(cls.matches_field(a) for a in inner) + return False + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.SLOT, OpaqueValueBundle({n: getattr(owner, n) for n in self.names}))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + if self.SLOT not in args: + return + meta = args[self.SLOT] + for n in self.names: + kwargs[n] = meta[n] + + +class _UnknownBucket(_Bucket): + """Fallback for fields no other bucket claims. + + Emits no slot; pack rejects non-trivial values (anything other than + ``None`` / all-``None`` sequence); unpack restores the field as ``None``. + """ + + 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 pack(self, owner: Any) -> List[Tuple[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, or Quantizer) " + "and carries a non-trivial value; add a matching bucket in " + "dynamo.py to handle it." + ) + return [] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = None + + +# Buckets, in priority order, owning ``try_build`` for a single field. +_FIELD_BUCKETS: Tuple[type, ...] = ( + _UniversalTensorBucket, + _TensorBucket, + _ReferenceOpaqueBucket, + _QuantizerBucket, +) + + +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_buckets(cls: type) -> List[_Bucket]: + """Build the bucket 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)." + ) + buckets: List[_Bucket] = [] + simple_names: List[str] = [] + for name, annot in _resolved_field_annotations(cls): + built: Optional[_Bucket] = None + for bucket_cls in _FIELD_BUCKETS: + built = bucket_cls.try_build(name, annot) + if built is not None: + break + if built is not None: + buckets.append(built) + elif _SimpleBundleBucket.matches_field(annot): + simple_names.append(name) + else: + buckets.append(_UnknownBucket(name, cls.__name__)) + if simple_names: + buckets.append(_SimpleBundleBucket(simple_names)) + return buckets + + +def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]: + """Return ``(schema_arg_str, slot_names)`` for a bucket list.""" + spec = [slot for b in buckets 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 _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: + """Build the op's flat ``{slot_name: value}`` argument dict from an args + dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every bucket's + packed slot(s). Inverse of :func:`_unpack`. + """ + out: Dict[str, Any] = {} + for bucket in buckets: + for name, value in bucket.pack(obj): + out[name] = value + return out + + +def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: + """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the + op's flat slot ``args`` dict, by letting every bucket restore its field(s). + Inverse of :func:`_pack`. + """ + kwargs: Dict[str, Any] = {} + for bucket in buckets: + bucket.unpack(args, kwargs) + obj = cls.__new__(cls) + for k, v in kwargs.items(): + object.__setattr__(obj, k, v) + return obj + +# --------------------------------------------------------------------------- # +# 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; the fake impl returns actual tensors whose __tensor_flatten__ +# provides the template for reassembly. +# --------------------------------------------------------------------------- # + + +def _value_to_flat_tensors( + value: Optional[Union[torch.Tensor, QuantizedTensorStorage]], +) -> List[torch.Tensor]: + """Return the flat ``Tensor[]`` slots that represent one op output ``value``. + + Inverse of :func:`_maybe_reassemble_tensor_subclass`; the slot count matches + :func:`_slot_count`. + """ + if value is None: + return [_encode_none(None)] + if isinstance(value, torch.Tensor): + if type(value) is not torch.Tensor and hasattr( # pylint: disable=unidiomatic-typecheck + value, "__tensor_flatten__" + ): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + return [_encode_none(value)] + if hasattr(value, "__tensor_flatten__"): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + raise TypeError( + f"unsupported value type {type(value).__name__}; expected None / " + "torch.Tensor / tensor subclass / bare storage." + ) + + +# Trailing slots in every fwd-impl return: ``tensors_to_save, tensor_objects, +# ctx_attrs``. User-output count is ``len(result) - this``. +_FWD_TRAILING_SLOTS = 3 + + +def _format_fwd_result(result: Any) -> List[torch.Tensor]: + """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. + + User outputs first, then saved-for-backward tensors in declaration order. + """ + num_outputs = len(result) - _FWD_TRAILING_SLOTS + flat: List[torch.Tensor] = [] + for value in result[:num_outputs]: + flat.extend(_value_to_flat_tensors(value)) + saved = result[num_outputs] or () + for value in saved: + flat.extend(_value_to_flat_tensors(value)) + return flat + + +def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: + """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. + + Each grad occupies exactly one slot (validated against ``num_grad_inputs``). + """ + grads = list(grads) + if len(grads) != num_grad_inputs: + raise RuntimeError( + f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads " + f"(one per input_tensors_for_grad entry), got {len(grads)}" + ) + return [_encode_none(g) for g in grads] + + +def _split_fwd_fake_result( + result: Tuple[Any, ...], +) -> Tuple[List[Any], List[Any], Dict[str, Any]]: + """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" + num_outputs = len(result) - _FWD_TRAILING_SLOTS + saved = result[num_outputs] + ctx_attrs = result[num_outputs + 2] + 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_buckets: List[_Bucket], + fwd_arg_type: type, + input_tensors_for_grad: List[str], +) -> Tuple[List[Any], List[int]]: + """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. + + Returns ``(fwd_slot_defaults, grad_targets)``: the per-slot no-grad template + (``[]`` for ``Tensor[]`` slots, ``None`` otherwise) and, for each requested + input name, the schema-slot index its gradient maps to. + """ + fwd_slot_defaults: List[Any] = [] + name_to_slot: Dict[str, int] = {} + slot_offset = 0 + for bucket in fwd_buckets: + slots = bucket.schema_slots() + for _, type_str in slots: + fwd_slot_defaults.append([] if type_str.endswith("[]") else None) + grad_slot = bucket.grad_slot() + if grad_slot is not None: + name_to_slot[bucket.name] = slot_offset + grad_slot + slot_offset += len(slots) + + unknown = [n for n in input_tensors_for_grad if n not in name_to_slot] + if unknown: + raise ValueError( + f"input_tensors_for_grad contains names not in {fwd_arg_type.__name__} " + f"schema: {unknown}" + ) + grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] + return fwd_slot_defaults, grad_targets + + +def _register_kernel( + *, + op_name: str, + schema_str: str, + arg_type: type, + arg_names: List[str], + buckets: List[_Bucket], + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + format_result: Callable[[Any], List[torch.Tensor]], +) -> Any: + """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the + ``fake_impl``, returning the ``CustomOpDef``. + + The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel + runs the fake impl directly on the unpacked object. Both go through + ``format_result``. + """ + + def _impl(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _unpack(arg_type, kwargs, buckets) + return format_result(impl(obj)) + + def _fake(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _unpack(arg_type, kwargs, buckets) + return format_result(fake_impl(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_name: str, + fwd_arg_type: type, + fwd_arg_names: List[str], + fwd_buckets: List[_Bucket], + bwd_arg_names: List[str], + bwd_buckets: List[_Bucket], + fwd_slot_defaults: List[Any], + grad_targets: List[int], + setup_context_user: Callable[..., Any], + backward_obj_type: type, + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> None: + """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op_name``. + + ``setup_context`` re-runs the fwd fake impl to recover output / saved + templates, reassembles each flat output chunk, and hands the saved tuple + + ``ctx_attrs`` to the module's ``setup_context``. + """ + + def _setup_context(ctx, inputs, output): + ctx._te_fwd_tensor_list_lengths = { + i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + } + kwargs = dict(zip(fwd_arg_names, inputs)) + fwd_obj = _unpack(fwd_arg_type, kwargs, fwd_buckets) + + user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(fwd_obj)) + + cursor = 0 + user_outputs: List[Any] = [] + for template in user_fakes: + n = _slot_count(template) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + user_outputs.append(_maybe_reassemble_tensor_subclass(template, chunk)) + + saved_list: List[Any] = [] + for template in saved_fakes: + n = _slot_count(template) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + saved_list.append(_maybe_reassemble_tensor_subclass(template, chunk)) + + bwd_obj = backward_obj_type() + tensors_to_save_from_setup = setup_context_user( + bwd_obj, + fwd_obj, + user_outputs[0] if len(user_fakes) == 1 else tuple(user_outputs), + ctx_attrs, + tuple(saved_list), + ) + tensors_to_save, tensor_objects = prepare_for_saving( + *(tensors_to_save_from_setup or ()) + ) + ctx.tensor_objects = tensor_objects + ctx.save_for_backward(*tensors_to_save) + ctx.bwd_obj = bwd_obj + + def _autograd_backward(ctx, *grad_outputs): + bwd_obj = ctx.bwd_obj + if hasattr(bwd_obj, "setup_saved_tensors"): + bwd_obj.setup_saved_tensors(ctx) + ctx.tensor_objects = None + per_output_grads = grad_outputs[0] + bwd_obj.grad_output = _decode_none(per_output_grads[0]) + kwargs = _pack(bwd_obj, bwd_buckets) + bwd_args_flat = [kwargs[name] for name in bwd_arg_names] + bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name) + grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + out: List[Any] = list(fwd_slot_defaults) + tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {}) + for pos, length in tensor_list_lengths.items(): + if isinstance(out[pos], list): + out[pos] = [None] * length + for pos, g in zip(grad_targets, grads): + out[pos] = g + return tuple(out) + + fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) + + +def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]: + """Start index of each ``_UniversalTensorBucket`` group in the flat args.""" + offsets: List[int] = [] + pos = 0 + for bucket in buckets: + if isinstance(bucket, _UniversalTensorBucket): + offsets.append(pos) + pos += len(bucket.schema_slots()) + return offsets + + +def _flatten_subclass_into_slots( + new_args: List[Any], slot_offsets: List[int], subclass: type +) -> None: + """Rewrite each universal-bucket 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 val is None or not isinstance(val, subclass): + continue + meta, tensors = _storage_flatten(val) + meta._data[_UniversalTensorBucket.KIND_KEY] = _UniversalKind.STORAGE + new_args[offset] = None + new_args[offset + 1] = list(tensors) + new_args[offset + 2] = meta + + +def _register_outer_forwarder( + *, + outer_op_name: str, + schema_str: str, + inner_op_name: str, + buckets: Optional[List[_Bucket]] = None, + subclass_list: Optional[List[type]] = None, +) -> Any: + """Define the outer op via ``torch.library.custom_op``: forward to the inner + op, optionally flattening registered subclass inputs in place first. Returns + the ``CustomOpDef``. + """ + inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name) + input_flatten_enabled = bool(subclass_list) and buckets is not None + slot_offsets = _collect_universal_slot_offsets(buckets) if input_flatten_enabled else [] + + def _forward(*flat: Any) -> List[torch.Tensor]: + if not input_flatten_enabled: + return inner_op(*flat) + new_args = list(flat) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, slot_offsets, sub) + return inner_op(*new_args) + + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, mutates_args=(), schema=schema_str + ) + op.register_fake(_forward) + return op + + +def _all_quantized_tensor_subclasses() -> List[type]: + """Return every imported ``QuantizedTensor`` wrapper subclass.""" + import transformer_engine.pytorch.tensor # noqa: F401 pylint: disable=import-outside-toplevel,unused-import + return [cls for cls in _STORAGE_REGISTRY.values() if issubclass(cls, QuantizedTensor)] + + +def register_custom_op( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_obj: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Optional[Callable[..., Any]]: + """Register a TE module's forward + backward as torch custom ops. + + Always two-tier: an inner ``_base`` op carries the real schema / + autograd, and an outer ```` op forwards to it, flattening any + quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an + empty subclass list simply makes the outer 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 outer op and returns the user-facing outputs. + + Registration touches experimental ``torch.library`` / opaque-object APIs + that may be missing on older PyTorch. If it fails, this warns once and + returns ``None`` instead of raising, so callers can fall back to eager under + ``torch.compile`` (a graph break) rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + input_tensors_for_grad=input_tensors_for_grad, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + setup_context=setup_context, + backward_arg_type=backward_arg_type, + backward_obj=backward_obj, + backward_impl=backward_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_fake_impl=bwd_fake_impl, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warnings.warn( + f"Could not register the torch.compile custom op '{op_name}' " + f"({type(e).__name__}: {e}); modules using it will fall back to eager " + "execution under torch.compile (a graph break, incompatible with " + "fullgraph=True)." + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_obj: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Callable[..., Any]: + """Body of :func:`register_custom_op`; see it for semantics.""" + outer_fwd_name = op_name + outer_bwd_name = f"{op_name}_backward" + inner_fwd_name = f"{op_name}_base" + inner_bwd_name = f"{outer_bwd_name}_base" + subclass_list = _all_quantized_tensor_subclasses() + + fwd_buckets = _get_buckets(fwd_arg_type) + bwd_buckets = _get_buckets(backward_arg_type) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_buckets) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets) + + num_grad_inputs = len(input_tensors_for_grad) + fwd_slot_defaults, grad_targets = _resolve_grad_targets( + fwd_buckets, fwd_arg_type, input_tensors_for_grad + ) + + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}" + + inner_fwd_def = _register_kernel( + op_name=inner_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + buckets=fwd_buckets, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=inner_bwd_name, + schema_str=bwd_schema, + arg_type=backward_arg_type, + arg_names=bwd_arg_names, + buckets=bwd_buckets, + impl=backward_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname), + ) + + outer_fwd_def = _register_outer_forwarder( + outer_op_name=outer_fwd_name, + schema_str=fwd_schema, + inner_op_name=inner_fwd_name, + buckets=fwd_buckets, + subclass_list=list(subclass_list), + ) + outer_bwd_def = _register_outer_forwarder( + outer_op_name=outer_bwd_name, schema_str=bwd_schema, inner_op_name=inner_bwd_name + ) + + autograd_common = { + "fwd_arg_type": fwd_arg_type, + "fwd_arg_names": fwd_arg_names, + "fwd_buckets": fwd_buckets, + "bwd_arg_names": bwd_arg_names, + "bwd_buckets": bwd_buckets, + "fwd_slot_defaults": fwd_slot_defaults, + "grad_targets": grad_targets, + "setup_context_user": setup_context, + "backward_obj_type": backward_obj, + "fwd_fake_impl": fwd_fake_impl, + } + _register_autograd_for_op( + fwd_op=inner_fwd_def, bwd_op_name=inner_bwd_name, **autograd_common + ) + _register_autograd_for_op( + fwd_op=outer_fwd_def, bwd_op_name=outer_bwd_name, **autograd_common + ) + + inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name) + inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name) + outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name) + outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name) + + fwd_slot_offsets = _collect_universal_slot_offsets(fwd_buckets) + bwd_slot_offsets = _collect_universal_slot_offsets(bwd_buckets) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return inner_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return inner_bwd_op(*new_args) + + for sub in subclass_list: + outer_fwd_def.register_torch_dispatch(sub, _fwd_rule) + outer_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + _quantized_tensor_passthrough_ops.add(outer_fwd_op.default) + _quantized_tensor_passthrough_ops.add(outer_bwd_op.default) + _quantized_tensor_passthrough_ops.add(inner_fwd_op.default) + _quantized_tensor_passthrough_ops.add(inner_bwd_op.default) + + def forward_fn(fwd_args): + user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(fwd_args)) + kwargs = _pack(fwd_args, fwd_buckets) + flat_in = [kwargs[name] for name in fwd_arg_names] + result = outer_fwd_op(*flat_in) + + cursor = 0 + outputs: List[Any] = [] + for template in user_fakes: + n = _slot_count(template) + chunk = [_decode_none(t) for t in result[cursor : cursor + n]] + cursor += n + outputs.append(_maybe_reassemble_tensor_subclass(template, chunk)) + + if len(outputs) == 1: + return outputs[0] + return tuple(outputs) + + return forward_fn diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py new file mode 100644 index 0000000000..8b8b3caa69 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -0,0 +1,116 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Value-opaque quantizers for torch.compile.""" + +from __future__ import annotations +from typing import Any, Dict, Tuple + +from ..constants import DType + + +# Registration marks the class with this attribute rather than recording it in a +# module-level set. It looks odd but is a deliberate workaround: the check must +# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a +# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate +# ``type(q) in some_set`` (no equality/hash rules for the opaque class object), +# which would graph-break under ``fullgraph=True``. +_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" + + +def is_value_opaque_quantizer(quantizer: Any) -> bool: + """Whether *quantizer*'s class is registered as a torch.compile value-opaque + type.""" + return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) + + +def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: + """Rebuild a tensorless quantizer of type *cls* from its value items. + + Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the + generated FX code calls this to materialize the quantizer constant. + """ + # Bypass ``__init__`` and restore the value attributes directly: the value + # items already capture every value-defining field (including derived ones), + # and the constructors have heterogeneous signatures / side effects. + obj = cls.__new__(cls) + field_names = set() + for name, value in items: + if name == "dtype": + value = DType.cast(value) + object.__setattr__(obj, name, value) + field_names.add(name) + # The deprecated amax-reduction group is not a value field; initialize it to + # None so attribute access keeps working on the rebuilt quantizer. + if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): + object.__setattr__(obj, "amax_reduction_group", None) + # Restore non-value derived state that ``__init__`` would normally build but + # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). + finalize = getattr(obj, "_rebuild_derived_state", None) + if finalize is not None: + finalize() + return obj + + +def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: + """``__fx_repr__`` for value-opaque quantizers (attached at registration). + + Returns an evaluable expression that rebuilds the quantizer via + :func:`_rebuild_quantizer`, capturing both the helper and the quantizer + class itself in the FX globals so codegen can resolve them with no global + registry and no qualname collisions. + + Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer + stores a process group (e.g. a non-``None`` deprecated + ``amax_reduction_group``): live distributed state must never be baked into + the graph as a constant. Pass the reduction group per quantize call instead + of storing it on the quantizer. + """ + cls = type(self) + items = self._value_key()[1] + return ( + f"_rebuild_quantizer({cls.__name__}, {items!r})", + {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls}, + ) + + +def register_value_opaque_quantizer(cls: type) -> None: + """Register a tensorless quantizer class as a torch.compile value opaque type. + + Attaches ``__fx_repr__`` and registers the class with + ``torch._library.opaque_object``. Safe to call on any PyTorch build: on + versions without the opaque-object API it only attaches ``__fx_repr__`` + (harmless), so Transformer Engine keeps importing and running in eager mode. + + The quantizer class must already provide value ``__eq__`` / ``__hash__`` and + a non-``None`` ``_value_fields`` (see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). + """ + # Stamp the class so it can be recognized as value-opaque in dynamo-traced + # code (used to fall back to eager for unregistered quantizers). + setattr(cls, _VALUE_OPAQUE_FLAG, True) + + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # class, so attach it before registering. + if "__fx_repr__" not in cls.__dict__: + cls.__fx_repr__ = _quantizer_fx_repr + + try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + register_opaque_type, + is_opaque_value_type, + ) + except (ImportError, AttributeError): + # Older PyTorch without the opaque-object API: eager value semantics + # still work; torch.compile specialization on the quantizer does not. + return + + try: + if not is_opaque_value_type(cls): + register_opaque_type(cls, typ="value") + except (RuntimeError, TypeError): + # 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. + pass diff --git a/transformer_engine/pytorch/dynamo/traceable_utils.py b/transformer_engine/pytorch/dynamo/traceable_utils.py new file mode 100644 index 0000000000..560599518a --- /dev/null +++ b/transformer_engine/pytorch/dynamo/traceable_utils.py @@ -0,0 +1,133 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pure-Python, torch.compile-traceable quantized tensor allocation and reassembly.""" + +from __future__ import annotations +import copy as _copy +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch + + +def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: + """Row-major (contiguous) stride for ``shape``.""" + stride: list = [] + acc = 1 + for dim in reversed(shape): + stride.append(acc) + acc *= dim + return tuple(reversed(stride)) + + +def make_empty_traceable( + quantizer, + shape: Tuple[int, ...], + *, + dtype: torch.dtype = torch.float32, + device: Optional[Union[torch.device, str]] = None, + requires_grad: bool = False, +) -> Any: + """Allocate a tensor purely in Python (traceable under torch.compile). + + When ``quantizer`` is not None, produces a quantized tensor via + ``alloc_tensors`` + ``__tensor_unflatten__`` (the compile-friendly + equivalent of ``Quantizer.make_empty``). The quantizer is copied first + so the caller's instance is never mutated. + + When ``quantizer`` is None, falls back to ``torch.empty`` for a plain tensor. + + Stashed metadata (``_te_flat_names``, ``_te_flat_ctx``) + --------------------------------------------------------- + The resulting quantized tensor has these attributes stashed so that + ``forward_fn`` (in custom_op.py) can read them at Dynamo trace time. + + Why: ``forward_fn`` runs inside torch.compile's trace. It needs the flat + buffer names and unflatten context to decode the custom op's flat Tensor[] + return back into a structured QuantizedTensor. Calling + ``__tensor_flatten__()`` for this would cause a graph break (it returns + non-Tensor Python objects -- List[str] and Dict -- that Dynamo cannot + represent as graph nodes). Accessing ``t._quantizer`` and calling + ``_describe_buffers()`` on it also fails: while Dynamo treats the retrieved + quantizer as constant metadata, it wraps it in a generic VariableTracker + that does not support method calls (unlike a closure-captured quantizer + which is recognized as a value-opaque constant). Stashing the buffer names + and context as plain attributes sidesteps both issues -- Dynamo reads them + as constants without needing to call any methods. + + Allocation cost: when ``forward_fn`` calls the fake impl to obtain these + templates, the ``torch.empty`` calls appear as nodes in the initial Dynamo + FX graph. However, because the tensors themselves are never used (only + the stashed metadata is read), AOT autograd's dead-code elimination removes + them before any kernel code is generated. They do not appear in the final + compiled graph. + """ + device = torch.device(device if device is not None else "cuda") + shape = tuple(shape) + if quantizer is None: + return torch.empty(shape, dtype=dtype, device=device, requires_grad=requires_grad) + from ..quantized_tensor import _STORAGE_REGISTRY # pylint: disable=import-outside-toplevel + + # Copy so the caller's quantizer is not mutated by alloc_tensors internals. + # The caller is expected to have already called set_usage() on the quantizer + # before passing it here -- Dynamo tracks those mutations as explicit setattr + # nodes in the graph, so the copy captures the post-mutation state and the + # stashed _te_flat_names reflects the correct buffer layout. + q = quantizer.copy() if hasattr(quantizer, "copy") else _copy.copy(quantizer) + ctx = q.create_metadata(shape, dtype=dtype, requires_grad=requires_grad) + inner = q.alloc_tensors(shape, device=device) + storage_cls = _STORAGE_REGISTRY[ctx["cls"]] + result = storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape)) + if requires_grad and hasattr(result, "requires_grad_"): + result.requires_grad_(True) + # TODO: understand why Dynamo does not recognize the quantizer retrieved via + # t._quantizer as the same value-opaque type it would if captured from a + # closure. If that is fixed upstream, the stashed attributes become + # unnecessary and we could compute slot counts directly from the quantizer. + result._te_flat_names = tuple(inner.keys()) + result._te_flat_ctx = ctx + return result + + +# --------------------------------------------------------------------------- # +# Slot counting and reassembly for the custom-op flat Tensor[] protocol. +# --------------------------------------------------------------------------- # + + +def _slot_count(value: Any) -> int: + """Number of flat tensor slots a value occupies in the op's Tensor[] return. + + Reads ``_te_flat_names`` stashed by :func:`make_empty_traceable`, which is + safe to access at Dynamo trace time (treated as constant metadata on a + traceable wrapper subclass). Plain tensors (no stashed names) occupy 1 slot. + """ + if value is None: + return 1 + names = getattr(value, "_te_flat_names", None) + if names is not None: + return len(names) + return 1 + + +def _maybe_reassemble_tensor_subclass( + template: Any, + chunk: List[Optional[torch.Tensor]], +) -> Optional[Union[torch.Tensor, Any]]: + """Rebuild a value from its flat tensors using ``template`` for geometry. + + ``template`` is a tensor produced by the fake impl via + :func:`make_empty_traceable`. Uses the stashed ``_te_flat_names`` / + ``_te_flat_ctx`` attributes for reassembly (trace-safe). For plain tensors + (no stashed metadata), returns the single chunk element directly. + """ + if template is None: + return None + names = getattr(template, "_te_flat_names", None) + ctx = getattr(template, "_te_flat_ctx", None) + if names is None or ctx is None: + return chunk[0] + inner_dict = dict(zip(names, chunk)) + shape = tuple(template.shape) if hasattr(template, "shape") else tuple(template.size()) + stride = _contiguous_stride(shape) + return type(template).__tensor_unflatten__(inner_dict, ctx, shape, stride) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index fed367bce5..d1dd5f362b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,10 +3,12 @@ # See LICENSE for license information. """Linear API""" + from dataclasses import dataclass from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op +import math import warnings import weakref @@ -38,10 +40,10 @@ divide, init_method_constant, needs_quantized_gemm, - assert_dim_for_fp8_exec, nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, + warn_compile_eager_fallback, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -58,9 +60,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, @@ -68,6 +71,8 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import register_custom_op, is_value_opaque_quantizer +from ..dynamo.traceable_utils import make_empty_traceable from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -92,11 +97,28 @@ class LinearFwdArgs: # --- Differentiable tensors (also passed positionally to autograd) --- weight: TensorOrQuantized - inp: torch.Tensor + inp: TensorOrQuantized 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] + + # --- CUDA-graph workspace pinning (torch.compile / reduce-overhead) --- + # Fetched in the *traced* module forward and threaded in as op inputs purely so + # the process-global, lru_cached cuBLAS / NVFP4-RHT workspaces become graph + # inputs: they are then allocated at trace time in the normal allocator (external + # to the cudagraph private pool) instead of being created inside the op during + # capture, where a persistent allocation trips check_memory_pool ("tensor not + # tracked as outputs"). The op body never reads these; general_gemm / the + # quantizer fetch the same lru_cached globals by address. Pinning the workspace + # in the forward also covers the backward GEMM, which reuses the same cached + # global. None on eager / non-compiled paths. + cublas_workspace: Optional[torch.Tensor] + rht_matrix: Optional[torch.Tensor] # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool @@ -130,7 +152,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 @@ -158,6 +182,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: @@ -195,7 +252,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 @@ -295,13 +353,19 @@ def _linear_forward_impl( if ub_name is not None: nvtx_label = f"{nvtx_label}.{ub_name}" - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - assert inp.shape[-1] == in_features, "GEMM not possible" + out_features = weight.shape[0] # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad + # Use the requires-grad flags captured into ``args`` at op-call time rather + # than the live tensors': the fake impl (``_linear_forward_impl_fake``) keys + # the number of FP8 inner buffers it emits off ``args.*_requires_grad``, so + # the real impl must agree to keep the custom-op output arity stable. Under + # ``torch.compile`` with CUDA-graph trees (``mode="reduce-overhead"``) the + # static graph inputs are detached during capture, so live + # ``weight.requires_grad`` / ``inp.requires_grad`` flip to False mid-capture + # and would otherwise diverge from the fake (schema/arity mismatch). + backward_needs_input = is_grad_enabled and args.weight_requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -331,7 +395,6 @@ def _linear_forward_impl( inputmat_total = None # Input tensor to pass to GEMM (gathered) own_quantized_input = False if fp8: - assert_dim_for_fp8_exec(inputmat, weight) if save_original_input: assert not isinstance( input_quantizer, Float8Quantizer @@ -418,7 +481,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: @@ -598,12 +661,24 @@ def _linear_forward_impl( if is_fsdp2 and weightmat is not 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. Needed because a custom op may + # not return a tensor that aliases an input or another return: the cached + # FP8 weight is the same tensor as ``new_weight_workspace`` (a return, on a + # cache miss) or ``weight_workspace`` (an input, on a cache hit). + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - "weight" if wt_save is weight else None, + wt_alias, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -622,10 +697,209 @@ def _linear_forward_impl( return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs +def _linear_forward_impl_fake( + args: LinearFwdArgs, +) -> Tuple[Any, Optional[Any], Optional[Tuple[Any, ...]], None, Optional[Dict]]: + """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile, + returning traceable tensors for the outputs and saved tensors instead of + allocating real data via C++ kernels.""" + 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 + + 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 isinstance(inp, QuantizedTensorStorage): + # 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 isinstance(weight, QuantizedTensorStorage) 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 isinstance(weight, QuantizedTensorStorage): + weight_quantizer = weight._quantizer + + if isinstance(weight, QuantizedTensorStorage): + # Primary-quantized weight: the impl reuses it as ``weightmat``. + weightmat = weight + weightmat_is_storage = True + weightmat_aliases_weight = True + else: + weightmat = make_empty_traceable( + weight_quantizer, tuple(weight.shape), + dtype=activation_dtype, device=weight.device, + ) + weightmat_is_storage = True + update_ws = args.is_first_microbatch is None or args.is_first_microbatch + if args.cache_weight and update_ws and args.weight_workspace is None: + new_weight_workspace = make_empty_traceable( + weight_quantizer, tuple(weight.shape), + dtype=activation_dtype, device=weight.device, + ) + else: + weightmat_aliases_weight = weight.dtype == activation_dtype + weightmat = make_empty_traceable( + None, tuple(weight.shape), dtype=activation_dtype, device=weight.device, + ) + + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). + # ------------------------------------------------------ + out_leading = inp.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + out_leading = out_leading * args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + out_leading = out_leading // args.tp_size + out = make_empty_traceable( + output_quantizer, + (out_leading, *tuple(inp.shape[1:-1]), out_features), + dtype=activation_dtype, + 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: + # 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. + # Copy the quantizer so the shared instance on fwd_args is not mutated. + save_q = input_quantizer.copy() if hasattr(input_quantizer, "copy") else input_quantizer + if own_quantized_input and not save_original_input: + if args.backward_override is not None: + save_q.set_usage(rowwise=True, columnwise=False) + elif ( + args.backward_input_needs_gather + and weight_quantizer is not None + and weight_quantizer.supports_only_rowwise_all_gather() + ): + save_q.set_usage(rowwise=True, columnwise=False) + else: + save_q.set_usage(rowwise=False, columnwise=True) + saved_inputmat = make_empty_traceable( + save_q, tuple(inp.shape), + dtype=activation_dtype, device=inp.device, + ) + else: + saved_inputmat = make_empty_traceable( + None, tuple(inp.shape), dtype=activation_dtype, device=inp.device, + ) + + # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached + # FP8 weight is shared with ``new_weight_workspace`` (a return, on a cache + # miss) or the ``weight_workspace`` input (on a cache hit), so it is + # reconstructed in ``_linear_setup_ctx`` rather than saved twice. + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered weight in backward. + elif weightmat_is_storage and new_weight_workspace is not None: + wt_alias = "new_workspace" + elif weightmat_is_storage and args.weight_workspace is not None: + wt_alias = "weight_workspace" + elif weightmat_is_storage: + wt_save = weightmat + else: + wt_save = make_empty_traceable( + None, 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, None, 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, ...]: @@ -638,7 +912,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`` are the op's user outputs ``(out, new_weight_workspace)``; + # only ``new_weight_workspace`` is needed here, to rebuild the deduped weight. inp = fwd_args.inp weight = fwd_args.weight @@ -661,7 +936,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 @@ -728,6 +1006,10 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight + elif wt_save_alias == "new_workspace": + wt_save = fwd_outputs[1] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace if saved_weight_alias == "weight": saved_weight = weight if bias_alias == "bias": @@ -805,6 +1087,19 @@ 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: + _w = saved_weight if saved_weight is not None else weight_fp8 + in_features = _w.shape[-1] + go_leading = grad_output.shape[0] + if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: + inp_leading = go_leading // bwd_args.tp_size + elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: + inp_leading = go_leading * bwd_args.tp_size + else: + inp_leading = go_leading + bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) + # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1311,6 +1606,92 @@ def wgrad_gemm( ) +def _linear_backward_impl_fake( + args: LinearBwdArgs, +) -> Tuple[Optional[Any], Optional[Any], Optional[Any]]: + """Traceable fake of :func:`_linear_backward`. + + The saved-tensor fields of ``args`` carry real tensors (fake-backed under + tracing). Returns ``(wgrad, dgrad, grad_bias)`` as traceable tensors, + mirroring the real backward's return contract without running C++ kernels. + + 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 if args.saved_weight is not None else args.weight_fp8 + out_dtype = args.activation_dtype + out_features, in_features = weight.shape + + # Mirror ``_linear_backward``: ``set_usage`` on ``grad_input_quantizer`` + # influences ``dgrad``'s buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + # 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. + _in_features = weight.shape[-1] + _go_leading = args.grad_output.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + _dgrad_leading = _go_leading // args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + _dgrad_leading = _go_leading * args.tp_size + else: + _dgrad_leading = _go_leading + dgrad = make_empty_traceable( + args.grad_input_quantizer, + (_dgrad_leading, *args.grad_output.shape[1:-1], _in_features), + dtype=out_dtype, + 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 = make_empty_traceable( + args.grad_weight_quantizer, + (out_features, in_features), + dtype=out_dtype, + device=weight.device, + ) + + grad_bias = None + if args.use_bias and args.requires_wgrad: + grad_bias = make_empty_traceable( + None, (out_features,), dtype=out_dtype, device=args.grad_output.device, + ) + + return wgrad, dgrad, grad_bias + + +# Custom op used under ``torch.compile``. +_linear_op = register_custom_op( + op_name="linear", + input_tensors_for_grad=["weight", "inp", "bias"], + fwd_arg_type=LinearFwdArgs, + fwd_impl=_linear_forward_impl, + fwd_fake_impl=_linear_forward_impl_fake, + setup_context=_linear_setup_ctx, + backward_arg_type=LinearBwdArgs, + backward_obj=LinearBwdArgs, + backward_impl=_linear_backward, + bwd_fake_impl=_linear_backward_impl_fake, +) + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1350,7 +1731,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, ) @@ -1395,6 +1776,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` @@ -1792,7 +2187,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, @@ -1867,12 +2261,7 @@ def forward( grad_output_quantizer, ) = quantizers - 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 = ( @@ -1912,16 +2301,58 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad + torch._check( + inp.shape[-1] == weight_tensor.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + if self.fp8: + torch._check( + math.prod(inp.shape[:-1]) % 8 == 0, + lambda: ( + "FP8 execution requires the product of all input dimensions except" + " the last to be divisible by 8" + ), + ) + torch._check( + inp.shape[-1] % 16 == 0, + lambda: "FP8 execution requires the input last dimension to be divisible by 16", + ) + torch._check( + weight_tensor.shape[0] % 16 == 0, + lambda: "FP8 execution requires out_features to be divisible by 16", + ) + torch._check( + weight_tensor.shape[1] % 16 == 0, + lambda: "FP8 execution requires in_features to be divisible by 16", + ) + linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + + # Pin the lazily-cached cuBLAS (and NVFP4-RHT) workspaces as op inputs so + # they are materialized at trace time (external to the cudagraph pool) + # rather than inside the op during capture. See LinearFwdArgs for details. + cublas_workspace = None + rht_matrix = None + if use_compiled_op: + cublas_workspace = get_cublas_workspace(inp.device.index, False, False) + from ..tensor.nvfp4_tensor import NVFP4Quantizer, get_rht_matrix + + if isinstance(input_quantizer, NVFP4Quantizer): + rht_matrix = get_rht_matrix( + input_quantizer._with_random_sign_mask, inp.device.index + ) + fwd_args = LinearFwdArgs( # tensors weight=weight_tensor, inp=inp, bias=linear_bias_tensor, weight_workspace=weight_workspace, + cublas_workspace=cublas_workspace, + rht_matrix=rht_matrix, # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, @@ -1976,13 +2407,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/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..2af15ca7b0 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -16,6 +16,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -23,12 +24,29 @@ ) +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False + + # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. _quantized_tensor_passthrough_ops: set = set() +#: Maps storage / wrapper class qualname -> class object, for ``__tensor_unflatten__``. +_STORAGE_REGISTRY: Dict[str, type] = {} + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -132,6 +150,66 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: f"{self.__class__.__name__} class does not implement copy_from_storage function" ) + # ----- PyTorch subclass flatten protocol (torch.compile / traceable allocation) ----- + + # Subclasses declare their tensor buffers once, as ``(attribute_name, + # constructor_kwarg)`` pairs in flatten order; everything else returned by + # :meth:`get_metadata` is treated as non-tensor context. + _FLATTEN_TENSOR_BUFFERS: Tuple[Tuple[str, str], ...] = () + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + # Register every storage / wrapper class so ``__tensor_unflatten__`` can + # resolve the concrete class from its qualname inside an FX graph. + _STORAGE_REGISTRY[cls.__qualname__] = cls + + def _flatten_nontensor_kwargs(self) -> Dict[str, Any]: + """Non-tensor constructor kwargs (scalars, dtype, quantizer).""" + tensor_kwargs = {kwarg for _, kwarg in self._FLATTEN_TENSOR_BUFFERS} + return {k: v for k, v in self.get_metadata().items() if k not in tensor_kwargs} + + def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: + """Return ``(inner_tensor_attr_names, context)``; see class comment.""" + present = [ + attr for attr, _ in self._FLATTEN_TENSOR_BUFFERS if getattr(self, attr) is not None + ] + ctx = { + "cls": type(self).__qualname__, + "is_tensor": isinstance(self, QuantizedTensor), + "requires_grad": ( + bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False + ), + "nontensor_kwargs": self._flatten_nontensor_kwargs(), + } + return present, ctx + + @staticmethod + def __tensor_unflatten__( + inner_tensors: Dict[str, torch.Tensor], + ctx: Dict[str, Any], + outer_size: Iterable[int], + outer_stride: Optional[Iterable[int]], + ) -> QuantizedTensorStorage: + """Rebuild a storage / wrapper from flat tensors + context.""" + cls = _STORAGE_REGISTRY[ctx["cls"]] + kwargs: Dict[str, Any] = dict(ctx["nontensor_kwargs"]) + # Map each declared buffer back to its constructor kwarg (absent -> None). + for attr, kwarg in cls._FLATTEN_TENSOR_BUFFERS: + kwargs[kwarg] = inner_tensors.get(attr) + if not ctx["is_tensor"]: + return cls(**kwargs) + # Wrapper subclass: it also needs outer shape / dtype / device / stride. + fake_dtype = kwargs.get("fake_dtype") + device = next((t.device for t in inner_tensors.values() if t is not None), None) + return cls( + shape=tuple(outer_size), + dtype=fake_dtype, + requires_grad=ctx["requires_grad"], + device=device, + stride=tuple(outer_stride) if outer_stride is not None else None, + **kwargs, + ) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], @@ -349,6 +427,71 @@ def make_empty( result.requires_grad_(True) return result + # ----- Data-free buffer/metadata primitives backing make_empty_traceable ----- + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + """Return ``{attr_name: (buffer_shape, buffer_dtype)}`` for the buffers + this quantizer would allocate for a logical tensor of ``shape``. + + Keys must match the buffer attribute names declared in the storage's + ``_FLATTEN_TENSOR_BUFFERS`` and respect the quantizer's usage flags. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement _describe_buffers; " + "it cannot be used with traceable allocation" + ) + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + """Non-tensor context for the produced storage. + + Returns ``{"cls": , "nontensor_kwargs": {...}}`` where ``cls`` is + the concrete class to instantiate (wrapper subclass for user-visible + tensors, bare storage class for ``internal`` quantizers) and + ``nontensor_kwargs`` are its non-tensor constructor kwargs (e.g. + ``fp8_dtype``, ``quantizer``, ``fake_dtype``). + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement _storage_metadata; " + "it cannot be used with traceable allocation" + ) + + def alloc_tensors( + self, + shape: Iterable[int], + *, + device: Optional[Union[torch.device, str]] = None, + ) -> Dict[str, torch.Tensor]: + """Allocate (uninitialized) the flat buffers for ``shape``. + + Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors`` + argument of the storage's ``__tensor_unflatten__``. + """ + device = torch.device(device if device is not None else "cuda") + return { + attr: torch.empty(buf_shape, dtype=buf_dtype, device=device) + for attr, (buf_shape, buf_dtype) in self._describe_buffers(tuple(shape)).items() + } + + def create_metadata( + self, + _shape: Iterable[int], + *, + dtype: torch.dtype, + requires_grad: bool = False, + ) -> Dict[str, Any]: + """Build the data-free ``__tensor_unflatten__`` context describing the + quantized tensor this quantizer would produce for ``shape`` / ``dtype``. + """ + meta = self._storage_metadata(dtype) + return { + "cls": meta["cls"].__qualname__, + "is_tensor": not self.internal, + "requires_grad": requires_grad, + "nontensor_kwargs": meta["nontensor_kwargs"], + } + def calibrate(self, tensor: torch.Tensor) -> None: """Calibrate quantizer state @@ -408,6 +551,78 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } + #: Attributes shared by every quantizer that take part in value identity. + _BASE_VALUE_FIELDS: Tuple[str, ...] = ( + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + ) + + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Subclass-specific value-defining attribute names, or ``None``. + + Returning ``None`` (the default) means the quantizer cannot be represented as + a value opaque object and keeps identity-based equality/hashing. + This also means that passing such a quantizer as an argument to a custom op + causes a graph break under torch.compile, since it cannot be baked into the + FX graph as a constant. + """ + return None + + def _check_value_has_no_process_group(self) -> None: + # A value quantizer is baked into the FX graph as a constant via its + # value key, which cannot carry live distributed state. Enforced here -- + # the single point every value-materialization path (``__eq__`` / + # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom + # ``__fx_repr__`` cannot bypass it. Reject any field holding a + # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than + # silently dropping it; pass the reduction group per quantize call. + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) + + def _value_key(self) -> Tuple[Any, ...]: + """Hashable, reproducible key identifying this quantizer's value. + + Only valid for value quantizers (``_value_fields()`` is not ``None``). + """ + fields = self._value_fields() # pylint: disable=assignment-from-none + assert fields is not None, f"{type(self).__name__} is not a value quantizer" + self._check_value_has_no_process_group() + items = [] + for name in self._BASE_VALUE_FIELDS + tuple(fields): + value = getattr(self, name) + if name == "dtype": + # ``DType`` is an ``IntEnum``; store the int so the key stays + # plain: hashable and ``repr``-reproducible for FX codegen. + value = int(value) + items.append((name, value)) + return (type(self).__qualname__, tuple(items)) + + def __eq__(self, other: object) -> Any: + # Value quantizers compare by configuration; everything else keeps the + # default identity semantics (returning ``NotImplemented`` makes Python + # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. + if self is other: + return True + if self._value_fields() is None or type(self) is not type(other): + return NotImplemented + if other._value_fields() is None: + return NotImplemented + return self._value_key() == other._value_key() + + def __hash__(self) -> int: + if self._value_fields() is None: + return object.__hash__(self) + return hash(self._value_key()) + class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ba46508d74..04e5281f65 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -7,13 +7,14 @@ from collections.abc import Iterable import math import warnings -from typing import Any, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -69,6 +70,42 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") + + # ----- traceable allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "is_2D_scaled": self.block_scaling_dim == 2, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Blockwise FP8 scales are FP32; columnwise data is stored transposed. + if self.rowwise_usage: + buffers["_rowwise_data"] = (shape, torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.float32, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = (tuple(self.get_columnwise_shape(shape)), torch.uint8) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.float32, + ) + return buffers + def update_quantized( self, src: torch.Tensor, @@ -211,6 +248,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return Float8BlockScaling +register_value_opaque_quantizer(Float8BlockQuantizer) + + class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): """Tensor class with FP8 data quantized via NxN blocks or 1xN blocks. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e26abf7df0..423b055f15 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -4,7 +4,7 @@ """Tensor class with FP8 data""" from __future__ import annotations -from typing import Any, Optional, Tuple, Iterable, Union +from typing import Any, Dict, Optional, Tuple, Iterable, Union import warnings import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState @@ -15,9 +15,10 @@ Float8CurrentScaling, Recipe, ) -from ..utils import canonicalize_process_group, devices_match +from ..utils import canonicalize_process_group, devices_match, is_non_tn_fp8_gemm_supported from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import dist_group_type, DType @@ -386,6 +387,45 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. + return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") + + # ----- traceable allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8TensorStorage if self.internal else Float8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Mirror the C++ quantizer allocation (csrc/quantizer.cpp): on non-TN-capable + # archs (Blackwell+) a single ``_data`` buffer backs both row- and column-wise + # usage and no separate transpose is materialized. This must match what the + # real kernel produces so the torch.compile fake layout lines up slot-for-slot. + non_tn = is_non_tn_fp8_gemm_supported() + if self.rowwise_usage or non_tn: + buffers["_data"] = (shape, torch.uint8) + if self.columnwise_usage and not non_tn: + buffers["_transpose"] = ((shape[-1], *shape[:-1]), torch.uint8) + # Per-tensor scale-inv is always present for current scaling. + buffers["_scale_inv"] = ((1,), torch.float32) + return buffers + + +register_value_opaque_quantizer(Float8CurrentScalingQuantizer) + class Float8Tensor(Float8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d759aaf5c4..bd51a88aaf 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable import math -from typing import Optional, Tuple, Union, Any +from typing import Optional, Tuple, Union, Any, Dict import warnings import torch @@ -18,6 +18,7 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -57,6 +58,41 @@ def copy(self) -> MXFP8Quantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype",) + + # ----- traceable allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + if self.rowwise_usage: + buffers["_rowwise_data"] = (shape, torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = (shape, torch.uint8) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + return buffers + def update_quantized( self, src: torch.Tensor, @@ -1058,3 +1094,6 @@ def backward( ) return dgrad, None return grad.view(ctx.shape), None + + +register_value_opaque_quantizer(MXFP8Quantizer) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index aa92be004f..d22423353c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -7,7 +7,7 @@ from collections.abc import Iterable import math import warnings -from typing import Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import functools import torch @@ -23,6 +23,7 @@ from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -173,6 +174,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") + self._with_random_sign_mask = with_random_sign_mask self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -184,6 +186,16 @@ def __getstate__(self): state["amax_reduction_group"] = None return state + def _rebuild_derived_state(self) -> None: + """Restore the derived ``rht_matrix`` after value-key reconstruction. + + ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask`` + and the device, so it cannot be part of the (hashable) value key. + ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit. + """ + self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) + def update_quantized( self, src: torch.Tensor, @@ -233,8 +245,14 @@ def copy(self) -> NVFP4Quantizer: ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm - quantizer.rht_matrix = self.rht_matrix quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t + if not torch.compiler.is_compiling(): + # Under Dynamo tracing rht_matrix is a FakeTensor on an opaque script + # object; accessing it triggers SourcelessBuilder which cannot wrap + # FakeTensor. The fake impl never runs real quantization so the matrix + # is unnecessary -- it will be rebuilt lazily via _rebuild_derived_state + # if the quantizer is later used outside tracing. + quantizer.rht_matrix = self.rht_matrix return quantizer @@ -333,6 +351,79 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group, not a value (``_value_key`` rejects a stored group). + # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived + # from ``_with_random_sign_mask``; kept in the key so the rebuilt + # quantizer carries it without recomputation. + return ( + "dtype", + "with_rht", + "with_post_rht_amax", + "with_2d_quantization", + "stochastic_rounding", + "row_scaled_nvfp4", + "nvfp4_use_4over6", + "nvfp4_e4m3_max", + "nvfp4_4over6_err_mode", + "_with_random_sign_mask", + "rht_matrix_random_sign_mask_t", + "with_amax_reduction", + ) + + # ----- traceable allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, + "nontensor_kwargs": { + "fp4_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored + # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar). + # Order matches NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (the canonical + # __tensor_flatten__ order): data + scale_inv per usage first, amax last. + # Workaround: call @staticmethods via the class, not the instance -- + # instance access breaks torch.compile guard generation (pytorch #182741). + if self.rowwise_usage: + buffers["_rowwise_data"] = (type(self).convert_shape_for_fp4(shape), torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = ( + type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)), + torch.uint8, + ) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + if self.rowwise_usage: + amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,) + buffers["_amax_rowwise"] = (amax_rowwise_shape, torch.float32) + if self.columnwise_usage: + buffers["_amax_columnwise"] = ((1,), torch.float32) + return buffers + + +register_value_opaque_quantizer(NVFP4Quantizer) + class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): """Quantized tensor class with FP4 data diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index f7a3dae70b..ec2c40ef9a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -35,6 +35,15 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a97162f91c..1419a02559 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -75,6 +75,14 @@ class Float8TensorStorage(QuantizedTensorStorage): _transpose: Optional[torch.Tensor] _transpose_invalid: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_data", "data"), + ("_transpose", "data_transpose"), + ("_scale_inv", "fp8_scale_inv"), + ) + def __new__( cls, *args, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index ea592cd989..a2a3bf2f4c 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -82,6 +82,15 @@ class MXFP8TensorStorage(QuantizedTensorStorage): # GEMM _with_gemm_swizzled_scales: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 53bb5e7c11..5ed6d1d641 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -108,6 +108,17 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Global E4M3 scale bound used by this NVFP4 tensor _nvfp4_e4m3_max: int + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ("_amax_rowwise", "amax_rowwise"), + ("_amax_columnwise", "amax_columnwise"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index ffbfbc1fdd..8eeaf2417d 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,6 +26,20 @@ ] +def warn_compile_eager_fallback(reason: str) -> None: + """Warn that a TE module is running eagerly under ``torch.compile``. + + Emitted when ``reason`` is unsupported on the module's compiled custom-op + path. Python's default warning filter dedups identical messages, so each + distinct ``reason`` is surfaced once. + """ + warnings.warn( + f"Falling back to eager execution under torch.compile: {reason} is " + "unsupported on the compiled path (graph-breaks under fullgraph=True).", + stacklevel=2, + ) + + @functools.lru_cache(maxsize=None) def get_cached_ones_tensor( num_elements: int,