diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 46795415e5..0b9b0aa13c 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -18,6 +18,11 @@ import torch import torch.distributed as dist +try: + from torch._dynamo.utils import counters as dynamo_counters +except ImportError: # pragma: no cover + dynamo_counters = None + import transformer_engine.pytorch as te from transformer_engine.common.recipe import ( DelayedScaling, @@ -200,6 +205,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." ) @@ -285,6 +303,13 @@ def _parse_args(argv=None, namespace=None): ) args = parser.parse_args(argv, namespace) + if args.compile and args.use_cuda_graphs: + parser.error( + "--compile and --use-cuda-graphs are mutually exclusive; to test" + " torch.compile with CUDA graphs use --compile --compile-mode" + " reduce-overhead." + ) + if args.use_cuda_graphs and args.layer_type in [te.MultiheadAttention, te.TransformerLayer]: warnings.warn(f"{args.layer_type.__name__} does not support CUDA Graphs!") args.use_cuda_graphs = False @@ -535,6 +560,14 @@ def run_fwd_bwd(model, x): loss.backward() return out + if opts.compile: + for i, layer in enumerate(test_model.layers): + test_model.layers[i] = torch.compile(layer, fullgraph=True, mode=opts.compile_mode) + dist_print( + f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...", + debug=True, + ) + 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: @@ -545,7 +578,18 @@ def run_fwd_bwd(model, x): if not opts.benchmark: del test_graph else: + if opts.compile and opts.compile_mode == "reduce-overhead": + # Warm up so the measured run below replays captured CUDA graphs. + for _ in range(2): + torch.compiler.cudagraph_mark_step_begin() + run_fwd_bwd(test_model, test_x) + test_model.zero_grad(set_to_none=True) + test_x.grad = None + torch.compiler.cudagraph_mark_step_begin() test_out = run_fwd_bwd(test_model, test_x) + if opts.compile and opts.compile_mode == "reduce-overhead" and dynamo_counters is not None: + skips = dynamo_counters["inductor"]["cudagraph_skips"] + assert skips == 0, f"reduce-overhead fell back to eager: {skips} cudagraph skip(s)" test_grads = [test_out, test_x.grad] names = ["output", "input.grad"] for test_name, test_param in test_model.named_parameters(): diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index fe02f990b4..7d331ccb55 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -292,6 +292,15 @@ def _check_gradients(model_distributed, model_single, main_grad_check=False): assert not bool(numerics_failed.item()) +def _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel): + grad_d = input_distributed.grad + if parallel_mode == "row": + grad_d = _gather(grad_d, dim=1) + elif sequence_parallel: + grad_d = _gather(grad_d, dim=0) + _check_outputs(input_single_node.grad, grad_d) + + def _copy_params(model_distributed, model_single): for dist_param, single_param in zip(model_distributed.parameters(), model_single.parameters()): with torch.no_grad(): @@ -310,22 +319,35 @@ 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: + # Reset the compile cache so parametrized cases don't trip recompile_limit. + torch._dynamo.reset() + forward_single_node = torch.compile(model_single_node, fullgraph=True, mode=compile_mode) + forward_distributed = torch.compile(model_distributed, fullgraph=True, mode=compile_mode) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), ): - 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 +663,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 +726,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: @@ -728,6 +763,8 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): main_grad_check=("fuse_wgrad_accumulation" in kwargs), ) + _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel) + def test_linear(): """Run linear layer tests with various configurations.""" @@ -740,12 +777,20 @@ 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: if kwargs.get("save_original_input", False) and QUANTIZATION == "fp8": continue - if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: + # use_compile: debug instrumentation forces the eager fallback, so + # compile is a no-op there. + if NVTE_TEST_NVINSPECT_ENABLED and ( + kwargs.get("delay_wgrad_compute", False) or kwargs.get("use_compile", False) + ): + continue + if kwargs.get("use_compile", False) and QUANTIZATION == "fp8": continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6b1ad870e9..12d3747c3a 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, + use_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 use_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,45 @@ def test_layers_with_overlap_bf16( ) +@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize( + "quantization", + [None, "fp8_current_scaling", "mxfp8"], + ids=["bf16", "fp8_current_scaling", "mxfp8"], +) +@pytest.mark.parametrize( + "linear_parallel_mode,overlap_rs_dgrad", + [ + ("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, quantization, compile_mode +): + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile, + checked numerically against the eager, non-overlap reference.""" + if quantization is not None and linear_parallel_mode == "row": + pytest.skip( + "FP8 row-parallel UB forces differentiable fp8_output, unsupported under compile." + ) + _run_layer_with_overlap( + te.Linear.__name__, + linear_parallel_mode, + overlap_rs_dgrad, + quantization is not None, + quantization, + use_compile=True, + compile_mode=compile_mode, + ) + + @pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 06cda114f0..e5aaf814a3 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -496,7 +496,7 @@ def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self): ``gather_along_first_dim`` cannot operate on a columnwise-only NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather`` must return True in this case so ``_linear_forward_impl`` / - ``_linear_backward`` preserve rowwise data (which NVFP4 can + ``_linear_backward_impl`` preserve rowwise data (which NVFP4 can dequantize) instead. """ hq = HybridQuantizer( diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6ebc408dcb..8c136dd797 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,10 +4,18 @@ import abc import contextlib +import os +import re +import sys import warnings import pytest import torch + +try: + from torch._dynamo.utils import counters +except ImportError: # pragma: no cover + counters = None from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: @@ -22,16 +30,16 @@ except ImportError: _opaque_available = False -from torch._dynamo.utils import counters - import transformer_engine.pytorch as te import transformer_engine_torch as tex 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 FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec @@ -40,11 +48,14 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, - Float8BlockQuantizer, - MXFP8Quantizer, - NVFP4Quantizer, ) + +# Import from the local utils.py by explicit path: importing cutedsl makes a +# top-level ``utils`` package visible that would shadow it. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id + +sys.path.pop(0) from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) @@ -57,6 +68,15 @@ nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True) +@pytest.fixture(autouse=True) +def _reset_fp8_global_state(): + """Pending FP8 global state (e.g. delayed-scaling amax reductions) must not + leak between tests: a leftover buffer makes a later autocast __exit__ call + raw tex bindings, which graph-breaks fullgraph=True tests.""" + yield + FP8GlobalStateManager.reset() + + def nvfp4_row_scaled(): nvfp4_recipe = recipe.NVFP4BlockScaling( disable_rht=True, @@ -96,6 +116,69 @@ def nvfp4_4over6(): _all_recipes.append(nvfp4_row_scaled()) +# Modes exercised by the te.Linear tests; "reduce-overhead" = CUDA-graph trees. +_compile_modes = ["default", "reduce-overhead"] + + +def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: + """One eager iteration so lazily-initialized TE state (fp8 meta, workspaces) + is allocated before any CUDA-graph capture.""" + out = fn(inp) + if backward: + out.sum().backward() + + +def _dynamo_counter(group: str, key: str): + """Read a torch._dynamo counter; None (with a warning) if the private + counters API is gone, so CI degrades instead of failing.""" + try: + return counters[group][key] + except Exception: # pylint: disable=broad-except + warnings.warn(f"torch._dynamo.utils.counters[{group!r}][{key!r}] unavailable") + return None + + +@contextlib.contextmanager +def _assert_no_cudagraph_skips(enabled: bool): + """Assert reduce-overhead really captured CUDA graphs: inductor may skip + capture and silently fall back to eager, which ``fullgraph=True`` does not + catch. No-op when ``enabled`` is False.""" + before = _dynamo_counter("inductor", "cudagraph_skips") + yield + if enabled and before is not None: + skipped = _dynamo_counter("inductor", "cudagraph_skips") - before + assert skipped == 0, ( + f"reduce-overhead fell back to eager: {skipped} cudagraph skip(s); " + "see the 'skipping cudagraphs due to ...' log for the reason" + ) + + +# All compute runs inside the op and the loss grad is ones, so bit-exact. +_EAGER_ATOL, _EAGER_RTOL = 0.0, 0.0 + + +def _assert_close_eager_compiled(fn, compiled, model, base): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the + forward output and the input / weight gradients match.""" + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = fn(inp_eager) + 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) @@ -1330,12 +1413,18 @@ def _hw_available(quantizer): return fp8_available # Float8CurrentScalingQuantizer -# (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ pytest.param(_mxfp8, id="mxfp8"), pytest.param(_blockwise, id="float8_blockwise"), pytest.param(_current_scaling, id="float8_current_scaling"), - pytest.param(_nvfp4, id="nvfp4"), + pytest.param( + _nvfp4, + id="nvfp4", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), ] @@ -1702,9 +1791,13 @@ def test_tensor_spec_matches_primitives(factory, shape): # Metadata matches the quantizer's. assert spec.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16) - # inner_names + create_inner_tensors match inner_tensor_specs. + # inner_names follows the storage's canonical __tensor_flatten__ order (the + # order the real op flattens its outputs to), while create_inner_tensors + # matches the inner_tensor_specs geometry (a name->shape/dtype mapping). specs = q.inner_tensor_specs(shape) - names = tuple(specs) + direct = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(direct.__tensor_flatten__()[0]) + assert set(names) == set(specs) assert spec.inner_names() == names inner_tensors = spec.create_inner_tensors() assert len(inner_tensors) == len(names) @@ -1714,7 +1807,6 @@ def test_tensor_spec_matches_primitives(factory, shape): assert inner.dtype == exp_dtype # The assembled tensor matches one built directly from the primitives. - direct = _build_from_primitives(q, shape, torch.bfloat16) assert _signature(spec.create_tensor(), names) == _signature(direct, names) @@ -1781,3 +1873,446 @@ def test_to_tensor_spec_quantized(factory, shape): assert _signature(spec.create_tensor(), spec.inner_names()) == _signature( tensor, spec.inner_names() ) + + +# --------------------------------------------------------------------------- +# te.Linear +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize( + "fp8_recipe", + [None, *_all_recipes], + ids=lambda r: "bf16" if r is None else type(r).__name__, +) +def test_te_linear_compiles(fp8_recipe, compile_mode): + """ + torch.compile(fullgraph=True) of ``te.Linear`` under every built-in + recipe (plus the bf16-only baseline with no autocast), for both the default + backend and ``mode="reduce-overhead"`` (CUDA-graph trees). + """ + 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) + + # Iterate a few times so reduce-overhead actually replays a captured graph. + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + 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 of Linear with the weight initialized as an FP8 tensor + (exercises the wrapper op's ``register_torch_dispatch`` input flattening).""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + 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)`` under no_grad: + forward must return a working :class:`Float8Tensor` (exercises the output + rewrap path). The differentiable case falls back to eager, so it is not + covered here.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + 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" + deq = out.dequantize() + assert deq.shape == (32, 32) + assert deq.dtype == dtype + torch.testing.assert_close( + deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + +# Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a +# single-GPU unit test can construct. Distributed-only reasons (fsdp_group, +# DistributedWeight) and CPU offloading need machinery this file doesn't have; +# delayed scaling is a hard error (check_recipe_support), tested separately. +# Modes: "bwd" = fwd+bwd vs eager; "fwd_grad" = grad-enabled forward only +# (differentiable fp8_output backward hits a PyTorch limitation: the Float8 +# output crossing the graph-break boundary gets a plain-tensor tangent); +# "no_grad" = forward under no_grad. +_FALLBACK_CASES = [ + "fp8_output_differentiable", + "fuse_wgrad_accumulation", + "delayed_wgrad", + "quantized_input", +] + + +def _fallback_case(case, dtype, device): + """Build ``(model, fn, mode, post_backward, reason)`` for one case.""" + model_kwargs = {} + if case == "fuse_wgrad_accumulation": + model_kwargs["fuse_wgrad_accumulation"] = True + elif case == "delayed_wgrad": + model_kwargs["delay_wgrad_compute"] = True + model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs) + + if case == "fp8_output_differentiable": + fp8_recipe = recipe.Float8CurrentScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, fp8_output=True).dequantize() + + return model, fn, "fwd_grad", None, "differentiable fp8_output=True" + if case == "fuse_wgrad_accumulation": + model.weight.main_grad = torch.zeros_like(model.weight, dtype=torch.float32) + return model, model, "bwd", None, "fuse_wgrad_accumulation" + if case == "delayed_wgrad": + return model, model, "bwd", model.backward_dw, "delayed wgrad compute" + if case == "quantized_input": + fp8_recipe = recipe.Float8CurrentScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + return model, fn, "no_grad", None, "a quantized input tensor" + raise ValueError(case) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("case", _FALLBACK_CASES) +def test_te_linear_compile_eager_fallback(case): + """Configs unsupported on the compiled custom-op path must fall back to + eager under ``torch.compile`` -- warning + numerics identical to eager -- + and graph-break with the explicit reason under ``fullgraph=True``.""" + dtype, device = torch.bfloat16, "cuda" + torch.manual_seed(0) + model_ref, fn_ref, mode, post_bwd_ref, _ = _fallback_case(case, dtype, device) + torch.manual_seed(0) + _model, fn, _, post_bwd, reason = _fallback_case(case, dtype, device) + + def make_inp(): + torch.manual_seed(1) + x = torch.randn(32, 64, dtype=dtype, device=device) + if case == "quantized_input": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, device=device + ) + return quantizer(x) + return x.requires_grad_(mode != "no_grad") + + torch._dynamo.reset() + compiled = torch.compile(fn) + grad_ctx = torch.no_grad() if mode == "no_grad" else contextlib.nullcontext() + + inp_ref, inp = make_inp(), make_inp() + with grad_ctx: + out_ref = fn_ref(inp_ref) + with pytest.warns(UserWarning, match="Falling back to eager execution under torch.compile"): + out = compiled(inp) + if mode == "bwd": + out_ref.sum().backward() + if post_bwd_ref is not None: + post_bwd_ref() + out.sum().backward() + if post_bwd is not None: + post_bwd() + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + torch._dynamo.reset() + compiled_fg = torch.compile(fn, fullgraph=True) + with pytest.raises(Exception, match=re.escape(reason)): + with grad_ctx: + compiled_fg(make_inp()) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_linear_compile_delayed_scaling_raises(): + """Delayed scaling is rejected under torch.compile with a hard error + (``check_recipe_support`` in ``te.autocast.__enter__``), not a fallback. + Without fullgraph the raising frame is skipped and re-run eagerly (where + the guard passes), so only ``fullgraph=True`` surfaces the error.""" + dtype, device = torch.bfloat16, "cuda" + model = te.Linear(64, 32, params_dtype=dtype, device=device) + fp8_recipe = recipe.DelayedScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + inp = torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True) + with pytest.raises(Exception, match="DelayedScaling is not supported under torch.compile"): + torch.compile(fn, fullgraph=True)(inp) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_is_first_microbatch(compile_mode): + """torch.compile of ``te.Linear`` across a microbatch schedule: + ``is_first_microbatch=True`` caches the FP8 weight, later steps must reuse + it and stay numerically aligned with eager. The eager reference runs on a + separate module so it cannot mask a corrupted or rebuilt cache.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + ref_model = te.Linear(64, 32, params_dtype=dtype, device=device) + with torch.no_grad(): + ref_model.weight.copy_(model.weight) + ref_model.bias.copy_(model.bias) + + schedule = [True, False, False] + is_first = schedule[0] # rebound each step; closed over by the fns. + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=is_first) + + def ref_fn(inp): + with te.autocast(recipe=fp8_recipe): + return ref_model(inp, is_first_microbatch=is_first) + + # Eager priming: FP8 state must exist before tracing (creating quantizers + # in-graph breaks later recompiles; upstream Dynamo bug). + is_first = None + fn(torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True)) + is_first = schedule[0] + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + 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) + + cached_workspace = None + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for step, is_first in enumerate(schedule): + base = torch.randn(32, 64, dtype=dtype, device=device) + + inp_ref = base.detach().clone().requires_grad_(True) + ref_model.zero_grad(set_to_none=True) + out_ref = ref_fn(inp_ref) + out_ref.sum().backward() + + inp = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out = compiled(inp).clone() + out.sum().backward() + + torch.testing.assert_close(out, out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close( + model.weight.grad, ref_model.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + workspace = model._fp8_workspaces.get("weight") + assert workspace is not None, f"no cached FP8 weight after step {step}" + if step == 0: + cached_workspace = workspace + else: + assert workspace is cached_workspace, f"cache rebuilt at step {step}" + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.xfail( + reason=( + "value-opaque module state comes back as None on recompile" + " (pytorch/pytorch#187041; fixed by #187057 (cold compile, merged)" + " + #193190 (FX-graph-cache hit, in review))" + ), + strict=False, +) +def test_te_linear_compile_train_eval_switch(): + """train -> eval -> train on the same compiled ``te.Linear``, vs eager.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=True) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + def train_step(): + inp = torch.randn(16, 64, dtype=dtype, device=device, requires_grad=True) + out = compiled(inp) + out.sum().backward() + inp_ref = inp.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_ref = fn(inp_ref) + out_ref.sum().backward() + torch.testing.assert_close( + out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + model.zero_grad(set_to_none=True) + + train_step() + + model.eval() + x = torch.randn(16, 64, dtype=dtype, device=device) + with torch.no_grad(): + out_eval = compiled(x) + ref_eval = fn(x) + torch.testing.assert_close(out_eval, ref_eval, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + model.train() + train_step() + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_linear_dynamic_shapes(): + """torch.compile of ``te.Linear`` with a ``mark_dynamic`` batch dimension: + one graph must serve all batch sizes -- no recompiles -- and match eager + numerically. + + Only the leading (batch/sequence) dims may be dynamic; the last dim is + fixed by the weight's ``in_features``. + """ + dtype = torch.bfloat16 + device = "cuda" + 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] + + # Two warmup calls: the second absorbs the one-time recompile from module + # attributes lazily created during call one (e.g. the cached ``is_fsdp2``). + for _ in range(2): + warm = torch.randn(batch_sizes[0], in_features, dtype=dtype, device=device) + torch._dynamo.mark_dynamic(warm, 0) + compiled(warm.requires_grad_(True)).sum().backward() + model.zero_grad(set_to_none=True) + unique_graphs_baseline = _dynamo_counter("stats", "unique_graphs") + if not unique_graphs_baseline: + warnings.warn("unique_graphs counter is stale; skipping the recompile check") + + for batch in batch_sizes: + inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) + # 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 unique_graphs_baseline: + unique_graphs_after = _dynamo_counter("stats", "unique_graphs") + assert unique_graphs_after == unique_graphs_baseline, ( + "Unexpected recompilation(s) across different batch sizes: " + f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" + ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 18451976ab..092622f23c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -31,6 +31,7 @@ "general_gemm", "general_grouped_gemm", "general_grouped_gemm_for_grouped_tensor", + "get_cublas_workspace", ] @@ -49,6 +50,10 @@ def get_cublas_workspace_size_bytes() -> None: def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor: """Returns workspace for cublas GEMM.""" assert not (ub and grouped_gemm), "UB is unsupported for grouped GEMM." + assert not torch.cuda.is_current_stream_capturing(), ( + "cuBLAS workspace would be first allocated during CUDA-graph capture and would live in" + " the graph's memory pool; run a warmup iteration outside the graph first." + ) if ub: return torch.empty( diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 4d5c76e9ce..e42eb8f9f6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,10 +6,13 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec +from .custom_op import register_custom_op, TensorOrQuantized __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", "TensorSpec", "to_tensor_spec", + "register_custom_op", + "TensorOrQuantized", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py new file mode 100644 index 0000000000..1378e65acb --- /dev/null +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -0,0 +1,1319 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile custom-op framework for Transformer Engine. + +Registers TE modules' eager forward/backward as ``torch.library`` custom ops so +``torch.compile(fullgraph=True)`` traces them as single graph nodes. +``register_custom_op`` is the entry point; ``module/linear.py`` is the first user. + +A TE forward/backward implementation takes one dataclass argument +(``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix +tensors, quantized tensors, quantizers, process groups and plain Python values. + +A ``torch.library`` custom op is narrower: it only accepts flat schema slots +(tensors plus opaque objects) and returns a flat ``Tensor[]``. + +Bridging the two takes three parts (below): a parsed per-op *arg plan* maps the +args dataclass onto the op's input slots; a per-trace *output plan*, parsed from +the data-free fake impl's result, maps the logical outputs / saved tensors / +grads onto ranges of the flat return; and a *two-tier op* lets a +quantized-tensor subclass be an op input. + +Field <-> slot mapping. ``_parse_arg_type`` parses the dataclass's field +annotations once, at registration, into an immutable ``_ArgPlan``: per field a +``_FieldPlan`` (its ``_FieldKind`` plus the schema slots it occupies), and the +derived layout -- schema string, slot order, gradient placement and +tensor-or-quantized offsets. ``_ArgPlan.pack`` / ``unpack`` interpret the plan +on each call. The kinds -- and how each represents its field as op inputs: + + * ``TENSOR`` -- a plain ``Tensor`` / ``Optional[Tensor]``: one tensor slot. + * ``TENSOR_OR_QUANTIZED`` -- a field that may be a plain tensor, a bare + quantized storage, or ``None``: three slots (the tensor, its flat inner + buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. + * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes, + quantizers -- value-opaque constants baked into the graph -- and nested + collections of them), gathered into one shared ``OpaqueValueBundle`` slot. + * ``PROCESS_GROUP`` -- rides in the shared bundle too, as its c10d registry + name, re-resolved inside the op. + * ``UNSUPPORTED`` -- a field no kind can encode; emits no slot and is allowed + only when its value is trivial (``None`` / all-``None``) at call time. + +What runs where. Each op registers a data-free fake (``register_fake``) so it +traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward +call through it: + + * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see + ``tensor_spec.py``) and parses its result into an ``_OutputPlan`` -- the + outputs' geometry and their ranges in the flat payload, in pure Python; + * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat + ``Tensor[]`` payload; + * rebuilds the structured user outputs from that payload per the plan + (``_OutputPlan.user_outputs``; ``_flatten_value`` is the pack-side inverse). + +Autograd, registered on the op, drives backward: + + * ``setup_context`` (run when the forward is taped) re-runs ``fwd_fake_impl``, + parses the ``_OutputPlan``, reassembles the saved tensors from the op's flat + output, then calls the user ``setup_context`` to fill the backward args from + forward state + ``ctx_attrs`` (e.g. saved-tensor aliases) and return the + tensors to persist; the plan is stashed on ``ctx``; + * on ``backward()`` the incoming flat grads are sliced per user output from the + stashed plan (a ``grad_outputs`` field on the backward args receives the + whole tuple; otherwise ``grad_output`` receives the first output's grad), + the container's optional ``setup_saved_tensors`` hook restores the saved + tensors, then the *backward op* runs the real ``bwd_impl`` and returns the + flat grads (``bwd_fake_impl`` is its data-free fake). + +Two-tier op (``base`` + ``wrapper``), so a ``QuantizedTensor`` subclass can be an +op *input*. The ``_base`` op carries the real schema + autograd; a custom op +can't take a tensor-subclass input directly, so the ```` wrapper intercepts +those via ``register_torch_dispatch`` and flattens each into the base op's slots +(``_flatten_subclass_into_slots``) before forwarding. An empty subclass list makes +the wrapper a pass-through (plain / bf16 calls go straight through). +""" + +from __future__ import annotations +import dataclasses +import math +import types as _types # aliased: torch_dispatch rules take a ``types`` param +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + get_args, + get_origin, + get_type_hints, +) + +import torch + +from torch._prims_common import make_contiguous_strides_for + +from .tensor_spec import TensorSpec, to_tensor_spec +from ..quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, + _quantized_tensor_passthrough_ops, + prepare_for_saving, +) +from ..utils import record_compile_disabled + +_TE_OP_NAMESPACE = "transformer_engine_compile" + +# Annotation for an op arg field that may hold a plain tensor, a quantized +# tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer +# optimization). Matched exactly by ``_TensorOrQuantizedAdapter``. +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +# ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a +# 0-element sentinel tensor: a non-nullable ``Tensor[]`` schema is required for +# ``register_autograd`` to attach a ``grad_fn`` to the outputs. The sentinel is +# recognized by (numel == 0, dtype), so its dtype must be one no real payload +# tensor can have -- complex is never a TE payload (quantized data / scales are +# uint8 / fp32, outputs are float), while e.g. a uint8 sentinel would collide +# with a genuinely empty FP8 data buffer (empty batch). +_NONE_SENTINEL_DTYPE = torch.complex32 + + +def _encode_none(t: Optional[torch.Tensor]) -> torch.Tensor: + """Replace ``None`` with a 0-element sentinel tensor.""" + if t is None: + return torch.empty(0, dtype=_NONE_SENTINEL_DTYPE) + return t + + +def _decode_none(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Inverse of :func:`_encode_none`.""" + if t is None: + return None + if t.numel() == 0 and t.dtype == _NONE_SENTINEL_DTYPE: + return None + return t + + +# --------------------------------------------------------------------------- # +# OpaqueValueBundle: bundle of simple / value-opaque Python values +# --------------------------------------------------------------------------- # + + +class OpaqueValueBundle: + """Opaque value-type bundle of simple Python values. + + Wraps a ``{name: value}`` dict so many small non-Tensor args pass through a + single custom-op input; registered as a torch.compile *value* opaque type + (Dynamo specializes the graph on its contents). Allowed values: primitives + in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, classes, + any registered value-opaque type (e.g. TE quantizers), plus nested tuples / + lists / dicts thereof (so a bundle can carry a ``__tensor_flatten__`` + context verbatim -- including its ``cls`` entry). + """ + + PRIMITIVE_TYPES: Tuple[type, ...] = ( + type(None), + bool, + int, + float, + str, + torch.dtype, + torch.device, + torch.Size, + ) + + @classmethod + def is_simple_value(cls, value: Any) -> bool: + """Whether ``value`` may be stored inside an instance (recursive).""" + if isinstance(value, cls.PRIMITIVE_TYPES): + return True + if isinstance(value, Enum): + return True + if isinstance(value, type): + return True + if _is_opaque_value_type is not None and _is_opaque_value_type(type(value)): + return True + if isinstance(value, dict): + return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()) + 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: + # Tag with the concrete type so e.g. [1] / (1,) / Size([1]) or True / 1 + # stay distinct under __eq__ / __hash__ (graph guards compare bundles). + if isinstance(value, dict): + return ("dict", tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items()))) + if isinstance(value, (list, tuple)): # incl. torch.Size + return (type(value).__name__, tuple(cls._to_hashable(v) for v in value)) + return (type(value).__name__, value) + + @classmethod + def _fmt_simple(cls, value: Any) -> str: + """Repr for a value, evaluable in a context with ``torch`` globals.""" + if isinstance(value, torch.dtype): + return f"__import__('torch').{str(value).split('.')[-1]}" + if isinstance(value, torch.device): + return f"__import__('torch').device({str(value)!r})" + if isinstance(value, torch.Size): + return f"__import__('torch').Size({list(value)!r})" + # Enum before primitives: IntEnum is also ``int`` but must render as + # ``EnumName.MEMBER`` (the Enum class is added to globals by ``_collect``). + if isinstance(value, Enum): + return f"{type(value).__name__}.{value.name}" + # Class objects (e.g. the flatten context's ``cls``) render by name; the + # class itself is added to globals by ``_collect``. + if isinstance(value, type): + return value.__name__ + if isinstance(value, dict): + body = ", ".join(f"{k!r}: {cls._fmt_simple(v)}" for k, v in value.items()) + return f"{{{body}}}" + if isinstance(value, list): + return "[" + ", ".join(cls._fmt_simple(v) for v in value) + "]" + if isinstance(value, tuple): + body = ", ".join(cls._fmt_simple(v) for v in value) + return f"({body},)" if len(value) == 1 else f"({body})" + if _is_opaque_value_type(type(value)): + return value.__fx_repr__()[0] + # repr(float('inf')) is 'inf', which is not an evaluable literal. + if isinstance(value, float) and not math.isfinite(value): + return f"float({str(value)!r})" + return repr(value) + + def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: + 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()) + ) + # Precomputed: Dynamo guards hash bundles on every compiled call. + self._hash: int = hash(self._frozen) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + 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 self._hash + + def __fx_repr__(self) -> Tuple[str, Dict[str, Any]]: + items = ", ".join( + f"{k!r}: {OpaqueValueBundle._fmt_simple(v)}" for k, v in self._data.items() + ) + globals_: Dict[str, Any] = {"OpaqueValueBundle": OpaqueValueBundle} + + def _collect(value: Any) -> None: + if isinstance(value, dict): + for v in value.values(): + _collect(v) + return + if isinstance(value, (list, tuple)): + for v in value: + _collect(v) + return + if isinstance(value, Enum): + globals_[type(value).__name__] = type(value) + return + if isinstance(value, type): + globals_[value.__name__] = value + return + if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES): + return + if _is_opaque_value_type(type(value)): + _, extra = value.__fx_repr__() + globals_.update(extra) + + for v in self._data.values(): + _collect(v) + return (f"OpaqueValueBundle({{{items}}})", globals_) + + +try: + from torch._library.opaque_object import ( + get_opaque_type_name, + is_opaque_value_type as _is_opaque_value_type, + register_opaque_type, + ) + + register_opaque_type(OpaqueValueBundle, typ="value") + _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) +# Older torch without opaque_object support. +except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover + record_compile_disabled( + f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build" + ) + _is_opaque_value_type = None + _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None + +try: + from torch._C._distributed_c10d import ProcessGroup as _PROCESS_GROUP_TYPE + from torch._C._distributed_c10d import _resolve_process_group +except ImportError: # pragma: no cover + _PROCESS_GROUP_TYPE = None + _resolve_process_group = None + + +# --------------------------------------------------------------------------- # +# Storage flatten / unflatten (value-opaque quantizer; no ProcessGroup) +# --------------------------------------------------------------------------- # + + +def _storage_flatten( + value: Any, extra_meta: Optional[Dict[str, Any]] = None +) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: + """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``. + + The flatten context (embedding the value-opaque quantizer) plus inner names + and -- for a wrapper subclass -- the outer geometry are stashed in the bundle + so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. + ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` + hash key stays consistent) -- used to tag the tensor-or-quantized slot ``__kind__``. + """ + inner_names, ctx = value.__tensor_flatten__() + meta = dict(ctx) + meta["_inner_names"] = list(inner_names) + if isinstance(value, torch.Tensor): + meta["_outer_shape"] = torch.Size(value.shape) + if extra_meta: + meta.update(extra_meta) + tensors = [getattr(value, name) for name in inner_names] + return OpaqueValueBundle(meta), tensors + + +def _storage_unflatten(meta: "OpaqueValueBundle", tensors: List[torch.Tensor]) -> Any: + """Inverse of :func:`_storage_flatten`.""" + meta_dict = meta.as_dict() + inner_names = meta_dict["_inner_names"] + inner = dict(zip(inner_names, tensors)) + outer_shape = meta_dict.get("_outer_shape") + stride = make_contiguous_strides_for(tuple(outer_shape)) if outer_shape is not None else None + return QuantizedTensorStorage.__tensor_unflatten__(inner, meta_dict, outer_shape, stride) + + +# --------------------------------------------------------------------------- # +# Arg plans: dataclass annotations are parsed once, at registration, into an +# immutable per-op plan (schema string, slot order, gradient placement, +# tensor-or-quantized offsets); packing / unpacking interpret that plan on +# each call. +# --------------------------------------------------------------------------- # + + +def _is_union(annot: Any) -> bool: + """True for both ``typing.Union[...]`` / ``Optional[...]`` and PEP 604 ``X | Y``. + + ``get_origin`` returns ``typing.Union`` for the former but ``types.UnionType`` + for the latter, so the two syntaxes must be checked separately. + """ + origin = get_origin(annot) + return origin is Union or origin is _types.UnionType + + +def _strip_optional(annot: Any) -> Tuple[Any, bool]: + """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``.""" + if _is_union(annot): + args = get_args(annot) + if type(None) in args: + non_none = [a for a in args if a is not type(None)] + if len(non_none) == 1: + return non_none[0], True + return annot, False + + +class _FieldKind(Enum): + """How one dataclass field crosses the custom-op boundary.""" + + TENSOR = "tensor" # one ``Tensor`` / ``Tensor?`` slot + TENSOR_OR_QUANTIZED = "tensor_or_quantized" # 3 slots: tensor / inner / meta + PROCESS_GROUP = "process_group" # c10d group name inside the shared bundle + SIMPLE = "simple" # value carried verbatim inside the shared bundle + UNSUPPORTED = "unsupported" # no slots; only a trivial value may cross + + +class _TensorOrQuantizedKind(Enum): + """What a tensor-or-quantized slot group carries, tagged in its ``__meta``.""" + + NONE = "none" + TENSOR = "tensor" + STORAGE = "storage" + + +_TQ_KIND_KEY = "__kind__" +_SIMPLE_META_SLOT = "_simple_meta" + +# Matched by exact member set, so a bare quantized annotation or an accidental +# extra union member is rejected rather than silently taken as tensor-or-quantized. +_TQ_MEMBERS = frozenset(get_args(TensorOrQuantized)) + + +@dataclasses.dataclass(frozen=True) +class _SlotSpec: + """One schema slot: its name and torch.library type string.""" + + name: str + type_str: str + + +@dataclasses.dataclass(frozen=True) +class _FieldPlan: + """Parsed record for one dataclass field. + + ``slots`` are the schema slots the field occupies (empty for the kinds that + ride in the shared simple bundle, or cross nothing). + """ + + name: str + kind: _FieldKind + slots: Tuple[_SlotSpec, ...] + + +def _is_tensor_storage_union(annot: Any) -> bool: + """Whether ``annot`` is exactly the tensor-or-quantized union.""" + if not _is_union(annot): + return False + members = frozenset(a for a in get_args(annot) if a is not type(None)) + return members == _TQ_MEMBERS + + +def _is_process_group_annot(annot: Any) -> bool: + """Whether the field annotation is (Optional) ProcessGroup.""" + if _PROCESS_GROUP_TYPE is None: + return False + stripped, _ = _strip_optional(annot) + return stripped is _PROCESS_GROUP_TYPE + + +def _is_simple_annot(annot: Any) -> bool: + """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" + annot, _ = _strip_optional(annot) + if annot in OpaqueValueBundle.PRIMITIVE_TYPES: + return True + if isinstance(annot, type) and issubclass(annot, Enum): + return True + # Quantizers are value-opaque constants; the abstract ``Quantizer`` + # annotation itself is not a registered opaque type, so match by base. + if isinstance(annot, type) and issubclass(annot, Quantizer): + return True + if ( + isinstance(annot, type) + and _is_opaque_value_type is not None + and _is_opaque_value_type(annot) + ): + return True + if get_origin(annot) in (tuple, list): + inner = [a for a in get_args(annot) if a is not Ellipsis] + return bool(inner) and all(_is_simple_annot(a) for a in inner) + return False + + +def _parse_field(name: str, annot: Any) -> _FieldPlan: + """Parse one field's annotation into its :class:`_FieldPlan`.""" + if _is_tensor_storage_union(annot): + slots = ( + _SlotSpec(name, "Tensor?"), + _SlotSpec(name + "__tensors", "Tensor[]"), + _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + ) + return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots) + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Tensor: + slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor") + return _FieldPlan(name, _FieldKind.TENSOR, (slot,)) + # A union mixing tensor types with anything else is a malformed signature + # (e.g. a bare quantized-storage Optional, or Tensor | int): reject it at + # registration instead of silently degrading to an unsupported field. + if _is_union(annot): + members = [a for a in get_args(annot) if a is not type(None)] + if any( + isinstance(m, type) and issubclass(m, (torch.Tensor, QuantizedTensorStorage)) + for m in members + ): + raise TypeError( + f"field {name!r}: union {annot!r} is not a supported tensor " + "signature; use Tensor, Optional[Tensor], or TensorOrQuantized." + ) + if _is_process_group_annot(annot): + return _FieldPlan(name, _FieldKind.PROCESS_GROUP, ()) + if _is_simple_annot(annot): + return _FieldPlan(name, _FieldKind.SIMPLE, ()) + return _FieldPlan(name, _FieldKind.UNSUPPORTED, ()) + + +def _is_trivial(value: Any) -> bool: + """Whether an unsupported field's runtime value carries nothing.""" + if value is None: + return True + if isinstance(value, (list, tuple)): + return all(v is None for v in value) + return False + + +def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, Any]) -> None: + """Fill a tensor-or-quantized field's three slots from its runtime value.""" + tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots) + if value is None: + slots[tensor_slot] = None + slots[inner_slot] = [] + slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE}) + elif isinstance(value, torch.Tensor): + # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the + # ``Tensor?`` slot; subclass flattening (if any) is done by the + # wrapper op's ``register_torch_dispatch`` rule. + slots[tensor_slot] = value + slots[inner_slot] = [] + slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.TENSOR}) + elif isinstance(value, QuantizedTensorStorage): + meta, tensors = _storage_flatten(value, {_TQ_KIND_KEY: _TensorOrQuantizedKind.STORAGE}) + slots[tensor_slot] = None + slots[inner_slot] = tensors + slots[meta_slot] = meta + else: + raise TypeError( + f"field {field.name!r} expected None, torch.Tensor, or " + f"QuantizedTensorStorage, got {type(value).__name__}" + ) + + +def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any: + """Inverse of :func:`_pack_tensor_or_quantized`.""" + tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots) + meta = slots[meta_slot] + kind = meta[_TQ_KIND_KEY] + if kind == _TensorOrQuantizedKind.NONE: + return None + if kind == _TensorOrQuantizedKind.TENSOR: + return slots[tensor_slot] + return _storage_unflatten(meta, slots[inner_slot]) + + +@dataclasses.dataclass(frozen=True) +class _ArgPlan: + """The parsed plan for one args dataclass: the single source of truth for + the schema string, slot order, packing / unpacking, gradient placement and + the tensor-or-quantized slot offsets used for subclass flattening. + + Built once per registration by :func:`_parse_arg_type`; :meth:`pack` and + :meth:`unpack` interpret it on each call. ProcessGroup fields ride in the + shared bundle as their c10d registry *name* -- mirroring traceable + functional collectives -- and the op re-resolves the very group the caller + passed; groups created outside the c10d registry fail the resolve loudly. + """ + + arg_type: type + fields: Tuple[_FieldPlan, ...] + slot_names: Tuple[str, ...] + schema_str: str + simple_slot: Optional[str] + tensor_field_names: Tuple[str, ...] + tq_offsets: Tuple[int, ...] + + @property + def slot_count(self) -> int: + """Total number of input schema slots.""" + return len(self.slot_names) + + def resolve_grad_targets(self, input_tensors_for_grad: Sequence[str]) -> List[int]: + """Absolute schema-slot index receiving each requested field's gradient. + + Derived from ``fields`` on demand -- called once per registration, so + the plan doesn't cache the mapping. + """ + index: Dict[str, int] = {} + offset = 0 + for field in self.fields: + if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): + # The gradient flows to the group's first slot -- for + # tensor-or-quantized that is the ``Tensor?`` slot, the one + # autograd sees the (subclass) tensor in. + index[field.name] = offset + offset += len(field.slots) + non_differentiable = [n for n in input_tensors_for_grad if n not in index] + if non_differentiable: + raise ValueError( + f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" + ) + return [index[n] for n in input_tensors_for_grad] + + def pack(self, obj: Any) -> Dict[str, Any]: + """Flatten an ``arg_type`` instance into the op's ``{slot: value}`` dict. + + Inverse of :meth:`unpack`. + """ + slots: Dict[str, Any] = {} + simple: Dict[str, Any] = {} + for field in self.fields: + value = getattr(obj, field.name, None) + match field.kind: + case _FieldKind.TENSOR: + slots[field.slots[0].name] = value + case _FieldKind.TENSOR_OR_QUANTIZED: + _pack_tensor_or_quantized(field, value, slots) + case _FieldKind.PROCESS_GROUP: + simple[field.name] = None if value is None else value.group_name + case _FieldKind.SIMPLE: + simple[field.name] = value + case _FieldKind.UNSUPPORTED: + # Annotation alone (e.g. Optional[Any]) can't decide; only + # the runtime value can, so the check runs at pack time. + if not _is_trivial(value): + raise TypeError( + f"{self.arg_type.__name__} field {field.name!r} has a type not " + "supported by torch.compile (not Tensor, simple, Quantizer, or " + "ProcessGroup) and carries a " + "non-trivial value; add a matching field kind in custom_op.py " + "to handle it." + ) + if self.simple_slot is not None: + slots[self.simple_slot] = OpaqueValueBundle(simple) + return slots + + def unpack(self, slots: Dict[str, Any]) -> Any: + """Rebuild a fresh ``arg_type`` instance from the op's flat slot dict. + + Inverse of :meth:`pack`. + """ + kwargs: Dict[str, Any] = {} + bundle = slots.get(self.simple_slot) if self.simple_slot is not None else None + for field in self.fields: + match field.kind: + case _FieldKind.TENSOR: + kwargs[field.name] = slots[field.slots[0].name] + case _FieldKind.TENSOR_OR_QUANTIZED: + kwargs[field.name] = _unpack_tensor_or_quantized(field, slots) + case _FieldKind.PROCESS_GROUP: + if bundle is not None: + name = bundle[field.name] + kwargs[field.name] = None if name is None else _resolve_process_group(name) + case _FieldKind.SIMPLE: + if bundle is not None: + kwargs[field.name] = bundle[field.name] + case _FieldKind.UNSUPPORTED: + kwargs[field.name] = None + obj = self.arg_type.__new__(self.arg_type) + for k, v in kwargs.items(): + object.__setattr__(obj, k, v) + return obj + + +def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: + """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 _parse_arg_type(cls: type) -> _ArgPlan: + """Parse an args ``@dataclass`` into its immutable :class:`_ArgPlan`. + + The layout pass assigns absolute slot positions (the shared simple bundle, + if any field needs it, takes the last slot) and validates the result: + duplicate slot names are rejected here, before any op is registered. + """ + if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None: + raise RuntimeError( + f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle " + "is not registered as a torch._library value-opaque type (PyTorch build " + "without opaque-object support)." + ) + fields = tuple(_parse_field(name, annot) for name, annot in _resolved_field_annotations(cls)) + + slot_specs: List[_SlotSpec] = [] + tq_offsets: List[int] = [] + tensor_field_names: List[str] = [] + for field in fields: + if field.kind is _FieldKind.TENSOR_OR_QUANTIZED: + tq_offsets.append(len(slot_specs)) + if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): + tensor_field_names.append(field.name) + slot_specs.extend(field.slots) + if any(f.kind in (_FieldKind.SIMPLE, _FieldKind.PROCESS_GROUP) for f in fields): + simple_slot: Optional[str] = _SIMPLE_META_SLOT + slot_specs.append(_SlotSpec(_SIMPLE_META_SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)) + else: + simple_slot = None + + slot_names = tuple(s.name for s in slot_specs) + if len(set(slot_names)) != len(slot_names): + dupes = sorted(n for n in set(slot_names) if slot_names.count(n) > 1) + raise ValueError(f"{cls.__name__}: duplicate schema slot names: {dupes}") + schema_str = "(" + ", ".join(f"{s.type_str} {s.name}" for s in slot_specs) + ")" + + return _ArgPlan( + arg_type=cls, + fields=fields, + slot_names=slot_names, + schema_str=schema_str, + simple_slot=simple_slot, + tensor_field_names=tuple(tensor_field_names), + tq_offsets=tuple(tq_offsets), + ) + + +def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any: + """Copy of dataclass ``obj`` with each tensor field replaced by a :class:`TensorSpec`. + + Only tensor fields have a ``TensorSpec`` equivalent, so quantizer / scalar + fields are simply carried over unchanged; the fake impl works purely on + geometry. Built with :func:`dataclasses.replace` (the only such construction + Dynamo can trace). + """ + overrides: Dict[str, Any] = {} + for name in tensor_field_names: + value = getattr(obj, name, None) + if value is not None and not isinstance(value, TensorSpec): + overrides[name] = to_tensor_spec(value) + if not overrides: + return obj + return dataclasses.replace(obj, **overrides) + + +# --------------------------------------------------------------------------- # +# Op outputs <-> flat ``Tensor[]`` payload: this is how an op returns / saves +# quantized tensors (and wrapper subclasses). Outputs are flattened to their +# inner buffers on the way out and rebuilt via ``__tensor_unflatten__`` on the +# way back; on the fake side a TensorSpec supplies the geometry. +# --------------------------------------------------------------------------- # + + +def _spec_slot_count(spec: Optional[TensorSpec]) -> int: + """Flat ``Tensor[]`` slots the value for ``spec`` occupies.""" + if spec is None: + return 1 + return len(spec.inner_names()) + + +def _flatten_value( + value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorSpec]], +) -> List[torch.Tensor]: + """Return the flat ``Tensor[]`` slots that represent one op output ``value``. + + Pack-side inverse of :meth:`_OutputPlan.user_outputs`; the slot count + matches :func:`_spec_slot_count`. + """ + if value is None: + return [_encode_none(None)] + if isinstance(value, TensorSpec): + return [_encode_none(t) for t in value.create_inner_tensors()] + if hasattr(value, "__tensor_flatten__"): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + if isinstance(value, torch.Tensor): + return [_encode_none(value)] + raise TypeError( + f"unsupported value type {type(value).__name__}; expected None / " + "torch.Tensor / tensor subclass / bare storage / TensorSpec." + ) + + +# Trailing slots in every fwd-impl return: ``tensors_to_save, ctx_attrs``. +# User-output count is ``len(result) - this``. +_FWD_TRAILING_SLOTS = 2 + + +def _check_fwd_result(result: Any) -> None: + """Validate a fwd-impl return against the + ``(*user_outputs, tensors_to_save, ctx_attrs)`` contract, with a clear + message for op authors (user-output *types* are checked later, by + :func:`_flatten_value`). + + Only called on the fake path (:meth:`_OutputPlan.parse`), which runs at + trace/compile time -- so this is a compile-time check with no per-call cost. + The real impl must return the same shape as the fake, so validating the fake + covers both. + """ + if not isinstance(result, tuple) or len(result) < _FWD_TRAILING_SLOTS: + raise TypeError( + f"fwd impl must return a tuple of >= {_FWD_TRAILING_SLOTS} elements " + "(*user_outputs, tensors_to_save, ctx_attrs); " + f"got {type(result).__name__}" + ) + tensors_to_save, ctx_attrs = result[-2], result[-1] + if tensors_to_save is not None and not isinstance(tensors_to_save, (list, tuple)): + raise TypeError("fwd impl 'tensors_to_save' slot must be a list/tuple or None") + if ctx_attrs is not None and not isinstance(ctx_attrs, dict): + raise TypeError("fwd impl 'ctx_attrs' slot must be a dict or None") + + +def _pack_fwd_result(result: Any) -> List[torch.Tensor]: + """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. + + User outputs first, then saved-for-backward tensors in declaration order. + """ + num_outputs = len(result) - _FWD_TRAILING_SLOTS + flat: List[torch.Tensor] = [] + for value in result[:num_outputs]: + flat.extend(_flatten_value(value)) + saved = result[num_outputs] + if saved is not None: + for value in saved: + flat.extend(_flatten_value(value)) + return flat + + +def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: + """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. + + Each grad occupies exactly one slot (validated against ``num_grad_inputs``); + a :class:`TensorSpec` grad is materialized into a single tensor. + """ + grads = list(grads) + if len(grads) != num_grad_inputs: + raise RuntimeError( + f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads " + f"(one per input_tensors_for_grad entry), got {len(grads)}" + ) + out: List[torch.Tensor] = [] + for g in grads: + if isinstance(g, TensorSpec): + out.append(_encode_none(g.create_tensor())) + else: + out.append(_encode_none(g)) + return out + + +@dataclasses.dataclass(frozen=True) +class _OutputPlan: + """Per-trace layout of an op's flat ``Tensor[]`` return. + + Parsed from a fwd fake-impl result: the logical user outputs and the + saved-for-backward tensors, each with its range in the flat payload. The + single source of truth for rebuilding forward outputs and saved tensors and + for slicing backward grad_outputs per user output. Per-trace rather than + per-registration because a quantized output's inner-tensor count is only + known from the fake result. + """ + + user_specs: Tuple[Optional[TensorSpec], ...] + saved_specs: Tuple[Optional[TensorSpec], ...] + ctx_attrs: Dict[str, Any] + user_ranges: Tuple[Tuple[int, int], ...] + saved_start: int + + @classmethod + def parse(cls, result: Tuple[Any, ...]) -> "_OutputPlan": + """Slice a fwd fake-impl return into the plan (validating the contract).""" + _check_fwd_result(result) + num_outputs = len(result) - _FWD_TRAILING_SLOTS + user_specs = tuple(result[:num_outputs]) + saved = result[num_outputs] + ctx_attrs = result[num_outputs + 1] + cursor = 0 + user_ranges: List[Tuple[int, int]] = [] + for spec in user_specs: + n = _spec_slot_count(spec) + user_ranges.append((cursor, cursor + n)) + cursor += n + return cls( + user_specs=user_specs, + saved_specs=tuple(saved) if saved is not None else (), + ctx_attrs=dict(ctx_attrs) if ctx_attrs else {}, + user_ranges=tuple(user_ranges), + saved_start=cursor, + ) + + @staticmethod + def _assemble( + spec: Optional[TensorSpec], flat: Sequence[Optional[torch.Tensor]], start: int, stop: int + ) -> Any: + chunk = [_decode_none(t) for t in flat[start:stop]] + # ``spec is None`` is the op-boundary sentinel for an absent output. + return spec.assemble(chunk) if spec is not None else None + + def user_outputs(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Rebuild the structured user outputs from the op's flat return.""" + return [ + self._assemble(spec, flat, start, stop) + for spec, (start, stop) in zip(self.user_specs, self.user_ranges) + ] + + def saved_tensors(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Rebuild the saved-for-backward tensors from the op's flat return.""" + values: List[Any] = [] + cursor = self.saved_start + for spec in self.saved_specs: + n = _spec_slot_count(spec) + values.append(self._assemble(spec, flat, cursor, cursor + n)) + cursor += n + return values + + def user_grads(self, flat_grads: Sequence[Optional[torch.Tensor]]) -> List[Any]: + """Gradient of each user output, sliced from the op's flat grad list. + + A single-slot output yields its tensor grad; a flattened quantized + output yields the tuple of its inner-buffer grads. + """ + grads: List[Any] = [] + for start, stop in self.user_ranges: + chunk = [_decode_none(g) for g in flat_grads[start:stop]] + grads.append(chunk[0] if stop - start == 1 else tuple(chunk)) + return grads + + +# --------------------------------------------------------------------------- # +# Op registration +# --------------------------------------------------------------------------- # + + +def _register_base_op( + *, + op_name: str, + schema_str: str, + plan: _ArgPlan, + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + pack_result: Callable[[Any], List[torch.Tensor]], +) -> Any: + """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the + ``fake_impl`` (spec), returning the ``CustomOpDef``. + + The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel + runs the spec fake impl on the :func:`_spec_view`. Both go through + ``pack_result``. + """ + + def _impl(*flat: Any) -> List[torch.Tensor]: + obj = plan.unpack(dict(zip(plan.slot_names, flat))) + return pack_result(impl(obj)) + + def _fake(*flat: Any) -> List[torch.Tensor]: + obj = plan.unpack(dict(zip(plan.slot_names, flat))) + spec_obj = _spec_view(obj, plan.tensor_field_names) + return pack_result(fake_impl(spec_obj)) + + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str + ) + op.register_fake(_fake) + return op + + +def _register_autograd_for_op( + *, + fwd_op: Any, + bwd_op: Any, + fwd_plan: _ArgPlan, + bwd_plan: _ArgPlan, + grad_targets: List[int], + setup_context_user: Callable[..., Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> None: + """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. + + ``setup_context`` re-runs the spec fwd fake impl to parse the + :class:`_OutputPlan`, reassembles the outputs / saved tensors from it, hands + the saved tuple + ``ctx_attrs`` to the module's ``setup_context`` and stashes + the plan on ``ctx`` so backward can slice its grads per user output. + """ + bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields) + + def _setup_context(ctx, inputs, output): + ctx.fwd_tensor_list_lengths = { + i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + } + fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs))) + spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names) + + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + user_outputs = out_plan.user_outputs(output) + saved_list = out_plan.saved_tensors(output) + + bwd_obj = bwd_plan.arg_type() + tensors_to_save_from_setup = setup_context_user( + bwd_obj, + fwd_obj, + user_outputs[0] if len(user_outputs) == 1 else tuple(user_outputs), + out_plan.ctx_attrs, + tuple(saved_list), + ) + tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ())) + ctx.tensor_objects = tensor_objects + ctx.save_for_backward(*tensors_to_save) + ctx.backward_objects = bwd_obj + ctx.output_plan = out_plan + + def _autograd_backward(ctx, *grad_outputs): + bwd_obj = ctx.backward_objects + if hasattr(bwd_obj, "setup_saved_tensors"): + bwd_obj.setup_saved_tensors(ctx) + ctx.tensor_objects = None + user_grads = ctx.output_plan.user_grads(grad_outputs[0]) + ctx.output_plan = None + if bwd_takes_grad_tuple: + bwd_obj.grad_outputs = tuple(user_grads) + else: + bwd_obj.grad_output = user_grads[0] + kwargs = bwd_plan.pack(bwd_obj) + bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names] + grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + ctx.backward_objects = None + # One grad per input schema slot: default None, but a ``Tensor[]`` slot + # (always recorded in ``fwd_tensor_list_lengths``) needs a + # list-shaped no-grad of matching length. + out: List[Any] = [None] * fwd_plan.slot_count + for pos, length in ctx.fwd_tensor_list_lengths.items(): + out[pos] = [None] * length + for pos, g in zip(grad_targets, grads): + out[pos] = g + return tuple(out) + + fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) + + +def _flatten_subclass_into_slots( + new_args: List[Any], slot_offsets: Sequence[int], subclass: type +) -> None: + """Rewrite each tensor-or-quantized slot group whose ``Tensor?`` slot holds an + instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). + """ + for offset in slot_offsets: + val = new_args[offset] + if not isinstance(val, subclass): + continue + meta, tensors = _storage_flatten(val, {_TQ_KIND_KEY: _TensorOrQuantizedKind.STORAGE}) + new_args[offset] = None + new_args[offset + 1] = tensors + new_args[offset + 2] = meta + + +def _make_slot_forwarder( + base_op: Any, slot_offsets: Sequence[int], subclasses: Sequence[type] +) -> Callable[[Sequence[Any]], List[torch.Tensor]]: + """Return ``call(args)`` forwarding to ``base_op``, first flattening any + ``subclasses`` instance sitting in the tensor-or-quantized slot groups at + ``slot_offsets``. + + A ``torch.library`` op cannot take a tensor subclass directly, so the wrapper + op body and its ``register_torch_dispatch`` rules all funnel through this one + path -- see the two-tier op note in the module docstring. With no slots or no + subclasses to flatten it is a plain pass-through. + """ + enabled = bool(slot_offsets) and bool(subclasses) + + def call(args: Sequence[Any]) -> List[torch.Tensor]: + if not enabled: + return base_op(*args) + new_args = list(args) + for sub in subclasses: + _flatten_subclass_into_slots(new_args, slot_offsets, sub) + return base_op(*new_args) + + return call + + +def _make_dispatch_rule( + forward: Callable[[Sequence[Any]], List[torch.Tensor]], +) -> Callable[..., Any]: + """Adapt a slot forwarder to the ``register_torch_dispatch`` signature.""" + + def _rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + return forward(args) + + return _rule + + +def _register_wrapper_op( + *, + wrapper_op_name: str, + schema_str: str, + base_op: Any, + slot_offsets: Sequence[int] = (), + subclasses: Sequence[type] = (), +) -> Any: + """Define the wrapper op via ``torch.library.custom_op``: forward to the base + op through :func:`_make_slot_forwarder`. Returns the ``CustomOpDef``. + """ + forward = _make_slot_forwarder(base_op, slot_offsets, subclasses) + + def _forward(*flat: Any) -> List[torch.Tensor]: + return forward(flat) + + op_def = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + ) + op_def.register_fake(_forward) + return op_def + + +def _all_quantized_tensor_subclasses() -> List[type]: + """Return every imported ``QuantizedTensor`` wrapper subclass.""" + import transformer_engine.pytorch.tensor # noqa: F401 pylint: disable=import-outside-toplevel,unused-import + + found: List[type] = [] + stack = list(QuantizedTensor.__subclasses__()) + while stack: + cls = stack.pop() + if cls not in found: + found.append(cls) + stack.extend(cls.__subclasses__()) + return found + + +def register_custom_op( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Optional[Callable[..., Any]]: + """Register a TE module's forward + backward as torch custom ops. + + Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for + ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches + through the op and returns the user-facing outputs. + + ``fwd_arg_type`` and ``bwd_arg_type`` are ``@dataclass``es whose *field + annotations* define the op schema (see the module docstring for the + field <-> slot mapping). The caller builds a ``fwd_arg_type`` instance and + passes it to ``forward_fn``. ``input_tensors_for_grad`` lists the + ``fwd_arg_type`` fields that receive gradients and fixes the backward grad + order. ``bwd_arg_type`` is also instantiated by the framework + (``bwd_arg_type()``), so it must be constructible with no arguments. + + Callable contracts: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` -- the + real forward. ``user_outputs``: op outputs (tensor / quantized / ``None``); + ``tensors_to_save``: list/tuple (or ``None``) of tensors for backward; + ``ctx_attrs``: dict (or ``None``) of plain metadata for ``setup_context``. + The trailing two slots are fixed (``_FWD_TRAILING_SLOTS``); everything + before them is a user output. + * ``fwd_fake_impl(fwd_args)`` -- data-free traceable twin of ``fwd_impl``: + same return shape, but tensor outputs are :class:`TensorSpec`. Must match + ``fwd_impl``'s shape (checked at compile time by ``_check_fwd_result``). + * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved) + -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the + tensors to persist across the boundary. + * ``bwd_impl(bwd_args) -> grads`` -- exactly one grad per + ``input_tensors_for_grad`` entry, in that order (``None`` for a + non-differentiable input). + * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning + :class:`TensorSpec` grads. + * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if + absent. + + How the backward container is populated: ``setup_context`` fills the + ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from + forward state and returns the tensors to persist; the framework saves them + via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework + restores them into the container's tensor fields through the + ``setup_saved_tensors`` hook and sets the incoming gradient directly -- + into a ``grad_outputs`` field (tuple, one grad per user output) if + ``bwd_arg_type`` declares one, else into ``grad_output`` (the first user + output's grad) -- so ``bwd_impl`` receives a fully-populated + ``bwd_arg_type``. + + Registration touches experimental ``torch.library`` / opaque-object APIs + that may be missing on older PyTorch. If it fails, this warns once and + returns ``None`` instead of raising, so callers can fall back to eager under + ``torch.compile`` (a graph break) rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + input_tensors_for_grad=input_tensors_for_grad, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + setup_context=setup_context, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_fake_impl=bwd_fake_impl, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + record_compile_disabled( + f"could not register the custom op '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Callable[..., Any]: + """Body of :func:`register_custom_op`; see it for semantics.""" + # Existence check at the API boundary: every ``input_tensors_for_grad`` name + # must be an actual field of ``fwd_arg_type`` (differentiability -- whether + # that field can carry a gradient -- is checked later, in + # :meth:`_ArgPlan.resolve_grad_targets`). + fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} + missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] + if missing: + raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + subclass_list = _all_quantized_tensor_subclasses() + + fwd_plan = _parse_arg_type(fwd_arg_type) + bwd_plan = _parse_arg_type(bwd_arg_type) + + num_grad_inputs = len(input_tensors_for_grad) + grad_targets = fwd_plan.resolve_grad_targets(input_tensors_for_grad) + + fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]" + bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]" + + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_base_op( + op_name=base_fwd_name, + schema_str=fwd_schema, + plan=fwd_plan, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + pack_result=_pack_fwd_result, + ) + _register_base_op( + op_name=base_bwd_name, + schema_str=bwd_schema, + plan=bwd_plan, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + fwd_slot_offsets = fwd_plan.tq_offsets + bwd_slot_offsets = bwd_plan.tq_offsets + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + slot_offsets=fwd_slot_offsets, + subclasses=subclass_list, + ) + # Pass-through: a subclass input reaches the base op through the dispatch + # rule below, never through the wrapper body. + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + ) + + autograd_common = { + "fwd_plan": fwd_plan, + "bwd_plan": bwd_plan, + "grad_targets": grad_targets, + "setup_context_user": setup_context, + "fwd_fake_impl": fwd_fake_impl, + } + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) + _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) + + _fwd_rule = _make_dispatch_rule( + _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list) + ) + _bwd_rule = _make_dispatch_rule( + _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list) + ) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) + _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) + _quantized_tensor_passthrough_ops.add(base_fwd_op.default) + _quantized_tensor_passthrough_ops.add(base_bwd_op.default) + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + kwargs = fwd_plan.pack(fwd_args) + flat_in = [kwargs[name] for name in fwd_plan.slot_names] + result = wrapper_fwd_op(*flat_in) + + outputs = out_plan.user_outputs(result) + if len(outputs) == 1: + return outputs[0] + return tuple(outputs) + + return forward_fn diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index a342dd1e6c..dc689258d1 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -9,6 +9,7 @@ from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType +from ..utils import record_compile_disabled # Qualnames of the registered quantizer classes. The set holds strings rather @@ -117,18 +118,22 @@ def register_value_opaque_quantizer(cls: type) -> None: register_opaque_type, is_opaque_value_type, ) - except (ImportError, AttributeError): + except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. + record_compile_disabled( + f"this PyTorch build has no opaque-object API ({e}); use a newer build" + ) return try: if not is_opaque_value_type(cls): register_opaque_type(cls, typ="value") - except (RuntimeError, TypeError): + except (RuntimeError, TypeError) as e: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. + record_compile_disabled(f"could not register {cls.__name__} as an opaque type ({e})") return _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6..2330a8c314 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,7 +3,8 @@ # See LICENSE for license information. """Linear API""" -from dataclasses import dataclass + +from dataclasses import dataclass, replace as dataclass_replace from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -48,6 +49,9 @@ nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, + check_gemm_dims, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -69,10 +73,11 @@ ) from ..cpp_extensions import ( general_gemm, + get_cublas_workspace, ) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type -from ..jit import no_torch_dynamo from ..graph import is_graph_capturing +from ..jit import no_torch_dynamo from ..quantized_tensor import ( QuantizedTensor, QuantizedTensorStorage, @@ -80,6 +85,12 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op, + is_value_opaque_quantizer, +) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -95,9 +106,6 @@ __all__ = ["Linear"] -TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] - - @dataclass(slots=True) class LinearFwdArgs: """Single-argument bag for the forward path of :class:`_Linear`.""" @@ -108,7 +116,8 @@ class LinearFwdArgs: bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- - weight_workspace: Optional[torch.Tensor] + # TensorOrQuantized so a cached quantized workspace can cross the op boundary. + weight_workspace: Optional[TensorOrQuantized] # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool @@ -142,7 +151,8 @@ class LinearFwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] - tp_group: Optional[Any] + # Crosses the op boundary as its c10d registry name, re-resolved in the op. + tp_group: Optional[dist_group_type] tp_size: int tensor_parallel: bool sequence_parallel: bool @@ -170,6 +180,43 @@ 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 is_distributed_weight(self.weight): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)): + return "a quantized input tensor" + if self.fsdp_group is not None: + return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" + if ( + self.fp8_output + and self.is_grad_enabled + and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad) + ): + return "differentiable fp8_output=True" + if self.cpu_offloading: + return "CPU activation offloading" + if self.wgrad_store is not None: + # Non-None only when delayed wgrad compute is on (see Linear.forward). + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + for quantizer in ( + self.input_quantizer, + self.weight_quantizer, + self.output_quantizer, + self.grad_input_quantizer, + self.grad_weight_quantizer, + self.grad_output_quantizer, + ): + # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe + # quantizers are not value-opaque and can't cross the custom-op boundary. + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + @dataclass(slots=True) class LinearBwdArgs: @@ -207,7 +254,8 @@ class LinearBwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] = None - tp_group: Optional[Any] = None + # See LinearFwdArgs.tp_group. + tp_group: Optional[dist_group_type] = None tp_size: int = 1 tensor_parallel: bool = False sequence_parallel: bool = False @@ -240,7 +288,7 @@ class LinearBwdArgs: cpu_offloading: bool = False owns_input: bool = False - # --- Per-backward scratch state (populated inside _linear_backward) --- + # --- Per-backward scratch state (populated inside _linear_backward_impl) --- ub_obj_gradout: Optional[Any] = None def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: @@ -265,15 +313,49 @@ def _check_fp8_reduce_and_update(): return result +def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Output's leading (sequence) dim from the input's: sequence parallelism + gathers it (column-parallel) or scatters it (row-parallel).""" + if not args.sequence_parallel: + return leading + if args.parallel_mode == "column": + return leading * args.tp_size + if args.parallel_mode == "row": + return leading // args.tp_size + return leading + + +def _inp_leading_from_out(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: + """Inverse of :func:`_out_leading_from_inp`.""" + if not args.sequence_parallel: + return leading + if args.parallel_mode == "column": + return leading // args.tp_size + if args.parallel_mode == "row": + return leading * args.tp_size + return leading + + +def _fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer]) -> bool: + """Spec-level mirror of ``_is_weight_workspace_valid``: the cached workspace + must already hold every inner buffer the quantizer's current usage needs.""" + if quantizer is None: + return True + required = TensorSpec( + shape=workspace.shape, dtype=workspace.dtype, quantizer=quantizer, device=workspace.device + ).inner_names() + return set(required) <= set(workspace.inner_names()) + + def _linear_forward_impl( args: LinearFwdArgs, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], None, Optional[Dict]]: +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: """Forward implementation for the linear layer. - Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, None, + Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight workspace (returned alongside ``out`` so the caller can refresh its - cache). The last three are ``None`` when gradients are disabled. + cache). The last two are ``None`` when gradients are disabled. """ weight = args.weight @@ -300,7 +382,7 @@ def _linear_forward_impl( debug = args.debug backward_override = args.backward_override is_fsdp2 = args.is_fsdp2 - backward_needs_input = is_grad_enabled and weight.requires_grad + backward_needs_input = is_grad_enabled and args.weight_requires_grad if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": @@ -464,7 +546,7 @@ def _linear_forward_impl( # No need to set the quantizer states if weight is already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug): - columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 + columnwise_usage = is_grad_enabled and args.input_requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -647,12 +729,24 @@ def _linear_forward_impl( if is_dist_weight: wt_save = None - # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` - # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. - # Needed for torch.compile to work correctly. + # Dedup save slots that alias forward inputs or other op returns; + # ``_linear_setup_ctx`` rebuilds the refs. A custom op may not return a + # tensor aliasing an input or another return, and the cached FP8 weight + # is the same object as ``new_weight_workspace`` (cache miss) or + # ``weight_workspace`` (cache hit). + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_weight_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - "weight" if wt_save is weight else None, + wt_alias, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -668,13 +762,227 @@ def _linear_forward_impl( "saved_tensor_aliases": saved_tensor_aliases, } - return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _linear_forward_fake( + args: LinearFwdArgs, +) -> Tuple[TensorSpec, Optional[TensorSpec], Optional[Tuple[Any, ...]], Optional[Dict]]: + """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile, + returning ``TensorSpec`` descriptors for the outputs and saved tensors instead + of allocating real data.""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time Linear forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.weight + inp = args.inp + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + fp8 = args.fp8 + debug = args.debug + fp8_or_debug = fp8 or debug + is_grad_enabled = args.is_grad_enabled + activation_dtype = args.activation_dtype + save_original_input = args.save_original_input + if args.backward_override == "high_precision": + save_original_input = True + elif args.backward_override == "dequantized": + save_original_input = False + + out_features, _ = weight.shape + backward_needs_input = is_grad_enabled and args.weight_requires_grad + if ( + args.backward_override is None + and save_original_input + and backward_needs_input + and input_quantizer is not None + ): + if not isinstance( + input_quantizer, Float8Quantizer + ) and not can_reconstruct_wgrad_input_from_original(input_quantizer): + save_original_input = False + + own_quantized_input = False + inputmat_is_storage = False + inputmat_aliases_inp = False + if fp8_or_debug: + if inp.is_quantized: + # Primary-quantized input reused as-is. + inputmat_is_storage = True + inputmat_aliases_inp = True + else: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + input_quantizer.set_usage( + rowwise=True, + columnwise=( + backward_needs_input + and not save_original_input + and args.backward_override is None + ), + ) + own_quantized_input = True + inputmat_is_storage = True + else: + inputmat_aliases_inp = inp.dtype == activation_dtype + + if save_original_input: + inputmat_aliases_inp = True + inputmat_is_storage = False + + # ------------------------------------------------------ + # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ------------------------------------------------------ + new_weight_workspace = None + workspace = None # args.weight_workspace after validation + weightmat = None + weightmat_is_storage = False + weightmat_aliases_weight = False + if fp8_or_debug: + if weight_quantizer is not None and (not weight.is_quantized or debug): + columnwise_usage = is_grad_enabled and args.input_requires_grad and not args.is_fsdp2 + if args.backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() + and not in_fp8_activation_recompute_phase() + ) + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif weight.is_quantized: + weight_quantizer = weight.quantizer + + if weight.is_quantized: + # Primary-quantized weight: the impl reuses it as ``weightmat``. + weightmat = weight + weightmat_is_storage = True + weightmat_aliases_weight = True + else: + weightmat_is_storage = True + workspace = args.weight_workspace + if workspace is not None and not _fake_workspace_valid(workspace, weight_quantizer): + # quantize_weight drops a stale workspace and builds a new one. + workspace = None + if workspace is not None: + # Copy, so the ``update_usage`` below stays off the input spec. + weightmat = dataclass_replace(workspace) + else: + weightmat = TensorSpec( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=weight_quantizer, + device=weight.device, + ) + if args.cache_weight: + # Persistent cache entries are wrappers, not bare storages. + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_weight_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) + else: + weightmat_aliases_weight = weight.dtype == activation_dtype + weightmat = TensorSpec( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). + # ------------------------------------------------------ + out_leading = _out_leading_from_inp(inp.shape[0], args) + out = TensorSpec( + shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), + dtype=activation_dtype, + quantizer=output_quantizer, + requires_grad=is_grad_enabled + and (args.input_requires_grad or args.weight_requires_grad or args.bias_requires_grad), + device=inp.device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout + # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing. + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if is_grad_enabled: + # Slot 0 -- ``saved_inputmat``. + inputmat_alias = None + saved_inputmat = None + if backward_needs_input: + if inputmat_aliases_inp: + inputmat_alias = "inp" + elif inputmat_is_storage: + saved_inputmat = TensorSpec( + shape=tuple(inp.shape), + dtype=activation_dtype, + quantizer=input_quantizer, + device=inp.device, + ) + # Mirror the impl's post-quantization ``update_usage``. + if own_quantized_input and not save_original_input: + if args.backward_override is not None: + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + elif ( + args.backward_input_needs_gather + and weight_quantizer is not None + and weight_quantizer.supports_only_rowwise_all_gather() + ): + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + saved_inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + else: + saved_inputmat = TensorSpec( + shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device + ) + + # Slot 1 -- ``wt_save``, with the impl's alias dedup (rebuilt in + # ``_linear_setup_ctx`` instead of being saved twice). + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered weight in backward. + elif weightmat_is_storage and new_weight_workspace is not None: + wt_alias = "new_weight_workspace" + elif weightmat_is_storage and workspace is not None: + wt_alias = "weight_workspace" + elif weightmat_is_storage: + wt_save = weightmat + else: + wt_save = TensorSpec( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + # Slot 2 -- ``saved_weight`` (always aliased to ``weight``). + # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent). + saved_tensor_aliases = ( + inputmat_alias, + wt_alias, + "weight", + "bias" if bias is not None else None, + ) + tensors_to_save_from_forward = (saved_inputmat, wt_save, None, None) + ctx_attrs = { + "fsdp_shapes": [], + "saved_tensor_aliases": saved_tensor_aliases, + } + + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs def _linear_setup_ctx( bwd_args: LinearBwdArgs, fwd_args: LinearFwdArgs, - out: torch.Tensor, + fwd_outputs: Tuple[Any, ...], ctx_attrs: Dict, tensors_to_save_from_forward: Tuple[Any, ...], ) -> Tuple[Any, ...]: @@ -687,7 +995,8 @@ def _linear_setup_ctx( for FSDP2 re-quantization) without having to mutate the structured metadata returned by ``prepare_for_saving``. """ - del out # No-op; kept for symmetry with the compile-time helper signature. + # ``fwd_outputs`` is ``(out, new_weight_workspace)``; only the latter is used, + # to rebuild the deduped weight save slot. inp = fwd_args.inp weight = fwd_args.weight @@ -710,7 +1019,10 @@ def _linear_setup_ctx( bwd_args.use_bias = bias is not None bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad - bwd_args.inp_shape = inp.shape + # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) + # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle. + # The backward reconstructs inp_shape from grad_output + weight + SP config. + bwd_args.inp_shape = None # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype @@ -779,6 +1091,10 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight + elif wt_save_alias == "new_weight_workspace": + wt_save = fwd_outputs[1] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace if saved_weight_alias == "weight": saved_weight = weight if bias_alias == "bias": @@ -786,7 +1102,7 @@ def _linear_setup_ctx( return (saved_inputmat, wt_save, saved_weight, saved_bias) -def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: +def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: """Backward implementation for the linear layer. Caller must have populated ``args.grad_output`` and run @@ -858,6 +1174,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") + # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). + if bwd_args.inp_shape is None: + in_features = saved_weight.shape[-1] + inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args) + bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) + # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1387,6 +1709,80 @@ def wgrad_gemm( ) +def _linear_backward_fake( + args: LinearBwdArgs, +) -> Tuple[Optional[TensorSpec], Optional[TensorSpec], Optional[TensorSpec]]: + """Allocation-free fake of :func:`_linear_backward_impl` on ``TensorSpec``. + + Returns ``(wgrad, dgrad, grad_bias)`` specs. TP/SP gather/scatter happens + inside the eager op, so the specs carry rank-local shapes. + """ + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake Linear backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.saved_weight + out_dtype = args.activation_dtype + out_features, in_features = weight.shape + + # Mirrors the impl; affects dgrad's buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + # Input shape rederived from grad_output + SP config (inp_shape is not + # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). + dgrad_leading = _inp_leading_from_out(args.grad_output.shape[0], args) + # Under UB reduce-scatter overlap the returned dgrad is the plain + # reduce-scatter output; the quantizer only feeds the comm buffer. + dgrad_quantizer = None if args.ub_overlap_rs_dgrad else args.grad_input_quantizer + dgrad = TensorSpec( + shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), + dtype=out_dtype, + quantizer=dgrad_quantizer, + device=args.grad_output.device, + ) + + wgrad = None + # Under fuse_wgrad_accumulation the grad goes into main_grad in place. + if args.requires_wgrad and not args.fuse_wgrad_accumulation: + wgrad = TensorSpec( + shape=(out_features, in_features), + dtype=out_dtype, + quantizer=args.grad_weight_quantizer, + device=weight.device, + ) + + grad_bias = None + # FP8 backward computes bgrad in grad_output_preprocess whenever bias is + # used; in high precision it is fused into the wgrad GEMM, so it only + # exists when wgrad runs. + fp8_bwd = args.fp8 and args.backward_override is None + if args.use_bias and (args.requires_wgrad or fp8_bwd): + grad_bias = TensorSpec( + shape=(out_features,), dtype=out_dtype, device=args.grad_output.device + ) + + return wgrad, dgrad, grad_bias + + +# Custom op used under ``torch.compile``. +_linear_op = register_custom_op( + op_name="linear", + input_tensors_for_grad=["weight", "inp", "bias"], + fwd_arg_type=LinearFwdArgs, + fwd_impl=_linear_forward_impl, + fwd_fake_impl=_linear_forward_fake, + setup_context=_linear_setup_ctx, + bwd_arg_type=LinearBwdArgs, + bwd_impl=_linear_backward_impl, + bwd_fake_impl=_linear_backward_fake, +) + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1418,7 +1814,6 @@ def forward( out, new_weight_workspace, tensors_to_save_from_forward, - _, ctx_attrs, ) = _linear_forward_impl(fwd_args) if ctx is not None: @@ -1426,7 +1821,7 @@ def forward( tensors_to_save_from_setup = _linear_setup_ctx( bwd_args, fwd_args, - out, + (out, new_weight_workspace), ctx_attrs, tensors_to_save_from_forward, ) @@ -1458,7 +1853,7 @@ def backward( nvtx_label = "transformer_engine._Linear.backward" if bwd_args.ub_name is not None: nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" - result = _linear_backward(bwd_args) + (None,) # fwd_args grad slot + result = _linear_backward_impl(bwd_args) + (None,) # fwd_args grad slot reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, # main_grad closure) so they don't outlive backward via ctx under retain_graph. @@ -1471,6 +1866,20 @@ def backward( return result +@no_torch_dynamo() +def _linear_eager( + weight_tensor: torch.Tensor, + inp: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LinearFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run ``_Linear`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _Linear.apply(weight_tensor, inp, bias, fwd_args) + return _Linear.forward(None, weight_tensor, inp, bias, fwd_args) + + class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` @@ -1869,7 +2278,15 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - @no_torch_dynamo() + # Allocate the process-global cuBLAS workspaces eagerly: under + # torch.compile the first GEMM can run inside CUDA-graph capture, + # and a workspace first allocated there would live in the graph pool. + device = getattr(self, self.weight_names[0]).device + if device.type == "cuda": + get_cublas_workspace(device.index, False, False) + if self.ub_name is not None: + get_cublas_workspace(device.index, True, False) + def forward( self, inp: torch.Tensor, @@ -1948,12 +2365,9 @@ def forward( weight_quantizer, weight_tensor ) - if is_grad_enabled: - linear_fn = _Linear.apply - autograd_ctx = [] - else: - linear_fn = _Linear.forward - autograd_ctx = [None] + use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None + if _linear_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( @@ -1997,6 +2411,7 @@ def forward( 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 + fwd_args = LinearFwdArgs( # tensors weight=weight_tensor, @@ -2057,13 +2472,26 @@ 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 first: the break below makes Dynamo skip this frame, + # so anything after it is never traced. Explicit break so + # fullgraph=True errors show the reason. + warn_compile_eager_fallback(fallback_reason) + torch._dynamo.graph_break( + msg=f"te.Linear falling back to eager: {fallback_reason}" + ) + use_compiled_op = False + + if use_compiled_op: + check_gemm_dims(inp, weight_tensor, self.fp8) + out, new_weight_workspace = _linear_op(fwd_args) + else: + out, new_weight_workspace = _linear_eager( + weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled + ) if new_weight_workspace is not None and cache_name is not None: if isinstance(new_weight_workspace, torch.Tensor): diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 5c31022123..2799ec3edf 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -473,6 +473,11 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): amax_reduction_group: Optional[dist_group_type] = None def __repr__(self, *, tensor_contents=None): + # Data-free repr for a fake/meta scale_inv (exact class check: fake / + # functional tensors subclass Tensor); materializing it under tracing + # would leak an unbacked symbol into the ShapeEnv. + if self._scale_inv.__class__ is not torch.Tensor or self._scale_inv.is_meta: + return safe_quantized_repr(self, "Float8Tensor") try: return ( "Float8Tensor(" diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 16bba391c4..89b6c2cfa9 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -253,6 +253,11 @@ def view(self, shape: torch.Size): ) def __repr__(self): + # Data-free repr for a fake/meta scale_inv (exact class check: fake / + # functional tensors subclass Tensor); materializing it under tracing + # would leak an unbacked symbol into the ShapeEnv. + if self._scale_inv.__class__ is not torch.Tensor or self._scale_inv.is_meta: + return safe_quantized_repr(self, "Float8TensorStorage") try: return ( "Float8TensorStorage(" diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 39162f8311..c4c6024f47 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -5,6 +5,7 @@ """Utility functions for Transformer Engine modules""" from __future__ import annotations import functools +import logging import math import os import warnings @@ -26,6 +27,76 @@ ] +_compile_disabled_reason: Optional[str] = None +_compile_disabled_warned = False + +try: + from torch._dynamo.comptime import comptime as _comptime +except ImportError: # pragma: no cover + _comptime = None + + +def _compile_safe_warn(msg: str) -> None: + """``warnings.warn`` that also works from code being traced by Dynamo. + + Dynamo silently drops a traced ``warnings.warn``, and TE forwards wrap + everything in try/finally, so a graph break there makes Dynamo skip the + whole frame and re-run it with ``is_compiling() == False`` -- a runtime + warning branch is never reached. Under compilation ``comptime`` runs for + real inside the compiler instead, so the warning fires once per + compilation; the message is read back via ``get_local`` (a traced closure + would capture a VariableTracker, not the value). In eager this is a plain + ``warnings.warn``. + """ + if torch.compiler.is_compiling() and _comptime is not None: + _comptime(lambda ctx: warnings.warn(ctx.get_local("msg").as_python_constant())) + else: + warnings.warn(msg, stacklevel=3) + + +def record_compile_disabled(reason: str) -> None: + """Record why TE's torch.compile custom-op path is off; the warning is + emitted only when a compiled TE module actually runs (see + :func:`warn_if_compile_disabled`), so a plain import stays silent. + The first recorded reason wins. Distinct from + :func:`warn_compile_eager_fallback`, which reports a single *configuration* + falling back while the path itself is available. + """ + global _compile_disabled_reason # pylint: disable=global-statement + if _compile_disabled_reason is None: + _compile_disabled_reason = reason + logging.getLogger("TransformerEngine").info( + "torch.compile custom-op path disabled: %s", reason + ) + + +def warn_if_compile_disabled() -> None: + """Warn once, at the first compile attempt, that the path is off.""" + global _compile_disabled_warned # pylint: disable=global-statement + if _compile_disabled_warned: + return + _compile_disabled_warned = True + msg = ( + "Transformer Engine torch.compile support is disabled: " + f"{_compile_disabled_reason or 'custom-op registration unavailable'}. " + "Modules will fall back to eager execution under torch.compile, i.e. " + "a graph break, which is incompatible with fullgraph=True." + ) + _compile_safe_warn(msg) + + +def warn_compile_eager_fallback(reason: str) -> None: + """Warn that a TE module is running eagerly under ``torch.compile``. + + Emitted when ``reason`` is unsupported on the module's compiled custom-op + path -- once per compilation (see :func:`_compile_safe_warn`). + """ + _compile_safe_warn( + f"Falling back to eager execution under torch.compile: {reason} is " + "unsupported on the compiled path (graph-breaks under fullgraph=True)." + ) + + @functools.lru_cache(maxsize=None) def get_cached_ones_tensor( num_elements: int, @@ -624,6 +695,28 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: ) +def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: + """Emit the TN GEMM (``y = x @ w^T``) dim constraints as ``torch._check`` + guards at trace time. torch.compile path only; eager validation lives in + the op impl. Messages are constant: Dynamo forbids tensor closures here. + """ + # pylint: disable=protected-access + torch._check( + inp.shape[-1] == weight.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + if not fp8: + return + for tensor, name in ((inp, "input"), (weight, "weight")): + torch._check( + math.prod(tensor.shape[:-1]) % 8 == 0 and tensor.shape[-1] % 16 == 0, + lambda n=name: ( + f"FP8 execution requires the {n}'s product of all dimensions except the" + " last to be divisible by 8 and its last dimension to be divisible by 16" + ), + ) + + def is_bf16_compatible() -> bool: """Replaces torch.cuda.is_bf16_compatible() with an explicit check on device compute capability to enforce sm_80 or higher.