From 3e2b79d40477c8703b5af13e68611b206ca25796 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 14:39:39 +0200 Subject: [PATCH] [PyTorch] [torch.compile] Prepare the custom-op framework and LayerNorm modules for compile wiring Behavior-neutral groundwork for registering LayerNormLinear and LayerNormMLP as torch.compile custom ops (follows the Linear split #2967 / #3053). dynamo/custom_op.py: - A bwd args dataclass may declare GRAD_OUTPUT_FIELDS naming one grad field per user output; the LN modules have two differentiable outputs (out and the returned norm output), so a single grad_output slot is not enough. - Dict[str, ] annotations are bundle-simple (activation_params). module/_common.py: sp_out_leading / sp_inp_leading / fake_workspace_valid shared by Linear and the LN modules (moved out of linear.py). layernorm_linear.py / layernorm_mlp.py: - GRAD_OUTPUT_FIELDS on the bwd args; fp8_output carried in the fwd args. - inp_shape is rederived from grad_output in backward instead of being stored on the bwd args (SymInt dims are not hashable in the value bundle). - LayerNormMLP: the recipe object no longer rides on the backward args; the properties the backward needs (float8_block_scaling, custom, dbias-dact fusion availability) are bools computed in Module.forward. The activation tables are split into per-activation (act, dact) pairs plus the fused dbias kernels so the backward looks them up without a recipe. - The returned-norm-output grad is tolerated as None in backward. No functional change. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 28 ++- transformer_engine/pytorch/module/_common.py | 37 ++++ .../pytorch/module/layernorm_linear.py | 22 ++- .../pytorch/module/layernorm_mlp.py | 179 ++++++++++-------- transformer_engine/pytorch/module/linear.py | 45 +---- 5 files changed, 180 insertions(+), 131 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 00846d615a..38941ec8b0 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -61,8 +61,9 @@ forward state + ``ctx_attrs`` (e.g. saved-tensor aliases) and return the tensors to persist; the plan's output ranges are 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), + stashed plan (the backward args' ``GRAD_OUTPUT_FIELDS`` class attribute names + one field per user output; else a ``grad_outputs`` field 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). @@ -468,6 +469,9 @@ def _is_simple_annot(annot: Any) -> bool: 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) + if get_origin(annot) is dict: + key, value = get_args(annot) + return key is str and _is_simple_annot(value) return False @@ -1009,6 +1013,10 @@ def _register_autograd_for_op( 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. """ + # Where the incoming grads land on the backward args: the fields named by + # ``GRAD_OUTPUT_FIELDS`` (one per user output, in order), else a + # ``grad_outputs`` tuple field, else ``grad_output`` (first output only). + grad_output_fields = getattr(bwd_plan.arg_type, "GRAD_OUTPUT_FIELDS", None) bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields) def _setup_context(ctx, inputs, output): @@ -1049,7 +1057,10 @@ def _autograd_backward(ctx, *grad_outputs): ctx.tensor_objects = None user_grads = _slice_user_grads(ctx.output_ranges, grad_outputs[0]) ctx.output_ranges = None - if bwd_takes_grad_tuple: + if grad_output_fields is not None: + for name, grad in zip(grad_output_fields, user_grads): + setattr(bwd_obj, name, grad) + elif bwd_takes_grad_tuple: bwd_obj.grad_outputs = tuple(user_grads) else: bwd_obj.grad_output = user_grads[0] @@ -1219,11 +1230,12 @@ def register_custom_op( forward state and returns the tensors to persist; the framework saves them via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework restores them into the container's tensor fields through the - ``setup_saved_tensors`` hook and sets 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``. + ``setup_saved_tensors`` hook and sets the incoming gradients directly -- + into the fields named by a ``GRAD_OUTPUT_FIELDS`` class attribute (one per + user output, in order) if ``bwd_arg_type`` declares one, else into a + ``grad_outputs`` field (tuple, one grad per user output), 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 diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 98262bd99e..d497fe5910 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -13,8 +13,10 @@ from .. import cpp_extensions as tex from ..constants import TE_DType from ..distributed import in_fp8_activation_recompute_phase +from ..dynamo import TensorSpec from ..export import is_in_onnx_export_mode from ..quantization import FP8GlobalStateManager +from ..quantized_tensor import Quantizer from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -336,3 +338,38 @@ def check_fp8_reduce_and_update(restore_first_module: bool = False) -> bool: if restore_first_module or in_fp8_activation_recompute_phase(): qstate.is_first_fp8_module = first_fp8_module return result + + +def sp_out_leading(leading: int, args: Any) -> int: + """Output's leading (sequence) dim from the input's: sequence parallelism + gathers it (column-parallel) or scatters it (row-parallel). ``args`` carries + ``sequence_parallel`` / ``parallel_mode`` / ``tp_size``.""" + if not args.sequence_parallel: + return leading + if args.parallel_mode == "column": + return leading * args.tp_size + if args.parallel_mode == "row": + return leading // args.tp_size + return leading + + +def sp_inp_leading(leading: int, args: Any) -> int: + """Inverse of :func:`sp_out_leading`.""" + 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()) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index b0af7f05b2..f6db139ec1 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -7,7 +7,7 @@ import warnings import weakref from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Tuple, Union, List +from typing import Any, Callable, ClassVar, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -68,6 +68,7 @@ apply_normalization, check_fp8_reduce_and_update, noop_cat, + sp_inp_leading, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -143,6 +144,7 @@ class LayerNormLinearFwdArgs: activation_dtype: torch.dtype fp8: bool fp8_calibration: bool + fp8_output: bool backward_override: Optional[str] dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool @@ -199,6 +201,9 @@ def any_requires_grad(self) -> bool: class LayerNormLinearBwdArgs: """Single-argument bag for the backward path of :class:`_LayerNormLinear`.""" + # One field per user output of the forward op, in order (see custom_op.py). + GRAD_OUTPUT_FIELDS: ClassVar[Tuple[str, ...]] = ("grad_output", "grad_ln_out") + # --- Incoming gradients (populated at backward entry) --- grad_output: Optional[torch.Tensor] = None grad_ln_out: Optional[torch.Tensor] = None @@ -780,7 +785,9 @@ def _layernorm_linear_setup_ctx( bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad bwd_args.ln_out_needs_gather = ctx_attrs["ln_out_needs_gather"] - bwd_args.inp_shape = inp.shape + # Not stored (SymInt dims are not hashable in OpaqueValueBundle under + # torch.compile(dynamic=True)); backward rederives it from grad_output. + bwd_args.inp_shape = None # Normalization bwd_args.normalization = fwd_args.normalization @@ -908,6 +915,10 @@ def _layernorm_linear_backward_impl( """ grad_output = args.grad_output assert grad_output is not None + if args.inp_shape is None: + in_features = args.saved_weight.shape[-1] + inp_leading = sp_inp_leading(grad_output.shape[0], args) + args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) # NVTX label for profiling nvtx_label = "transformer_engine._LayerNormLinear.backward" @@ -1397,7 +1408,11 @@ def wgrad_gemm( # Residual gradient dgrad = dgrad.view(inputmat.shape) - if args.return_layernorm_output and not args.return_layernorm_output_gathered: + if ( + args.return_layernorm_output + and not args.return_layernorm_output_gathered + and args.grad_ln_out is not None + ): dgrad = dgrad + args.grad_ln_out.view_as(dgrad) # Norm gradient @@ -2177,6 +2192,7 @@ def forward( activation_dtype=self.activation_dtype, fp8=self.fp8, fp8_calibration=self.fp8_calibration, + fp8_output=fp8_output, backward_override=backward_override, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3be30fa83e..59ffa9a16c 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -7,7 +7,7 @@ import warnings from dataclasses import dataclass, replace as dataclass_replace import weakref -from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, List +from typing import Any, Callable, ClassVar, Dict, Optional, Sequence, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -100,76 +100,54 @@ __all__ = ["LayerNormMLP"] -def _get_act_func_supported_list(recipe: Optional[Recipe] = None): - if recipe is None: - # bf16 (recipe is None): - return { - "gelu": (tex.gelu, tex.dgelu, None), - "geglu": (tex.geglu, tex.dgeglu, None), - "glu": (tex.glu, tex.dglu, None), - "qgelu": (tex.qgelu, tex.dqgelu, None), - "qgeglu": (tex.qgeglu, tex.dqgeglu, None), - "relu": (tex.relu, tex.drelu, None), - "reglu": (tex.reglu, tex.dreglu, None), - "srelu": (tex.srelu, tex.dsrelu, None), - "sreglu": (tex.sreglu, tex.dsreglu, None), - "silu": (tex.silu, tex.dsilu, None), - "swiglu": (tex.swiglu, tex.dswiglu, None), - "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), - } - if recipe.delayed() or recipe.mxfp8(): - # Delayed scaling, fusion supported list: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] - # MXFP8: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] - return { - "gelu": (tex.gelu, tex.dgelu, tex.dbias_dgelu), - "geglu": (tex.geglu, tex.dgeglu, None), - "glu": (tex.glu, tex.dglu, None), - "qgelu": (tex.qgelu, tex.dqgelu, tex.dbias_dqgelu), - "qgeglu": (tex.qgeglu, tex.dqgeglu, None), - "relu": (tex.relu, tex.drelu, tex.dbias_drelu), - "reglu": (tex.reglu, tex.dreglu, None), - "srelu": (tex.srelu, tex.dsrelu, tex.dbias_dsrelu), - "sreglu": (tex.sreglu, tex.dsreglu, None), - "silu": (tex.silu, tex.dsilu, tex.dbias_dsilu), - "swiglu": (tex.swiglu, tex.dswiglu, None), - "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), - } - # no activation fusion written yet - # Per-tensor current scaling or fp8 blockwise scaling or custom quantization: [] - # TODO(ksivaman): Fuse nvfp4 act once kernel is available. - if ( - recipe.float8_current_scaling() - or recipe.float8_block_scaling() - or recipe.nvfp4() - or recipe.custom() - ): - return { - "gelu": (tex.gelu, tex.dgelu, None), - "geglu": (tex.geglu, tex.dgeglu, None), - "glu": (tex.glu, tex.dglu, None), - "qgelu": (tex.qgelu, tex.dqgelu, None), - "qgeglu": (tex.qgeglu, tex.dqgeglu, None), - "relu": (tex.relu, tex.drelu, None), - "reglu": (tex.reglu, tex.dreglu, None), - "srelu": (tex.srelu, tex.dsrelu, None), - "sreglu": (tex.sreglu, tex.dsreglu, None), - "silu": (tex.silu, tex.dsilu, None), - "swiglu": (tex.swiglu, tex.dswiglu, None), - "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), - } - raise NotImplementedError(f"Unhandled recipe type {recipe}") +_ACT_FUNCS = { + "gelu": (tex.gelu, tex.dgelu), + "geglu": (tex.geglu, tex.dgeglu), + "glu": (tex.glu, tex.dglu), + "qgelu": (tex.qgelu, tex.dqgelu), + "qgeglu": (tex.qgeglu, tex.dqgeglu), + "relu": (tex.relu, tex.drelu), + "reglu": (tex.reglu, tex.dreglu), + "srelu": (tex.srelu, tex.dsrelu), + "sreglu": (tex.sreglu, tex.dsreglu), + "silu": (tex.silu, tex.dsilu), + "swiglu": (tex.swiglu, tex.dswiglu), + "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu), +} + +# Fused dbias + dact + quantize kernels; only delayed scaling and MXFP8 have them. +_DBIAS_DACT_FUNCS = { + "gelu": tex.dbias_dgelu, + "qgelu": tex.dbias_dqgelu, + "relu": tex.dbias_drelu, + "srelu": tex.dbias_dsrelu, + "silu": tex.dbias_dsilu, +} + +# Activations whose output halves the last dim (gated linear units). +_GATED_ACTIVATIONS = frozenset( + {"geglu", "glu", "qgeglu", "reglu", "sreglu", "swiglu", "clamped_swiglu"} +) + + +def _recipe_has_dbias_dact_fusion(recipe: Optional[Recipe]) -> bool: + return recipe is not None and (recipe.delayed() or recipe.mxfp8()) -def _act_func(activation: str, recipe: Optional[Recipe] = None): - # based on each quantization mode, we have different kernel fusion supported: - # bf16 (recipe is None): [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] - # Delayed scaling, fusion supported list: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] - # MXFP8: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] - # Per-tensor current scaling or fp8 blockwise scaling: [] - funcs = _get_act_func_supported_list(recipe) - if activation not in funcs: +def _act_func( + activation: str, recipe: Optional[Recipe] = None, dbias_fusion: Optional[bool] = None +): + """``(act, dact, dbias_dact_quantize or None)`` for ``activation``. + + The fused dbias kernel is available for delayed scaling and MXFP8 only; + pass ``dbias_fusion`` to decide without a recipe object. + """ + if activation not in _ACT_FUNCS: raise NotImplementedError("Activation type " + activation + " is not supported!") - return funcs[activation] + if dbias_fusion is None: + dbias_fusion = _recipe_has_dbias_dact_fusion(recipe) + act, dact = _ACT_FUNCS[activation] + return act, dact, _DBIAS_DACT_FUNCS.get(activation) if dbias_fusion else None @dataclass(slots=True) @@ -234,6 +212,10 @@ class LayerNormMLPFwdArgs: backward_override: Optional[str] dgrad_use_split_accumulator: bool wgrad_use_split_accumulator: bool + # Recipe properties the backward needs (the recipe itself can't cross the op boundary). + recipe_float8_block_scaling: bool + recipe_custom: bool + recipe_dbias_dact_fusion: bool debug: bool # --- Weight-workspace caching --- @@ -292,6 +274,9 @@ def any_requires_grad(self) -> bool: class LayerNormMLPBwdArgs: """Single-argument bag for the backward path of :class:`_LayerNormMLP`.""" + # One field per user output of the forward op, in order (see custom_op.py). + GRAD_OUTPUT_FIELDS: ClassVar[Tuple[str, ...]] = ("grad_output", "grad_ln_out") + # --- Incoming gradients (populated at backward entry) --- grad_output: Optional[torch.Tensor] = None grad_ln_out: Optional[torch.Tensor] = None @@ -354,7 +339,9 @@ class LayerNormMLPBwdArgs: # --- Numerical / dtype config --- activation_dtype: Optional[torch.dtype] = None fp8: bool = False - fp8_recipe: Optional[Any] = None + recipe_float8_block_scaling: bool = False + recipe_custom: bool = False + recipe_dbias_dact_fusion: bool = False dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD backward_override: Optional[str] = None @@ -1157,7 +1144,9 @@ def _layernorm_mlp_setup_ctx( bwd_args.fc1_weight_requires_grad = fc1_weight_requires_grad bwd_args.fc1_bias_requires_grad = fwd_args.fc1_bias_requires_grad bwd_args.fc2_weight_requires_grad = fc2_weight_requires_grad - bwd_args.inp_shape = inp.shape + # Not stored (SymInt dims are not hashable in OpaqueValueBundle under + # torch.compile(dynamic=True)); backward rederives it from grad_output. + bwd_args.inp_shape = None # Normalization bwd_args.normalization = fwd_args.normalization @@ -1176,7 +1165,9 @@ def _layernorm_mlp_setup_ctx( # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype bwd_args.fp8 = fp8 - bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.recipe_float8_block_scaling = fwd_args.recipe_float8_block_scaling + bwd_args.recipe_custom = fwd_args.recipe_custom + bwd_args.recipe_dbias_dact_fusion = fwd_args.recipe_dbias_dact_fusion bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator bwd_args.backward_override = fwd_args.backward_override @@ -1325,6 +1316,14 @@ def _layernorm_mlp_backward_impl( the saved-tensor fields before invocation. Returns ``(dgrad, dgamma, dbeta, fc1_wgrad, fc1_bias_grad, fc2_wgrad, fc2_bias_grad)``. """ + if args.inp_shape is None: + in_features = args.ln_weight.shape[-1] + inp_leading = args.grad_output.shape[0] + if args.sequence_parallel and not args.set_parallel_mode: + # FC1's input was all-gathered but FC2's output was not reduce-scattered. + inp_leading = inp_leading // args.tp_size + args.inp_shape = torch.Size([inp_leading, *args.grad_output.shape[1:-1], in_features]) + with get_nvtx_range_context("_LayerNormMLP_backward"): inputmat = args.inputmat ln_weight = args.ln_weight @@ -1610,7 +1609,7 @@ def _layernorm_mlp_backward_impl( # Whether to set grad arg in general_gemm grad_arg = True - if args.fp8 and args.fp8_recipe.float8_block_scaling(): + if args.fp8 and args.recipe_float8_block_scaling: grad_arg = False # Arguments to include in wgrad GEMM closure @@ -1656,7 +1655,7 @@ def fc2_wgrad_gemm( # Update grad bias if needed if fc2_bias_grad is None: - if args.fp8 and args.fp8_recipe.float8_block_scaling() and fc2_bias is not None: + if args.fp8 and args.recipe_float8_block_scaling and fc2_bias is not None: # BGRAD not fused with GEMM for float8 blockwise gemm. fc2_bias_grad_ = act_out.view(-1, act_out.shape[-1]).sum(dim=0) fc2_bias_grad = fc2_bias_grad_ @@ -1689,14 +1688,15 @@ def fc2_wgrad_gemm( fc1_bias_grad = dact.sum(dim=0) dact = args.fc1_grad_output_quantizer(dact) elif ( - _act_func(args.activation, args.fp8_recipe if args.fp8 else None)[2] is not None + _act_func(args.activation, dbias_fusion=args.fp8 and args.recipe_dbias_dact_fusion)[2] + is not None and args.fp8 ): # Fusion: gemm, bias + gelu + quantize dbias_dact_quantize_func = _act_func( - args.activation, args.fp8_recipe if args.fp8 else None + args.activation, dbias_fusion=args.fp8 and args.recipe_dbias_dact_fusion )[2] - fc1_bias_grad, dact = dbias_dact_quantize_func( + fc1_bias_grad, dact = dbias_dact_quantize_func( # pylint: disable=not-callable fc2_dgrad, fc1_out.to(args.activation_dtype), args.fc1_grad_output_quantizer, @@ -1705,9 +1705,7 @@ def fc2_wgrad_gemm( else: # Fusion: gemm + gelu, if not fc2_dgrad_gemm_gelu_fusion: - activation_func_bwd = _act_func( - args.activation, args.fp8_recipe if args.fp8 else None - )[1] + activation_func_bwd = _act_func(args.activation)[1] dact = activation_func_bwd( fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params ) # activation in high precision @@ -1719,7 +1717,7 @@ def fc2_wgrad_gemm( args.fc1_grad_output_quantizer, (Float8BlockQuantizer, IdentityQuantizer), ) - or args.fp8_recipe.custom() + or args.recipe_custom ): fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) dact = args.fc1_grad_output_quantizer(dact) @@ -1825,7 +1823,11 @@ def fc2_wgrad_gemm( elif args.set_parallel_mode and not ub_bulk_wgrad: fc1_dgrad = gemm_out if args.sequence_parallel: - if args.return_layernorm_output and args.return_layernorm_output_gathered: + if ( + args.return_layernorm_output + and args.return_layernorm_output_gathered + and args.grad_ln_out is not None + ): fc1_dgrad = fc1_dgrad + args.grad_ln_out.view_as(fc1_dgrad) fc1_dgrad, fc1_dgrad_work = reduce_scatter_along_first_dim( fc1_dgrad, @@ -1966,7 +1968,11 @@ def fc1_wgrad_gemm( # Residual gradient dgrad = fc1_dgrad.view(inputmat.shape) - if args.return_layernorm_output and not args.return_layernorm_output_gathered: + if ( + args.return_layernorm_output + and not args.return_layernorm_output_gathered + and args.grad_ln_out is not None + ): dgrad = dgrad + args.grad_ln_out.view_as(dgrad) # Norm gradient @@ -2711,6 +2717,9 @@ def forward( dgrad_use_split_accumulator = _2X_ACC_DGRAD wgrad_use_split_accumulator = _2X_ACC_WGRAD + recipe_float8_block_scaling = False + recipe_custom = False + recipe_dbias_dact_fusion = False if self.fp8: _recipe = FP8GlobalStateManager.get_fp8_recipe() backward_override = _recipe.backward_override @@ -2718,6 +2727,9 @@ def forward( dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator if hasattr(_recipe, "fp8_gemm_wgrad"): wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + recipe_float8_block_scaling = _recipe.float8_block_scaling() + recipe_custom = _recipe.custom() + recipe_dbias_dact_fusion = _recipe_has_dbias_dact_fusion(_recipe) else: backward_override = None @@ -2799,6 +2811,9 @@ def forward( backward_override=backward_override, dgrad_use_split_accumulator=dgrad_use_split_accumulator, wgrad_use_split_accumulator=wgrad_use_split_accumulator, + recipe_float8_block_scaling=recipe_float8_block_scaling, + recipe_custom=recipe_custom, + recipe_dbias_dact_fusion=recipe_dbias_dact_fusion, debug=debug, # weight-workspace caching is_first_microbatch=is_first_microbatch, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 43e252f5f8..a3c4589a4c 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -34,6 +34,9 @@ from ._common import ( can_reconstruct_wgrad_input_from_original, check_fp8_reduce_and_update, + fake_workspace_valid, + sp_inp_leading, + sp_out_leading, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -318,40 +321,6 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: ) # pylint: disable=unbalanced-tuple-unpacking -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], Optional[Dict]]: @@ -870,7 +839,7 @@ def _linear_forward_fake( else: weightmat_is_storage = True workspace = args.weight_workspace - if workspace is not None and not _fake_workspace_valid(workspace, weight_quantizer): + 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: @@ -903,7 +872,7 @@ def _linear_forward_fake( # ------------------------------------------------------ # A rank-1 input is viewed to (1, in_features), so the output leads with 1. inp_leading = inp.shape[0] if len(inp.shape) > 1 else 1 - out_leading = _out_leading_from_inp(inp_leading, args) + out_leading = sp_out_leading(inp_leading, args) out = TensorSpec( shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), dtype=activation_dtype, @@ -1184,7 +1153,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). if bwd_args.inp_shape is None: in_features = saved_weight.shape[-1] - inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args) + inp_leading = sp_inp_leading(grad_output.shape[0], bwd_args) bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) # Configure Userbuffers communication (comm+GEMM overlap) @@ -1742,7 +1711,7 @@ def _linear_backward_fake( if args.requires_dgrad: # Input shape rederived from grad_output + SP config (inp_shape is not # stored: torch.Size with SymInt cannot cross in OpaqueValueBundle). - dgrad_leading = _inp_leading_from_out(args.grad_output.shape[0], args) + dgrad_leading = sp_inp_leading(args.grad_output.shape[0], args) # Under UB reduce-scatter or bulk-wgrad overlap the returned dgrad is a # plain tensor; the quantizer only feeds the comm buffer. dgrad_quantizer = (