[PyTorch] [torch.compile] Prepare the custom-op framework and LayerNorm modules for compile wiring

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

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())

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,