diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..b86daa394d 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -38,6 +38,8 @@ 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.module.layernorm_linear import LayerNormLinearFwdArgs +from transformer_engine.pytorch.module.layernorm_mlp import LayerNormMLPFwdArgs 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 @@ -2339,3 +2341,335 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" ) + + +# --------------------------------------------------------------------------- +# te.LayerNormLinear / te.LayerNormMLP +# --------------------------------------------------------------------------- + + +def _flatten_outputs(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + +def _assert_module_close_eager_compiled(fn, compiled, model, base): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert every + forward output and the input / parameter gradients match.""" + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + outs_eager = _flatten_outputs(fn(inp_eager)) + sum(o.sum() for o in outs_eager).backward() + ref_outs = [o.detach().clone() for o in outs_eager] + ref_igrad = inp_eager.grad.detach().clone() + ref_pgrads = { + name: p.grad.detach().clone() for name, p in model.named_parameters() if p.grad is not None + } + + 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. + outs_compiled = [o.clone() for o in _flatten_outputs(compiled(inp_compiled))] + sum(o.sum() for o in outs_compiled).backward() + + assert len(outs_compiled) == len(ref_outs) + for out, ref in zip(outs_compiled, ref_outs): + torch.testing.assert_close(out, ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + pgrads = {name: p.grad for name, p in model.named_parameters() if p.grad is not None} + assert pgrads.keys() == ref_pgrads.keys() + for name, ref in ref_pgrads.items(): + torch.testing.assert_close( + pgrads[name], ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL, msg=f"grad mismatch: {name}" + ) + + +def _run_module_compile_test(model, fp8_recipe, compile_mode, in_features, backward=True): + dtype, device = torch.bfloat16, "cuda" + + 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": + warm = torch.randn(32, in_features, dtype=dtype, device=device, requires_grad=True) + outs = _flatten_outputs(fn(warm)) + sum(o.sum() for o in outs).backward() + 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, in_features, dtype=dtype, device=device) + _assert_module_close_eager_compiled(fn, compiled, model, base) + + +@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_layernorm_linear_compiles(fp8_recipe, compile_mode): + """torch.compile(fullgraph=True) of ``te.LayerNormLinear`` under every + built-in recipe (plus the bf16 baseline), default and reduce-overhead.""" + model = te.LayerNormLinear(64, 32, params_dtype=torch.bfloat16, device="cuda") + _run_module_compile_test(model, fp8_recipe, compile_mode, 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("return_layernorm_output", [False, True]) +@pytest.mark.parametrize("zero_centered_gamma", [False, True]) +def test_te_layernorm_linear_compile_variants( + normalization, return_layernorm_output, zero_centered_gamma +): + """Norm-type / returned-norm-output / zero-centered-gamma variants of + ``te.LayerNormLinear`` under torch.compile, bf16 and FP8 current scaling.""" + model = te.LayerNormLinear( + 64, + 32, + params_dtype=torch.bfloat16, + device="cuda", + normalization=normalization, + return_layernorm_output=return_layernorm_output, + zero_centered_gamma=zero_centered_gamma, + ) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_layernorm_linear_compile_with_quantized_fp8_weight(): + """``te.LayerNormLinear`` with an FP8 primary weight under torch.compile.""" + fp8_recipe = recipe.Float8CurrentScaling() + with te.quantized_model_init(enabled=True, recipe=fp8_recipe): + model = te.LayerNormLinear(64, 32, params_dtype=torch.bfloat16, device="cuda") + assert isinstance(model.weight, te.Float8Tensor) + _run_module_compile_test(model, fp8_recipe, "default", 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("module", ["LayerNormLinear", "LayerNormMLP"]) +def test_te_layernorm_module_dynamic_shapes(module): + """LayerNorm modules with a ``mark_dynamic`` batch dim: one graph for all + batch sizes, numerics matching eager.""" + dtype, device = torch.bfloat16, "cuda" + if module == "LayerNormLinear": + model = te.LayerNormLinear(64, 32, params_dtype=dtype, device=device) + weight = model.weight + else: + model = te.LayerNormMLP(64, 128, params_dtype=dtype, device=device) + # The fused bias-gelu helpers are torch.compile'd on their own inside the + # op and recompile per shape; keep them out of the graph count. + model.bias_gelu_nvfusion = False + weight = model.fc1_weight + + def fn(inp): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + for _ in range(2): + warm = torch.randn(16, 64, 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") + + for batch in (16, 32, 48): + base = torch.randn(batch, 64, dtype=dtype, device=device) + inp = base.detach().clone().requires_grad_(True) + torch._dynamo.mark_dynamic(inp, 0) + model.zero_grad(set_to_none=True) + out = compiled(inp) + out.sum().backward() + igrad, wgrad = inp.grad.clone(), weight.grad.clone() + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = fn(inp_eager) + out_eager.sum().backward() + torch.testing.assert_close(out.detach(), out_eager.detach(), atol=0.0, rtol=0.0) + torch.testing.assert_close(igrad, inp_eager.grad, atol=0.0, rtol=0.0) + torch.testing.assert_close(wgrad, weight.grad, atol=0.0, rtol=0.0) + + if unique_graphs_baseline: + assert _dynamo_counter("stats", "unique_graphs") == unique_graphs_baseline + + +@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", ["fuse_wgrad_accumulation", "delayed_wgrad", "is_first_microbatch"] +) +def test_te_layernorm_linear_compile_eager_fallback(case): + """Unsupported configs fall back to eager with a warning (numerics identical) + and graph-break with the explicit reason under ``fullgraph=True``.""" + dtype, device = torch.bfloat16, "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + kwargs = {} + if case == "fuse_wgrad_accumulation": + kwargs["fuse_wgrad_accumulation"] = True + reason = "fuse_wgrad_accumulation" + elif case == "delayed_wgrad": + kwargs["delay_wgrad_compute"] = True + reason = "delayed wgrad compute" + else: + reason = "FP8 weight caching" + torch.manual_seed(0) + model_ref = te.LayerNormLinear(64, 32, params_dtype=dtype, device=device, **kwargs) + torch.manual_seed(0) + model = te.LayerNormLinear(64, 32, params_dtype=dtype, device=device, **kwargs) + for m in (model_ref, model): + if case == "fuse_wgrad_accumulation": + m.weight.main_grad = torch.zeros_like(m.weight, dtype=torch.float32) + + def make_fn(m): + def fn(inp): + with te.autocast(recipe=fp8_recipe): + if case == "is_first_microbatch": + return m(inp, is_first_microbatch=True) + return m(inp) + + return fn + + fn_ref, fn = make_fn(model_ref), make_fn(model) + torch._dynamo.reset() + compiled = torch.compile(fn) + base = torch.randn(32, 64, dtype=dtype, device=device) + inp_ref = base.clone().requires_grad_(True) + inp = base.clone().requires_grad_(True) + out_ref = fn_ref(inp_ref) + with pytest.warns(UserWarning, match="Falling back to eager execution under torch.compile"): + out = compiled(inp) + out_ref.sum().backward() + out.sum().backward() + if case == "delayed_wgrad": + model_ref.backward_dw() + model.backward_dw() + torch.testing.assert_close(out.detach(), out_ref.detach(), atol=0.0, rtol=0.0) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=0.0, rtol=0.0) + if case == "fuse_wgrad_accumulation": + torch.testing.assert_close(model.weight.main_grad, model_ref.weight.main_grad) + else: + torch.testing.assert_close(model.weight.grad, model_ref.weight.grad, atol=0.0, rtol=0.0) + + torch._dynamo.reset() + with pytest.raises(Exception, match=re.escape(reason)): + torch.compile(fn, fullgraph=True)(base.clone().requires_grad_(True)) + + +@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_layernorm_mlp_compiles(fp8_recipe, compile_mode): + """torch.compile(fullgraph=True) of ``te.LayerNormMLP`` under every + built-in recipe (plus the bf16 baseline), default and reduce-overhead.""" + model = te.LayerNormMLP(64, 128, params_dtype=torch.bfloat16, device="cuda") + _run_module_compile_test(model, fp8_recipe, compile_mode, 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("activation", ["gelu", "swiglu", "relu", "qgeglu"]) +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("return_layernorm_output", [False, True]) +def test_te_layernorm_mlp_compile_variants(activation, normalization, return_layernorm_output): + """Activation / norm-type / returned-norm-output variants of + ``te.LayerNormMLP`` under torch.compile, bf16 and FP8 current scaling.""" + model = te.LayerNormMLP( + 64, + 128, + params_dtype=torch.bfloat16, + device="cuda", + activation=activation, + normalization=normalization, + return_layernorm_output=return_layernorm_output, + ) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("bias", [True, False]) +def test_te_layernorm_mlp_compile_no_bias_or_frozen(bias): + """``te.LayerNormMLP`` without biases, and with frozen FC1 weight (exercises + the ``None`` saved-tensor / grad slots).""" + model = te.LayerNormMLP(64, 128, params_dtype=torch.bfloat16, device="cuda", bias=bias) + _run_module_compile_test(model, None, "default", 64) + model.fc1_weight.requires_grad_(False) + _run_module_compile_test(model, None, "default", 64) + if fp8_available: + _run_module_compile_test(model, recipe.Float8CurrentScaling(), "default", 64) + + +@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", ["checkpoint", "fuse_wgrad_accumulation", "delayed_wgrad"]) +def test_te_layernorm_mlp_compile_eager_fallback(case): + """Unsupported configs fall back to eager with a warning (numerics identical) + and graph-break with the explicit reason under ``fullgraph=True``.""" + dtype, device = torch.bfloat16, "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + kwargs = {} + if case == "checkpoint": + kwargs["checkpoint"] = True + reason = "activation checkpointing" + elif case == "fuse_wgrad_accumulation": + kwargs["fuse_wgrad_accumulation"] = True + reason = "fuse_wgrad_accumulation" + else: + kwargs["delay_wgrad_compute"] = True + reason = "delayed wgrad compute" + torch.manual_seed(0) + model_ref = te.LayerNormMLP(64, 128, params_dtype=dtype, device=device, **kwargs) + torch.manual_seed(0) + model = te.LayerNormMLP(64, 128, params_dtype=dtype, device=device, **kwargs) + for m in (model_ref, model): + if case == "fuse_wgrad_accumulation": + for w in (m.fc1_weight, m.fc2_weight): + w.main_grad = torch.zeros_like(w, dtype=torch.float32) + + def make_fn(m): + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return m(inp) + + return fn + + fn_ref, fn = make_fn(model_ref), make_fn(model) + torch._dynamo.reset() + compiled = torch.compile(fn) + base = torch.randn(32, 64, dtype=dtype, device=device) + inp_ref = base.clone().requires_grad_(True) + inp = base.clone().requires_grad_(True) + out_ref = fn_ref(inp_ref) + with pytest.warns(UserWarning, match="Falling back to eager execution under torch.compile"): + out = compiled(inp) + out_ref.sum().backward() + out.sum().backward() + if case == "delayed_wgrad": + model_ref.backward_dw() + model.backward_dw() + torch.testing.assert_close(out.detach(), out_ref.detach(), atol=0.0, rtol=0.0) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=0.0, rtol=0.0) + if case == "fuse_wgrad_accumulation": + torch.testing.assert_close(model.fc1_weight.main_grad, model_ref.fc1_weight.main_grad) + else: + torch.testing.assert_close( + model.fc1_weight.grad, model_ref.fc1_weight.grad, atol=0.0, rtol=0.0 + ) + + torch._dynamo.reset() + with pytest.raises(Exception, match=re.escape(reason)): + torch.compile(fn, fullgraph=True)(base.clone().requires_grad_(True)) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index f6db139ec1..ea04c8e9f2 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -6,7 +6,7 @@ import os import warnings import weakref -from dataclasses import dataclass +from dataclasses import dataclass, replace as dataclass_replace from typing import Any, Callable, ClassVar, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -44,6 +44,9 @@ nvtx_range_push, needs_quantized_gemm, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, + check_gemm_dims, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -67,8 +70,10 @@ from ._common import ( apply_normalization, check_fp8_reduce_and_update, + fake_workspace_valid, noop_cat, sp_inp_leading, + sp_out_leading, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -80,9 +85,15 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorOrQuantized +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op, + is_value_opaque_quantizer, +) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.identity_tensor import IdentityQuantizer from ..cpu_offload import ( @@ -196,6 +207,53 @@ def any_requires_grad(self) -> 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.any_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 LayerNormLinear.forward). + return "delayed wgrad compute (wgrad_store)" + if ( + self.grad_input_quantizer is not None + and self.is_grad_enabled + and self.input_requires_grad + and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad) + ): + # A quantized dgrad can't cross the op boundary: grads are packed + # one plain Tensor[] slot each (_pack_bwd_result). + return "a quantized input grad (fp8_grad=True)" + if self.cache_weight and self.fp8: + # The cached workspace is updated in place on the first microbatch, + # which the functional op (mutates_args=()) can't express. Without + # FP8 no workspace exists, so is_first_microbatch is inert. + return "FP8 weight caching (is_first_microbatch)" + 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 LayerNormLinearBwdArgs: @@ -1474,6 +1532,347 @@ def wgrad_gemm( ) +def _layernorm_linear_forward_fake( + args: LayerNormLinearFwdArgs, +) -> Tuple[ + TensorSpec, + Optional[TensorSpec], + Optional[TensorSpec], + Optional[Tuple[Any, ...]], + Optional[Dict], +]: + """Shape/metadata-only twin of :func:`_layernorm_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 LayerNormLinear forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight_quantizer = args.weight_quantizer + fp8_or_debug = args.fp8 or args.debug + with_input_all_gather = args.parallel_mode == "column" and args.sequence_parallel + backward_needs_input = args.is_grad_enabled and args.weight_requires_grad + device = args.inp.device + + out_features, in_features = args.weight.shape + # The impl views the input as (-1, in_features). + rows = reduce(multiply_op, args.inp.shape[:-1], 1) + inp_leading = args.inp.shape[0] if len(args.inp.shape) > 1 else 1 + inputmat_aliases_inp = args.inp.dtype == args.activation_dtype + ln_weight_aliases = args.ln_weight.dtype == args.activation_dtype + + # Norm output quantizer usage -- mirrors the impl; the ln_out spec must be + # taken now, before the later rowwise-only usage for the all-gather. + if args.fp8: + if args.input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + args.input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and args.backward_override is None, + ) + if with_input_all_gather and args.input_quantizer.supports_only_rowwise_all_gather(): + args.input_quantizer.set_usage(columnwise=False) + custom = is_custom(args.input_quantizer) + hybrid = isinstance(args.input_quantizer, HybridQuantizer) + identity = isinstance(args.input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + args.fp8 + and not args.debug + and not args.return_layernorm_output + and not args.return_layernorm_output_gathered + and args.backward_override is None + and not custom + and not hybrid + and not identity + ) + # A custom quantizer's ln_out stays in high precision on the plain + # all-gather path (only the gathered copy is quantized). + ln_out_quantized = fp8_or_debug and not ( + with_input_all_gather and not args.return_layernorm_output_gathered and custom + ) + ln_out_spec = TensorSpec( + shape=(rows, in_features), + dtype=args.activation_dtype, + quantizer=args.input_quantizer if ln_out_quantized else None, + device=device, + ) + if with_input_all_gather and fp8_or_debug: + args.input_quantizer.set_usage(rowwise=True, columnwise=False) + # ``ln_out_return`` is the high-precision norm output (or its all-gathered copy). + ln_out_return_is_total = with_input_all_gather and args.return_layernorm_output_gathered + # Whether the quantized ln_out is a different object from the returned one. + ln_out_rebound = ln_out_quantized and not with_quantized_norm + mu = ( + TensorSpec(shape=(rows,), dtype=torch.float32, device=device) + if args.normalization == "LayerNorm" + else None + ) + rsigma = TensorSpec(shape=(rows,), dtype=torch.float32, device=device) + + # ------------------------------------------------------ + # 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 + is_weight_param_quantized = False + if fp8_or_debug: + is_weight_param_quantized = args.weight.is_quantized + if is_weight_param_quantized and not args.debug: + weight_quantizer = args.weight.quantizer + elif weight_quantizer is not None: + weight_quantizer.set_usage( + rowwise=True, + columnwise=args.is_grad_enabled + and not args.is_fsdp2 + and args.backward_override is None, + ) + + if args.weight.is_quantized: + # Primary-quantized args.weight: the impl reuses it as ``weightmat``. + weightmat = args.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(args.weight.shape), + dtype=args.activation_dtype, + quantizer=weight_quantizer, + device=args.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 = args.weight.dtype == args.activation_dtype + weightmat = TensorSpec( + shape=tuple(args.weight.shape), dtype=args.activation_dtype, device=args.weight.device + ) + + # Bias cast: cuBLAS has no FP8 GEMM with FP32 args.bias. + bias_dtype = args.activation_dtype + if fp8_or_debug and args.activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 + bias_aliases = args.bias is not None and args.bias.dtype == bias_dtype + + if args.output_quantizer is not None: + args.output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Outputs: y = norm(x) @ w^T and the optional norm output. + # ------------------------------------------------------ + requires_grad = args.is_grad_enabled and args.any_requires_grad() + out = TensorSpec( + shape=(sp_out_leading(inp_leading, args), *tuple(args.inp.shape[1:-1]), out_features), + dtype=args.activation_dtype, + quantizer=args.output_quantizer, + requires_grad=requires_grad, + device=device, + ) + ln_out_for_return = None + if args.return_layernorm_output: + ln_leading = inp_leading + if args.return_layernorm_output_gathered and with_input_all_gather: + ln_leading = inp_leading * args.tp_size + ln_out_for_return = TensorSpec( + shape=( + (ln_leading, *tuple(args.inp.shape[1:])) + if len(args.inp.shape) > 1 + else (in_features,) + ), + dtype=args.activation_dtype, + requires_grad=requires_grad, + device=device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout + # (inputmat, wt_save, saved_weight, args.bias, args.ln_weight, ln_out, mu, rsigma). + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if args.is_grad_enabled: + ln_out_needs_gather = args.weight_requires_grad and with_input_all_gather + + # Slot 5 -- ``ln_out_to_save``. + ln_out_alias = None + ln_out_to_save = None + if args.backward_override == "high_precision": + # ``ln_out_hp`` is taken before ln_out may be dropped, so it is always saved. + ln_out_to_save = TensorSpec( + shape=(rows, in_features), dtype=args.activation_dtype, device=device + ) + if args.return_layernorm_output and not ln_out_return_is_total: + ln_out_alias = "ln_out" + elif args.weight_requires_grad or args.return_layernorm_output: + ln_out_to_save = ln_out_spec + if ( + backward_needs_input + and args.backward_override is None + and ln_out_quantized + and ( + isinstance(args.input_quantizer, (MXFP8Quantizer, Float8BlockQuantizer)) + or not ln_out_needs_gather + ) + ): + ln_out_to_save.update_usage(rowwise_usage=False) + if args.return_layernorm_output and not ln_out_return_is_total and not ln_out_rebound: + ln_out_alias = "ln_out" + + # Slot 1 -- ``wt_save``, with the impl's alias dedup. + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered args.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(args.weight.shape), + dtype=args.activation_dtype, + device=args.weight.device, + ) + + saved_tensor_aliases = ( + "inp" if inputmat_aliases_inp else None, + wt_alias, + "weight", + "bias" if bias_aliases else None, + "ln_weight" if ln_weight_aliases else None, + ln_out_alias, + None, + None, + ) + tensors_to_save_from_forward = ( + ( + None + if inputmat_aliases_inp + else TensorSpec( + shape=(rows, in_features), dtype=args.activation_dtype, device=device + ) + ), + wt_save, + None, + ( + None + if (args.bias is None or bias_aliases) + else TensorSpec(shape=tuple(args.bias.shape), dtype=bias_dtype, device=device) + ), + ( + None + if ln_weight_aliases + else TensorSpec( + shape=tuple(args.ln_weight.shape), dtype=args.activation_dtype, device=device + ) + ), + None if ln_out_alias is not None else ln_out_to_save, + mu, + rsigma, + ) + ctx_attrs = { + "fsdp_shapes": [], + "saved_tensor_aliases": saved_tensor_aliases, + "is_weight_param_quantized": is_weight_param_quantized, + "ln_out_needs_gather": ln_out_needs_gather, + } + + return out, ln_out_for_return, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _layernorm_linear_backward_fake( + args: LayerNormLinearBwdArgs, +) -> Tuple[Optional[TensorSpec], ...]: + """Allocation-free fake of :func:`_layernorm_linear_backward_impl` on + ``TensorSpec``. Returns ``(dgrad, dgamma, dbeta, wgrad, 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 LayerNormLinear 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 + device = args.grad_output.device + + # Mirrors the impl; affects dgrad's buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + dgrad_leading = sp_inp_leading(args.grad_output.shape[0], args) + dgrad = TensorSpec( + shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), + dtype=out_dtype, + device=device, + ) + + # The norm backward always runs; its grads take the saved (cast) ln_weight dtype. + dgamma = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + dbeta = None + if args.normalization == "LayerNorm": + dbeta = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=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=device) + + return dgrad, dgamma, dbeta, wgrad, grad_bias + + +# Custom op used under ``torch.compile``. +_layernorm_linear_op = register_custom_op( + op_name="layernorm_linear", + input_tensors_for_grad=["inp", "ln_weight", "ln_bias", "weight", "bias"], + fwd_arg_type=LayerNormLinearFwdArgs, + fwd_impl=_layernorm_linear_forward_impl, + fwd_fake_impl=_layernorm_linear_forward_fake, + setup_context=_layernorm_linear_setup_ctx, + bwd_arg_type=LayerNormLinearBwdArgs, + bwd_impl=_layernorm_linear_backward_impl, + bwd_fake_impl=_layernorm_linear_backward_fake, +) + + class _LayerNormLinear(torch.autograd.Function): """LayerNormLinear semi-top level module Calls custom cuda extensions. @@ -1564,6 +1963,22 @@ def backward( ) +@no_torch_dynamo() +def _layernorm_linear_eager( + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + weight: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LayerNormLinearFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run ``_LayerNormLinear`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _LayerNormLinear.apply(inp, ln_weight, ln_bias, weight, bias, fwd_args) + return _LayerNormLinear.forward(None, inp, ln_weight, ln_bias, weight, bias, fwd_args) + + class LayerNormLinear(TransformerEngineBaseModule): r""" Applies layer normalization followed by linear transformation to the incoming data. @@ -2031,7 +2446,6 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2083,6 +2497,16 @@ def forward( if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True + if torch.compiler.is_compiling() and _layernorm_linear_op is not None: + reason = self._compile_eager_fallback_reason( + inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug + ) + if reason is not None: + # A break inside the try/finally below would skip the whole frame. + warn_compile_eager_fallback(reason) + torch._dynamo.graph_break(msg=f"te.LayerNormLinear falling back to eager: {reason}") + return self._forward_eager_fallback(inp, is_first_microbatch, fp8_output, fp8_grad) + inp = self.prepare_forward( inp, allow_non_contiguous=False # removed .contiguous from inside the layer ) @@ -2114,6 +2538,15 @@ def forward( weight_quantizer, weight_tensor ) + use_compiled_op = torch.compiler.is_compiling() and _layernorm_linear_op is not None + if _layernorm_linear_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() + if use_compiled_op: + # Process groups cross the op boundary separately from quantizers. + for quantizer in (input_quantizer, grad_output_quantizer): + if getattr(quantizer, "amax_reduction_group", None) is not None: + set_quantizer_amax_reduction_group(quantizer, None) + cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( self._fp8_workspaces.get(cache_name) if cache_name is not None else None @@ -2227,24 +2660,28 @@ def forward( is_grad_enabled=is_grad_enabled, ) - if is_grad_enabled: - out, ln_out, new_weight_workspace = _LayerNormLinear.apply( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - weight_tensor, - linear_bias_tensor, - fwd_args, - ) + if use_compiled_op: + # Safety net for quantizer-dependent conditions only. + fallback_reason = fwd_args.compile_unsupported_reason() + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + torch._dynamo.graph_break( + msg=f"te.LayerNormLinear falling back to eager: {fallback_reason}" + ) + use_compiled_op = False + + if use_compiled_op: + check_gemm_dims(inp, weight_tensor, self.fp8) + out, ln_out, new_weight_workspace = _layernorm_linear_op(fwd_args) else: - out, ln_out, new_weight_workspace = _LayerNormLinear.forward( - None, + out, ln_out, new_weight_workspace = _layernorm_linear_eager( inp, self.layer_norm_weight, self.layer_norm_bias, weight_tensor, linear_bias_tensor, fwd_args, + is_grad_enabled, ) if new_weight_workspace is not None and cache_name is not None: @@ -2311,6 +2748,71 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): for name, q in zip(names, original_quantizers) ) + def _compile_eager_fallback_reason( + self, + inp: torch.Tensor, + is_first_microbatch: Optional[bool], + fp8_output: bool, + fp8_grad: bool, + is_grad_enabled: bool, + debug: bool, + ) -> Optional[str]: + """Why this call can't use the compiled op (else None), decided before + prepare_forward. Quantizer checks stay in compile_unsupported_reason.""" + if debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + if is_distributed_weight(weight_tensor): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if isinstance(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" + any_requires_grad = ( + inp.requires_grad + or self.layer_norm_weight.requires_grad + or (self.layer_norm_bias is not None and self.layer_norm_bias.requires_grad) + or weight_tensor.requires_grad + or (bias_tensor is not None and bias_tensor.requires_grad) + ) + if fp8_output and is_grad_enabled and any_requires_grad: + return "differentiable fp8_output=True" + if is_cpu_offload_enabled(): + return "CPU activation offloading" + if self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute(): + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + fp8 = FP8GlobalStateManager.is_fp8_enabled() + needs_dgrad = is_grad_enabled and inp.requires_grad + if ( + fp8 + and fp8_grad + and needs_dgrad + and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad) + ): + return "a quantized input grad (fp8_grad=True)" + if fp8 and is_first_microbatch is not None and not self.is_fsdp2: + return "FP8 weight caching (is_first_microbatch)" + return None + + @torch._dynamo.disable + def _forward_eager_fallback( + self, + inp: torch.Tensor, + is_first_microbatch: Optional[bool], + fp8_output: bool, + fp8_grad: bool, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """Re-run forward outside Dynamo (unsupported-config fallback).""" + return LayerNormLinear.forward( + self, + inp, + is_first_microbatch=is_first_microbatch, + fp8_output=fp8_output, + fp8_grad=fp8_grad, + ) + def _get_weight_and_bias_tensors(self): # Get concatenated weight and bias tensors unfused_weights = self._get_weight_tensors() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 59ffa9a16c..d1d88af7dd 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -49,6 +49,9 @@ clear_tensor_data, needs_quantized_gemm, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, + check_gemm_dims, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -74,6 +77,7 @@ from ._common import ( apply_normalization, check_fp8_reduce_and_update, + fake_workspace_valid, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -85,12 +89,18 @@ mark_activation_offload, ) from ..quantized_tensor import ( + QuantizedTensor, QuantizedTensorStorage, Quantizer, prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorOrQuantized +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op, + is_value_opaque_quantizer, +) from ..cpp_extensions import ( general_gemm, ) @@ -269,6 +279,53 @@ def any_requires_grad(self) -> bool: ) ) + def compile_unsupported_reason(self) -> Optional[str]: + """Reason this config can't use the torch.compile custom-op path (else None).""" + if self.debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if self.checkpoint and self.is_grad_enabled: + return "activation checkpointing (checkpoint=True)" + 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.fc2_output_quantizer is not None + and self.is_grad_enabled + and self.any_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 LayerNormMLP.forward). + return "delayed wgrad compute (wgrad_store)" + if self.cache_weight and self.fp8: + # The cached workspaces are updated in place on the first microbatch, + # which the functional op (mutates_args=()) can't express. + return "FP8 weight caching (is_first_microbatch)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + if self.fp8 and self.gemm_gelu_fusion and self.activation == "gelu": + return "gemm_gelu_fusion with FP8" + for quantizer in ( + self.fc1_input_quantizer, + self.fc1_weight_quantizer, + self.fc1_output_quantizer, + self.fc1_grad_input_quantizer, + self.fc1_grad_weight_quantizer, + self.fc1_grad_output_quantizer, + self.fc2_input_quantizer, + self.fc2_weight_quantizer, + self.fc2_output_quantizer, + self.fc2_grad_input_quantizer, + self.fc2_grad_weight_quantizer, + self.fc2_grad_output_quantizer, + ): + 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 LayerNormMLPBwdArgs: @@ -2069,6 +2126,434 @@ def fc1_wgrad_gemm( ) +def _fake_quantized_weight( + weight: TensorSpec, + quantizer: Optional[Quantizer], + workspace: Optional[TensorSpec], + *, + activation_dtype: torch.dtype, + cache_weight: bool, +) -> Tuple[TensorSpec, Optional[TensorSpec], Optional[str]]: + """Spec-level mirror of ``quantize_weight``: ``(weightmat, new_workspace, + save_alias)`` where ``save_alias`` names the forward input / output the + saved weight aliases (``None`` when it is a fresh tensor).""" + if weight.is_quantized: + return weight, None, "weight" + if workspace is not None and not fake_workspace_valid(workspace, quantizer): + workspace = None + if workspace is not None: + weightmat = dataclass_replace(workspace) + weightmat.update_usage(rowwise_usage=True) + return weightmat, None, "weight_workspace" + weightmat = TensorSpec( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=quantizer, + device=weight.device, + ) + new_workspace = None + if cache_weight: + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) + return weightmat, new_workspace, "new_weight_workspace" if cache_weight else None + + +def _layernorm_mlp_forward_fake( + args: LayerNormMLPFwdArgs, +) -> Tuple[ + TensorSpec, + Optional[TensorSpec], + Optional[TensorSpec], + Optional[TensorSpec], + Optional[Tuple[Any, ...]], + Optional[Dict], +]: + """Shape/metadata-only twin of :func:`_layernorm_mlp_forward_impl` for + torch.compile (no args.activation checkpointing: that path falls back to eager).""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time LayerNormMLP forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + if args.checkpoint and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time LayerNormMLP forward does not support args.activation checkpointing" + ) + + fc1_weight_quantizer = args.fc1_weight_quantizer + fc2_weight_quantizer = args.fc2_weight_quantizer + fp8_or_debug = args.fp8 or args.debug + bias_gelu_fusion = args.bias_gelu_fusion + gemm_gelu_fusion = args.gemm_gelu_fusion + device = args.inp.device + + in_features = args.ln_weight.shape[0] + fc1_out_features = args.fc1_weight.shape[0] + act_features = ( + fc1_out_features // 2 if args.activation in _GATED_ACTIVATIONS else fc1_out_features + ) + fc2_out_features = args.fc2_weight.shape[0] + # The impl views the input as (-1, in_features); FC1 consumes the (all-gathered) rows. + rows = reduce(multiply_op, args.inp.shape[:-1], 1) + rows_total = rows * args.tp_size if args.sequence_parallel else rows + inp_leading = args.inp.shape[0] if len(args.inp.shape) > 1 else 1 + inputmat_aliases_inp = args.inp.dtype == args.activation_dtype + ln_weight_aliases = args.ln_weight.dtype == args.activation_dtype + + backwards_needs_fc1_input = args.fc1_weight_requires_grad and args.is_grad_enabled + + # Norm output quantizer usage -- mirrors the impl; the ln_out spec must be + # taken now, before the later rowwise-only usage for the all-gather. + if args.fp8: + if args.fc1_input_quantizer is None: + raise ValueError("Missing quantizer for FC1 input tensor") + args.fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) + if args.sequence_parallel and args.fc1_input_quantizer.supports_only_rowwise_all_gather(): + args.fc1_input_quantizer.set_usage(columnwise=False) + custom = is_custom(args.fc1_input_quantizer) + hybrid = isinstance(args.fc1_input_quantizer, HybridQuantizer) + identity = isinstance(args.fc1_input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + args.fp8 + and not args.debug + and not args.return_layernorm_output + and not args.return_layernorm_output_gathered + and not custom + and not hybrid + and not identity + ) + # A custom quantizer's ln_out stays in high precision on the plain + # all-gather path (only the gathered copy is quantized). + ln_out_quantized = fp8_or_debug and not ( + args.sequence_parallel and not args.return_layernorm_output_gathered and custom + ) + ln_out = TensorSpec( + shape=(rows, in_features), + dtype=args.activation_dtype, + quantizer=args.fc1_input_quantizer if ln_out_quantized else None, + device=device, + ) + if args.sequence_parallel and fp8_or_debug: + args.fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_return_is_total = args.sequence_parallel and args.return_layernorm_output_gathered + ln_out_rebound = ln_out_quantized and not with_quantized_norm + mu = ( + TensorSpec(shape=(rows,), dtype=torch.float32, device=device) + if args.normalization == "LayerNorm" + else None + ) + rsigma = TensorSpec(shape=(rows,), dtype=torch.float32, device=device) + + # ------------------------------------------------------ + # Weights -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ------------------------------------------------------ + new_fc1_weight_workspace = None + new_fc2_weight_workspace = None + if fp8_or_debug: + if args.fc1_weight.is_quantized and not args.debug: + fc1_weight_quantizer = args.fc1_weight.quantizer + elif fc1_weight_quantizer is not None: + fc1_weight_quantizer.set_usage( + rowwise=True, columnwise=args.is_grad_enabled and not args.is_fsdp2 + ) + if args.fc2_weight.is_quantized and not args.debug: + fc2_weight_quantizer = args.fc2_weight.quantizer + elif fc2_weight_quantizer is not None: + fc2_weight_quantizer.set_usage( + rowwise=True, columnwise=args.is_grad_enabled and not args.is_fsdp2 + ) + fc1_weight_final, new_fc1_weight_workspace, fc1_wt_alias = _fake_quantized_weight( + args.fc1_weight, + fc1_weight_quantizer, + args.fc1_weight_workspace, + activation_dtype=args.activation_dtype, + cache_weight=args.cache_weight, + ) + fc2_weight_final, new_fc2_weight_workspace, fc2_wt_alias = _fake_quantized_weight( + args.fc2_weight, + fc2_weight_quantizer, + args.fc2_weight_workspace, + activation_dtype=args.activation_dtype, + cache_weight=args.cache_weight, + ) + fc1_wt_alias = ( + None if fc1_wt_alias is None else fc1_wt_alias.replace("weight", "fc1_weight") + ) + fc2_wt_alias = ( + None if fc2_wt_alias is None else fc2_wt_alias.replace("weight", "fc2_weight") + ) + else: + fc1_weight_final = TensorSpec( + shape=tuple(args.fc1_weight.shape), + dtype=args.activation_dtype, + device=args.fc1_weight.device, + ) + fc2_weight_final = TensorSpec( + shape=tuple(args.fc2_weight.shape), + dtype=args.activation_dtype, + device=args.fc2_weight.device, + ) + fc1_wt_alias = "fc1_weight" if args.fc1_weight.dtype == args.activation_dtype else None + fc2_wt_alias = "fc2_weight" if args.fc2_weight.dtype == args.activation_dtype else None + + # Bias cast: cuBLAS has no FP8 GEMM with FP32 bias. + bias_dtype = args.activation_dtype + if fp8_or_debug and args.activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 + fc1_bias_aliases = args.fc1_bias is not None and args.fc1_bias.dtype == bias_dtype + fc2_bias_aliases = args.fc2_bias is not None and args.fc2_bias.dtype == bias_dtype + + # ------------------------------------------------------ + # FC1 GEMM + args.activation (see the impl for the fusion rules). + # ------------------------------------------------------ + if args.activation != "gelu": + gemm_gelu_fusion = bias_gelu_fusion = False + else: + if not args.fp8: + gemm_gelu_fusion = True + if gemm_gelu_fusion and bias_gelu_fusion: + gemm_gelu_fusion = False + if args.debug: + gemm_gelu_fusion = False + fc1_out = None + fc1_out_without_bias = None + if bias_gelu_fusion: + fc1_out_without_bias = TensorSpec( + shape=(rows_total, fc1_out_features), dtype=args.activation_dtype, device=device + ) + else: + fc1_out = TensorSpec( + shape=(rows_total, fc1_out_features), dtype=args.activation_dtype, device=device + ) + act_out = TensorSpec( + shape=(rows_total, act_features), + dtype=args.activation_dtype, + quantizer=args.fc2_input_quantizer if fp8_or_debug else None, + device=device, + ) + if args.fc2_output_quantizer is not None: + args.fc2_output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Outputs. + # ------------------------------------------------------ + requires_grad = args.is_grad_enabled and args.any_requires_grad() + out_leading = inp_leading + if args.sequence_parallel and not args.set_parallel_mode: + out_leading = inp_leading * args.tp_size + fc2_out = TensorSpec( + shape=(out_leading, *tuple(args.inp.shape[1:-1]), fc2_out_features), + dtype=args.activation_dtype, + quantizer=args.fc2_output_quantizer, + requires_grad=requires_grad, + device=device, + ) + ln_out_for_return = None + if args.return_layernorm_output: + ln_leading = inp_leading + if ( + args.return_layernorm_output_gathered + and args.sequence_parallel + and args.set_parallel_mode + ): + ln_leading = inp_leading * args.tp_size + ln_out_for_return = TensorSpec( + shape=( + (ln_leading, *tuple(args.inp.shape[1:])) + if len(args.inp.shape) > 1 + else (in_features,) + ), + dtype=args.activation_dtype, + requires_grad=requires_grad, + device=device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout (see the impl). + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if args.is_grad_enabled: + if not args.fc1_weight_requires_grad: + ln_out = None + if not args.fc2_weight_requires_grad: + act_out = None + fc1_wt_save = fc1_weight_final + fc2_wt_save = fc2_weight_final + if args.is_fsdp2: + if fc1_wt_alias != "fc1_weight": + fc1_wt_save, fc1_wt_alias = None, None + if fc2_wt_alias != "fc2_weight": + fc2_wt_save, fc2_wt_alias = None, None + ln_out_alias = ( + "ln_out" + if args.return_layernorm_output + and ln_out is not None + and not ln_out_return_is_total + and not ln_out_rebound + else None + ) + saved_tensor_aliases = ( + "inp" if inputmat_aliases_inp else None, + "ln_weight" if ln_weight_aliases else None, + ln_out_alias, + fc1_wt_alias, + "fc1_weight", + "fc1_bias" if fc1_bias_aliases else None, + None, + None, + None, + fc2_wt_alias, + "fc2_weight", + "fc2_bias" if fc2_bias_aliases else None, + None, + None, + ) + saved = ( + TensorSpec(shape=(rows, in_features), dtype=args.activation_dtype, device=device), + TensorSpec( + shape=tuple(args.ln_weight.shape), dtype=args.activation_dtype, device=device + ), + ln_out, + fc1_wt_save, + None, + ( + None + if args.fc1_bias is None + else TensorSpec(shape=tuple(args.fc1_bias.shape), dtype=bias_dtype, device=device) + ), + fc1_out, + fc1_out_without_bias, + act_out, + fc2_wt_save, + None, + ( + None + if args.fc2_bias is None + else TensorSpec(shape=tuple(args.fc2_bias.shape), dtype=bias_dtype, device=device) + ), + mu, + rsigma, + ) + tensors_to_save_from_forward = tuple( + None if alias is not None else spec for alias, spec in zip(saved_tensor_aliases, saved) + ) + ctx_attrs = { + "saved_tensor_aliases": saved_tensor_aliases, + "fsdp_shapes": None, + "is_recomputation": False, + } + + return ( + fc2_out, + ln_out_for_return, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) + + +def _layernorm_mlp_backward_fake( + args: LayerNormMLPBwdArgs, +) -> Tuple[Optional[TensorSpec], ...]: + """Allocation-free fake of :func:`_layernorm_mlp_backward_impl` on + ``TensorSpec``. Returns ``(dgrad, dgamma, dbeta, fc1_wgrad, fc1_bias_grad, + fc2_wgrad, fc2_bias_grad)`` specs with rank-local shapes.""" + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake LayerNormMLP backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + out_dtype = args.activation_dtype + device = args.grad_output.device + in_features = args.ln_weight.shape[-1] + fc1_out_features = args.fc1_weight.shape[0] + fc2_out_features, act_features = args.fc2_weight.shape + + if args.fc1_grad_output_quantizer is not None: + args.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + + dgrad = None + if args.requires_dgrad: + inp_leading = args.grad_output.shape[0] + if args.sequence_parallel and not args.set_parallel_mode: + inp_leading = inp_leading // args.tp_size + dgrad = TensorSpec( + shape=(inp_leading, *args.grad_output.shape[1:-1], in_features), + dtype=out_dtype, + device=device, + ) + + # The norm backward always runs; its grads take the saved (cast) ln_weight dtype. + dgamma = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + dbeta = None + if args.normalization == "LayerNorm": + dbeta = TensorSpec(shape=(in_features,), dtype=args.ln_weight.dtype, device=device) + + fc1_wgrad = None + if args.fc1_weight_requires_grad and not args.fuse_wgrad_accumulation: + fc1_wgrad = TensorSpec( + shape=(fc1_out_features, in_features), + dtype=out_dtype, + quantizer=args.fc1_grad_weight_quantizer, + device=device, + ) + fc1_bias_grad = None + # FP8, debug and the fused bias-gelu path always produce it; in high + # precision it is fused into the FC1 wgrad GEMM or summed when only the + # bias needs a grad. + if args.fc1_bias is not None and any( + ( + args.fp8, + args.debug, + args.bias_gelu_fusion, + args.fc1_weight_requires_grad, + args.fc1_bias_requires_grad, + ) + ): + fc1_bias_grad = TensorSpec(shape=(fc1_out_features,), dtype=out_dtype, device=device) + + fc2_wgrad = None + if args.fc2_weight_requires_grad and not args.fuse_wgrad_accumulation: + fc2_wgrad = TensorSpec( + shape=(fc2_out_features, act_features), + dtype=out_dtype, + quantizer=args.fc2_grad_weight_quantizer, + device=device, + ) + fc2_bias_grad = None + fp8_bwd = args.fp8 and args.backward_override is None + if args.use_bias and (args.fc2_weight_requires_grad or fp8_bwd): + fc2_bias_grad = TensorSpec(shape=(fc2_out_features,), dtype=out_dtype, device=device) + + return dgrad, dgamma, dbeta, fc1_wgrad, fc1_bias_grad, fc2_wgrad, fc2_bias_grad + + +# Custom op used under ``torch.compile``. +_layernorm_mlp_op = register_custom_op( + op_name="layernorm_mlp", + input_tensors_for_grad=[ + "inp", + "ln_weight", + "ln_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", + ], + fwd_arg_type=LayerNormMLPFwdArgs, + fwd_impl=_layernorm_mlp_forward_impl, + fwd_fake_impl=_layernorm_mlp_forward_fake, + setup_context=_layernorm_mlp_setup_ctx, + bwd_arg_type=LayerNormMLPBwdArgs, + bwd_impl=_layernorm_mlp_backward_impl, + bwd_fake_impl=_layernorm_mlp_backward_fake, +) + + class _LayerNormMLP(torch.autograd.Function): """LayerNormMLP semi-top level module Calls custom cuda extensions. @@ -2171,6 +2656,18 @@ def backward( ) +@no_torch_dynamo() +def _layernorm_mlp_eager( + tensors: Tuple[Optional[torch.Tensor], ...], + fwd_args: LayerNormMLPFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run ``_LayerNormMLP`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _LayerNormMLP.apply(*tensors, fwd_args) + return _LayerNormMLP.forward(None, *tensors, fwd_args) + + class LayerNormMLP(TransformerEngineBaseModule): r""" Applies layer normalization on the input followed by the MLP module, consisting of @@ -2601,8 +3098,6 @@ def reset_parameters(self, defer_init=False): if self.set_parallel_mode: setattr(self.fc2_bias, "sequence_parallel", self.sequence_parallel) - @no_torch_dynamo() - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2650,6 +3145,16 @@ def forward( if get_ub_is_fp8("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True + if torch.compiler.is_compiling() and _layernorm_mlp_op is not None: + reason = self._compile_eager_fallback_reason( + inp, is_first_microbatch, fp8_output, is_grad_enabled, debug + ) + if reason is not None: + # A break inside the try/finally below would skip the whole frame. + warn_compile_eager_fallback(reason) + torch._dynamo.graph_break(msg=f"te.LayerNormMLP falling back to eager: {reason}") + return self._forward_eager_fallback(inp, is_first_microbatch) + inp = self.prepare_forward(inp, num_gemms=2) try: @@ -2702,6 +3207,15 @@ def forward( if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): self.fast_setattr("bias_gelu_nvfusion", False) + use_compiled_op = torch.compiler.is_compiling() and _layernorm_mlp_op is not None + if _layernorm_mlp_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() + if use_compiled_op: + # Process groups cross the op boundary separately from quantizers. + for quantizer in (fc1_input_quantizer, fc2_grad_output_quantizer): + if getattr(quantizer, "amax_reduction_group", None) is not None: + set_quantizer_amax_reduction_group(quantizer, None) + cache_name_fc1 = ( None if (is_first_microbatch is None or self.is_fsdp2) else "fc1_weight" ) @@ -2847,28 +3361,32 @@ def forward( is_grad_enabled=is_grad_enabled, ) - if is_grad_enabled: - out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.apply( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias_tensor, - fwd_args, - ) + if use_compiled_op: + # Safety net for quantizer-dependent conditions only. + fallback_reason = fwd_args.compile_unsupported_reason() + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + torch._dynamo.graph_break( + msg=f"te.LayerNormMLP falling back to eager: {fallback_reason}" + ) + use_compiled_op = False + + if use_compiled_op: + check_gemm_dims(inp, fc1_weight, self.fp8) + out, ln_out, new_fc1_ws, new_fc2_ws = _layernorm_mlp_op(fwd_args) else: - out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.forward( - None, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias_tensor, + out, ln_out, new_fc1_ws, new_fc2_ws = _layernorm_mlp_eager( + ( + inp, + self.layer_norm_weight, + self.layer_norm_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias_tensor, + ), fwd_args, + is_grad_enabled, ) if new_fc1_ws is not None and cache_name_fc1 is not None: @@ -2894,6 +3412,52 @@ def forward( return out, ln_out return out + def _compile_eager_fallback_reason( + self, + inp: torch.Tensor, + is_first_microbatch: Optional[bool], + fp8_output: bool, + is_grad_enabled: bool, + debug: bool, + ) -> Optional[str]: + """Why this call can't use the compiled op (else None), decided before + prepare_forward. Quantizer checks stay in compile_unsupported_reason.""" + if debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if self.checkpoint and is_grad_enabled: + return "activation checkpointing (checkpoint=True)" + if isinstance(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 fp8_output and is_grad_enabled and (inp.requires_grad or self.requires_grad_params()): + return "differentiable fp8_output=True" + if is_cpu_offload_enabled(): + return "CPU activation offloading" + if self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute(): + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + fp8 = FP8GlobalStateManager.is_fp8_enabled() + if fp8 and is_first_microbatch is not None and not self.is_fsdp2: + return "FP8 weight caching (is_first_microbatch)" + if fp8 and self.gemm_gelu_fusion and self.activation == "gelu": + return "gemm_gelu_fusion with FP8" + return None + + def requires_grad_params(self) -> bool: + """Whether any parameter of this module requires a gradient.""" + return any(p.requires_grad for p in self.parameters()) + + @torch._dynamo.disable + def _forward_eager_fallback( + self, + inp: torch.Tensor, + is_first_microbatch: Optional[bool], + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """Re-run forward outside Dynamo (unsupported-config fallback).""" + return LayerNormMLP.forward(self, inp, is_first_microbatch=is_first_microbatch) + def _get_quantizers(self, fp8_output, is_grad_enabled): if self.fp8: self._warn_missing_output_quantizer_role(fp8_output, False)