From 92fc8ff2200d2475c5704542a284ff2291b78922 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:58:30 +0200 Subject: [PATCH] [PyTorch] [torch.compile] Prepare LayerNormLinear and LayerNormMLP for torch.compile Restructure the two modules the same way Linear was restructured in #2967, so the follow-up custom-op registration is a pure wiring change: - LayerNormLinearFwdArgs / LayerNormLinearBwdArgs and LayerNormMLPFwdArgs / LayerNormMLPBwdArgs dataclasses replace the positional non_tensor_args tuple and the ad-hoc ctx. plumbing. - The autograd Functions shrink to thin wrappers around module-level _layernorm_{linear,mlp}_forward_impl / _setup_ctx / _backward_impl that take one dataclass argument. Saved tensors that alias forward inputs or outputs are deduplicated by name and rebuilt in setup_ctx. - LayerNormMLP keeps activation checkpointing: the forward saves only its inputs, and _layernorm_mlp_recompute re-runs the forward impl in backward with recompute_for_bwd set, filling the backward args from the recomputed state. - The weight workspaces move from positional autograd inputs into the args bag and are returned as extra outputs; backward-only recipe knobs (split accumulators, backward_override, debug UB flags, delayed wgrad_store) are resolved in Module.forward. - check_fp8_reduce_and_update moves to module/_common.py and is shared with Linear. No functional change. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/_common.py | 16 + .../pytorch/module/layernorm_linear.py | 2586 +++++++----- .../pytorch/module/layernorm_mlp.py | 3666 +++++++++-------- transformer_engine/pytorch/module/linear.py | 13 +- 4 files changed, 3556 insertions(+), 2725 deletions(-) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 908e88f30d..98262bd99e 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -12,7 +12,9 @@ from .. import cpp_extensions as tex from ..constants import TE_DType +from ..distributed import in_fp8_activation_recompute_phase from ..export import is_in_onnx_export_mode +from ..quantization import FP8GlobalStateManager from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -320,3 +322,17 @@ def assert_empty(self): assert self.enabled is True, "delay_wgrad_compute is not enabled" rank = torch.distributed.get_rank() assert self.context.empty(), f"Queue is not empty. rank {rank}" + + +def check_fp8_reduce_and_update(restore_first_module: bool = False) -> bool: + """Whether this module's backward should reduce and update the FP8 scaling factors. + + Consumes the "first FP8 module" flag, restored when the forward is a + recomputation so the flag survives for the real forward's owner. + """ + qstate = FP8GlobalStateManager.quantization_state + first_fp8_module = qstate.is_first_fp8_module + result = FP8GlobalStateManager.is_first_fp8_module() + if restore_first_module or in_fp8_activation_recompute_phase(): + qstate.is_first_fp8_module = first_fp8_module + return result diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 561e813348..b0af7f05b2 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -6,7 +6,8 @@ import os import warnings import weakref -from typing import Callable, Dict, Optional, Tuple, Union, List +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -41,7 +42,6 @@ init_method_constant, nvtx_range_pop, nvtx_range_push, - requires_grad, needs_quantized_gemm, get_nvtx_range_context, ) @@ -52,7 +52,6 @@ symmetric_all_reduce, reduce_scatter_along_first_dim, gather_along_first_dim, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _fsdp_gather_tensors, ) @@ -67,6 +66,7 @@ from ..graph import is_graph_capturing from ._common import ( apply_normalization, + check_fp8_reduce_and_update, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -79,6 +79,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorOrQuantized from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.hybrid_tensor import HybridQuantizer @@ -100,1107 +101,1451 @@ __all__ = ["LayerNormLinear"] -class _LayerNormLinear(torch.autograd.Function): - """LayerNormLinear semi-top level module - Calls custom cuda extensions. - """ +@dataclass(slots=True) +class LayerNormLinearFwdArgs: + """Single-argument bag for the forward path of :class:`_LayerNormLinear`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + inp: torch.Tensor + ln_weight: torch.Tensor + ln_bias: Optional[torch.Tensor] + weight: TensorOrQuantized + bias: Optional[torch.Tensor] + + # --- Non-differentiable cached tensors --- + weight_workspace: Optional[TensorOrQuantized] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + ln_weight_requires_grad: bool + ln_bias_requires_grad: bool + weight_requires_grad: bool + bias_requires_grad: bool + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] + weight_quantizer: Optional[Quantizer] + output_quantizer: Optional[Quantizer] + grad_input_quantizer: Optional[Quantizer] + grad_weight_quantizer: Optional[Quantizer] + grad_output_quantizer: Optional[Quantizer] + + # --- Normalization --- + eps: float + normalization: str + zero_centered_gamma: bool + fwd_ln_sm_margin: int + bwd_ln_sm_margin: int + return_layernorm_output: bool + return_layernorm_output_gathered: bool + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + skip_fp8_weight_update: Optional[torch.Tensor] + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] + tp_group: Optional[dist_group_type] + tp_size: int + tensor_parallel: bool + sequence_parallel: bool + symmetric_ar_type: Optional[str] + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] + ub_overlap_ag_fprop: bool + ub_overlap_rs_fprop: bool + ub_overlap_ag_dgrad: bool + ub_overlap_rs_dgrad: bool + ub_bulk_dgrad: bool + ub_bulk_wgrad: bool + + # --- FSDP --- + fsdp_group: Optional[Any] + is_fsdp2: bool + + # --- Weight-grad scheduling --- + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + + # --- Misc --- + cpu_offloading: bool + is_grad_enabled: bool + + def any_requires_grad(self) -> bool: + """Whether any differentiable input requires a gradient.""" + return any( + ( + self.input_requires_grad, + self.ln_weight_requires_grad, + self.ln_bias_requires_grad, + self.weight_requires_grad, + self.bias_requires_grad, + ) + ) - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: Union[torch.Tensor, None], - weight: torch.Tensor, - weight_workspace: Optional[torch.Tensor], - bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - - # Reduce number of arguments to autograd function in order - # to reduce CPU overhead due to pytorch arg checking. + +@dataclass(slots=True) +class LayerNormLinearBwdArgs: + """Single-argument bag for the backward path of :class:`_LayerNormLinear`.""" + + # --- Incoming gradients (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + grad_ln_out: Optional[torch.Tensor] = None + + # --- Saved / restored tensors (populated at backward entry) --- + inputmat: Optional[torch.Tensor] = None + weight_fp8: Optional[TensorOrQuantized] = None + saved_weight: Optional[TensorOrQuantized] = None + bias: Optional[torch.Tensor] = None + ln_weight: Optional[torch.Tensor] = None + ln_out: Optional[TensorOrQuantized] = None + mu: Optional[torch.Tensor] = None + rsigma: Optional[torch.Tensor] = None + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] = None + weight_quantizer: Optional[Quantizer] = None + grad_input_quantizer: Optional[Quantizer] = None + grad_weight_quantizer: Optional[Quantizer] = None + grad_output_quantizer: Optional[Quantizer] = None + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + requires_wgrad: bool = False + ln_out_needs_gather: bool = False + inp_shape: Optional[torch.Size] = None + + # --- Normalization --- + normalization: str = "LayerNorm" + zero_centered_gamma: bool = False + bwd_ln_sm_margin: int = 0 + return_layernorm_output: bool = False + return_layernorm_output_gathered: bool = False + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + backward_override: Optional[str] = None + is_weight_param_quantized: bool = False + debug: bool = False + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] = None + tp_group: Optional[dist_group_type] = None + tp_size: int = 1 + tensor_parallel: bool = False + sequence_parallel: bool = False + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] = None + ub_overlap_ag: bool = False + ub_overlap_rs_dgrad: bool = False + ub_bulk_dgrad: bool = False + ub_bulk_wgrad: bool = False + + # --- FSDP --- + fsdp_group: Optional[Any] = None + fsdp_shapes: Any = None + is_fsdp2: bool = False + + # --- Weight-grad scheduling / accumulation --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + origin_weight_ref: Optional[Any] = None + origin_weight_overwrites_main_grad: bool = False + main_grad_func: Optional[Callable[[], torch.Tensor]] = None + + # --- FP8 reduce-and-update bookkeeping --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + + # --- Per-backward scratch state (populated inside the backward impl) --- + ub_obj_gradout: Optional[Any] = None + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" ( - eps, - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - fuse_wgrad_accumulation, + self.inputmat, + self.weight_fp8, + self.saved_weight, + self.bias, + self.ln_weight, + self.ln_out, + self.mu, + self.rsigma, + ) = restore_from_func_ctx( + ctx + ) # pylint: disable=unbalanced-tuple-unpacking + + +def _layernorm_linear_forward_impl( + args: LayerNormLinearFwdArgs, +) -> Tuple[ + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[Tuple], + Optional[Dict], +]: + """Forward implementation for the layernorm-linear layer. + + Returns ``(out, ln_out_return, new_weight_workspace, + tensors_to_save_from_forward, ctx_attrs)``. ``new_weight_workspace`` is + the freshly produced FP8 weight workspace (returned alongside ``out`` so + the caller can refresh its cache). The last two are ``None`` when + gradients are disabled. + """ + inp = args.inp + ln_weight = args.ln_weight + ln_bias = args.ln_bias + weight = args.weight + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + is_first_microbatch = args.is_first_microbatch + fp8 = args.fp8 + debug = args.debug + cpu_offloading = args.cpu_offloading + tp_group = args.tp_group + tp_size = args.tp_size + sequence_parallel = args.sequence_parallel + activation_dtype = args.activation_dtype + parallel_mode = args.parallel_mode + is_grad_enabled = args.is_grad_enabled + return_layernorm_output = args.return_layernorm_output + return_layernorm_output_gathered = args.return_layernorm_output_gathered + backward_override = args.backward_override + ub_name = args.ub_name + ub_overlap_ag_fprop = args.ub_overlap_ag_fprop + ub_overlap_rs_fprop = args.ub_overlap_rs_fprop + fsdp_group = args.fsdp_group + is_fsdp2 = args.is_fsdp2 + weight_requires_grad = args.weight_requires_grad + + # NVTX label for profiling + nvtx_label = "transformer_engine._LayerNormLinear.forward" + if ub_name is not None: + nvtx_label = f"{nvtx_label}.{ub_name}" + + with_input_all_gather = parallel_mode == "column" and sequence_parallel + + # Make sure input dimensions are compatible + out_features, in_features = weight.shape + inp_shape = inp.shape + assert inp_shape[-1] == in_features, "GEMM not possible" + inp = inp.view((-1, in_features)) + inputmat = inp + if fp8: + assert_dim_for_fp8_exec(inputmat, weight) + + # Cast for native AMP + nvtx_range_push(f"{nvtx_label}.norm_input_cast") + inputmat = cast_if_needed(inputmat, activation_dtype) + ln_weight_cast = cast_if_needed(ln_weight, activation_dtype) + if ln_bias is not None: + ln_bias = cast_if_needed(ln_bias, activation_dtype) + nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + + if is_cpu_offload_enabled(): + start_offload(inputmat) + + tp_world_size = get_distributed_world_size(tp_group) + + backward_needs_input = is_grad_enabled and weight_requires_grad + + # Configure Userbuffers communication (comm+GEMM overlap) + ub_obj = None + ub_type = None + ub_overlap_ag_fprop = ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output + if ub_overlap_rs_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.RS + elif ub_overlap_ag_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.AG + + # Configure quantizer for norm output + if fp8: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and backward_override is None, + ) + if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather(): + # All-gather is not supported with FP8 column-wise data + input_quantizer.set_usage(columnwise=False) + # Amax reduction group for the input quantizer (column-parallel sequence parallel) + set_quantizer_amax_reduction_group( input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - parallel_mode, - return_layernorm_output, - return_layernorm_output_gathered, - is_grad_enabled, - fwd_ln_sm_margin, - bwd_ln_sm_margin, - zero_centered_gamma, - normalization, - ub_overlap_ag_fprop, - ub_overlap_rs_fprop, - ub_overlap_ag_dgrad, - ub_overlap_rs_dgrad, - ub_bulk_wgrad, - ub_bulk_dgrad, - ub_name, - fsdp_group, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - debug, - is_fsdp2, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + tp_group if (sequence_parallel and parallel_mode == "column") else None, + ) + + # Avoid quantized norm kernel if norm output will be returned + # or if a gather of ln_out must be in high precision. + custom = is_custom(input_quantizer) + hybrid = isinstance(input_quantizer, HybridQuantizer) + identity = isinstance(input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + fp8 + and not debug + and not return_layernorm_output + and not return_layernorm_output_gathered + and backward_override is None + and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() + and not hybrid + and not identity + ) + + # Apply normalization + nvtx_range_push(f"{nvtx_label}.norm") + ln_out, mu, rsigma = apply_normalization( + inputmat, + None, # ln_out + ln_weight_cast, + ln_bias, + args.eps, + input_quantizer if with_quantized_norm else None, + inputmat.dtype, + args.normalization, + args.fwd_ln_sm_margin, + args.zero_centered_gamma, + ) + nvtx_range_pop(f"{nvtx_label}.norm") + + # Store unquantized layer norm output if we need to return it + ln_out_return = None + if return_layernorm_output or return_layernorm_output_gathered: + ln_out_return = ln_out + ln_out_hp = ln_out if backward_override == "high_precision" else None + + # ------------------------------------------------------ + # Prepare GEMM input tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm") + ln_out_total = None + if with_input_all_gather: + if return_layernorm_output_gathered: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) + ln_out_return = ln_out_total + if fp8 or debug: + ln_out = input_quantizer(ln_out) + input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = input_quantizer(ln_out_total) else: - backward_override = None + quantizer = None + if fp8 or debug: + quantizer = input_quantizer + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: + ln_out = quantizer(ln_out) + quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag_fprop: # Initialize Userbuffers all-gather + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj, + ln_out, + quantizer, + tp_group, + ) + else: # Perform NCCL all-gather + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=quantizer, + ) + else: + if (fp8 or debug) and not with_quantized_norm: + ln_out = input_quantizer(ln_out) + ln_out_total = ln_out + nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm") + # ------------------------------------------------------ + # GEMM input tensor is ready... + # ------------------------------------------------------ + + # ------------------------------------------------------ + # Prepare weight tensor + # ------------------------------------------------------ + is_dist_weight = is_distributed_weight(args.weight) + if is_dist_weight: + weight = materialize_weight_for_forward(weight)[0] + out_features = weight.shape[0] + new_weight_workspace = None + weightmat = weight + is_weight_param_quantized = False + if fp8 or debug: + is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) + + # Configure quantizer + # If weight is already quantized, weight._quantizer is its true quantizer. + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if is_weight_param_quantized and not debug: + weight_quantizer = weight._quantizer + elif weight_quantizer is not None: + # FSDP2: Skip columnwise/transpose creation during forward + # to avoid accumulating caches across layers. Backward's + # FSDP2 all-gather will recreate them. (Issue #2681) + weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not is_fsdp2 and backward_override is None, + ) - # NVTX label for profiling - nvtx_label = "transformer_engine._LayerNormLinear.forward" - if ub_name is not None: - nvtx_label = f"{nvtx_label}.{ub_name}" + # Get quantized weight + update_ws = is_first_microbatch is None or is_first_microbatch + weightmat, new_weight_workspace = quantize_weight( + tensor=weight, + quantizer=weight_quantizer, + workspace=args.weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, + ) - with_input_all_gather = parallel_mode == "column" and sequence_parallel + weightmat.update_usage(rowwise_usage=True) + + else: + weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP + # ------------------------------------------------------ + # Weight tensor is ready for GEMM... + # ------------------------------------------------------ + + # Cast bias to expected dtype + bias_dtype = activation_dtype + if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: + # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 + bias_dtype = torch.bfloat16 + bias_cast = cast_if_needed(bias, bias_dtype) if bias is not None else bias + + # Calibrate quantizers if needed + if not fp8 and args.fp8_calibration: + if input_quantizer is not None: + input_quantizer.calibrate(ln_out_total) + if weight_quantizer is not None: + weight_quantizer.calibrate(weight) - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - inp_shape = inp.shape - inp_requires_grad = inp.requires_grad - assert inp_shape[-1] == in_features, "GEMM not possible" - inp = inp.view((-1, in_features)) - inputmat = inp - if fp8: - assert_dim_for_fp8_exec(inputmat, weight) + # Choose whether to use GEMM kernel with split accumulator + use_split_accumulator = _2X_ACC_FPROP + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if hasattr(recipe, "fp8_gemm_fprop"): + use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + + # Configure output quantizer + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # Output buffer for Userbuffers reduce-scatter + reduce_scatter_out = None + if ub_overlap_rs_fprop: + out_shape = list(inp_shape) + out_shape[0] //= tp_world_size + out_shape[-1] = out_features + reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) + + # ------------------------------------------------------ + # Forward GEMM + # Note: y = x * w^T + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.gemm") + gemm_out, *_, reduce_scatter_out = general_gemm( + weightmat, + ln_out_total, + quantization_params=output_quantizer, + out_dtype=activation_dtype, + bias=bias_cast, + use_split_accumulator=use_split_accumulator, + ub=ub_obj, + ub_type=ub_type, + extra_output=reduce_scatter_out, + ) + nvtx_range_pop(f"{nvtx_label}.gemm") + # ------------------------------------------------------ + # Finished forward GEMM... + # ------------------------------------------------------ + + # Deallocate GEMM input tensor if no longer needed + if not weight_requires_grad and not return_layernorm_output: + clear_tensor_data(ln_out, ln_out_total) + ln_out = ln_out_total = None + elif with_input_all_gather and not return_layernorm_output_gathered: + # ln_out_total aliases ln_out for the cuBLASMp backend; skip the + # deallocation to avoid corrupting the backward-saved tensor. + if ln_out_total is not ln_out: + clear_tensor_data(ln_out_total) + ln_out_total = None - # Cast for native AMP - nvtx_range_push(f"{nvtx_label}.norm_input_cast") - inputmat = cast_if_needed(inputmat, activation_dtype) - ln_weight = cast_if_needed(ln_weight, activation_dtype) - if ln_bias is not None: - ln_bias = cast_if_needed(ln_bias, activation_dtype) - nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + # ------------------------------------------------------ + # Prepare output tensor + # Note: Perform tensor-parallel communication + # ------------------------------------------------------ + out = None + if ub_overlap_rs_fprop: + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out + elif parallel_mode == "row" and tp_size > 1: + nvtx_range_push(f"{nvtx_label}.row_parallel_comm") + out = gemm_out + if sequence_parallel: + out, _ = reduce_scatter_along_first_dim(out, tp_group) + elif args.tensor_parallel: + if args.symmetric_ar_type is not None: + out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=args.symmetric_ar_type) + else: + out, _ = allreduce(out, tp_group) + nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") + else: + out = gemm_out + out = out.view(-1, *inp_shape[1:-1], out_features) + # ------------------------------------------------------ + # Output tensor is ready to return... + # ------------------------------------------------------ + + # Prepare backward state + tensors_to_save_from_forward = None + ctx_attrs = None + + if is_grad_enabled: + ln_out_to_save = ln_out + if backward_override == "high_precision": + ln_out_to_save = ln_out_hp + ln_out_needs_gather = ( + weight_requires_grad and parallel_mode == "column" and sequence_parallel + ) - if is_cpu_offload_enabled(): - start_offload(inputmat) + # Input with column-wise usage is needed for wgrad GEMM. + if backward_needs_input and backward_override is None: + if isinstance(ln_out, QuantizedTensorStorage): + # For sequence parallel in vanilla FP8, rowwise data is + # to gather the input. For MXFP8, columnwise only data + # can be allgathered. + if ( + isinstance(ln_out, (MXFP8TensorStorage, Float8BlockwiseQTensorStorage)) + or not ln_out_needs_gather + ): + ln_out.update_usage(rowwise_usage=False) + + if cpu_offloading: + mark_activation_offload(inputmat, mu, rsigma, ln_out_to_save) + + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + nvtx_range_push(f"{nvtx_label}.fsdp_scatter") + fsdp_shapes = _fsdp_scatter_tensors( + fsdp_group, + mu, + rsigma, + weightmat if fp8 and not is_weight_param_quantized else None, + ln_out_to_save if weight_requires_grad else None, + ) + nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - tp_world_size = get_distributed_world_size(tp_group) + if cpu_offloading: + mark_not_offload( + weightmat, + weight, + bias_cast, + ln_weight_cast, + ln_bias, + ) + + # FSDP2: Don't save FP8 workspace for non-quantized weights. + # Backward will re-quantize from FSDP2 all-gathered weight. + # (Issue #2681) + wt_save = weightmat + if is_fsdp2 and weightmat is not weight: + wt_save = None + # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; + # backward re-gathers from the saved (sharded) weight and re-quantizes. + if is_dist_weight: + wt_save = None + + # Dedup save slots that alias forward inputs or other outputs; + # ``_layernorm_linear_setup_ctx`` rebuilds the refs. + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_weight_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None + saved_tensor_aliases = ( + "inp" if inputmat is inp else None, + wt_alias, + "weight", # ``saved_weight`` slot is always the weight parameter + "bias" if bias_cast is not None and bias_cast is bias else None, + "ln_weight" if ln_weight_cast is ln_weight else None, + ( + "ln_out" + if return_layernorm_output + and ln_out_to_save is not None + and ln_out_to_save is ln_out_return + else None + ), + None, + None, + ) + tensors_to_save_from_forward = ( + None if saved_tensor_aliases[0] is not None else inputmat, + None if saved_tensor_aliases[1] is not None else wt_save, + None, + None if saved_tensor_aliases[3] is not None else bias_cast, + None if saved_tensor_aliases[4] is not None else ln_weight_cast, + None if saved_tensor_aliases[5] is not None else ln_out_to_save, + mu, + rsigma, + ) - weight_requires_grad = weight.requires_grad - backward_needs_input = is_grad_enabled and weight_requires_grad + ctx_attrs = { + "fsdp_shapes": fsdp_shapes, + "saved_tensor_aliases": saved_tensor_aliases, + "is_weight_param_quantized": is_weight_param_quantized, + "ln_out_needs_gather": ln_out_needs_gather, + } + + ln_out_for_return = None + if return_layernorm_output: + if return_layernorm_output_gathered: + shape = list(inp_shape) + shape[0] *= tp_size if with_input_all_gather else 1 + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + return out, ln_out_for_return, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _layernorm_linear_setup_ctx( + bwd_args: LayerNormLinearBwdArgs, + fwd_args: LayerNormLinearFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state. + + Returns the tensors that should be passed through ``prepare_for_saving`` + by the caller. ``fwd_outputs`` is ``(out, ln_out_return, + new_weight_workspace)``; the last two rebuild the deduped save slots. + """ + inp = fwd_args.inp + weight = fwd_args.weight + bias = fwd_args.bias + ln_weight = fwd_args.ln_weight + + backward_override = fwd_args.backward_override + fp8 = fwd_args.fp8 + debug = fwd_args.debug + fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation + is_weight_param_quantized = ctx_attrs["is_weight_param_quantized"] + + # Quantizers + bwd_args.input_quantizer = fwd_args.input_quantizer + bwd_args.weight_quantizer = ( + weight._quantizer + if (is_weight_param_quantized and not debug and isinstance(weight, QuantizedTensorStorage)) + else fwd_args.weight_quantizer + ) + bwd_args.grad_input_quantizer = fwd_args.grad_input_quantizer + bwd_args.grad_weight_quantizer = fwd_args.grad_weight_quantizer + bwd_args.grad_output_quantizer = fwd_args.grad_output_quantizer + + # Differentiability summary + bwd_args.use_bias = bias is not None + bwd_args.requires_dgrad = fwd_args.input_requires_grad + bwd_args.requires_wgrad = fwd_args.weight_requires_grad + bwd_args.ln_out_needs_gather = ctx_attrs["ln_out_needs_gather"] + bwd_args.inp_shape = inp.shape + + # Normalization + bwd_args.normalization = fwd_args.normalization + bwd_args.zero_centered_gamma = fwd_args.zero_centered_gamma + bwd_args.bwd_ln_sm_margin = fwd_args.bwd_ln_sm_margin + bwd_args.return_layernorm_output = fwd_args.return_layernorm_output + bwd_args.return_layernorm_output_gathered = fwd_args.return_layernorm_output_gathered + + # Numerical / dtype config + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fp8 + 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 = backward_override + bwd_args.is_weight_param_quantized = is_weight_param_quantized + bwd_args.debug = debug + + # Tensor / sequence parallelism + bwd_args.parallel_mode = fwd_args.parallel_mode + bwd_args.tp_group = fwd_args.tp_group + bwd_args.tp_size = fwd_args.tp_size + bwd_args.tensor_parallel = fwd_args.tensor_parallel + bwd_args.sequence_parallel = fwd_args.sequence_parallel + + # Userbuffers + bwd_args.ub_name = fwd_args.ub_name + bwd_args.ub_overlap_ag = fwd_args.ub_overlap_ag_dgrad + bwd_args.ub_overlap_rs_dgrad = fwd_args.ub_overlap_rs_dgrad + bwd_args.ub_bulk_dgrad = fwd_args.ub_bulk_dgrad + bwd_args.ub_bulk_wgrad = fwd_args.ub_bulk_wgrad + + # FSDP + bwd_args.fsdp_group = fwd_args.fsdp_group + bwd_args.fsdp_shapes = ctx_attrs["fsdp_shapes"] + bwd_args.is_fsdp2 = fwd_args.is_fsdp2 + + # Weight-grad scheduling / accumulation + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + if fuse_wgrad_accumulation and fwd_args.weight_requires_grad: + # Keep weakref to weight to preserve attributes like main_grad + # when we need to modify the weight python object + bwd_args.origin_weight_ref = weakref.ref(weight) + bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) + # MCore FSDP creates main_grad lazily before backward, so don't touch it here + if hasattr(weight, "__fsdp_param__"): + bwd_args.main_grad_func = weight.get_main_grad + elif is_distributed_weight(weight): + bwd_args.main_grad_func = weight.grad_buffer + else: + bwd_args.main_grad_func = lambda: weight.main_grad + + # Misc + bwd_args.cpu_offloading = fwd_args.cpu_offloading + + if backward_override is not None: + bwd_args.fp8 = False + bwd_args.debug = False + bwd_args.ub_overlap_ag = False + bwd_args.ub_overlap_rs_dgrad = False + bwd_args.ub_bulk_dgrad = False + bwd_args.ub_bulk_wgrad = False + bwd_args.grad_input_quantizer = None + bwd_args.grad_weight_quantizer = None + bwd_args.grad_output_quantizer = None + + ( + saved_inputmat, + wt_save, + saved_weight, + saved_bias, + saved_ln_weight, + saved_ln_out, + mu, + rsigma, + ) = tensors_to_save_from_forward + ( + inputmat_alias, + wt_save_alias, + saved_weight_alias, + bias_alias, + ln_weight_alias, + ln_out_alias, + _, + _, + ) = ctx_attrs["saved_tensor_aliases"] + in_features = inp.shape[-1] + if inputmat_alias == "inp": + saved_inputmat = inp.view((-1, in_features)) + if wt_save_alias == "weight": + wt_save = weight + elif wt_save_alias == "new_weight_workspace": + wt_save = fwd_outputs[2] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace + if saved_weight_alias == "weight": + saved_weight = weight + if bias_alias == "bias": + saved_bias = bias + if ln_weight_alias == "ln_weight": + saved_ln_weight = ln_weight + if ln_out_alias == "ln_out": + saved_ln_out = fwd_outputs[1].view((-1, in_features)) + return ( + saved_inputmat, + wt_save, + saved_weight, + saved_bias, + saved_ln_weight, + saved_ln_out, + mu, + rsigma, + ) + + +def _layernorm_linear_backward_impl( + args: LayerNormLinearBwdArgs, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward implementation for the layernorm-linear layer. + + Caller must have populated ``args.grad_output`` / ``args.grad_ln_out`` and + run ``args.setup_saved_tensors(ctx)`` before invocation. Returns + ``(dgrad, dgamma, dbeta, wgrad, grad_bias)``. + """ + grad_output = args.grad_output + assert grad_output is not None + + # NVTX label for profiling + nvtx_label = "transformer_engine._LayerNormLinear.backward" + if args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{args.ub_name}" + + with get_nvtx_range_context("_LayerNormLinear_backward"): + inputmat = args.inputmat + weight = args.weight_fp8 + saved_weight = args.saved_weight + bias = args.bias + ln_weight = args.ln_weight + ln_out = args.ln_out + mu = args.mu + rsigma = args.rsigma + + is_dist_weight = is_distributed_weight(saved_weight) + if is_dist_weight: + weight = materialize_weight_for_backward(saved_weight)[0] + # Restore from weakref to get original weight python object + # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) + # Only needed when fuse_wgrad_accumulation is enabled. + origin_weight = None + origin_weight_overwrites_main_grad = args.origin_weight_overwrites_main_grad + main_grad = None + if args.fuse_wgrad_accumulation and args.requires_wgrad: + origin_weight_ref = args.origin_weight_ref + args.origin_weight_ref = None + origin_weight = origin_weight_ref() if origin_weight_ref is not None else None + assert ( + origin_weight is not None + ), "weight was removed while fuse_wgrad_accumulation=True" + # Since main_grad can be modified inplace, it should not be a part of saved_tensors + main_grad = args.main_grad_func() if weight is not None else None + if main_grad is not None and not is_dist_weight: + origin_weight.main_grad = main_grad + + # Gather intermediate/activation tensors if needed + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + nvtx_range_push(f"{nvtx_label}.fsdp_gather") + _fsdp_gather_tensors( + args.fsdp_group, + args.fsdp_shapes, + mu, + rsigma, + weight if args.fp8 and not args.is_weight_param_quantized else None, + ln_out, + ) + nvtx_range_pop(f"{nvtx_label}.fsdp_gather") # Configure Userbuffers communication (comm+GEMM overlap) - if debug: # turn off userbuffers in debug mode - ub_overlap_ag_fprop = False - ub_overlap_rs_fprop = False - ub_overlap_ag_dgrad = False - ub_overlap_rs_dgrad = False - ub_bulk_wgrad = False - ub_bulk_dgrad = False - ub_obj = None - ub_type = None - ub_overlap_ag_fprop = ( - ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output - ) - if ub_overlap_rs_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.RS - elif ub_overlap_ag_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.AG - - # Configure quantizer for norm output - if fp8: - if input_quantizer is None: - raise ValueError("Missing quantizer for input tensor") - input_quantizer.set_usage( - rowwise=True, - columnwise=backward_needs_input and backward_override is None, - ) - if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather(): - # All-gather is not supported with FP8 column-wise data - input_quantizer.set_usage(columnwise=False) - # Amax reduction group for the input quantizer (column-parallel sequence parallel) + args.ub_obj_gradout = None + ub_obj_dgrad = None + ub_obj_wgrad = None + ub_type_dgrad = None + ub_type_wgrad = None + dgrad_shape = [reduce(multiply_op, args.inp_shape[:-1]), args.inp_shape[-1]] + if args.ub_overlap_ag: + # Overlap grad_output all-gather with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.AG + elif args.ub_overlap_rs_dgrad: + # Overlap dgrad reduce-scatter with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.RS + else: + if args.ub_bulk_dgrad: + # Overlap inputmat all-gather with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.AG + if args.ub_bulk_wgrad: + # Overlap dgrad reduce-scatter with wgrad compute + ub_obj_wgrad = get_ub(args.ub_name + "_wgrad", args.fp8) + ub_type_wgrad = tex.CommOverlapType.RS + + # -------------------------------------------------- + # Prepare grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # -------------------------------------------------- + + # Configure quantizer for grad output tensor + # Note: dgrad GEMM requires row-wise usage, wgrad GEMM + # requires column-wise usage + if args.grad_output_quantizer is not None: + quantizer = args.grad_output_quantizer + quantizer.set_usage(rowwise=True, columnwise=True) + if args.ub_overlap_ag: + # Userbuffers only supports communication for one + # tensor usage at a time. Configure quantizer with + # usage for only dgrad GEMM. + quantizer.set_usage(columnwise=False) + # Amax reduction group for grad output (row-parallel sequence parallel) set_quantizer_amax_reduction_group( - input_quantizer, - tp_group if (sequence_parallel and parallel_mode == "column") else None, + quantizer, + ( + args.tp_group + if (args.sequence_parallel and args.parallel_mode == "row") + else None + ), ) - # Avoid quantized norm kernel if norm output will be returned - # or if a gather of ln_out must be in high precision. - custom = is_custom(input_quantizer) - hybrid = isinstance(input_quantizer, HybridQuantizer) - identity = isinstance(input_quantizer, IdentityQuantizer) - with_quantized_norm = ( - fp8 - and not debug - and not return_layernorm_output - and not return_layernorm_output_gathered - and backward_override is None - and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() - and not hybrid - and not identity - ) - - # Apply normalization - nvtx_range_push(f"{nvtx_label}.norm") - ln_out, mu, rsigma = apply_normalization( - inputmat, - None, # ln_out - ln_weight, - ln_bias, - eps, - input_quantizer if with_quantized_norm else None, - inputmat.dtype, - normalization, - fwd_ln_sm_margin, - zero_centered_gamma, + # Prepare grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") + ( + grad_output, + grad_bias, + ) = TransformerEngineBaseModule.grad_output_preprocess( + args, + grad_output, + args.parallel_mode == "row", + args.grad_output_quantizer, ) - nvtx_range_pop(f"{nvtx_label}.norm") + nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - # Store unquantized layer norm output if we need to return it - ln_out_return = None - if return_layernorm_output or return_layernorm_output_gathered: - ln_out_return = ln_out - ln_out_hp = ln_out if backward_override == "high_precision" else None + # -------------------------------------------------- + # Grad output tensor is ready for computing grad input... + # -------------------------------------------------- - # ------------------------------------------------------ + # -------------------------------------------------- # Prepare GEMM input tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm") + # Note: Input tensor is needed for wgrad GEMM. + # Tensor-parallel communication is overlapped with dgrad + # GEMM. + # -------------------------------------------------- ln_out_total = None - if with_input_all_gather: - if return_layernorm_output_gathered: - # Perform all-gather in high precision if gathered - # norm output will be returned - ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) - ln_out_return = ln_out_total - if fp8 or debug: - ln_out = input_quantizer(ln_out) - input_quantizer.set_usage(rowwise=True, columnwise=False) - ln_out_total = input_quantizer(ln_out_total) + ln_out_total_work = None + if args.backward_override == "dequantized": + if isinstance(ln_out, QuantizedTensorStorage): + ln_out = ln_out.dequantize(dtype=args.activation_dtype) else: - quantizer = None - if fp8 or debug: - quantizer = input_quantizer - # custom recipe doesn't need to support quantized AG - if not with_quantized_norm and not custom: - ln_out = quantizer(ln_out) - quantizer.set_usage(rowwise=True, columnwise=False) - if ub_overlap_ag_fprop: # Initialize Userbuffers all-gather - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj, - ln_out, - quantizer, - tp_group, - ) - else: # Perform NCCL all-gather - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=quantizer, - ) - else: - if (fp8 or debug) and not with_quantized_norm: - ln_out = input_quantizer(ln_out) - ln_out_total = ln_out - nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm") - # ------------------------------------------------------ - # GEMM input tensor is ready... - # ------------------------------------------------------ - - # ------------------------------------------------------ - # Prepare weight tensor - # ------------------------------------------------------ - origin_weight = weight - is_dist_weight = is_distributed_weight(origin_weight) - if is_dist_weight: - weight = materialize_weight_for_forward(weight)[0] - out_features = weight.shape[0] - new_weight_workspace = None - weightmat = weight - is_weight_param_quantized = False - if fp8 or debug: - is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) - - # Configure quantizer - # If weight is already quantized, weight._quantizer is its true quantizer. - # for debug mode we create quantizer every iteration, thus we need to set the quantizer states - if is_weight_param_quantized and not debug: - weight_quantizer = weight._quantizer - elif weight_quantizer is not None: - # FSDP2: Skip columnwise/transpose creation during forward - # to avoid accumulating caches across layers. Backward's - # FSDP2 all-gather will recreate them. (Issue #2681) - weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not is_fsdp2 and backward_override is None, + ln_out = cast_if_needed(ln_out, args.activation_dtype) + if args.ln_out_needs_gather: + quantizer = None + if args.input_quantizer is not None and args.fp8: + quantizer = args.input_quantizer + set_quantizer_usage_for_wgrad_all_gather(quantizer) + if args.ub_bulk_dgrad: + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_dgrad, + ln_out, + quantizer, + args.tp_group, ) - - # Get quantized weight - update_ws = is_first_microbatch is None or is_first_microbatch - weightmat, new_weight_workspace = quantize_weight( - tensor=weight, - quantizer=weight_quantizer, - workspace=weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - - weightmat.update_usage(rowwise_usage=True) - + else: + nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") + ln_out_total, ln_out_total_work = gather_along_first_dim( + ln_out, + args.tp_group, + async_op=True, + quantizer=quantizer, + ) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") else: - weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP - # ------------------------------------------------------ - # Weight tensor is ready for GEMM... - # ------------------------------------------------------ - - # Cast bias to expected dtype - bias_dtype = activation_dtype - if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: - # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 - bias_dtype = torch.bfloat16 - bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias - - # Calibrate quantizers if needed - if not fp8 and fp8_calibration: - if input_quantizer is not None: - input_quantizer.calibrate(ln_out_total) - if weight_quantizer is not None: - weight_quantizer.calibrate(weight) + ln_out_total = ln_out + # -------------------------------------------------- + # Input tensor is ready for computing grad weight... + # -------------------------------------------------- + + # -------------------------------------------------- + # Compute grad input tensor + # Note: Gradient w.r.t. GEMM input (i.e. norm output). + # -------------------------------------------------- + + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + # Use saved_weight (the original weight parameter) since + # origin_weight is only set when fuse_wgrad_accumulation=True. + if weight is None: + if isinstance(saved_weight, QuantizedTensorStorage): + # saved weight is already set to right usages by + # fsdp2 quantized-tensor hooks when workspace was not saved. + weight = saved_weight + elif args.weight_quantizer is not None: + args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = args.weight_quantizer(saved_weight) + + # Make sure required data is available + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(rowwise_usage=True) + if ( + args.fp8 + and args.weight_quantizer is not None + and isinstance(weight, QuantizedTensorStorage) + ): + weight.update_usage(columnwise_usage=True) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + use_split_accumulator = args.dgrad_use_split_accumulator - # Configure output quantizer - if output_quantizer is not None: - output_quantizer.set_usage(rowwise=True, columnwise=False) + # Update grad input quantizer + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - # Output buffer for Userbuffers reduce-scatter + # Output buffers for Userbuffers reduce-scatter + gemm_out = None reduce_scatter_out = None - if ub_overlap_rs_fprop: - out_shape = list(inp_shape) - out_shape[0] //= tp_world_size - out_shape[-1] = out_features - reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) - - # ------------------------------------------------------ - # Forward GEMM - # Note: y = x * w^T - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.gemm") + if args.ub_overlap_rs_dgrad: + reduce_scatter_out = torch.empty( + dgrad_shape, dtype=args.activation_dtype, device=args.grad_output.device + ) + elif args.ub_bulk_wgrad: + gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) + + # dgrad GEMM + # Note: dx = dy * w + nvtx_range_push(f"{nvtx_label}.dgrad_gemm") + weight_for_dgrad = weight + if args.backward_override == "dequantized": + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=args.activation_dtype) + else: + weight_for_dgrad = cast_if_needed(weight_for_dgrad, args.activation_dtype) + elif args.backward_override == "high_precision": + weight_for_dgrad = saved_weight + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=args.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( - weightmat, - ln_out_total, - quantization_params=output_quantizer, - out_dtype=activation_dtype, - bias=bias, + weight_for_dgrad, + grad_output, + layout="NN", + grad=True, + quantization_params=args.grad_input_quantizer, + out=gemm_out, + out_dtype=args.activation_dtype, use_split_accumulator=use_split_accumulator, - ub=ub_obj, - ub_type=ub_type, + ub=ub_obj_dgrad, + ub_type=ub_type_dgrad, extra_output=reduce_scatter_out, + bulk_overlap=args.ub_bulk_dgrad, ) - nvtx_range_pop(f"{nvtx_label}.gemm") - # ------------------------------------------------------ - # Finished forward GEMM... - # ------------------------------------------------------ - - # Deallocate GEMM input tensor if no longer needed - if not weight.requires_grad and not return_layernorm_output: - clear_tensor_data(ln_out, ln_out_total) - ln_out = ln_out_total = None - elif with_input_all_gather and not return_layernorm_output_gathered: - # ln_out_total aliases ln_out for the cuBLASMp backend; skip the - # deallocation to avoid corrupting the backward-saved tensor. - if ln_out_total is not ln_out: - clear_tensor_data(ln_out_total) - ln_out_total = None - - # ------------------------------------------------------ - # Prepare output tensor - # Note: Perform tensor-parallel communication - # ------------------------------------------------------ - out = None - if ub_overlap_rs_fprop: - # cuBLASMp writes the reduce-scattered output directly into the - # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out - elif parallel_mode == "row" and tp_size > 1: - nvtx_range_push(f"{nvtx_label}.row_parallel_comm") - out = gemm_out - if sequence_parallel: - out, _ = reduce_scatter_along_first_dim(out, tp_group) - elif tensor_parallel: - if symmetric_ar_type is not None: - out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type) - else: - out, _ = allreduce(out, tp_group) - nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") - else: - out = gemm_out - out = out.view(-1, *inp_shape[1:-1], out_features) - # ------------------------------------------------------ - # Output tensor is ready to return... - # ------------------------------------------------------ + nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") - # ------------------------------------------------------ - # Cache state for backward pass - # ------------------------------------------------------ - - if is_grad_enabled: - ln_out_to_save = ln_out - if backward_override == "high_precision": - ln_out_to_save = ln_out_hp - ctx.weight_quantizer = weight_quantizer - ctx.ln_out_needs_gather = ( - weight.requires_grad and parallel_mode == "column" and sequence_parallel - ) + # FSDP2 only handles deallocation all-gathered weights that it allocates. + # Columnwise data is derived from rowwise data after allgather for fp8 + # and 2d block-scaled weights in TE managed memory. So we need to clear + # it here. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(weight, QuantizedTensorStorage): + clear_columnwise_cache(weight) - # Input with column-wise usage is needed for wgrad GEMM. - if backward_needs_input and backward_override is None: - if isinstance(ln_out, QuantizedTensorStorage): - # For sequence parallel in vanilla FP8, rowwise data is - # to gather the input. For MXFP8, columnwise only data - # can be allgathered. - if ( - isinstance(ln_out, (MXFP8TensorStorage, Float8BlockwiseQTensorStorage)) - or not ctx.ln_out_needs_gather - ): - ln_out.update_usage(rowwise_usage=False) - - if cpu_offloading: - mark_activation_offload(inputmat, mu, rsigma, ln_out_to_save) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - nvtx_range_push(f"{nvtx_label}.fsdp_scatter") - ctx.fsdp_group = fsdp_group - ctx.fsdp_shapes = _fsdp_scatter_tensors( - fsdp_group, - mu, - rsigma, - weightmat if fp8 and not is_weight_param_quantized else None, - ln_out_to_save if weight.requires_grad else None, + # Prepare grad input tensor + # Note: Perform tensor-parallel communication + dgrad = None + dgrad_work = None + if args.ub_overlap_rs_dgrad: + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + dgrad = ( + gemm_out + if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() + else reduce_scatter_out ) - nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - - if cpu_offloading: - mark_not_offload( - weightmat, - weight, - bias, - ln_weight, - ln_bias, + elif args.ub_bulk_wgrad: + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) + elif args.parallel_mode == "column" and args.tp_size > 1: + nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") + dgrad = gemm_out + if args.sequence_parallel: + dgrad, dgrad_work = reduce_scatter_along_first_dim( + dgrad, + args.tp_group, + async_op=True, ) + else: + dgrad, dgrad_work = allreduce(dgrad, args.tp_group, async_op=True) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") + else: + dgrad = gemm_out + + # -------------------------------------------------- + # Grad input tensor has been computed... + # -------------------------------------------------- + + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for wgrad. Userbuffers leaves the gathered + # tensor in its persistent buffer; cuBLASMp does not, so we gather + # here. Route through the same FP8-aware all-gather as the + # non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Columnwise data needed for wgrad is + # produced by ``update_usage(columnwise_usage=True)`` further below. + if ( + args.requires_wgrad + and args.ub_overlap_ag + and args.ub_obj_gradout is not None + and args.ub_obj_gradout.with_cublasmp() + ): + if args.grad_output_quantizer is not None: + set_quantizer_usage_for_wgrad_all_gather(args.grad_output_quantizer) + grad_output, _ = gather_along_first_dim( + grad_output, + args.tp_group, + quantizer=args.grad_output_quantizer, + ) - # FSDP2: Don't save FP8 workspace for non-quantized weights. - # Backward will re-quantize from FSDP2 all-gathered weight. - # (Issue #2681) - wt_save = weightmat - if is_fsdp2 and weightmat is not weight: - wt_save = None - # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; - # backward re-gathers from the saved (sharded) weight and re-quantizes. - if is_dist_weight: - wt_save = None + # -------------------------------------------------- + # Compute grad weight + # -------------------------------------------------- - tensors_to_save, tensor_objects = prepare_for_saving( - inputmat, - wt_save, - origin_weight, - bias, - ln_weight, - ln_out_to_save, - mu, - rsigma, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.requires_dgrad = inp_requires_grad - ctx.requires_wgrad = weight.requires_grad - ctx.is_weight_param_quantized = is_weight_param_quantized - ctx.is_fsdp2 = is_fsdp2 - if fuse_wgrad_accumulation and weight.requires_grad: - # Keep weakref to weight to preserve attributes like main_grad - # when we need to modify the weight python object - ctx.origin_weight_ref = weakref.ref(weight) - # Save overwrite_main_grad flag now while we have access to weight object - ctx.origin_weight_overwrites_main_grad = getattr( - weight, "overwrite_main_grad", False + wgrad = None + if args.requires_wgrad: + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.ub_overlap_ag and isinstance(args.grad_output_quantizer, MXFP8Quantizer): + # UB does not support pipelined overlapping grad output + # all-gather with wgrad GEMM. Also, we can't + # convert row-scaled MXFP8 to column-scaled, so we + # can't reuse the grad output that was gathered + # for the dgrad GEMM. We work around by explicitly + # overlapping the AG operation with the dgrad GEMM. + + # Get the communication stream from the dgrad GEMM to use for the AG + dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream() + + # This object is separate from the ub_obj_wgrad object which is passed to the GEMM + ub_obj_overlap_wgrad = get_ub(args.ub_name + "_wgrad", args.fp8) + + args.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + + # We use the send stream to copy into the userbuffers. + # This is the same stream that we will use to access the data in the AG, + # so we dont need to add any syncs yet. + with torch.cuda.stream(dgrad_send_stream): + grad_output, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_overlap_wgrad, + args.grad_output, + args.grad_output_quantizer, + args.tp_group, + ) + + # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm + tex.bulk_overlap_ag_with_external_gemm( + ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - if hasattr(weight, "__fsdp_param__"): - # MCore FSDP creates main_grad lazily before backward - ctx.main_grad_func = weight.get_main_grad - elif is_dist_weight: - ctx.main_grad_func = origin_weight.grad_buffer - else: - ctx.main_grad_func = lambda: weight.main_grad - ctx.grad_input_quantizer = grad_input_quantizer - ctx.grad_weight_quantizer = grad_weight_quantizer - ctx.grad_output_quantizer = grad_output_quantizer - ctx.input_quantizer = input_quantizer - ctx.owns_input = inputmat is not inp - ctx.weight = weight - ctx.activation_dtype = activation_dtype - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp_shape - ctx.parallel_mode = parallel_mode - ctx.tp_group = tp_group - ctx.tp_size = tp_size - ctx.return_layernorm_output = return_layernorm_output - ctx.return_layernorm_output_gathered = return_layernorm_output_gathered - ctx.bwd_ln_sm_margin = bwd_ln_sm_margin - ctx.zero_centered_gamma = zero_centered_gamma - ctx.ub_overlap_ag = ub_overlap_ag_dgrad - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_bulk_wgrad = ub_bulk_wgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_name = ub_name - ctx.requires_dgrad = inp_requires_grad - ctx.normalization = normalization - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - ctx.wgrad_store = wgrad_store - ctx.debug = debug - - # backward overrides - if backward_override is not None: - ctx.fp8 = False - ctx.debug = False - ctx.ub_overlap_ag = False - ctx.ub_overlap_rs_dgrad = False - ctx.ub_bulk_dgrad = False - ctx.ub_bulk_wgrad = False - ctx.grad_input_quantizer = None - ctx.grad_weight_quantizer = None - ctx.grad_output_quantizer = None - ctx.reduce_and_update_bwd_fp8_tensors = False - - # ------------------------------------------------------ - # Cached state for backward pass is ready... - # ------------------------------------------------------ - - ln_out_for_return = None - if return_layernorm_output: - if return_layernorm_output_gathered: - shape = list(inp_shape) - shape[0] *= tp_size if with_input_all_gather else 1 - ln_out_for_return = ln_out_return.view(shape) - else: - ln_out_for_return = ln_out_return.view(inp_shape) - return out, ln_out_for_return, new_weight_workspace - @staticmethod - def backward( - ctx, *grad_outputs: Tuple[torch.Tensor, ...] - ) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if args.fp8 or args.debug: + if isinstance(ln_out_total, QuantizedTensorStorage): + ln_out_total.update_usage(columnwise_usage=True) + else: + args.input_quantizer.set_usage(rowwise=False, columnwise=True) + ln_out_total = args.input_quantizer(ln_out_total) - # NVTX label for profiling - nvtx_label = "transformer_engine._LayerNormLinear.backward" - if ctx.ub_name is not None: - nvtx_label = f"{nvtx_label}.{ctx.ub_name}" + if args.fp8 or args.debug: + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(columnwise_usage=True) + else: + args.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + grad_output = args.grad_output_quantizer(grad_output) - with get_nvtx_range_context("_LayerNormLinear_backward"): - ( # pylint: disable=unbalanced-tuple-unpacking - inputmat, - weight, - saved_weight, - bias, - ln_weight, - ln_out, - mu, - rsigma, - ) = restore_from_func_ctx(ctx) + # Figure out whether to use split accumulator + use_split_accumulator = args.wgrad_use_split_accumulator - is_dist_weight = is_distributed_weight(saved_weight) + # Figure out whether to output wgrad GEMM directly into main grad if is_dist_weight: - weight = materialize_weight_for_backward(saved_weight)[0] - # Restore from weakref to get original weight python object - # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) - # Only needed when fuse_wgrad_accumulation is enabled. - origin_weight = None - origin_weight_overwrites_main_grad = getattr( - ctx, "origin_weight_overwrites_main_grad", False - ) - main_grad = None - if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: - origin_weight_ref = ctx.origin_weight_ref - ctx.origin_weight_ref = None - origin_weight = origin_weight_ref() if origin_weight_ref is not None else None - assert ( - origin_weight is not None - ), "weight was removed while fuse_wgrad_accumulation=True" - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - main_grad = ctx.main_grad_func() if weight is not None else None - if main_grad is not None and not is_dist_weight: - origin_weight.main_grad = main_grad - - # Gather intermediate/activation tensors if needed - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - nvtx_range_push(f"{nvtx_label}.fsdp_gather") - _fsdp_gather_tensors( - ctx.fsdp_group, - ctx.fsdp_shapes, - mu, - rsigma, - weight if ctx.fp8 and not ctx.is_weight_param_quantized else None, - ln_out, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - - # Configure Userbuffers communication (comm+GEMM overlap) - ctx.ub_obj_gradout = None - ub_obj_dgrad = None - ub_obj_wgrad = None - ub_type_dgrad = None - ub_type_wgrad = None - dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]] - if ctx.ub_overlap_ag: - # Overlap grad_output all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - elif ctx.ub_overlap_rs_dgrad: - # Overlap dgrad reduce-scatter with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.RS - else: - if ctx.ub_bulk_dgrad: - # Overlap inputmat all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: - # Overlap dgrad reduce-scatter with wgrad compute - ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) - ub_type_wgrad = tex.CommOverlapType.RS - - # -------------------------------------------------- - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # -------------------------------------------------- - - # Configure quantizer for grad output tensor - # Note: dgrad GEMM requires row-wise usage, wgrad GEMM - # requires column-wise usage - if ctx.grad_output_quantizer is not None: - quantizer = ctx.grad_output_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: - # Userbuffers only supports communication for one - # tensor usage at a time. Configure quantizer with - # usage for only dgrad GEMM. - quantizer.set_usage(columnwise=False) - # Amax reduction group for grad output (row-parallel sequence parallel) - set_quantizer_amax_reduction_group( - quantizer, - ( - ctx.tp_group - if (ctx.sequence_parallel and ctx.parallel_mode == "row") - else None - ), + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif args.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + args.fuse_wgrad_accumulation and not args.is_first_microbatch ) - - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") - ( - grad_output, - grad_bias, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, - grad_outputs[0], - ctx.parallel_mode == "row", - ctx.grad_output_quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - - # -------------------------------------------------- - # Grad output tensor is ready for computing grad input... - # -------------------------------------------------- - - # -------------------------------------------------- - # Prepare GEMM input tensor - # Note: Input tensor is needed for wgrad GEMM. - # Tensor-parallel communication is overlapped with dgrad - # GEMM. - # -------------------------------------------------- - ln_out_total = None - ln_out_total_work = None - if ctx.backward_override == "dequantized": - if isinstance(ln_out, QuantizedTensorStorage): - ln_out = ln_out.dequantize(dtype=ctx.activation_dtype) - else: - ln_out = cast_if_needed(ln_out, ctx.activation_dtype) - if ctx.ln_out_needs_gather: - quantizer = None - if ctx.input_quantizer is not None and ctx.fp8: - quantizer = ctx.input_quantizer - set_quantizer_usage_for_wgrad_all_gather(quantizer) - if ctx.ub_bulk_dgrad: - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_dgrad, - ln_out, - quantizer, - ctx.tp_group, - ) - else: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") - ln_out_total, ln_out_total_work = gather_along_first_dim( - ln_out, - ctx.tp_group, - async_op=True, - quantizer=quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") else: - ln_out_total = ln_out - # -------------------------------------------------- - # Input tensor is ready for computing grad weight... - # -------------------------------------------------- - - # -------------------------------------------------- - # Compute grad input tensor - # Note: Gradient w.r.t. GEMM input (i.e. norm output). - # -------------------------------------------------- - - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved. (Issue #2681) - # Use saved_weight (the original weight parameter) since - # origin_weight is only set when fuse_wgrad_accumulation=True. - if weight is None: - if isinstance(saved_weight, QuantizedTensorStorage): - # saved weight is already set to right usages by - # fsdp2 quantized-tensor hooks when workspace was not saved. - weight = saved_weight - elif ctx.weight_quantizer is not None: - ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) - weight = ctx.weight_quantizer(saved_weight) - - # Make sure required data is available - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(rowwise_usage=True) - if ( - ctx.fp8 - and ctx.weight_quantizer is not None - and isinstance(weight, QuantizedTensorStorage) - ): - weight.update_usage(columnwise_usage=True) - - # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - - # Update grad input quantizer - if ctx.grad_input_quantizer is not None: - ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - - # Output buffers for Userbuffers reduce-scatter - gemm_out = None + accumulate_wgrad_into_param_main_grad = args.fuse_wgrad_accumulation + + # Output buffer for overlapping FP8 grad input + # reduce-scatter with wgrad GEMM reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: + if args.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device - ) - elif ctx.ub_bulk_wgrad: - gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) - - # dgrad GEMM - # Note: dx = dy * w - nvtx_range_push(f"{nvtx_label}.dgrad_gemm") - weight_for_dgrad = weight - if ctx.backward_override == "dequantized": - if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) - else: - weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) - elif ctx.backward_override == "high_precision": - weight_for_dgrad = saved_weight - if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) - gemm_out, *_, reduce_scatter_out = general_gemm( - weight_for_dgrad, - grad_output, - layout="NN", - grad=True, - quantization_params=ctx.grad_input_quantizer, - out=gemm_out, - out_dtype=ctx.activation_dtype, - use_split_accumulator=use_split_accumulator, - ub=ub_obj_dgrad, - ub_type=ub_type_dgrad, - extra_output=reduce_scatter_out, - bulk_overlap=ctx.ub_bulk_dgrad, - ) - nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") - - # FSDP2 only handles deallocation all-gathered weights that it allocates. - # Columnwise data is derived from rowwise data after allgather for fp8 - # and 2d block-scaled weights in TE managed memory. So we need to clear - # it here. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(weight, QuantizedTensorStorage): - clear_columnwise_cache(weight) - - # Prepare grad input tensor - # Note: Perform tensor-parallel communication - dgrad = None - dgrad_work = None - if ctx.ub_overlap_rs_dgrad: - # cuBLASMp writes the reduce-scattered dgrad directly into the - # GEMM output tensor; Userbuffers uses the extra-output buffer. - dgrad = ( - gemm_out - if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() - else reduce_scatter_out + dgrad_shape, dtype=args.activation_dtype, device=args.grad_output.device ) - elif ctx.ub_bulk_wgrad: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) - elif ctx.parallel_mode == "column" and ctx.tp_size > 1: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") - dgrad = gemm_out - if ctx.sequence_parallel: - dgrad, dgrad_work = reduce_scatter_along_first_dim( - dgrad, - ctx.tp_group, - async_op=True, + + # Arguments to include in wgrad GEMM closure + wgrad_gemm_kwargs = { + "out_dtype": ( + main_grad.dtype if args.fuse_wgrad_accumulation else args.activation_dtype + ), + "quantization_params": args.grad_weight_quantizer, + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not origin_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": main_grad if args.fuse_wgrad_accumulation else None, + "bias": (bias if (grad_bias is None and not args.fp8) else None), + "use_split_accumulator": use_split_accumulator, + "grad": True, + "ub": ub_obj_wgrad, + "ub_type": ub_type_wgrad, + "extra_output": reduce_scatter_out, + "bulk_overlap": args.ub_bulk_wgrad, + } + + def wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform wgrad GEMM: dw = dy^T * x + + May be fused with bgrad computation. + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + nvtx_range_push(f"{nvtx_label}.wgrad_gemm") + dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") + return dw, db + + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + if ( + wgrad_gemm_kwargs["ub"] is not None + or wgrad_gemm_kwargs["ub_type"] is not None + or wgrad_gemm_kwargs["extra_output"] is not None + or wgrad_gemm_kwargs["bulk_overlap"] + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported " + "with Userbuffers (tensor-parallel communication overlapping)" ) - else: - dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") + args.wgrad_store.put([ln_out_total, grad_output], wgrad_gemm) else: - dgrad = gemm_out - - # -------------------------------------------------- - # Grad input tensor has been computed... - # -------------------------------------------------- - - # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and - # does not preserve it for wgrad. Userbuffers leaves the gathered - # tensor in its persistent buffer; cuBLASMp does not, so we gather - # here. Route through the same FP8-aware all-gather as the - # non-overlap path in - # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing - # the grad_output quantizer. Columnwise data needed for wgrad is - # produced by ``update_usage(columnwise_usage=True)`` further below. - if ( - ctx.requires_wgrad - and ctx.ub_overlap_ag - and ctx.ub_obj_gradout is not None - and ctx.ub_obj_gradout.with_cublasmp() - ): - if ctx.grad_output_quantizer is not None: - set_quantizer_usage_for_wgrad_all_gather(ctx.grad_output_quantizer) - grad_output, _ = gather_along_first_dim( - grad_output, - ctx.tp_group, - quantizer=ctx.grad_output_quantizer, - ) - - # -------------------------------------------------- - # Compute grad weight - # -------------------------------------------------- - wgrad = None - if ctx.requires_wgrad: - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.ub_overlap_ag and isinstance(ctx.grad_output_quantizer, MXFP8Quantizer): - # UB does not support pipelined overlapping grad output - # all-gather with wgrad GEMM. Also, we can't - # convert row-scaled MXFP8 to column-scaled, so we - # can't reuse the grad output that was gathered - # for the dgrad GEMM. We work around by explicitly - # overlapping the AG operation with the dgrad GEMM. - - # Get the communication stream from the dgrad GEMM to use for the AG - dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream() - - # This object is separate from the ub_obj_wgrad object which is passed to the GEMM - ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) - - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - - # We use the send stream to copy into the userbuffers. - # This is the same stream that we will use to access the data in the AG, - # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): - grad_output, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_overlap_wgrad, - grad_outputs[0], - ctx.grad_output_quantizer, - ctx.tp_group, - ) - - # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm - tex.bulk_overlap_ag_with_external_gemm( - ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream - ) + # Call wgrad GEMM now + wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if ctx.fp8 or ctx.debug: - if isinstance(ln_out_total, QuantizedTensorStorage): - ln_out_total.update_usage(columnwise_usage=True) - else: - ctx.input_quantizer.set_usage(rowwise=False, columnwise=True) - ln_out_total = ctx.input_quantizer(ln_out_total) - - if ctx.fp8 or ctx.debug: - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(columnwise_usage=True) - else: - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - grad_output = ctx.grad_output_quantizer(grad_output) - - # Figure out whether to use split accumulator - use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator - - # Figure out whether to output wgrad GEMM directly into main grad if is_dist_weight: - # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. - accumulate_wgrad_into_param_main_grad = False - elif ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch - ) + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + + # Update grad bias if needed + if grad_bias is None: + grad_bias = grad_bias_ + del grad_bias_ + + # Deallocate input tensors if permitted + if not args.return_layernorm_output and not args.return_layernorm_output_gathered: + # Input tensors have not been exposed externally + clear_tensor_data(ln_out) + elif args.ln_out_needs_gather and args.return_layernorm_output_gathered: + # Non-gathered input has not been exposed externally + clear_tensor_data(ln_out) + if args.ln_out_needs_gather: + # Gathered input is internal + clear_tensor_data(ln_out_total) + if args.sequence_parallel and ( + args.parallel_mode == "row" or (args.parallel_mode == "column" and args.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal + clear_tensor_data(grad_output) + + # Update grad input if overlapping reduce-scatter with wgrad GEMM + if args.ub_bulk_wgrad: + if ub_obj_wgrad.is_fp8_ubuf(): + dgrad = reduce_scatter_out else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - # Output buffer for overlapping FP8 grad input - # reduce-scatter with wgrad GEMM - reduce_scatter_out = None - if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): - reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device - ) + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() - # Arguments to include in wgrad GEMM closure - wgrad_gemm_kwargs = { - "out_dtype": ( - main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype - ), - "quantization_params": ctx.grad_weight_quantizer, - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not origin_weight_overwrites_main_grad - else False - ), - "layout": "NT", - "out": main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": (bias if (grad_bias is None and not ctx.fp8) else None), - "use_split_accumulator": use_split_accumulator, - "grad": True, - "ub": ub_obj_wgrad, - "ub_type": ub_type_wgrad, - "extra_output": reduce_scatter_out, - "bulk_overlap": ctx.ub_bulk_wgrad, - } - - def wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform wgrad GEMM: dw = dy^T * x - - May be fused with bgrad computation. - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - nvtx_range_push(f"{nvtx_label}.wgrad_gemm") - dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) - nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - if ( - wgrad_gemm_kwargs["ub"] is not None - or wgrad_gemm_kwargs["ub_type"] is not None - or wgrad_gemm_kwargs["extra_output"] is not None - or wgrad_gemm_kwargs["bulk_overlap"] - ): - raise NotImplementedError( - "Delayed weight grad computation is not supported " - "with Userbuffers (tensor-parallel communication overlapping)" - ) - ctx.wgrad_store.put([ln_out_total, grad_output], wgrad_gemm) - else: + # -------------------------------------------------- + # Grad weight has been computed... + # -------------------------------------------------- - # Call wgrad GEMM now - wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - - if is_dist_weight: - wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] - - # Update grad bias if needed - if grad_bias is None: - grad_bias = grad_bias_ - del grad_bias_ - - # Deallocate input tensors if permitted - if not ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - # Input tensors have not been exposed externally - clear_tensor_data(ln_out) - elif ctx.ln_out_needs_gather and ctx.return_layernorm_output_gathered: - # Non-gathered input has not been exposed externally - clear_tensor_data(ln_out) - if ctx.ln_out_needs_gather: - # Gathered input is internal - clear_tensor_data(ln_out_total) - if ctx.sequence_parallel and ( - ctx.parallel_mode == "row" or (ctx.parallel_mode == "column" and ctx.fp8) - ): - # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal - clear_tensor_data(grad_output) - - # Update grad input if overlapping reduce-scatter with wgrad GEMM - if ctx.ub_bulk_wgrad: - if ub_obj_wgrad.is_fp8_ubuf(): - dgrad = reduce_scatter_out - else: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() - - # -------------------------------------------------- - # Grad weight has been computed... - # -------------------------------------------------- - - # Don't return grad bias if not needed - if not ctx.use_bias: - grad_bias = None - - # Synchronize tensor parallel communication - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if dgrad_work is not None: - dgrad_work.wait() - dgrad_work = None + # Don't return grad bias if not needed + if not args.use_bias: + grad_bias = None + + # Synchronize tensor parallel communication + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if dgrad_work is not None: + dgrad_work.wait() + dgrad_work = None - # Residual gradient - dgrad = dgrad.view(inputmat.shape) - if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - dgrad = dgrad + grad_outputs[1].view_as(dgrad) + # Residual gradient + dgrad = dgrad.view(inputmat.shape) + if args.return_layernorm_output and not args.return_layernorm_output_gathered: + dgrad = dgrad + args.grad_ln_out.view_as(dgrad) - # Norm gradient - dgamma = None + # Norm gradient + dgamma = None + dbeta = None + nvtx_range_push(f"{nvtx_label}.norm") + if args.normalization == "LayerNorm": + dgrad, dgamma, dbeta = tex.layernorm_bwd( + dgrad, + inputmat, + mu, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + dgrad = dgrad.reshape(inputmat.size()) + elif args.normalization == "RMSNorm": + dgrad, dgamma = tex.rmsnorm_bwd( + dgrad, + inputmat, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + dgrad = dgrad.reshape(inputmat.size()) dbeta = None - nvtx_range_push(f"{nvtx_label}.norm") - if ctx.normalization == "LayerNorm": - dgrad, dgamma, dbeta = tex.layernorm_bwd( - dgrad, - inputmat, - mu, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + nvtx_range_pop(f"{nvtx_label}.norm") + clear_tensor_data(mu) + clear_tensor_data(rsigma) + + if args.requires_wgrad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): + origin_weight.grad_added_to_main_grad = True + if getattr(origin_weight, "zero_out_wgrad", False): + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight.dtype, + zero=True, ) - dgrad = dgrad.reshape(inputmat.size()) - elif ctx.normalization == "RMSNorm": - dgrad, dgamma = tex.rmsnorm_bwd( - dgrad, - inputmat, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + else: + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight.dtype, ) - dgrad = dgrad.reshape(inputmat.size()) - dbeta = None - nvtx_range_pop(f"{nvtx_label}.norm") - clear_tensor_data(mu) - clear_tensor_data(rsigma) - - if ctx.requires_wgrad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): - origin_weight.grad_added_to_main_grad = True - if getattr(origin_weight, "zero_out_wgrad", False): - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight.dtype, - zero=True, - ) - else: - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight.dtype, - ) - elif ctx.fuse_wgrad_accumulation: - wgrad = None - else: + elif args.fuse_wgrad_accumulation: wgrad = None + else: + wgrad = None + + return ( + dgrad.view(args.inp_shape) if args.requires_dgrad else None, + dgamma, + dbeta, + wgrad, + grad_bias, + ) + + +class _LayerNormLinear(torch.autograd.Function): + """LayerNormLinear semi-top level module + Calls custom cuda extensions. + """ - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + @staticmethod + def forward( + ctx, + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + weight: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LayerNormLinearFwdArgs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Forward pass: compute the output and set up the autograd context. + + The tensors are positional so autograd tracks them; they are + re-attached to ``fwd_args`` so every downstream helper takes a single + argument. ``weight_workspace`` is a non-differentiable cached tensor + passed in via ``fwd_args`` and the freshly produced workspace is + returned as a separate output so the module can refresh its cache. + """ + fwd_args.inp = inp + fwd_args.ln_weight = ln_weight + fwd_args.ln_bias = ln_bias + fwd_args.weight = weight + fwd_args.bias = bias + ( + out, + ln_out_return, + new_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_linear_forward_impl(fwd_args) + if ctx is not None: + bwd_args = LayerNormLinearBwdArgs() + tensors_to_save_from_setup = _layernorm_linear_setup_ctx( + bwd_args, + fwd_args, + (out, ln_out_return, new_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + ctx.backward_objects = bwd_args + if fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() + if fwd_args.backward_override is not None: + bwd_args.reduce_and_update_bwd_fp8_tensors = False + + return out, ln_out_return, new_weight_workspace + + @staticmethod + def backward( + ctx, + grad_output: torch.Tensor, + grad_ln_out: Optional[torch.Tensor], + _grad_weight_workspace, + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward pass: compute gradients and reduce FP8 scaling factors.""" + bwd_args: LayerNormLinearBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.grad_ln_out = grad_ln_out + bwd_args.setup_saved_tensors(ctx) + nvtx_label = "transformer_engine._LayerNormLinear.backward" + if bwd_args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" + dgrad, dgamma, dbeta, wgrad, grad_bias = _layernorm_linear_backward_impl(bwd_args) + reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, + # main_grad closure) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") - - # Scatter fp8 weight buffers - # if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage): - # _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) - return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + dgrad, dgamma, dbeta, wgrad, - None, # weight_workspace grad_bias, - None, + None, # fwd_args ) @@ -1754,69 +2099,138 @@ def forward( weight_quantizer, weight_tensor ) - if is_grad_enabled: - fwd_fn = _LayerNormLinear.apply - autograd_ctx = [] - else: - fwd_fn = _LayerNormLinear.forward - autograd_ctx = [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 ) - non_tensor_args = ( - self.eps, - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - self.fuse_wgrad_accumulation, - input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, - is_cpu_offload_enabled(), - self.tp_group, - self.tp_size, - self.sequence_parallel, - self.tp_size > 1, - self.activation_dtype, - self.parallel_mode, - self.return_layernorm_output, - self.return_layernorm_output_gathered, - is_grad_enabled, - self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, - self.bwd_ln_sm_margin, - self.zero_centered_gamma, - self.normalization, - self.ub_overlap_ag_fprop, - self.ub_overlap_rs_fprop, - self.ub_overlap_ag_dgrad, - self.ub_overlap_rs_dgrad, - self.ub_bulk_wgrad, - self.ub_bulk_dgrad, - self.ub_name, - self.fsdp_group, - cache_name is not None, - skip_fp8_weight_update, - self.symmetric_ar_type, - debug, - self.is_fsdp2, + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD + if self.fp8: + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + 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 + else: + backward_override = None + + if debug: # turn off userbuffers in debug mode + ub_overlap_ag_fprop = False + ub_overlap_rs_fprop = False + ub_overlap_ag_dgrad = False + ub_overlap_rs_dgrad = False + ub_bulk_wgrad = False + ub_bulk_dgrad = False + else: + ub_overlap_ag_fprop = self.ub_overlap_ag_fprop + ub_overlap_rs_fprop = self.ub_overlap_rs_fprop + ub_overlap_ag_dgrad = self.ub_overlap_ag_dgrad + ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad + ub_bulk_wgrad = self.ub_bulk_wgrad + ub_bulk_dgrad = self.ub_bulk_dgrad + + linear_bias_tensor = ( + bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) - out, ln_out, new_weight_workspace = fwd_fn( - *autograd_ctx, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - weight_tensor, - weight_workspace, - bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None, - non_tensor_args, + wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + + fwd_args = LayerNormLinearFwdArgs( + # tensors + inp=inp, + ln_weight=self.layer_norm_weight, + ln_bias=self.layer_norm_bias, + weight=weight_tensor, + bias=linear_bias_tensor, + weight_workspace=weight_workspace, + # requires_grad flags + input_requires_grad=inp.requires_grad, + ln_weight_requires_grad=self.layer_norm_weight.requires_grad, + ln_bias_requires_grad=( + self.layer_norm_bias.requires_grad + if self.layer_norm_bias is not None + else False + ), + weight_requires_grad=weight_tensor.requires_grad, + bias_requires_grad=( + linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False + ), + # quantizers + input_quantizer=input_quantizer, + weight_quantizer=weight_quantizer, + output_quantizer=output_quantizer, + grad_input_quantizer=grad_input_quantizer, + grad_weight_quantizer=grad_weight_quantizer, + grad_output_quantizer=grad_output_quantizer, + # normalization + eps=self.eps, + normalization=self.normalization, + zero_centered_gamma=self.zero_centered_gamma, + fwd_ln_sm_margin=( + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin + ), + bwd_ln_sm_margin=self.bwd_ln_sm_margin, + return_layernorm_output=self.return_layernorm_output, + return_layernorm_output_gathered=self.return_layernorm_output_gathered, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_name is not None, + skip_fp8_weight_update=skip_fp8_weight_update, + # tensor / sequence parallelism + parallel_mode=self.parallel_mode, + tp_group=self.tp_group, + tp_size=self.tp_size, + tensor_parallel=self.tp_size > 1, + sequence_parallel=self.sequence_parallel, + symmetric_ar_type=self.symmetric_ar_type, + # userbuffers + ub_name=self.ub_name, + ub_overlap_ag_fprop=ub_overlap_ag_fprop, + ub_overlap_rs_fprop=ub_overlap_rs_fprop, + ub_overlap_ag_dgrad=ub_overlap_ag_dgrad, + ub_overlap_rs_dgrad=ub_overlap_rs_dgrad, + ub_bulk_dgrad=ub_bulk_dgrad, + ub_bulk_wgrad=ub_bulk_wgrad, + # FSDP + fsdp_group=self.fsdp_group, + is_fsdp2=self.is_fsdp2, + # weight-grad scheduling + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + # misc + cpu_offloading=is_cpu_offload_enabled(), + 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, + ) + else: + out, ln_out, new_weight_workspace = _LayerNormLinear.forward( + None, + inp, + self.layer_norm_weight, + self.layer_norm_bias, + weight_tensor, + linear_bias_tensor, + fwd_args, + ) + if new_weight_workspace is not None and cache_name is not None: if isinstance(new_weight_workspace, torch.Tensor): new_weight_workspace = new_weight_workspace.detach() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3ee0cda50c..3be30fa83e 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -5,8 +5,9 @@ """LayerNormMLP API""" import os import warnings +from dataclasses import dataclass, replace as dataclass_replace import weakref -from typing import Callable, Optional, Tuple, Union, List +from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -46,7 +47,6 @@ cast_if_needed, assert_dim_for_fp8_exec, clear_tensor_data, - requires_grad, needs_quantized_gemm, get_nvtx_range_context, ) @@ -58,7 +58,6 @@ reduce_scatter_along_first_dim, gather_along_first_dim, use_reentrant_activation_recompute, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _get_cuda_rng_state, _set_cuda_rng_state, @@ -74,6 +73,7 @@ from ..tensor.identity_tensor import IdentityQuantizer from ._common import ( apply_normalization, + check_fp8_reduce_and_update, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -90,6 +90,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorOrQuantized from ..cpp_extensions import ( general_gemm, ) @@ -171,1663 +172,1996 @@ def _act_func(activation: str, recipe: Optional[Recipe] = None): return funcs[activation] -class _LayerNormMLP(torch.autograd.Function): - """LayerNormMLP semi-top level module - Calls custom cuda extensions. - """ - - @staticmethod - def _forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: torch.Tensor, - fc1_weight: torch.Tensor, - fc1_weight_workspace: Optional[torch.Tensor], - fc1_bias: torch.Tensor, - fc2_weight: torch.Tensor, - fc2_weight_workspace: Optional[torch.Tensor], - fc2_bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: - # pylint: disable=missing-function-docstring - - # Reduce number of arguments to autograd function in order - # to reduce CPU overhead due to pytorch arg checking. - ( - eps, - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - fuse_wgrad_accumulation, - fc1_input_quantizer, - fc1_weight_quantizer, - fc1_output_quantizer, - fc1_grad_input_quantizer, - fc1_grad_weight_quantizer, - fc1_grad_output_quantizer, - fc2_input_quantizer, - fc2_weight_quantizer, - fc2_output_quantizer, - fc2_grad_input_quantizer, - fc2_grad_weight_quantizer, - fc2_grad_output_quantizer, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - return_layernorm_output, - return_layernorm_output_gathered, - bias_gelu_fusion, - set_parallel_mode, - is_grad_enabled, - fwd_ln_sm_margin, - bwd_ln_sm_margin, - zero_centered_gamma, - activation, - activation_params, - normalization, - ub_overlap_ag, - ub_overlap_rs, - ub_overlap_rs_dgrad, - ub_bulk_wgrad, - ub_bulk_dgrad, - gemm_gelu_fusion, - fsdp_group, - fp8_meta, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - checkpoint, - debug, - is_fsdp2, - recompute_for_bwd, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None - assert backward_override is None, ( - "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." - " Replace LayerNormMLP with LayerNormLinear + Linear to enable" - " high_precision/dequantized backward." +@dataclass(slots=True) +class LayerNormMLPFwdArgs: + """Single-argument bag for the forward path of :class:`_LayerNormMLP`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + inp: torch.Tensor + ln_weight: torch.Tensor + ln_bias: Optional[torch.Tensor] + fc1_weight: TensorOrQuantized + fc1_bias: Optional[torch.Tensor] + fc2_weight: TensorOrQuantized + fc2_bias: Optional[torch.Tensor] + + # --- Non-differentiable cached tensors --- + fc1_weight_workspace: Optional[TensorOrQuantized] + fc2_weight_workspace: Optional[TensorOrQuantized] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + ln_weight_requires_grad: bool + ln_bias_requires_grad: bool + fc1_weight_requires_grad: bool + fc1_bias_requires_grad: bool + fc2_weight_requires_grad: bool + fc2_bias_requires_grad: bool + + # --- Quantizers --- + fc1_input_quantizer: Optional[Quantizer] + fc1_weight_quantizer: Optional[Quantizer] + fc1_output_quantizer: Optional[Quantizer] + fc1_grad_input_quantizer: Optional[Quantizer] + fc1_grad_weight_quantizer: Optional[Quantizer] + fc1_grad_output_quantizer: Optional[Quantizer] + fc2_input_quantizer: Optional[Quantizer] + fc2_weight_quantizer: Optional[Quantizer] + fc2_output_quantizer: Optional[Quantizer] + fc2_grad_input_quantizer: Optional[Quantizer] + fc2_grad_weight_quantizer: Optional[Quantizer] + fc2_grad_output_quantizer: Optional[Quantizer] + + # --- Normalization --- + eps: float + normalization: str + zero_centered_gamma: bool + fwd_ln_sm_margin: int + bwd_ln_sm_margin: int + return_layernorm_output: bool + return_layernorm_output_gathered: bool + + # --- Activation --- + activation: str + activation_params: Optional[Dict[str, Any]] + bias_gelu_fusion: bool + gemm_gelu_fusion: bool + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + skip_fp8_weight_update: Optional[torch.Tensor] + + # --- Tensor / sequence parallelism --- + set_parallel_mode: bool + tp_group: Optional[dist_group_type] + tp_size: int + tensor_parallel: bool + sequence_parallel: bool + symmetric_ar_type: Optional[str] + + # --- Userbuffers (comm + GEMM overlap) --- + ub_overlap_ag: bool + ub_overlap_rs: bool + ub_overlap_rs_dgrad: bool + ub_bulk_dgrad: bool + ub_bulk_wgrad: bool + + # --- FSDP --- + fsdp_group: Optional[Any] + is_fsdp2: bool + + # --- Weight-grad scheduling --- + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + + # --- Activation checkpointing (recompute in backward) --- + checkpoint: bool + fp8_meta: Optional[Any] + recompute_for_bwd: bool + + # --- Misc --- + cpu_offloading: bool + is_grad_enabled: bool + + def any_requires_grad(self) -> bool: + """Whether any differentiable input requires a gradient.""" + return any( + ( + self.input_requires_grad, + self.ln_weight_requires_grad, + self.ln_bias_requires_grad, + self.fc1_weight_requires_grad, + self.fc1_bias_requires_grad, + self.fc2_weight_requires_grad, + self.fc2_bias_requires_grad, + ) ) - # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take - if is_grad_enabled and not recompute_for_bwd: - ctx.checkpoint = checkpoint - if checkpoint: - # save the state of autocast and quantizers for recomputation - ctx.autocast_state = FP8GlobalStateManager.get_autocast_state() - if ( - fp8 - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ - == "DelayedScaling" - ): # only applicable for delayed scaling - FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( - fp8_meta - ) # to restore quantizers during recomputation - # save the rng states - ctx.cpu_rng_state = torch.get_rng_state() - ctx.cuda_rng_state = _get_cuda_rng_state() - - # whether to save activations regularly, or save inputs for recomputation in bwd - save_for_checkpoint = checkpoint and is_grad_enabled and not recompute_for_bwd - - # whether we are in the forward stage, or recomputing in the bwd stage (false if not checkpointing) - is_recomputation = checkpoint and is_grad_enabled and recompute_for_bwd - - # save the initial state for recomputation by bwd - if save_for_checkpoint: - - # save tensors - tensors_to_save, tensor_objects = prepare_for_saving( - inp, - ln_weight, - ln_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.other_args = { - "eps": eps, - "is_first_microbatch": is_first_microbatch, - "fp8": fp8, - "fp8_calibration": fp8_calibration, - "wgrad_store": wgrad_store, - "fuse_wgrad_accumulation": fuse_wgrad_accumulation, - "fc1_input_quantizer": fc1_input_quantizer, - "fc1_weight_quantizer": fc1_weight_quantizer, - "fc1_output_quantizer": fc1_output_quantizer, - "fc1_grad_input_quantizer": fc1_grad_input_quantizer, - "fc1_grad_weight_quantizer": fc1_grad_weight_quantizer, - "fc1_grad_output_quantizer": fc1_grad_output_quantizer, - "fc2_input_quantizer": fc2_input_quantizer, - "fc2_weight_quantizer": fc2_weight_quantizer, - "fc2_output_quantizer": fc2_output_quantizer, - "fc2_grad_input_quantizer": fc2_grad_input_quantizer, - "fc2_grad_weight_quantizer": fc2_grad_weight_quantizer, - "fc2_grad_output_quantizer": fc2_grad_output_quantizer, - "cpu_offloading": cpu_offloading, - "tp_group": tp_group, - "tp_size": tp_size, - "sequence_parallel": sequence_parallel, - "tensor_parallel": tensor_parallel, - "activation_dtype": activation_dtype, - "return_layernorm_output": return_layernorm_output, - "return_layernorm_output_gathered": return_layernorm_output_gathered, - "bias_gelu_fusion": bias_gelu_fusion, - "set_parallel_mode": set_parallel_mode, - "is_grad_enabled": is_grad_enabled, - "fwd_ln_sm_margin": fwd_ln_sm_margin, - "bwd_ln_sm_margin": bwd_ln_sm_margin, - "zero_centered_gamma": zero_centered_gamma, - "activation": activation, - "activation_params": activation_params, - "normalization": normalization, - "ub_overlap_ag": ub_overlap_ag, - "ub_overlap_rs": ub_overlap_rs, - "ub_overlap_rs_dgrad": ub_overlap_rs_dgrad, - "ub_bulk_wgrad": ub_bulk_wgrad, - "ub_bulk_dgrad": ub_bulk_dgrad, - "gemm_gelu_fusion": gemm_gelu_fusion, - "fsdp_group": fsdp_group, - "fp8_meta": fp8_meta, - "cache_weight": False, - "skip_fp8_weight_update": skip_fp8_weight_update, - "symmetric_ar_type": symmetric_ar_type, - "checkpoint": checkpoint, - "debug": debug, - "is_fsdp2": is_fsdp2, - "recompute_for_bwd": True, # set this to true for recomputation phase - } - # Make sure input dimensions are compatible - in_features, inp_shape = ln_weight.numel(), inp.shape - assert inp_shape[-1] == in_features, "GEMM not possible" - inputmat = inp.view((-1, in_features)) - if fp8: - assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) - - activation_func = _act_func( - activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - )[0] - - # Cast for native AMP - inputmat = cast_if_needed(inputmat, activation_dtype) - ln_weight = cast_if_needed(ln_weight, activation_dtype) - if ln_bias is not None: - ln_bias = cast_if_needed(ln_bias, activation_dtype) - if is_cpu_offload_enabled(): - start_offload(inputmat) - - tp_world_size = get_distributed_world_size(tp_group) - - # bwd needs fc1 input when grad is enabled, fc1 needs grad, and either - # 1) no checkpointing - # or 2) doing the recomputation with checkpointing - backwards_needs_fc1_input = fc1_weight.requires_grad and ( - (is_grad_enabled and not checkpoint) or is_recomputation - ) +@dataclass(slots=True) +class LayerNormMLPBwdArgs: + """Single-argument bag for the backward path of :class:`_LayerNormMLP`.""" + + # --- Incoming gradients (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + grad_ln_out: Optional[torch.Tensor] = None + + # --- Saved / restored tensors (populated at backward entry) --- + inputmat: Optional[torch.Tensor] = None + ln_weight: Optional[torch.Tensor] = None + ln_out: Optional[TensorOrQuantized] = None + fc1_weight_fp8: Optional[TensorOrQuantized] = None + fc1_weight: Optional[TensorOrQuantized] = None + fc1_bias: Optional[torch.Tensor] = None + fc1_out: Optional[TensorOrQuantized] = None + fc1_out_without_bias: Optional[torch.Tensor] = None + act_out: Optional[TensorOrQuantized] = None + fc2_weight_fp8: Optional[TensorOrQuantized] = None + fc2_weight: Optional[TensorOrQuantized] = None + fc2_bias: Optional[torch.Tensor] = None + mu: Optional[torch.Tensor] = None + rsigma: Optional[torch.Tensor] = None + + # --- Activation checkpointing (forward inputs saved, recomputed in backward) --- + checkpoint: bool = False + checkpoint_fwd_args: Optional[Any] = None + autocast_state: Optional[Any] = None + cpu_rng_state: Optional[Any] = None + cuda_rng_state: Optional[Any] = None + + # --- Quantizers --- + fc1_input_quantizer: Optional[Quantizer] = None + fc1_weight_quantizer: Optional[Quantizer] = None + fc1_grad_input_quantizer: Optional[Quantizer] = None + fc1_grad_weight_quantizer: Optional[Quantizer] = None + fc1_grad_output_quantizer: Optional[Quantizer] = None + fc2_input_quantizer: Optional[Quantizer] = None + fc2_weight_quantizer: Optional[Quantizer] = None + fc2_grad_input_quantizer: Optional[Quantizer] = None + fc2_grad_weight_quantizer: Optional[Quantizer] = None + fc2_grad_output_quantizer: Optional[Quantizer] = None + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + fc1_weight_requires_grad: bool = False + fc1_bias_requires_grad: bool = False + fc2_weight_requires_grad: bool = False + inp_shape: Optional[torch.Size] = None + + # --- Normalization --- + normalization: str = "LayerNorm" + zero_centered_gamma: bool = False + bwd_ln_sm_margin: int = 0 + return_layernorm_output: bool = False + return_layernorm_output_gathered: bool = False + + # --- Activation --- + activation: str = "gelu" + activation_params: Optional[Dict[str, Any]] = None + bias_gelu_fusion: bool = False + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + fp8_recipe: Optional[Any] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + backward_override: Optional[str] = None + debug: bool = False + + # --- Tensor / sequence parallelism --- + set_parallel_mode: bool = False + tp_group: Optional[dist_group_type] = None + tp_size: int = 1 + tensor_parallel: bool = False + sequence_parallel: bool = False + + # --- Userbuffers (comm + GEMM overlap) --- + ub_overlap_ag: bool = False + ub_overlap_rs_dgrad: bool = False + ub_bulk_dgrad: bool = False + ub_bulk_wgrad: bool = False + + # --- FSDP --- + fsdp_group: Optional[Any] = None + fsdp_shapes: Any = None + is_fsdp2: bool = False + + # --- Weight-grad scheduling / accumulation --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + fc1_weight_ref: Optional[Any] = None + fc2_weight_ref: Optional[Any] = None + fc1_weight_overwrites_main_grad: bool = False + fc2_weight_overwrites_main_grad: bool = False + fc1_main_grad_func: Optional[Callable[[], torch.Tensor]] = None + fc2_main_grad_func: Optional[Callable[[], torch.Tensor]] = None + + # --- FP8 reduce-and-update bookkeeping --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + + # --- Per-backward scratch state (populated inside the backward impl) --- + ub_obj_gradout: Optional[Any] = None + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" + self.set_saved_tensors(restore_from_func_ctx(ctx)) + + def set_saved_tensors(self, tensors: Sequence[Any]) -> None: + """Bind the (restored or recomputed) saved tensors to their fields.""" + ( + self.inputmat, + self.ln_weight, + self.ln_out, + self.fc1_weight_fp8, + self.fc1_weight, + self.fc1_bias, + self.fc1_out, + self.fc1_out_without_bias, + self.act_out, + self.fc2_weight_fp8, + self.fc2_weight, + self.fc2_bias, + self.mu, + self.rsigma, + ) = tensors + + +_CHECKPOINT_SAVED_ALIASES = ( + "inp", + "ln_weight", + "ln_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", +) - device = inp.device - # Configure Userbuffers communication (comm+GEMM overlap) - if debug: # turn off userbuffers in debug mode - ub_overlap_ag = False - ub_overlap_rs = False - ub_overlap_rs_dgrad = False - ub_bulk_wgrad = False - ub_bulk_dgrad = False - ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered - ub_overlap_rs = ub_overlap_rs and is_grad_enabled +def _layernorm_mlp_forward_impl( + args: LayerNormMLPFwdArgs, +) -> Tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[Tuple], + Optional[Dict], +]: + """Forward implementation for the layernorm-MLP layer. + + Returns ``(fc2_out, ln_out_return, new_fc1_weight_workspace, + new_fc2_weight_workspace, tensors_to_save_from_forward, ctx_attrs)``. The + new workspaces are the freshly produced FP8 weight workspaces (returned so + the caller can refresh its cache). The last two are ``None`` when + gradients are disabled. + + With ``args.checkpoint`` the forward saves only its inputs and the + backward recomputes the activations by calling this again with + ``args.recompute_for_bwd``; that call returns ``None`` user outputs and + the recomputed saved tensors. + """ + inp = args.inp + ln_weight = args.ln_weight + ln_bias = args.ln_bias + fc1_weight = args.fc1_weight + fc1_bias = args.fc1_bias + fc2_weight = args.fc2_weight + fc2_bias = args.fc2_bias + fc1_input_quantizer = args.fc1_input_quantizer + fc1_weight_quantizer = args.fc1_weight_quantizer + fc1_output_quantizer = args.fc1_output_quantizer + fc2_input_quantizer = args.fc2_input_quantizer + fc2_weight_quantizer = args.fc2_weight_quantizer + fc2_output_quantizer = args.fc2_output_quantizer + is_first_microbatch = args.is_first_microbatch + fp8 = args.fp8 + fp8_calibration = args.fp8_calibration + debug = args.debug + cpu_offloading = args.cpu_offloading + tp_group = args.tp_group + tp_size = args.tp_size + sequence_parallel = args.sequence_parallel + tensor_parallel = args.tensor_parallel + set_parallel_mode = args.set_parallel_mode + activation_dtype = args.activation_dtype + is_grad_enabled = args.is_grad_enabled + return_layernorm_output = args.return_layernorm_output + return_layernorm_output_gathered = args.return_layernorm_output_gathered + activation = args.activation + bias_gelu_fusion = args.bias_gelu_fusion + gemm_gelu_fusion = args.gemm_gelu_fusion + ub_overlap_ag = args.ub_overlap_ag + ub_overlap_rs = args.ub_overlap_rs + fsdp_group = args.fsdp_group + is_fsdp2 = args.is_fsdp2 + checkpoint = args.checkpoint + recompute_for_bwd = args.recompute_for_bwd + fc1_weight_requires_grad = args.fc1_weight_requires_grad + fc2_weight_requires_grad = args.fc2_weight_requires_grad + + assert args.backward_override is None, ( + "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." + " Replace LayerNormMLP with LayerNormLinear + Linear to enable" + " high_precision/dequantized backward." + ) + + ctx_attrs = None + # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take + if is_grad_enabled and not recompute_for_bwd: + ctx_attrs = {"checkpoint": checkpoint} + if checkpoint: + # save the state of autocast and quantizers for recomputation + ctx_attrs["autocast_state"] = FP8GlobalStateManager.get_autocast_state() + if ( + fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( + args.fp8_meta + ) # to restore quantizers during recomputation + # save the rng states + ctx_attrs["cpu_rng_state"] = torch.get_rng_state() + ctx_attrs["cuda_rng_state"] = _get_cuda_rng_state() + + # whether to save activations regularly, or save inputs for recomputation in bwd + save_for_checkpoint = checkpoint and is_grad_enabled and not recompute_for_bwd + + # whether we are in the forward stage, or recomputing in the bwd stage (false if not checkpointing) + is_recomputation = checkpoint and is_grad_enabled and recompute_for_bwd + + tensors_to_save_from_forward = None + # save the initial state for recomputation by bwd + if save_for_checkpoint: + tensors_to_save_from_forward = (None,) * len(_CHECKPOINT_SAVED_ALIASES) + ctx_attrs["saved_tensor_aliases"] = _CHECKPOINT_SAVED_ALIASES + + # Make sure input dimensions are compatible + in_features, inp_shape = ln_weight.numel(), inp.shape + assert inp_shape[-1] == in_features, "GEMM not possible" + inp = inp.view((-1, in_features)) + inputmat = inp + if fp8: + assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) + + activation_func = _act_func( + activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + )[0] + + # Cast for native AMP + inputmat = cast_if_needed(inputmat, activation_dtype) + ln_weight_cast = cast_if_needed(ln_weight, activation_dtype) + if ln_bias is not None: + ln_bias = cast_if_needed(ln_bias, activation_dtype) + if is_cpu_offload_enabled(): + start_offload(inputmat) + + tp_world_size = get_distributed_world_size(tp_group) + + # bwd needs fc1 input when grad is enabled, fc1 needs grad, and either + # 1) no checkpointing + # or 2) doing the recomputation with checkpointing + backwards_needs_fc1_input = fc1_weight_requires_grad and ( + (is_grad_enabled and not checkpoint) or is_recomputation + ) + + device = inp.device + + # Configure Userbuffers communication (comm+GEMM overlap) + ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered + ub_overlap_rs = ub_overlap_rs and is_grad_enabled + + # Choose whether to use GEMM kernel with split accumulator + use_split_accumulator = _2X_ACC_FPROP + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if hasattr(recipe, "fp8_gemm_fprop"): + use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + + # Configure quantizer for norm output + if fp8: + if fc1_input_quantizer is None: + raise ValueError("Missing quantizer for FC1 input tensor") + fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) + if sequence_parallel and fc1_input_quantizer.supports_only_rowwise_all_gather(): + # All-gather is not supported with FP8 column-wise data + fc1_input_quantizer.set_usage(columnwise=False) + # Amax reduction group for the FC1 input quantizer (column-parallel sequence parallel) + set_quantizer_amax_reduction_group( + fc1_input_quantizer, + tp_group if (sequence_parallel and set_parallel_mode) else None, + ) - # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + # for fp8 DelayedScaling: layernorm output = FP8 + # only output of the linear is returned + # for return_layernorm_output: layernorm output = High precision, then cast to FP8 + # high precision layernorm output and output of the linear are returned + # for debug: : layernorm output = High precision to enable processing of this norm + + custom = is_custom(fc1_input_quantizer) + hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) + identity = isinstance(fc1_input_quantizer, IdentityQuantizer) + with_quantized_norm = ( + fp8 + and not debug + and not return_layernorm_output + and not return_layernorm_output_gathered + and not custom + and not hybrid + and not identity + ) + + # Apply normalization + ln_out, mu, rsigma = apply_normalization( + inputmat, + None, # ln_out + ln_weight_cast, + ln_bias, + args.eps, + fc1_input_quantizer if with_quantized_norm else None, + inputmat.dtype, + args.normalization, + args.fwd_ln_sm_margin, + args.zero_centered_gamma, + ) + ln_out_return = None + + # do not return layernorm output unless 1) no checkpointing or 2) checkpointing but not recomputing + if (return_layernorm_output or return_layernorm_output_gathered) and not is_recomputation: + ln_out_return = ln_out + + # Prepare GEMM input + # Note: Cast to expected dtype and perform tensor-parallel communication + ln_out_total = None + ub_obj_lnout = None + if sequence_parallel: + + # do not return ln output if checkpointing and in recomputation, not necessary + if return_layernorm_output_gathered and not is_recomputation: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) + ln_out_return = ln_out_total + if fp8 or debug: + ln_out = fc1_input_quantizer(ln_out) + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = fc1_input_quantizer(ln_out_total) + else: + quantizer = None + if fp8 or debug: + quantizer = fc1_input_quantizer + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: + ln_out = fc1_input_quantizer(ln_out) + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag: + # Copy into Userbuffers buffer + ub_obj_lnout = get_ub("fc1_fprop", fp8) + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_lnout, + ln_out, + quantizer, + tp_group, + ) + else: + # All-gather with NCCL + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=quantizer, + ) + else: + if (fp8 or debug) and not with_quantized_norm: + ln_out = fc1_input_quantizer(ln_out) + ln_out_total = ln_out + + # Cast weights to expected dtype + new_fc1_weight_workspace = None + new_fc2_weight_workspace = None + fc1_weight_final = fc1_weight + fc2_weight_final = fc2_weight + # FSDP2: Skip columnwise/transpose creation during forward (not + # recompute) to avoid accumulating FP8 caches across layers. + # Backward's FSDP2 all-gather will recreate them. (Issue #2681) + fsdp2_skip_columnwise = is_fsdp2 and not is_recomputation + if fp8 or debug: + update_ws = is_first_microbatch is None or is_first_microbatch + # If weight is already quantized, weight._quantizer is its true quantizer. + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: + fc1_weight_quantizer = fc1_weight._quantizer + elif fc1_weight_quantizer is not None: + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, + ) - # Configure quantizer for norm output - if fp8: - if fc1_input_quantizer is None: - raise ValueError("Missing quantizer for FC1 input tensor") - fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) - if sequence_parallel and fc1_input_quantizer.supports_only_rowwise_all_gather(): - # All-gather is not supported with FP8 column-wise data - fc1_input_quantizer.set_usage(columnwise=False) - # Amax reduction group for the FC1 input quantizer (column-parallel sequence parallel) - set_quantizer_amax_reduction_group( - fc1_input_quantizer, - tp_group if (sequence_parallel and set_parallel_mode) else None, + if isinstance(fc2_weight, QuantizedTensorStorage) and not debug: + fc2_weight_quantizer = fc2_weight._quantizer + elif fc2_weight_quantizer is not None: + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, ) - # for fp8 DelayedScaling: layernorm output = FP8 - # only output of the linear is returned - # for return_layernorm_output: layernorm output = High precision, then cast to FP8 - # high precision layernorm output and output of the linear are returned - # for debug: : layernorm output = High precision to enable processing of this norm - - custom = is_custom(fc1_input_quantizer) - hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) - identity = isinstance(fc1_input_quantizer, IdentityQuantizer) - with_quantized_norm = ( - fp8 - and not debug - and not return_layernorm_output - and not return_layernorm_output_gathered - and not custom - and not hybrid - and not identity + fc1_weight_final, new_fc1_weight_workspace = quantize_weight( + tensor=fc1_weight, + quantizer=fc1_weight_quantizer, + workspace=args.fc1_weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, ) - - # Apply normalization - ln_out, mu, rsigma = apply_normalization( - inputmat, - None, # ln_out - ln_weight, - ln_bias, - eps, - fc1_input_quantizer if with_quantized_norm else None, - inputmat.dtype, - normalization, - fwd_ln_sm_margin, - zero_centered_gamma, + fc2_weight_final, new_fc2_weight_workspace = quantize_weight( + tensor=fc2_weight, + quantizer=fc2_weight_quantizer, + workspace=args.fc2_weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, ) - ln_out_return = None - - # do not return layernorm output unless 1) no checkpointing or 2) checkpointing but not recomputing - if (return_layernorm_output or return_layernorm_output_gathered) and not is_recomputation: - ln_out_return = ln_out - - # Prepare GEMM input - # Note: Cast to expected dtype and perform tensor-parallel communication - ln_out_total = None - ub_obj_lnout = None - if sequence_parallel: - - # do not return ln output if checkpointing and in recomputation, not necessary - if return_layernorm_output_gathered and not is_recomputation: - # Perform all-gather in high precision if gathered - # norm output will be returned - ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) - ln_out_return = ln_out_total - if fp8 or debug: - ln_out = fc1_input_quantizer(ln_out) - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - ln_out_total = fc1_input_quantizer(ln_out_total) + fc1_weight_final.update_usage(rowwise_usage=True) + fc2_weight_final.update_usage(rowwise_usage=True) + else: + fc1_weight_final = cast_if_needed(fc1_weight_final, activation_dtype) + fc2_weight_final = cast_if_needed(fc2_weight_final, activation_dtype) + + # Cast biases to expected dtype + bias_dtype = activation_dtype + if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: + # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 + bias_dtype = torch.bfloat16 + fc1_bias_cast = fc1_bias + fc2_bias_cast = fc2_bias + if fc1_bias is not None: + fc1_bias_cast = cast_if_needed(fc1_bias, bias_dtype) + if fc2_bias is not None: + fc2_bias_cast = cast_if_needed(fc2_bias, bias_dtype) + + # Calibrate quantizers if needed + if not fp8 and fp8_calibration: + if fc1_input_quantizer is not None: + fc1_input_quantizer.calibrate(ln_out_total) + if fc1_weight_quantizer is not None: + fc1_weight_quantizer.calibrate(fc1_weight) + + # ------------------------------------------------------ + # FC1 GEMM + # ------------------------------------------------------ + + # There are 2 fusions possible: + # - gemm_gelu_fusion - default for full precision, optional for fp8 - need to turn on gemm_gelu_fusion, + # - bias_gelu_fusion - only for full precision. + # If both gemm_gelu_fusion and bias_gelu_fusion are enabled, only bias_gelu_fusion will be performer + if activation != "gelu": + # blockwise scaled gemms don't support gemm_gelu_fusion in fwd. + gemm_gelu_fusion = bias_gelu_fusion = False + else: + if fp8: + assert not bias_gelu_fusion, "Bias gelu fusion is supported only for full precision" + else: + gemm_gelu_fusion = True + if gemm_gelu_fusion and bias_gelu_fusion: + gemm_gelu_fusion = False + if debug: + gemm_gelu_fusion = False + fc1_outputs = general_gemm( + fc1_weight_final, + ln_out_total, + quantization_params=( + fc2_input_quantizer + if gemm_gelu_fusion + else fc1_output_quantizer # fused gelu output is in fp8 + ), + out_dtype=activation_dtype, + bias=( + fc1_bias_cast if not bias_gelu_fusion else None + ), # otherwise bias is added later (fused with gelu) + gelu=gemm_gelu_fusion, + use_split_accumulator=use_split_accumulator, + ub=ub_obj_lnout, + ub_type=tex.CommOverlapType.AG if ub_overlap_ag else None, + ) + + # ------------------------------------------------------ + # Finished FC1 GEMM... + # ------------------------------------------------------ + + # Deallocate FC1 GEMM input tensor if no longer needed + # first part of if statement means that we only clear ln_out_total if + # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) + # 2) not checkpointing and grad disabled + # The `is not ln_out` guard avoids clearing the bwd-saved tensor when + # ln_out_total aliases ln_out (cuBLASMp AG-fprop path). + if ( + ((checkpoint and not is_recomputation) or not is_grad_enabled) + and ln_out_total is not ln_out_return + and ln_out_total is not ln_out + ): + clear_tensor_data(ln_out_total) + + # ACTIVATION - sometimes activation is fused with the GEMM above. + + fc1_out_without_bias = None + act_params = args.activation_params or {} + + if bias_gelu_fusion: + fc1_out = None + fc1_out_without_bias, *_ = fc1_outputs + act_out = bias_gelu_fused(fc1_out_without_bias, fc1_bias_cast) + elif gemm_gelu_fusion: + act_out, _, fc1_out, _ = fc1_outputs + elif debug: + fc1_out, *_ = fc1_outputs + act_out = activation_func(fc1_out, None, **act_params) + act_out = fc2_input_quantizer(act_out) + else: + fc1_out, *_ = fc1_outputs + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if recipe.float8_block_scaling(): + # tex.quantize does not support GELU fusion for blockwise + act_out = activation_func(fc1_out, None, **act_params) + act_out = tex.quantize(act_out, fc2_input_quantizer) + elif recipe.custom(): + # tex.quantize does not support custom quantizers + act_out = activation_func(fc1_out, None, **act_params) + act_out = fc2_input_quantizer(act_out) else: - quantizer = None - if fp8 or debug: - quantizer = fc1_input_quantizer - # custom recipe doesn't need to support quantized AG - if not with_quantized_norm and not custom: - ln_out = fc1_input_quantizer(ln_out) - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - if ub_overlap_ag: - # Copy into Userbuffers buffer - ub_obj_lnout = get_ub("fc1_fprop", fp8) - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_lnout, - ln_out, - quantizer, - tp_group, - ) - else: - # All-gather with NCCL - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=quantizer, - ) + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) else: - if (fp8 or debug) and not with_quantized_norm: - ln_out = fc1_input_quantizer(ln_out) - ln_out_total = ln_out - - # Cast weights to expected dtype - new_fc1_weight_workspace = None - new_fc2_weight_workspace = None - fc1_weight_final = fc1_weight - fc2_weight_final = fc2_weight - # FSDP2: Skip columnwise/transpose creation during forward (not - # recompute) to avoid accumulating FP8 caches across layers. - # Backward's FSDP2 all-gather will recreate them. (Issue #2681) - fsdp2_skip_columnwise = is_fsdp2 and not is_recomputation - if fp8 or debug: - update_ws = is_first_microbatch is None or is_first_microbatch - # If weight is already quantized, weight._quantizer is its true quantizer. - # for debug mode we create quantizer every iteration, thus we need to set the quantizer states - if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: - fc1_weight_quantizer = fc1_weight._quantizer - elif fc1_weight_quantizer is not None: - fc1_weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not fsdp2_skip_columnwise, - ) + if fp8_calibration: + act_out = activation_func(fc1_out, None, **act_params) + else: + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - if isinstance(fc2_weight, QuantizedTensorStorage) and not debug: - fc2_weight_quantizer = fc2_weight._quantizer - elif fc2_weight_quantizer is not None: - fc2_weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not fsdp2_skip_columnwise, - ) + if not fp8 and fp8_calibration: + if fc2_input_quantizer is not None: + fc2_input_quantizer.calibrate(act_out) + + # we want to skip fc2 computation if we are checkpointing and recomputing, + # otherwise we compute fc2 + fc2_out = None + if not (is_recomputation and checkpoint): + + # if we get to this point, we know this is not bwd recomputation + # so we must be in the fwd + # now is_grad_enabled can be true or false + # if false, can safely delete + # if true, we can only delete if checkpoint is true, since we will recompute anyways, + # otherwise, checkpoint is false, so cant delete + if checkpoint or not is_grad_enabled: # we can safely get rid of these if this is the case + clear_tensor_data(fc1_out) - fc1_weight_final, new_fc1_weight_workspace = quantize_weight( - tensor=fc1_weight, - quantizer=fc1_weight_quantizer, - workspace=fc1_weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - fc2_weight_final, new_fc2_weight_workspace = quantize_weight( - tensor=fc2_weight, - quantizer=fc2_weight_quantizer, - workspace=fc2_weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - fc1_weight_final.update_usage(rowwise_usage=True) - fc2_weight_final.update_usage(rowwise_usage=True) - else: - fc1_weight_final = cast_if_needed(fc1_weight_final, activation_dtype) - fc2_weight_final = cast_if_needed(fc2_weight_final, activation_dtype) - - # Cast biases to expected dtype - bias_dtype = activation_dtype - if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: - # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 - bias_dtype = torch.bfloat16 - if fc1_bias is not None: - fc1_bias = cast_if_needed(fc1_bias, bias_dtype) - if fc2_bias is not None: - fc2_bias = cast_if_needed(fc2_bias, bias_dtype) - - # Calibrate quantizers if needed if not fp8 and fp8_calibration: - if fc1_input_quantizer is not None: - fc1_input_quantizer.calibrate(ln_out_total) - if fc1_weight_quantizer is not None: - fc1_weight_quantizer.calibrate(fc1_weight) + + if fc2_weight_quantizer is not None: + fc2_weight_quantizer.calibrate(fc2_weight) + + # Configure Userbuffers reduce-scatter if needed + ub_obj_fc2out = None + reduce_scatter_out = None + if ub_overlap_rs: + ub_obj_fc2out = get_ub("fc2_fprop", fp8) + dim_size = list(act_out.size()) + dim_size[0] //= tp_world_size + dim_size[-1] = fc2_weight.size(0) + reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) # ------------------------------------------------------ - # FC1 GEMM + # FC2 GEMM # ------------------------------------------------------ - - # There are 2 fusions possible: - # - gemm_gelu_fusion - default for full precision, optional for fp8 - need to turn on gemm_gelu_fusion, - # - bias_gelu_fusion - only for full precision. - # If both gemm_gelu_fusion and bias_gelu_fusion are enabled, only bias_gelu_fusion will be performer - if activation != "gelu": - # blockwise scaled gemms don't support gemm_gelu_fusion in fwd. - gemm_gelu_fusion = bias_gelu_fusion = False - else: - if fp8: - assert not bias_gelu_fusion, "Bias gelu fusion is supported only for full precision" - else: - gemm_gelu_fusion = True - if gemm_gelu_fusion and bias_gelu_fusion: - gemm_gelu_fusion = False - if debug: - gemm_gelu_fusion = False - fc1_outputs = general_gemm( - fc1_weight_final, - ln_out_total, - quantization_params=( - fc2_input_quantizer - if gemm_gelu_fusion - else fc1_output_quantizer # fused gelu output is in fp8 - ), + gemm_out, *_, reduce_scatter_out = general_gemm( + fc2_weight_final, + act_out, out_dtype=activation_dtype, - bias=( - fc1_bias if not bias_gelu_fusion else None - ), # otherwise bias is added later (fused with gelu) - gelu=gemm_gelu_fusion, + bias=fc2_bias_cast, + quantization_params=fc2_output_quantizer, use_split_accumulator=use_split_accumulator, - ub=ub_obj_lnout, - ub_type=tex.CommOverlapType.AG if ub_overlap_ag else None, + ub=ub_obj_fc2out, + ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, + extra_output=reduce_scatter_out, ) - # ------------------------------------------------------ - # Finished FC1 GEMM... + # Finished FC2 GEMM... # ------------------------------------------------------ - # Deallocate FC1 GEMM input tensor if no longer needed - # first part of if statement means that we only clear ln_out_total if - # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) - # 2) not checkpointing and grad disabled - # The `is not ln_out` guard avoids clearing the bwd-saved tensor when - # ln_out_total aliases ln_out (cuBLASMp AG-fprop path). - if ( - ((checkpoint and not is_recomputation) or not is_grad_enabled) - and ln_out_total is not ln_out_return - and ln_out_total is not ln_out - ): - clear_tensor_data(ln_out_total) - - # ACTIVATION - sometimes activation is fused with the GEMM above. - - fc1_out_without_bias = None - act_params = activation_params or {} - - if bias_gelu_fusion: - fc1_out = None - fc1_out_without_bias, *_ = fc1_outputs - act_out = bias_gelu_fused(fc1_out_without_bias, fc1_bias) - elif gemm_gelu_fusion: - act_out, _, fc1_out, _ = fc1_outputs - elif debug: - fc1_out, *_ = fc1_outputs - act_out = activation_func(fc1_out, None, **act_params) - act_out = fc2_input_quantizer(act_out) - else: - fc1_out, *_ = fc1_outputs - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_block_scaling(): - # tex.quantize does not support GELU fusion for blockwise - act_out = activation_func(fc1_out, None, **act_params) - act_out = tex.quantize(act_out, fc2_input_quantizer) - elif recipe.custom(): - # tex.quantize does not support custom quantizers - act_out = activation_func(fc1_out, None, **act_params) - act_out = fc2_input_quantizer(act_out) - else: - act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - else: - if fp8_calibration: - act_out = activation_func(fc1_out, None, **act_params) - else: - act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - - if not fp8 and fp8_calibration: - if fc2_input_quantizer is not None: - fc2_input_quantizer.calibrate(act_out) - - # we want to skip fc2 computation if we are checkpointing and recomputing, - # otherwise we compute fc2 - if not (is_recomputation and checkpoint): - - # if we get to this point, we know this is not bwd recomputation - # so we must be in the fwd - # now is_grad_enabled can be true or false - # if false, can safely delete - # if true, we can only delete if checkpoint is true, since we will recompute anyways, - # otherwise, checkpoint is false, so cant delete - if ( - checkpoint or not is_grad_enabled - ): # we can safely get rid of these if this is the case - clear_tensor_data(fc1_out) - - if not fp8 and fp8_calibration: - - if fc2_weight_quantizer is not None: - fc2_weight_quantizer.calibrate(fc2_weight) - - # Configure Userbuffers reduce-scatter if needed - ub_obj_fc2out = None - reduce_scatter_out = None - if ub_overlap_rs: - ub_obj_fc2out = get_ub("fc2_fprop", fp8) - dim_size = list(act_out.size()) - dim_size[0] //= tp_world_size - dim_size[-1] = fc2_weight.size(0) - reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) - - # ------------------------------------------------------ - # FC2 GEMM - # ------------------------------------------------------ - gemm_out, *_, reduce_scatter_out = general_gemm( - fc2_weight_final, - act_out, - out_dtype=activation_dtype, - bias=fc2_bias, - quantization_params=fc2_output_quantizer, - use_split_accumulator=use_split_accumulator, - ub=ub_obj_fc2out, - ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, - extra_output=reduce_scatter_out, + # Deallocate tensors if no longer needed, again, can safely deallocate + if checkpoint or not is_grad_enabled: # same logic as last clear_tensor_data block + clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) + + # Prepare output tensor + # Note: Perform tensor-parallel communication if needed + if ub_overlap_rs: + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + fc2_out = ( + gemm_out + if ub_obj_fc2out is not None and ub_obj_fc2out.with_cublasmp() + else reduce_scatter_out ) - # ------------------------------------------------------ - # Finished FC2 GEMM... - # ------------------------------------------------------ - - # Deallocate tensors if no longer needed, again, can safely deallocate - if checkpoint or not is_grad_enabled: # same logic as last clear_tensor_data block - clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) - - # Prepare output tensor - # Note: Perform tensor-parallel communication if needed - fc2_out = None - if ub_overlap_rs: - # cuBLASMp writes the reduce-scattered output directly into the - # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - fc2_out = ( - gemm_out - if ub_obj_fc2out is not None and ub_obj_fc2out.with_cublasmp() - else reduce_scatter_out + elif set_parallel_mode and sequence_parallel: + fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) + elif set_parallel_mode and tensor_parallel: + if args.symmetric_ar_type is not None: + fc2_out, _ = symmetric_all_reduce( + gemm_out, tp_group, all_reduce_type=args.symmetric_ar_type ) - elif set_parallel_mode and sequence_parallel: - fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) - elif set_parallel_mode and tensor_parallel: - if symmetric_ar_type is not None: - fc2_out, _ = symmetric_all_reduce( - gemm_out, tp_group, all_reduce_type=symmetric_ar_type - ) - else: - fc2_out, _ = allreduce(gemm_out, tp_group) else: - fc2_out = gemm_out - fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) - - # now saving stuff for bwd: - # if we are using checkpointing, this information will be saved in the bwd recomputation stage, so can skip it in fwd - # if we are not checkpointing, then we must save this if grad is enabled - if is_grad_enabled and not save_for_checkpoint: - - ctx.fc1_weight_quantizer = fc1_weight_quantizer - ctx.fc2_weight_quantizer = fc2_weight_quantizer - - if not fc1_weight.requires_grad: - if not return_layernorm_output: - clear_tensor_data(ln_out) - ln_out = None - if not fc2_weight.requires_grad: - clear_tensor_data(act_out) - act_out = None - - if not checkpoint: # regular path, no selective activation checkpointing - - if cpu_offloading: - mark_activation_offload( - inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out - ) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - ctx.fsdp_group = fsdp_group - - ctx.fsdp_shapes = ( - _fsdp_scatter_tensors( # again, ony relevant if we have activations to save - fsdp_group, - mu, - rsigma, - ln_out, - fc1_out_without_bias if bias_gelu_fusion else fc1_out, - act_out, - ( - fc1_weight_final - if fp8 and not isinstance(fc1_weight, Float8Tensor) - else None - ), - ( - fc2_weight_final - if fp8 and not isinstance(fc2_weight, Float8Tensor) - else None - ), - ) + fc2_out, _ = allreduce(gemm_out, tp_group) + else: + fc2_out = gemm_out + fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) + + # now saving stuff for bwd: + # if we are using checkpointing, this information will be saved in the bwd recomputation stage, so can skip it in fwd + # if we are not checkpointing, then we must save this if grad is enabled + if is_grad_enabled and not save_for_checkpoint: + if ctx_attrs is None: + ctx_attrs = {} + + if not fc1_weight_requires_grad: + if not return_layernorm_output: + clear_tensor_data(ln_out) + ln_out = None + if not fc2_weight_requires_grad: + clear_tensor_data(act_out) + act_out = None + + fsdp_shapes = None + if not checkpoint: # regular path, no selective activation checkpointing + + if cpu_offloading: + mark_activation_offload( + inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out ) - if cpu_offloading: - mark_not_offload( - ln_weight, - ln_bias, - fc1_weight_final, - fc1_weight, - fc1_bias, - fc2_weight_final, - fc2_weight, - fc2_bias, - ) - # FSDP2: Don't save FP8 workspace copies for non-quantized - # weights. Backward will re-quantize from the FSDP2 - # all-gathered weight parameter. (Issue #2681) - fc1_wt_save = fc1_weight_final - fc2_wt_save = fc2_weight_final - if fsdp2_skip_columnwise: - if fc1_weight_final is not fc1_weight: - fc1_wt_save = None - if fc2_weight_final is not fc2_weight: - fc2_wt_save = None - tensors_to_save, tensor_objects = prepare_for_saving( - inputmat, - ln_weight, - ln_out, - fc1_wt_save, - fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, - fc2_wt_save, - fc2_weight, - fc2_bias, + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + fsdp_shapes = ( + _fsdp_scatter_tensors( # again, ony relevant if we have activations to save + fsdp_group, mu, rsigma, + ln_out, + fc1_out_without_bias if bias_gelu_fusion else fc1_out, + act_out, + ( + fc1_weight_final + if fp8 and not isinstance(fc1_weight, Float8Tensor) + else None + ), + ( + fc2_weight_final + if fp8 and not isinstance(fc2_weight, Float8Tensor) + else None + ), ) - - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - if fuse_wgrad_accumulation: - # Keep weakrefs to weights to preserve attributes like main_grad - # when we need to modify the weight python objects - ctx.fc1_weight_python_object_ref = ( - weakref.ref(fc1_weight) if fc1_weight.requires_grad else None - ) - ctx.fc2_weight_python_object_ref = ( - weakref.ref(fc2_weight) if fc2_weight.requires_grad else None - ) - ctx.fc1_weight_overwrites_main_grad = getattr( - fc1_weight, "overwrite_main_grad", False - ) - ctx.fc2_weight_overwrites_main_grad = getattr( - fc2_weight, "overwrite_main_grad", False - ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - if hasattr(fc1_weight, "__fsdp_param__") and hasattr(fc2_weight, "__fsdp_param__"): - # MCore FSDP creates main_grad lazily before backward - ctx.fc1_main_grad_func = ( - fc1_weight.get_main_grad if fc1_weight.requires_grad else lambda: None - ) - ctx.fc2_main_grad_func = ( - fc2_weight.get_main_grad if fc2_weight.requires_grad else lambda: None - ) - else: - ctx.fc1_main_grad_func = lambda: fc1_weight.main_grad - ctx.fc2_main_grad_func = lambda: fc2_weight.main_grad - - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.fc1_grad_input_quantizer = fc1_grad_input_quantizer - ctx.fc1_grad_weight_quantizer = fc1_grad_weight_quantizer - ctx.fc1_grad_output_quantizer = fc1_grad_output_quantizer - ctx.fc2_grad_input_quantizer = fc2_grad_input_quantizer - ctx.fc2_grad_weight_quantizer = fc2_grad_weight_quantizer - ctx.fc2_grad_output_quantizer = fc2_grad_output_quantizer - ctx.fc1_input_quantizer = fc1_input_quantizer - ctx.fc2_input_quantizer = fc2_input_quantizer - - ctx.fc1_weight_requires_grad = fc1_weight.requires_grad - ctx.fc2_weight_requires_grad = fc2_weight.requires_grad - ctx.fc1_weight = fc1_weight - ctx.fc2_weight = fc2_weight - ctx.fsdp2_skip_columnwise = fsdp2_skip_columnwise - # Store raw is_fsdp2 flag for backward cleanup — must not be - # gated on is_recomputation since backward cleanup runs after - # the real backward, not the recomputation forward. - ctx.is_fsdp2 = is_fsdp2 - - ctx.device = device - ctx.activation_dtype = activation_dtype - ctx.activation = activation - ctx.activation_params = activation_params - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = fc2_bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp_shape - ctx.tp_group = tp_group - ctx.tp_size = tp_size - ctx.bias_gelu_fusion = bias_gelu_fusion - ctx.return_layernorm_output = return_layernorm_output - ctx.return_layernorm_output_gathered = ( - return_layernorm_output_gathered and sequence_parallel ) - ctx.set_parallel_mode = set_parallel_mode - ctx.bwd_ln_sm_margin = bwd_ln_sm_margin - ctx.zero_centered_gamma = zero_centered_gamma - ctx.ub_bulk_wgrad = ub_bulk_wgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_overlap_ag = ub_overlap_ag - ctx.debug = debug - - ctx.requires_dgrad = ( - inp.requires_grad or ln_weight.requires_grad or ln_bias.requires_grad - ) - ctx.normalization = normalization - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad( - inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias - ): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase() or is_recomputation: - qstate.is_first_fp8_module = _first_fp8_module - - ctx.wgrad_store = wgrad_store - if is_recomputation: # return the recomputed tensors - return ( - ctx, - inputmat, - ln_weight, - ln_out, + + if cpu_offloading: + mark_not_offload( + ln_weight_cast, + ln_bias, fc1_weight_final, fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, + fc1_bias_cast, fc2_weight_final, fc2_weight, - fc2_bias, - mu, - rsigma, + fc2_bias_cast, ) - # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above - ln_out_for_return = None - if return_layernorm_output: - if return_layernorm_output_gathered: - shape = list(inp_shape) - shape[0] *= tp_size if (sequence_parallel and set_parallel_mode) else 1 - ln_out_for_return = ln_out_return.view(shape) - else: - ln_out_for_return = ln_out_return.view(inp_shape) - return fc2_out, ln_out_for_return, new_fc1_weight_workspace, new_fc2_weight_workspace - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: torch.Tensor, - fc1_weight: torch.Tensor, - fc1_weight_workspace: Optional[torch.Tensor], - fc1_bias: torch.Tensor, - fc2_weight: torch.Tensor, - fc2_weight_workspace: Optional[torch.Tensor], - fc2_bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - - # add recompute_for_bwd - non_tensor_args += (False,) - - return _LayerNormMLP._forward( - ctx, - inp, - ln_weight, - ln_bias, + # FSDP2: Don't save FP8 workspace copies for non-quantized + # weights. Backward will re-quantize from the FSDP2 + # all-gathered weight parameter. (Issue #2681) + fc1_wt_save = fc1_weight_final + fc2_wt_save = fc2_weight_final + if fsdp2_skip_columnwise: + if fc1_weight_final is not fc1_weight: + fc1_wt_save = None + if fc2_weight_final is not fc2_weight: + fc2_wt_save = None + + # Dedup save slots that alias forward inputs or other outputs; + # ``_layernorm_mlp_setup_ctx`` rebuilds the refs. + def _wt_alias(wt_save, weight, new_workspace, workspace, prefix): + if wt_save is None: + return None + if wt_save is weight: + return prefix + "_weight" + if new_workspace is not None and wt_save is new_workspace: + return "new_" + prefix + "_weight_workspace" + if workspace is not None and wt_save is workspace: + return prefix + "_weight_workspace" + return None + + saved_tensor_aliases = ( + "inp" if inputmat is inp else None, + "ln_weight" if ln_weight_cast is ln_weight else None, + ( + "ln_out" + if return_layernorm_output + and ln_out is not None + and ln_out_return is not None + and ln_out is ln_out_return + else None + ), + _wt_alias( + fc1_wt_save, + fc1_weight, + new_fc1_weight_workspace, + args.fc1_weight_workspace, + "fc1", + ), + "fc1_weight", + "fc1_bias" if fc1_bias_cast is not None and fc1_bias_cast is fc1_bias else None, + None, + None, + None, + _wt_alias( + fc2_wt_save, + fc2_weight, + new_fc2_weight_workspace, + args.fc2_weight_workspace, + "fc2", + ), + "fc2_weight", + "fc2_bias" if fc2_bias_cast is not None and fc2_bias_cast is fc2_bias else None, + None, + None, + ) + saved = ( + inputmat, + ln_weight_cast, + ln_out, + fc1_wt_save, fc1_weight, - fc1_weight_workspace, - fc1_bias, + fc1_bias_cast, + fc1_out, + fc1_out_without_bias, + act_out, + fc2_wt_save, fc2_weight, - fc2_weight_workspace, - fc2_bias, - non_tensor_args, + fc2_bias_cast, + mu, + rsigma, + ) + tensors_to_save_from_forward = tuple( + None if alias is not None else tensor + for alias, tensor in zip(saved_tensor_aliases, saved) + ) + ctx_attrs["saved_tensor_aliases"] = saved_tensor_aliases + ctx_attrs["fsdp_shapes"] = fsdp_shapes + ctx_attrs["is_recomputation"] = is_recomputation + + if is_recomputation: # return the recomputed tensors + return None, None, None, None, tensors_to_save_from_forward, ctx_attrs + + # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above + ln_out_for_return = None + if return_layernorm_output: + if return_layernorm_output_gathered: + shape = list(inp_shape) + shape[0] *= tp_size if (sequence_parallel and set_parallel_mode) else 1 + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + 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_setup_ctx( + bwd_args: LayerNormMLPBwdArgs, + fwd_args: LayerNormMLPFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state. + + Returns the tensors that should be passed through ``prepare_for_saving`` + by the caller. ``fwd_outputs`` is ``(fc2_out, ln_out_return, + new_fc1_weight_workspace, new_fc2_weight_workspace)``; the last three + rebuild the deduped save slots. + + With ``ctx_attrs["checkpoint"]`` and no recomputation yet, only the + forward inputs are saved; the rest of ``bwd_args`` is filled when the + backward recomputes the forward. + """ + checkpoint = ctx_attrs.get("checkpoint", False) + is_recomputation = ctx_attrs.get("is_recomputation", False) + if checkpoint and not is_recomputation: + bwd_args.checkpoint = True + bwd_args.autocast_state = ctx_attrs["autocast_state"] + bwd_args.cpu_rng_state = ctx_attrs["cpu_rng_state"] + bwd_args.cuda_rng_state = ctx_attrs["cuda_rng_state"] + # The saved inputs are re-bound to the args when the backward recomputes. + bwd_args.checkpoint_fwd_args = dataclass_replace( + fwd_args, + inp=None, + ln_weight=None, + ln_bias=None, + fc1_weight=None, + fc1_bias=None, + fc2_weight=None, + fc2_bias=None, + fc1_weight_workspace=None, + fc2_weight_workspace=None, + cache_weight=False, + recompute_for_bwd=True, + ) + return tuple(getattr(fwd_args, name) for name in ctx_attrs["saved_tensor_aliases"]) + + inp = fwd_args.inp + fc1_weight = fwd_args.fc1_weight + fc2_weight = fwd_args.fc2_weight + fp8 = fwd_args.fp8 + debug = fwd_args.debug + fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation + fc1_weight_requires_grad = fwd_args.fc1_weight_requires_grad + fc2_weight_requires_grad = fwd_args.fc2_weight_requires_grad + + # Quantizers + bwd_args.fc1_input_quantizer = fwd_args.fc1_input_quantizer + bwd_args.fc2_input_quantizer = fwd_args.fc2_input_quantizer + bwd_args.fc1_weight_quantizer = ( + fc1_weight._quantizer + if (fp8 or debug) and isinstance(fc1_weight, QuantizedTensorStorage) and not debug + else fwd_args.fc1_weight_quantizer + ) + bwd_args.fc2_weight_quantizer = ( + fc2_weight._quantizer + if (fp8 or debug) and isinstance(fc2_weight, QuantizedTensorStorage) and not debug + else fwd_args.fc2_weight_quantizer + ) + bwd_args.fc1_grad_input_quantizer = fwd_args.fc1_grad_input_quantizer + bwd_args.fc1_grad_weight_quantizer = fwd_args.fc1_grad_weight_quantizer + bwd_args.fc1_grad_output_quantizer = fwd_args.fc1_grad_output_quantizer + bwd_args.fc2_grad_input_quantizer = fwd_args.fc2_grad_input_quantizer + bwd_args.fc2_grad_weight_quantizer = fwd_args.fc2_grad_weight_quantizer + bwd_args.fc2_grad_output_quantizer = fwd_args.fc2_grad_output_quantizer + + # Differentiability summary + bwd_args.use_bias = fwd_args.fc2_bias is not None + bwd_args.requires_dgrad = ( + fwd_args.input_requires_grad + or fwd_args.ln_weight_requires_grad + or fwd_args.ln_bias_requires_grad + ) + 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 + + # Normalization + bwd_args.normalization = fwd_args.normalization + bwd_args.zero_centered_gamma = fwd_args.zero_centered_gamma + bwd_args.bwd_ln_sm_margin = fwd_args.bwd_ln_sm_margin + bwd_args.return_layernorm_output = fwd_args.return_layernorm_output + bwd_args.return_layernorm_output_gathered = ( + fwd_args.return_layernorm_output_gathered and fwd_args.sequence_parallel + ) + + # Activation + bwd_args.activation = fwd_args.activation + bwd_args.activation_params = fwd_args.activation_params + bwd_args.bias_gelu_fusion = fwd_args.bias_gelu_fusion + + # 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.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 + bwd_args.debug = debug + + # Tensor / sequence parallelism + bwd_args.set_parallel_mode = fwd_args.set_parallel_mode + bwd_args.tp_group = fwd_args.tp_group + bwd_args.tp_size = fwd_args.tp_size + bwd_args.tensor_parallel = fwd_args.tensor_parallel + bwd_args.sequence_parallel = fwd_args.sequence_parallel + + # Userbuffers + bwd_args.ub_overlap_ag = ( + fwd_args.ub_overlap_ag + and fwd_args.is_grad_enabled + and not fwd_args.return_layernorm_output_gathered + ) + bwd_args.ub_overlap_rs_dgrad = fwd_args.ub_overlap_rs_dgrad + bwd_args.ub_bulk_dgrad = fwd_args.ub_bulk_dgrad + bwd_args.ub_bulk_wgrad = fwd_args.ub_bulk_wgrad + + # FSDP + bwd_args.fsdp_group = fwd_args.fsdp_group + bwd_args.fsdp_shapes = ctx_attrs["fsdp_shapes"] + bwd_args.is_fsdp2 = fwd_args.is_fsdp2 + + # Weight-grad scheduling / accumulation + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + if fuse_wgrad_accumulation: + # Keep weakrefs to weights to preserve attributes like main_grad + # when we need to modify the weight python objects + bwd_args.fc1_weight_ref = weakref.ref(fc1_weight) if fc1_weight_requires_grad else None + bwd_args.fc2_weight_ref = weakref.ref(fc2_weight) if fc2_weight_requires_grad else None + bwd_args.fc1_weight_overwrites_main_grad = getattr(fc1_weight, "overwrite_main_grad", False) + bwd_args.fc2_weight_overwrites_main_grad = getattr(fc2_weight, "overwrite_main_grad", False) + # MCore FSDP creates main_grad lazily before backward, so don't touch it here + if hasattr(fc1_weight, "__fsdp_param__") and hasattr(fc2_weight, "__fsdp_param__"): + bwd_args.fc1_main_grad_func = ( + fc1_weight.get_main_grad if fc1_weight_requires_grad else lambda: None + ) + bwd_args.fc2_main_grad_func = ( + fc2_weight.get_main_grad if fc2_weight_requires_grad else lambda: None + ) + else: + bwd_args.fc1_main_grad_func = lambda: fc1_weight.main_grad + bwd_args.fc2_main_grad_func = lambda: fc2_weight.main_grad + + # Misc + bwd_args.cpu_offloading = fwd_args.cpu_offloading + + saved = list(tensors_to_save_from_forward) + aliases = ctx_attrs["saved_tensor_aliases"] + in_features = inp.shape[-1] + for i, alias in enumerate(aliases): + if alias is None: + continue + if alias == "inp": + saved[i] = inp.view((-1, in_features)) + elif alias == "ln_out": + saved[i] = fwd_outputs[1].view((-1, in_features)) + elif alias == "new_fc1_weight_workspace": + saved[i] = fwd_outputs[2] + elif alias == "new_fc2_weight_workspace": + saved[i] = fwd_outputs[3] + else: + saved[i] = getattr(fwd_args, alias) + return tuple(saved) + + +def _layernorm_mlp_recompute( + bwd_args: LayerNormMLPBwdArgs, ctx: torch.autograd.function.FunctionCtx +) -> None: + """Bind the saved tensors to ``bwd_args``, recomputing the forward when + the module ran with activation checkpointing.""" + tensors = restore_from_func_ctx(ctx) + if not bwd_args.checkpoint: + bwd_args.set_saved_tensors(tensors) + return + + # backward is not in autocast context, so we set the state here + # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) + fwd_args: LayerNormMLPFwdArgs = bwd_args.checkpoint_fwd_args + bwd_args.checkpoint_fwd_args = None + final_autocast_state = FP8GlobalStateManager.get_autocast_state() + FP8GlobalStateManager.set_autocast_state(bwd_args.autocast_state) + if ( + fwd_args.fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( + fwd_args.fp8_meta + ) # set old quantizer state + + # get current rng state + final_cpu_rng_state = torch.get_rng_state() + final_cuda_rng_state = _get_cuda_rng_state() + + # set rng state for fwd + torch.set_rng_state(bwd_args.cpu_rng_state) + _set_cuda_rng_state(bwd_args.cuda_rng_state) + + for name, tensor in zip(_CHECKPOINT_SAVED_ALIASES, tensors): + setattr(fwd_args, name, tensor) + ( + _, + _, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_mlp_forward_impl(fwd_args) + recomputed = _layernorm_mlp_setup_ctx( + bwd_args, + fwd_args, + (None, None, new_fc1_weight_workspace, new_fc2_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + bwd_args.set_saved_tensors(recomputed) + if fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update( + restore_first_module=True ) - @staticmethod - def _recompute(ctx): - # pylint: disable=missing-function-docstring - - tensors = restore_from_func_ctx(ctx) + FP8GlobalStateManager.set_autocast_state(final_autocast_state) + if ( + fwd_args.fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): + FP8GlobalStateManager.restore_fp8_meta_tensors(fwd_args.fp8_meta) # restore quantizers - if ctx.checkpoint: # do recomputation from the original args + # set rng state for fwd + torch.set_rng_state(final_cpu_rng_state) + _set_cuda_rng_state(final_cuda_rng_state) - # backward is not in autocast context, so we set the state here - # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) - final_autocast_state = FP8GlobalStateManager.get_autocast_state() - FP8GlobalStateManager.set_autocast_state(ctx.autocast_state) - if ( - ctx.other_args["fp8"] - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" - ): # only applicable for delayed scaling - FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( - ctx.other_args["fp8_meta"] - ) # set old quantizer state - # get current rng state - final_cpu_rng_state = torch.get_rng_state() - final_cuda_rng_state = _get_cuda_rng_state() +def _layernorm_mlp_backward_impl( + args: LayerNormMLPBwdArgs, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward implementation for the layernorm-MLP layer. - # set rng state for fwd - torch.set_rng_state(ctx.cpu_rng_state) - _set_cuda_rng_state(ctx.cuda_rng_state) + Caller must have populated ``args.grad_output`` / ``args.grad_ln_out`` and + the saved-tensor fields before invocation. Returns ``(dgrad, dgamma, + dbeta, fc1_wgrad, fc1_bias_grad, fc2_wgrad, fc2_bias_grad)``. + """ + with get_nvtx_range_context("_LayerNormMLP_backward"): + inputmat = args.inputmat + ln_weight = args.ln_weight + ln_out = args.ln_out + fc1_weight = args.fc1_weight_fp8 + origin_fc1_weight = args.fc1_weight + fc1_bias = args.fc1_bias + fc1_out = args.fc1_out + fc1_out_without_bias = args.fc1_out_without_bias + act_out = args.act_out + fc2_weight = args.fc2_weight_fp8 + origin_fc2_weight = args.fc2_weight + fc2_bias = args.fc2_bias + mu = args.mu + rsigma = args.rsigma + grad_output_arg = args.grad_output + + # Restore origin weights from weakrefs + # Only needed when fuse_wgrad_accumulation is enabled. + fc1_weight_python_object = None + fc2_weight_python_object = None + fc1_weight_main_grad = None + fc2_weight_main_grad = None + if args.fuse_wgrad_accumulation: + fc1_weight_ref = args.fc1_weight_ref + fc2_weight_ref = args.fc2_weight_ref + args.fc1_weight_ref = None + args.fc2_weight_ref = None + fc1_weight_python_object = fc1_weight_ref() if fc1_weight_ref is not None else None + fc2_weight_python_object = fc2_weight_ref() if fc2_weight_ref is not None else None + if args.fc1_weight_requires_grad: + assert ( + fc1_weight_python_object is not None + ), "fc1_weight was removed while fuse_wgrad_accumulation=True" + fc1_weight_main_grad = args.fc1_main_grad_func() + fc1_weight_python_object.main_grad = fc1_weight_main_grad + if args.fc2_weight_requires_grad: + assert ( + fc2_weight_python_object is not None + ), "fc2_weight was removed while fuse_wgrad_accumulation=True" + fc2_weight_main_grad = args.fc2_main_grad_func() + fc2_weight_python_object.main_grad = fc2_weight_main_grad + + # TODO: Fix this # pylint: disable=fixme + # Gather saved autograd context tensors when running with FSDP + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + # _fsdp_gather_tensors( + # args.fsdp_group, + # args.fsdp_shapes, + # mu, + # rsigma, + # ln_out, + # fc1_out_without_bias if bias_gelu_nvfusion else fc1_out,, + # gelu_out, + # fc1_weight_fp8 if args.fp8 and not isinstance(fc1_weight, Float8Tensor) else None, + # fc2_weight_fp8 if args.fp8 and not isinstance(fc2_weight, Float8Tensor) else None, + # ) - # Unpack saved tensors and pass None for weight workspaces (recomputed from scratch) - ( - inp_r, - ln_weight_r, - ln_bias_r, - fc1_weight_r, - fc1_bias_r, - fc2_weight_r, - fc2_bias_r, - ) = tensors - out = _LayerNormMLP._forward( # recompute - ctx, - inp_r, - ln_weight_r, - ln_bias_r, - fc1_weight_r, - None, - fc1_bias_r, - fc2_weight_r, - None, - fc2_bias_r, - tuple(ctx.other_args.values()), + # Choose whether to use GEMM kernel with split accumulator + dgrad_use_split_accumulator = args.dgrad_use_split_accumulator + wgrad_use_split_accumulator = args.wgrad_use_split_accumulator + + # No need to do bulk DGRAD/WGRAD overlap if WGRAD is not required + ub_bulk_dgrad = args.fc1_weight_requires_grad and args.ub_bulk_dgrad + ub_bulk_wgrad = args.fc1_weight_requires_grad and args.ub_bulk_wgrad + + # Configure quantizer for FC2 grad output tensor + # Note: dgrad GEMM requires row-wise usage, wgrad GEMM + # requires column-wise usage + if args.fc2_grad_output_quantizer is not None: + quantizer = args.fc2_grad_output_quantizer + quantizer.set_usage(rowwise=True, columnwise=True) + if args.ub_overlap_ag: + # Userbuffers only supports communication for one + # tensor usage at a time. Configure quantizer with + # usage for only dgrad GEMM. + quantizer.set_usage(columnwise=False) + # Amax reduction group for FC2 grad output (row-parallel sequence parallel) + set_quantizer_amax_reduction_group( + quantizer, + args.tp_group if (args.sequence_parallel and args.set_parallel_mode) else None, ) - FP8GlobalStateManager.set_autocast_state(final_autocast_state) - if ( - ctx.other_args["fp8"] - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" - ): - FP8GlobalStateManager.restore_fp8_meta_tensors( - ctx.other_args["fp8_meta"] - ) # restore quantizers - - # set rng state for fwd - torch.set_rng_state(final_cpu_rng_state) - _set_cuda_rng_state(final_cuda_rng_state) - - return out - - # load from saved (return ctx is just because the other branch does too) - return tuple([ctx] + tensors) + # Prepare FC2 grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + ub_obj_fc2_dgrad = None + if args.ub_overlap_ag: + ub_obj_fc2_dgrad = get_ub("fc2_dgrad", args.fp8) + args.ub_obj_gradout = ub_obj_fc2_dgrad + ( + grad_output, + fc2_bias_grad, + ) = TransformerEngineBaseModule.grad_output_preprocess( + args, grad_output_arg, True, args.fc2_grad_output_quantizer + ) - @staticmethod - def backward( - ctx, *grad_outputs: Tuple[torch.Tensor, ...] - ) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring - with get_nvtx_range_context("_LayerNormMLP_backward"): - ( # pylint: disable=unbalanced-tuple-unpacking - ctx, - inputmat, - ln_weight, - ln_out, - fc1_weight, - origin_fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, - fc2_weight, - origin_fc2_weight, - fc2_bias, - mu, - rsigma, - ) = _LayerNormMLP._recompute(ctx) - - # Restore origin weights from weakrefs - # Only needed when fuse_wgrad_accumulation is enabled. - fc1_weight_python_object = None - fc2_weight_python_object = None - fc1_weight_main_grad = None - fc2_weight_main_grad = None - if ctx.fuse_wgrad_accumulation: - fc1_weight_python_object_ref = getattr(ctx, "fc1_weight_python_object_ref", None) - fc2_weight_python_object_ref = getattr(ctx, "fc2_weight_python_object_ref", None) - ctx.fc1_weight_python_object_ref = None - ctx.fc2_weight_python_object_ref = None - fc1_weight_python_object = ( - fc1_weight_python_object_ref() - if fc1_weight_python_object_ref is not None - else None - ) - fc2_weight_python_object = ( - fc2_weight_python_object_ref() - if fc2_weight_python_object_ref is not None - else None - ) - if ctx.fc1_weight_requires_grad: - assert ( - fc1_weight_python_object is not None - ), "fc1_weight was removed while fuse_wgrad_accumulation=True" - fc1_weight_main_grad = ctx.fc1_main_grad_func() - fc1_weight_python_object.main_grad = fc1_weight_main_grad - if ctx.fc2_weight_requires_grad: - assert ( - fc2_weight_python_object is not None - ), "fc2_weight was removed while fuse_wgrad_accumulation=True" - fc2_weight_main_grad = ctx.fc2_main_grad_func() - fc2_weight_python_object.main_grad = fc2_weight_main_grad - - # TODO: Fix this # pylint: disable=fixme - # Gather saved autograd context tensors when running with FSDP - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - # _fsdp_gather_tensors( - # ctx.fsdp_group, - # ctx.fsdp_shapes, - # mu, - # rsigma, - # ln_out, - # fc1_out_without_bias if bias_gelu_nvfusion else fc1_out,, - # gelu_out, - # fc1_weight_fp8 if ctx.fp8 and not isinstance(fc1_weight, Float8Tensor) else None, - # fc2_weight_fp8 if ctx.fp8 and not isinstance(fc2_weight, Float8Tensor) else None, - # ) - - # Choose whether to use GEMM kernel with split accumulator - dgrad_use_split_accumulator = _2X_ACC_DGRAD - wgrad_use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - 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 - - # No need to do bulk DGRAD/WGRAD overlap if WGRAD is not required - ctx.ub_bulk_dgrad = ctx.fc1_weight_requires_grad and ctx.ub_bulk_dgrad - ctx.ub_bulk_wgrad = ctx.fc1_weight_requires_grad and ctx.ub_bulk_wgrad - - # Configure quantizer for FC2 grad output tensor - # Note: dgrad GEMM requires row-wise usage, wgrad GEMM - # requires column-wise usage - if ctx.fc2_grad_output_quantizer is not None: - quantizer = ctx.fc2_grad_output_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: - # Userbuffers only supports communication for one - # tensor usage at a time. Configure quantizer with - # usage for only dgrad GEMM. - quantizer.set_usage(columnwise=False) - # Amax reduction group for FC2 grad output (row-parallel sequence parallel) - set_quantizer_amax_reduction_group( + # Launch tensor-parallel communication for FC1 GEMM input + ln_out_total = None + ln_out_total_work = None + ub_obj_fc1_dgrad = None + if args.fc1_weight_requires_grad and args.tensor_parallel and args.sequence_parallel: + quantizer = None + if args.fp8 or args.debug: + quantizer = args.fc1_input_quantizer + set_quantizer_usage_for_wgrad_all_gather(quantizer) + if ub_bulk_dgrad: + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_fc1_dgrad, + ln_out, quantizer, - ctx.tp_group if (ctx.sequence_parallel and ctx.set_parallel_mode) else None, + args.tp_group, + ) + else: + ln_out_total, ln_out_total_work = gather_along_first_dim( + ln_out, + args.tp_group, + async_op=True, + quantizer=quantizer, ) + else: + ln_out_total = ln_out - # Prepare FC2 grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - ub_obj_fc2_dgrad = None - if ctx.ub_overlap_ag: - ub_obj_fc2_dgrad = get_ub("fc2_dgrad", ctx.fp8) - ctx.ub_obj_gradout = ub_obj_fc2_dgrad - ( + # Check whether to output wgrad GEMM directly into main grad + if args.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + args.fuse_wgrad_accumulation and not args.is_first_microbatch + ) + else: + accumulate_wgrad_into_param_main_grad = args.fuse_wgrad_accumulation + + # -------------------------------------------------- + # FC2 DGRAD + # -------------------------------------------------- + + # There are 6 possible fusion paths + # 1 high-precision bias_gelu_fusion: gemm, FC1_bias + gelu, + # 2 high-precision fc2_dgrad_gemm_gelu_fusion: gemm + gelu, FC1_bias + quantize + # 3 fp8 activation+bias+quantize fusion: gemm, activation + FC1_bias + quantize + # 4 fp8 bias+quantize fusion: gemm, activation, FC1_bias + quantize + # 5 high-precision unfused: gemm, activation, FC1_bias + FC1_gemm + # 6 fp8 unfused: gemm, activation, FC1_bias + FC1_gemm + fc2_dgrad_gemm_gelu_fusion = ( + not args.fp8 + and (args.activation == "gelu") + and (not args.bias_gelu_fusion) + and (not args.debug) + ) + + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved to avoid forward memory + # accumulation. (Issue #2681) + if fc2_weight is None: + if isinstance(origin_fc2_weight, QuantizedTensorStorage): + fc2_weight = origin_fc2_weight + elif args.fc2_weight_quantizer is not None: + args.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc2_weight = args.fc2_weight_quantizer(origin_fc2_weight) + + # Make sure required data is available + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(rowwise_usage=True) + if args.fc2_weight_quantizer is not None and isinstance(fc2_weight, QuantizedTensorStorage): + fc2_weight.update_usage(columnwise_usage=True) + + # Perform GEMM + gemm_output, *_ = general_gemm( + fc2_weight, + grad_output, + layout="NN", + grad=True, + quantization_params=( + args.fc1_grad_input_quantizer if fc2_dgrad_gemm_gelu_fusion or args.debug else None + ), # high precision to activation + out_dtype=args.activation_dtype, + gelu=fc2_dgrad_gemm_gelu_fusion, + gelu_in=fc1_out if fc2_dgrad_gemm_gelu_fusion else None, + use_split_accumulator=dgrad_use_split_accumulator, + ub=ub_obj_fc2_dgrad, + ub_type=tex.CommOverlapType.AG if args.ub_overlap_ag else None, + ) + + # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(fc2_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc2_weight) + + # Prepare input grad tensor + dact = None + fc2_dgrad = None + if fc2_dgrad_gemm_gelu_fusion: + dact = gemm_output + else: + fc2_dgrad = gemm_output + + # -------------------------------------------------- + # Finished FC2 DGRAD... + # -------------------------------------------------- + + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for fc2_wgrad. Userbuffers leaves the + # gathered tensor in its persistent buffer; cuBLASMp does not, so + # we gather here. Route through the same FP8-aware all-gather as + # the non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Columnwise data needed for fc2_wgrad + # is produced by ``update_usage(columnwise_usage=True)`` further + # below. + if ( + args.fc2_weight_requires_grad + and args.ub_overlap_ag + and args.ub_obj_gradout is not None + and args.ub_obj_gradout.with_cublasmp() + ): + if args.fc2_grad_output_quantizer is not None: + set_quantizer_usage_for_wgrad_all_gather(args.fc2_grad_output_quantizer) + grad_output, _ = gather_along_first_dim( grad_output, - fc2_bias_grad, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, grad_outputs[0], True, ctx.fc2_grad_output_quantizer + args.tp_group, + quantizer=args.fc2_grad_output_quantizer, ) - # Launch tensor-parallel communication for FC1 GEMM input - ln_out_total = None - ln_out_total_work = None - ub_obj_fc1_dgrad = None - if ctx.fc1_weight_requires_grad and ctx.tensor_parallel and ctx.sequence_parallel: - quantizer = None - if ctx.fp8 or ctx.debug: - quantizer = ctx.fc1_input_quantizer - set_quantizer_usage_for_wgrad_all_gather(quantizer) - if ctx.ub_bulk_dgrad: - ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_fc1_dgrad, - ln_out, - quantizer, - ctx.tp_group, - ) - else: - ln_out_total, ln_out_total_work = gather_along_first_dim( - ln_out, - ctx.tp_group, - async_op=True, - quantizer=quantizer, + # -------------------------------------------------- + # FC2 WGRAD + # -------------------------------------------------- + + fc2_wgrad = None + if args.fc2_weight_requires_grad: + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.ub_overlap_ag and isinstance(args.fc2_grad_output_quantizer, MXFP8Quantizer): + # UB does not support pipelined overlapping grad output + # all-gather with wgrad GEMM. Also, we can't + # convert row-scaled MXFP8 to column-scaled, so we + # can't reuse the grad output that was gathered + # for the dgrad GEMM. We work around by explicitly + # overlapping the AG operation with the dgrad GEMM. + + # Get the communication stream from the dgrad GEMM to use for the AG + dgrad_send_stream, dgrad_recv_stream = ub_obj_fc2_dgrad.get_communication_stream() + + ub_obj_fc2_wgrad = get_ub("fc2_wgrad", args.fp8) + + args.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + + # We use the send stream to copy into the userbuffers. + # This is the same stream that we will use to access the data in the AG, + # so we dont need to add any syncs yet. + with torch.cuda.stream(dgrad_send_stream): + grad_output, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_fc2_wgrad, + grad_output_arg, + args.fc2_grad_output_quantizer, + args.tp_group, ) - else: - ln_out_total = ln_out - # Check whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch + # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm + tex.bulk_overlap_ag_with_external_gemm( + ub_obj_fc2_wgrad, dgrad_send_stream, dgrad_recv_stream ) - else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - # -------------------------------------------------- - # FC2 DGRAD - # -------------------------------------------------- - - # There are 6 possible fusion paths - # 1 high-precision bias_gelu_fusion: gemm, FC1_bias + gelu, - # 2 high-precision fc2_dgrad_gemm_gelu_fusion: gemm + gelu, FC1_bias + quantize - # 3 fp8 activation+bias+quantize fusion: gemm, activation + FC1_bias + quantize - # 4 fp8 bias+quantize fusion: gemm, activation, FC1_bias + quantize - # 5 high-precision unfused: gemm, activation, FC1_bias + FC1_gemm - # 6 fp8 unfused: gemm, activation, FC1_bias + FC1_gemm - fc2_dgrad_gemm_gelu_fusion = ( - not ctx.fp8 - and (ctx.activation == "gelu") - and (not ctx.bias_gelu_fusion) - and (not ctx.debug) - ) - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved to avoid forward memory - # accumulation. (Issue #2681) - if fc2_weight is None: - if isinstance(origin_fc2_weight, QuantizedTensorStorage): - fc2_weight = origin_fc2_weight - elif ctx.fc2_weight_quantizer is not None: - ctx.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) - fc2_weight = ctx.fc2_weight_quantizer(origin_fc2_weight) - - # Make sure required data is available - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(rowwise_usage=True) - if ctx.fc2_weight_quantizer is not None and isinstance( - fc2_weight, QuantizedTensorStorage - ): - fc2_weight.update_usage(columnwise_usage=True) - - # Perform GEMM - gemm_output, *_ = general_gemm( - fc2_weight, - grad_output, - layout="NN", - grad=True, - quantization_params=( - ctx.fc1_grad_input_quantizer - if fc2_dgrad_gemm_gelu_fusion or ctx.debug - else None - ), # high precision to activation - out_dtype=ctx.activation_dtype, - gelu=fc2_dgrad_gemm_gelu_fusion, - gelu_in=fc1_out if fc2_dgrad_gemm_gelu_fusion else None, - use_split_accumulator=dgrad_use_split_accumulator, - ub=ub_obj_fc2_dgrad, - ub_type=tex.CommOverlapType.AG if ctx.ub_overlap_ag else None, - ) + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.fp8 or args.debug: + if isinstance(act_out, QuantizedTensorStorage): + act_out.update_usage(columnwise_usage=True) + else: + args.fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) + act_out = args.fc2_input_quantizer(act_out) + + if args.fp8 or args.debug: + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(columnwise_usage=True) + else: + args.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + grad_output = args.fc2_grad_output_quantizer(grad_output) + + # Whether to set grad arg in general_gemm + grad_arg = True + if args.fp8 and args.fp8_recipe.float8_block_scaling(): + grad_arg = False + + # Arguments to include in wgrad GEMM closure + fc2_wgrad_gemm_kwargs = { + "out_dtype": ( + fc2_weight_main_grad.dtype + if args.fuse_wgrad_accumulation + else args.activation_dtype + ), + "quantization_params": args.fc2_grad_weight_quantizer, # wgrad in high precision + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not args.fc2_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": fc2_weight_main_grad if args.fuse_wgrad_accumulation else None, + "bias": fc2_bias if fc2_bias is not None and fc2_bias_grad is None else None, + "use_split_accumulator": wgrad_use_split_accumulator, + "grad": grad_arg, + } + + def fc2_wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform FC2 WGRAD GEMM + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + dw, db, *_ = general_gemm(x, dy, **fc2_wgrad_gemm_kwargs) + return dw, db - # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM - # to prevent them from persisting on the all-gathered buffer. - # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs - # even when backward follows gradient-checkpoint recomputation. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(fc2_weight, QuantizedTensorStorage): - clear_columnwise_cache(fc2_weight) - - # Prepare input grad tensor - dact = None - fc2_dgrad = None - if fc2_dgrad_gemm_gelu_fusion: - dact = gemm_output + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + args.wgrad_store.put([act_out, grad_output], fc2_wgrad_gemm) else: - fc2_dgrad = gemm_output - - # -------------------------------------------------- - # Finished FC2 DGRAD... - # -------------------------------------------------- - - # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and - # does not preserve it for fc2_wgrad. Userbuffers leaves the - # gathered tensor in its persistent buffer; cuBLASMp does not, so - # we gather here. Route through the same FP8-aware all-gather as - # the non-overlap path in - # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing - # the grad_output quantizer. Columnwise data needed for fc2_wgrad - # is produced by ``update_usage(columnwise_usage=True)`` further - # below. - if ( - ctx.fc2_weight_requires_grad - and ctx.ub_overlap_ag - and ctx.ub_obj_gradout is not None - and ctx.ub_obj_gradout.with_cublasmp() - ): - if ctx.fc2_grad_output_quantizer is not None: - set_quantizer_usage_for_wgrad_all_gather(ctx.fc2_grad_output_quantizer) - grad_output, _ = gather_along_first_dim( - grad_output, - ctx.tp_group, - quantizer=ctx.fc2_grad_output_quantizer, - ) - # -------------------------------------------------- - # FC2 WGRAD - # -------------------------------------------------- + # Call wgrad GEMM now + fc2_wgrad, fc2_bias_grad_ = fc2_wgrad_gemm(act_out, grad_output) - fc2_wgrad = None - if ctx.fc2_weight_requires_grad: - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.ub_overlap_ag and isinstance(ctx.fc2_grad_output_quantizer, MXFP8Quantizer): - # UB does not support pipelined overlapping grad output - # all-gather with wgrad GEMM. Also, we can't - # convert row-scaled MXFP8 to column-scaled, so we - # can't reuse the grad output that was gathered - # for the dgrad GEMM. We work around by explicitly - # overlapping the AG operation with the dgrad GEMM. - - # Get the communication stream from the dgrad GEMM to use for the AG - dgrad_send_stream, dgrad_recv_stream = ( - ub_obj_fc2_dgrad.get_communication_stream() + # 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: + # 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_ + del fc2_bias_grad_ + + # Deallocate input tensor if permitted + if args.wgrad_store is not None and not args.wgrad_store.delay_wgrad_compute(): + clear_tensor_data(act_out) + + # -------------------------------------------------- + # Finished FC2 WGRAD... + # -------------------------------------------------- + + # bias computation + act_params = args.activation_params or {} + fc1_bias_grad = None + fuse_gemm_and_bias_fc1_wgrad = False + if args.fc1_grad_output_quantizer is not None: + args.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + if args.bias_gelu_fusion: + # Fusion: gemm, bias + gelu + assert args.activation == "gelu" + assert not args.fp8 + fc1_bias_grad, dact = bgrad_dgelu_fused(fc2_dgrad, fc1_out_without_bias, fc1_bias) + if args.fc1_grad_output_quantizer is not None: + dact = args.fc1_grad_output_quantizer(dact) + elif args.debug: + dact_func = _act_func(args.activation)[1] + dact = dact_func(fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params) + 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 + and args.fp8 + ): + # Fusion: gemm, bias + gelu + quantize + dbias_dact_quantize_func = _act_func( + args.activation, args.fp8_recipe if args.fp8 else None + )[2] + fc1_bias_grad, dact = dbias_dact_quantize_func( + fc2_dgrad, + fc1_out.to(args.activation_dtype), + args.fc1_grad_output_quantizer, + **act_params, + ) # quantize bgrad gelu fused + 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] + dact = activation_func_bwd( + fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params + ) # activation in high precision + + if args.fp8: + # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now + if ( + isinstance( + args.fc1_grad_output_quantizer, + (Float8BlockQuantizer, IdentityQuantizer), ) + or args.fp8_recipe.custom() + ): + fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) + dact = args.fc1_grad_output_quantizer(dact) + else: + fc1_bias_grad, dact = tex.bgrad_quantize(dact, args.fc1_grad_output_quantizer) + else: + fuse_gemm_and_bias_fc1_wgrad = ( + True # fc1_bias_grad is computed later, fused with wgrad gemm for the FC1 + ) + # it may not be calculated in case wgrad is not required. + if fc1_bias is not None: + if not args.fc1_weight_requires_grad and args.fc1_bias_requires_grad: + fc1_bias_grad = dact.sum(dim=0) + + # Overwrite data. Deleting the tensor does not release underlying memory. + clear_tensor_data(fc1_out, fc1_out_without_bias) + + # Set UB algo and UB obj for fc1_dgrad/wgrad bulk/pipelined overlap + ub_obj_fc1_dgrad = None + ub_obj_fc1_wgrad = None + ub_type_fc1_dgrad = None + ub_type_fc1_wgrad = None + fc1_dgrad_shape = [reduce(multiply_op, inputmat.shape[:-1]), inputmat.shape[-1]] + if args.ub_overlap_rs_dgrad: + # Overlap DGRAD+RS + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ub_type_fc1_dgrad = tex.CommOverlapType.RS + else: + if ub_bulk_dgrad: + # Overlap ln_out all-gather with DGRAD compute + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ub_type_fc1_dgrad = tex.CommOverlapType.AG + if ub_bulk_wgrad: + # Overlap FC1 DGRAD reduce-scatter with WGRAD compute + ub_obj_fc1_wgrad = get_ub("fc1_wgrad", args.fp8) + ub_type_fc1_wgrad = tex.CommOverlapType.RS + + # -------------------------------------------------- + # FC1 DGRAD + # -------------------------------------------------- + + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + if fc1_weight is None: + if isinstance(origin_fc1_weight, QuantizedTensorStorage): + fc1_weight = origin_fc1_weight + elif args.fc1_weight_quantizer is not None: + args.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc1_weight = args.fc1_weight_quantizer(origin_fc1_weight) + + # Make sure required data is available + if args.fc1_weight_quantizer is not None and isinstance(fc1_weight, QuantizedTensorStorage): + fc1_weight.update_usage(columnwise_usage=True) + + # Output buffers for Userbuffers reduce-scatter + gemm_out = None + reduce_scatter_out = None + if args.ub_overlap_rs_dgrad: + reduce_scatter_out = torch.empty( + fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" + ) + if ub_bulk_wgrad: + gemm_out = ub_obj_fc1_wgrad.get_buffer(local_chunk=False) - ub_obj_fc2_wgrad = get_ub("fc2_wgrad", ctx.fp8) + # dgrad GEMM + gemm_out, *_, reduce_scatter_out = general_gemm( + fc1_weight, + dact, + out=gemm_out, + out_dtype=args.activation_dtype, + quantization_params=args.fc1_grad_input_quantizer, + layout="NN", + grad=True, + use_split_accumulator=dgrad_use_split_accumulator, + ub=ub_obj_fc1_dgrad, + ub_type=ub_type_fc1_dgrad, + extra_output=reduce_scatter_out, + bulk_overlap=ub_bulk_dgrad, + ) - ctx.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(fc1_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc1_weight) + + # Prepare grad input tensor + # Note: Perform tensor-parallel communication + fc1_dgrad = None + fc1_dgrad_work = None + if args.ub_overlap_rs_dgrad: + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + fc1_dgrad = ( + gemm_out + if ub_obj_fc1_dgrad is not None and ub_obj_fc1_dgrad.with_cublasmp() + else reduce_scatter_out + ) + elif ub_bulk_wgrad: + fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True) + 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: + fc1_dgrad = fc1_dgrad + args.grad_ln_out.view_as(fc1_dgrad) + fc1_dgrad, fc1_dgrad_work = reduce_scatter_along_first_dim( + fc1_dgrad, + args.tp_group, + async_op=True, + ) + elif args.tensor_parallel: + fc1_dgrad, fc1_dgrad_work = allreduce(fc1_dgrad, args.tp_group, async_op=True) + else: + fc1_dgrad = gemm_out - # We use the send stream to copy into the userbuffers. - # This is the same stream that we will use to access the data in the AG, - # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): - grad_output, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_fc2_wgrad, - grad_outputs[0], - ctx.fc2_grad_output_quantizer, - ctx.tp_group, - ) + # -------------------------------------------------- + # Finished FC1 DGRAD... + # -------------------------------------------------- - # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm - tex.bulk_overlap_ag_with_external_gemm( - ub_obj_fc2_wgrad, dgrad_send_stream, dgrad_recv_stream - ) + # -------------------------------------------------- + # FC1 WGRAD + # -------------------------------------------------- + fc1_wgrad = None + if args.fc1_weight_requires_grad: - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.fp8 or ctx.debug: - if isinstance(act_out, QuantizedTensorStorage): - act_out.update_usage(columnwise_usage=True) - else: - ctx.fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) - act_out = ctx.fc2_input_quantizer(act_out) - - if ctx.fp8 or ctx.debug: - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(columnwise_usage=True) - else: - ctx.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - grad_output = ctx.fc2_grad_output_quantizer(grad_output) - - # Whether to set grad arg in general_gemm - grad_arg = True - if ctx.fp8 and ctx.fp8_recipe.float8_block_scaling(): - grad_arg = False - - # Arguments to include in wgrad GEMM closure - fc2_wgrad_gemm_kwargs = { - "out_dtype": ( - fc2_weight_main_grad.dtype - if ctx.fuse_wgrad_accumulation - else ctx.activation_dtype - ), - "quantization_params": ctx.fc2_grad_weight_quantizer, # wgrad in high precision - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "fc2_weight_overwrites_main_grad", False) - else False - ), - "layout": "NT", - "out": fc2_weight_main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": fc2_bias if fc2_bias is not None and fc2_bias_grad is None else None, - "use_split_accumulator": wgrad_use_split_accumulator, - "grad": grad_arg, - } - - def fc2_wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform FC2 WGRAD GEMM - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - dw, db, *_ = general_gemm(x, dy, **fc2_wgrad_gemm_kwargs) - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([act_out, grad_output], fc2_wgrad_gemm) + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if args.fp8 or args.debug: + if isinstance(ln_out_total, QuantizedTensorStorage): + ln_out_total.update_usage(columnwise_usage=True) else: - - # Call wgrad GEMM now - fc2_wgrad, fc2_bias_grad_ = fc2_wgrad_gemm(act_out, grad_output) - - # Update grad bias if needed - if fc2_bias_grad is None: - if ( - ctx.fp8 - and ctx.fp8_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_ - del fc2_bias_grad_ - - # Deallocate input tensor if permitted - if ctx.wgrad_store is not None and not ctx.wgrad_store.delay_wgrad_compute(): - clear_tensor_data(act_out) - - # -------------------------------------------------- - # Finished FC2 WGRAD... - # -------------------------------------------------- - - # bias computation - act_params = ctx.activation_params or {} - fc1_bias_grad = None - fuse_gemm_and_bias_fc1_wgrad = False - if ctx.fc1_grad_output_quantizer is not None: - ctx.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.bias_gelu_fusion: - # Fusion: gemm, bias + gelu - assert ctx.activation == "gelu" - assert not ctx.fp8 - fc1_bias_grad, dact = bgrad_dgelu_fused(fc2_dgrad, fc1_out_without_bias, fc1_bias) - if ctx.fc1_grad_output_quantizer is not None: - dact = ctx.fc1_grad_output_quantizer(dact) - elif ctx.debug: - dact_func = _act_func(ctx.activation)[1] - dact = dact_func(fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params) - fc1_bias_grad = dact.sum(dim=0) - dact = ctx.fc1_grad_output_quantizer(dact) - elif ( - _act_func(ctx.activation, ctx.fp8_recipe if ctx.fp8 else None)[2] is not None - and ctx.fp8 - ): - # Fusion: gemm, bias + gelu + quantize - dbias_dact_quantize_func = _act_func( - ctx.activation, ctx.fp8_recipe if ctx.fp8 else None - )[2] - fc1_bias_grad, dact = dbias_dact_quantize_func( - fc2_dgrad, - fc1_out.to(ctx.activation_dtype), - ctx.fc1_grad_output_quantizer, - **act_params, - ) # quantize bgrad gelu fused - else: - # Fusion: gemm + gelu, - if not fc2_dgrad_gemm_gelu_fusion: - activation_func_bwd = _act_func( - ctx.activation, ctx.fp8_recipe if ctx.fp8 else None - )[1] - dact = activation_func_bwd( - fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params - ) # activation in high precision - - if ctx.fp8: - # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now - if ( - isinstance( - ctx.fc1_grad_output_quantizer, - (Float8BlockQuantizer, IdentityQuantizer), - ) - or ctx.fp8_recipe.custom() - ): - fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) - dact = ctx.fc1_grad_output_quantizer(dact) - else: - fc1_bias_grad, dact = tex.bgrad_quantize( - dact, ctx.fc1_grad_output_quantizer - ) + args.fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) + ln_out_total = args.fc1_input_quantizer(ln_out_total) + + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.fp8 or args.debug: + if isinstance(dact, QuantizedTensorStorage): + dact.update_usage(columnwise_usage=True) else: - fuse_gemm_and_bias_fc1_wgrad = ( - True # fc1_bias_grad is computed later, fused with wgrad gemm for the FC1 - ) - # it may not be calculated in case wgrad is not required. - if fc1_bias is not None: - if not ctx.fc1_weight_requires_grad and fc1_bias.requires_grad: - fc1_bias_grad = dact.sum(dim=0) - - # Overwrite data. Deleting the tensor does not release underlying memory. - clear_tensor_data(fc1_out, fc1_out_without_bias) - - # Set UB algo and UB obj for fc1_dgrad/wgrad bulk/pipelined overlap - ub_obj_fc1_dgrad = None - ub_obj_fc1_wgrad = None - ub_type_fc1_dgrad = None - ub_type_fc1_wgrad = None - fc1_dgrad_shape = [reduce(multiply_op, inputmat.shape[:-1]), inputmat.shape[-1]] - if ctx.ub_overlap_rs_dgrad: - # Overlap DGRAD+RS - ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) - ub_type_fc1_dgrad = tex.CommOverlapType.RS - else: - if ctx.ub_bulk_dgrad: - # Overlap ln_out all-gather with DGRAD compute - ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) - ub_type_fc1_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: - # Overlap FC1 DGRAD reduce-scatter with WGRAD compute - ub_obj_fc1_wgrad = get_ub("fc1_wgrad", ctx.fp8) - ub_type_fc1_wgrad = tex.CommOverlapType.RS - - # -------------------------------------------------- - # FC1 DGRAD - # -------------------------------------------------- - - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved. (Issue #2681) - if fc1_weight is None: - if isinstance(origin_fc1_weight, QuantizedTensorStorage): - fc1_weight = origin_fc1_weight - elif ctx.fc1_weight_quantizer is not None: - ctx.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) - fc1_weight = ctx.fc1_weight_quantizer(origin_fc1_weight) - - # Make sure required data is available - if ctx.fc1_weight_quantizer is not None and isinstance( - fc1_weight, QuantizedTensorStorage - ): - fc1_weight.update_usage(columnwise_usage=True) - - # Output buffers for Userbuffers reduce-scatter - gemm_out = None + args.fc1_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + dact = args.fc1_grad_output_quantizer(dact) + + # Output buffer for overlapping grad input + # reduce-scatter with wgrad GEMM reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: + if ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" ) - if ctx.ub_bulk_wgrad: - gemm_out = ub_obj_fc1_wgrad.get_buffer(local_chunk=False) - # dgrad GEMM - gemm_out, *_, reduce_scatter_out = general_gemm( - fc1_weight, - dact, - out=gemm_out, - out_dtype=ctx.activation_dtype, - quantization_params=ctx.fc1_grad_input_quantizer, - layout="NN", - grad=True, - use_split_accumulator=dgrad_use_split_accumulator, - ub=ub_obj_fc1_dgrad, - ub_type=ub_type_fc1_dgrad, - extra_output=reduce_scatter_out, - bulk_overlap=ctx.ub_bulk_dgrad, - ) + # Arguments to include in wgrad GEMM closure + fc1_wgrad_gemm_kwargs = { + "out_dtype": ( + fc1_weight_main_grad.dtype + if args.fuse_wgrad_accumulation + else args.activation_dtype + ), + "quantization_params": args.fc1_grad_weight_quantizer, + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not args.fc1_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": fc1_weight_main_grad if args.fuse_wgrad_accumulation else None, + "bias": fc1_bias if fuse_gemm_and_bias_fc1_wgrad else None, + "use_split_accumulator": wgrad_use_split_accumulator, + "grad": fuse_gemm_and_bias_fc1_wgrad, + "ub": ub_obj_fc1_wgrad, + "ub_type": ub_type_fc1_wgrad, + "extra_output": reduce_scatter_out, + "bulk_overlap": ub_bulk_wgrad, + } - # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM - # to prevent them from persisting on the all-gathered buffer. - # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs - # even when backward follows gradient-checkpoint recomputation. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(fc1_weight, QuantizedTensorStorage): - clear_columnwise_cache(fc1_weight) - - # Prepare grad input tensor - # Note: Perform tensor-parallel communication - fc1_dgrad = None - fc1_dgrad_work = None - if ctx.ub_overlap_rs_dgrad: - # cuBLASMp writes the reduce-scattered dgrad directly into the - # GEMM output tensor; Userbuffers uses the extra-output buffer. - fc1_dgrad = ( - gemm_out - if ub_obj_fc1_dgrad is not None and ub_obj_fc1_dgrad.with_cublasmp() - else reduce_scatter_out - ) - elif ctx.ub_bulk_wgrad: - fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True) - elif ctx.set_parallel_mode and not ctx.ub_bulk_wgrad: - fc1_dgrad = gemm_out - if ctx.sequence_parallel: - if ctx.return_layernorm_output and ctx.return_layernorm_output_gathered: - fc1_dgrad = fc1_dgrad + grad_outputs[1].view_as(fc1_dgrad) - fc1_dgrad, fc1_dgrad_work = reduce_scatter_along_first_dim( - fc1_dgrad, - ctx.tp_group, - async_op=True, + def fc1_wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + _is_delayed: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform FC1 WGRAD GEMM + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + dw, db, *_ = general_gemm(x, dy, **fc1_wgrad_gemm_kwargs) + return dw, db + + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + if ( + fc1_wgrad_gemm_kwargs["ub"] is not None + or fc1_wgrad_gemm_kwargs["ub_type"] is not None + or fc1_wgrad_gemm_kwargs["extra_output"] is not None + or fc1_wgrad_gemm_kwargs["bulk_overlap"] + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported " + "with Userbuffers (tensor-parallel communication overlapping)" ) - elif ctx.tensor_parallel: - fc1_dgrad, fc1_dgrad_work = allreduce(fc1_dgrad, ctx.tp_group, async_op=True) + args.wgrad_store.put([ln_out_total, dact], fc1_wgrad_gemm) + if fuse_gemm_and_bias_fc1_wgrad: + fc1_bias_grad = None else: - fc1_dgrad = gemm_out - # -------------------------------------------------- - # Finished FC1 DGRAD... - # -------------------------------------------------- + # Call wgrad GEMM now + fc1_wgrad_outputs = fc1_wgrad_gemm(ln_out_total, dact) + if fuse_gemm_and_bias_fc1_wgrad: + fc1_wgrad, fc1_bias_grad = fc1_wgrad_outputs + else: + fc1_wgrad, _ = fc1_wgrad_outputs - # -------------------------------------------------- - # FC1 WGRAD - # -------------------------------------------------- - fc1_wgrad = None - if ctx.fc1_weight_requires_grad: - - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if ctx.fp8 or ctx.debug: - if isinstance(ln_out_total, QuantizedTensorStorage): - ln_out_total.update_usage(columnwise_usage=True) - else: - ctx.fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) - ln_out_total = ctx.fc1_input_quantizer(ln_out_total) - - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.fp8 or ctx.debug: - if isinstance(dact, QuantizedTensorStorage): - dact.update_usage(columnwise_usage=True) - else: - ctx.fc1_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - dact = ctx.fc1_grad_output_quantizer(dact) - - # Output buffer for overlapping grad input - # reduce-scatter with wgrad GEMM - reduce_scatter_out = None - if ctx.ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): - reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" - ) + # Deallocate tensors if permitted + clear_tensor_data(dact) + if not args.return_layernorm_output_gathered: + clear_tensor_data(ln_out_total) - # Arguments to include in wgrad GEMM closure - fc1_wgrad_gemm_kwargs = { - "out_dtype": ( - fc1_weight_main_grad.dtype - if ctx.fuse_wgrad_accumulation - else ctx.activation_dtype - ), - "quantization_params": ctx.fc1_grad_weight_quantizer, - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "fc1_weight_overwrites_main_grad", False) - else False - ), - "layout": "NT", - "out": fc1_weight_main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": fc1_bias if fuse_gemm_and_bias_fc1_wgrad else None, - "use_split_accumulator": wgrad_use_split_accumulator, - "grad": fuse_gemm_and_bias_fc1_wgrad, - "ub": ub_obj_fc1_wgrad, - "ub_type": ub_type_fc1_wgrad, - "extra_output": reduce_scatter_out, - "bulk_overlap": ctx.ub_bulk_wgrad, - } - - def fc1_wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - _is_delayed: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform FC1 WGRAD GEMM - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - dw, db, *_ = general_gemm(x, dy, **fc1_wgrad_gemm_kwargs) - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - if ( - fc1_wgrad_gemm_kwargs["ub"] is not None - or fc1_wgrad_gemm_kwargs["ub_type"] is not None - or fc1_wgrad_gemm_kwargs["extra_output"] is not None - or fc1_wgrad_gemm_kwargs["bulk_overlap"] - ): - raise NotImplementedError( - "Delayed weight grad computation is not supported " - "with Userbuffers (tensor-parallel communication overlapping)" - ) - ctx.wgrad_store.put([ln_out_total, dact], fc1_wgrad_gemm) - if fuse_gemm_and_bias_fc1_wgrad: - fc1_bias_grad = None + # Update grad input if overlapping reduce-scatter with wgrad GEMM + if ub_bulk_wgrad: + if ub_obj_fc1_wgrad.is_fp8_ubuf(): + fc1_dgrad = reduce_scatter_out else: + fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True).clone() - # Call wgrad GEMM now - fc1_wgrad_outputs = fc1_wgrad_gemm(ln_out_total, dact) - if fuse_gemm_and_bias_fc1_wgrad: - fc1_wgrad, fc1_bias_grad = fc1_wgrad_outputs - else: - fc1_wgrad, _ = fc1_wgrad_outputs - - # Deallocate tensors if permitted - clear_tensor_data(dact) - if not ctx.return_layernorm_output_gathered: - clear_tensor_data(ln_out_total) - - # Update grad input if overlapping reduce-scatter with wgrad GEMM - if ctx.ub_bulk_wgrad: - if ub_obj_fc1_wgrad.is_fp8_ubuf(): - fc1_dgrad = reduce_scatter_out - else: - fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True).clone() - - # -------------------------------------------------- - # Finished FC1 WGRAD... - # -------------------------------------------------- - - # Make sure all tensor-parallel communication is finished - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if fc1_dgrad_work is not None: - fc1_dgrad_work.wait() - fc1_dgrad_work = None + # -------------------------------------------------- + # Finished FC1 WGRAD... + # -------------------------------------------------- - # Residual gradient - dgrad = fc1_dgrad.view(inputmat.shape) - if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - dgrad = dgrad + grad_outputs[1].view_as(dgrad) + # Make sure all tensor-parallel communication is finished + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if fc1_dgrad_work is not None: + fc1_dgrad_work.wait() + fc1_dgrad_work = None - # Norm gradient - dgamma = None + # Residual gradient + dgrad = fc1_dgrad.view(inputmat.shape) + if args.return_layernorm_output and not args.return_layernorm_output_gathered: + dgrad = dgrad + args.grad_ln_out.view_as(dgrad) + + # Norm gradient + dgamma = None + dbeta = None + if args.normalization == "LayerNorm": + dgrad, dgamma, dbeta = tex.layernorm_bwd( + dgrad, + inputmat, + mu, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + elif args.normalization == "RMSNorm": + dgrad, dgamma = tex.rmsnorm_bwd( + dgrad, + inputmat, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) dbeta = None - if ctx.normalization == "LayerNorm": - dgrad, dgamma, dbeta = tex.layernorm_bwd( - dgrad, - inputmat, - mu, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + clear_tensor_data(mu, rsigma) + + if args.fc1_weight_requires_grad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr( + fc1_weight_python_object, "grad_added_to_main_grad" + ): + fc1_weight_python_object.grad_added_to_main_grad = True + if getattr(fc1_weight_python_object, "zero_out_wgrad", False): + fc1_wgrad = torch.zeros( + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, ) - elif ctx.normalization == "RMSNorm": - dgrad, dgamma = tex.rmsnorm_bwd( - dgrad, - inputmat, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + else: + fc1_wgrad = torch.empty( + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, ) - dbeta = None - clear_tensor_data(mu, rsigma) - - if ctx.fc1_weight_requires_grad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - fc1_weight_python_object, "grad_added_to_main_grad" - ): - fc1_weight_python_object.grad_added_to_main_grad = True - if getattr(fc1_weight_python_object, "zero_out_wgrad", False): - fc1_wgrad = torch.zeros( - fc1_weight_main_grad.shape, - dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - fc1_wgrad = torch.empty( - fc1_weight_main_grad.shape, - dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - fc1_wgrad = None - else: + elif args.fuse_wgrad_accumulation: fc1_wgrad = None + else: + fc1_wgrad = None - if ctx.fc2_weight_requires_grad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - fc2_weight_python_object, "grad_added_to_main_grad" - ): - fc2_weight_python_object.grad_added_to_main_grad = True - if getattr(fc2_weight_python_object, "zero_out_wgrad", False): - fc2_wgrad = torch.zeros( - fc2_weight_main_grad.shape, - dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - fc2_wgrad = torch.empty( - fc2_weight_main_grad.shape, - dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - fc2_wgrad = None - else: + if args.fc2_weight_requires_grad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr( + fc2_weight_python_object, "grad_added_to_main_grad" + ): + fc2_weight_python_object.grad_added_to_main_grad = True + if getattr(fc2_weight_python_object, "zero_out_wgrad", False): + fc2_wgrad = torch.zeros( + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) + else: + fc2_wgrad = torch.empty( + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) + elif args.fuse_wgrad_accumulation: fc2_wgrad = None + else: + fc2_wgrad = None + + # FIX THIS + # Scatter Fp8 tranposed-weight buffers + # if args.fp8: + # _fsdp_scatter_tensors( + # args.fsdp_group, + # fc1_weight_fp8 if not isinstance(fc1_weight, Float8Tensor) else None, + # fc2_weight_fp8 if not isinstance(fc2_weight, Float8Tensor) else None, + # ) + return ( + dgrad.view(args.inp_shape) if args.requires_dgrad else None, + dgamma, + dbeta, + fc1_wgrad, + fc1_bias_grad if fc1_bias is not None else None, + fc2_wgrad, # pylint: disable=possibly-used-before-assignment + fc2_bias_grad, + ) - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - # FIX THIS - # Scatter Fp8 tranposed-weight buffers - # if ctx.fp8: - # _fsdp_scatter_tensors( - # ctx.fsdp_group, - # fc1_weight_fp8 if not isinstance(fc1_weight, Float8Tensor) else None, - # fc2_weight_fp8 if not isinstance(fc2_weight, Float8Tensor) else None, - # ) +class _LayerNormMLP(torch.autograd.Function): + """LayerNormMLP semi-top level module + Calls custom cuda extensions. + """ + + @staticmethod + def forward( + ctx, + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + fc1_weight: torch.Tensor, + fc1_bias: Optional[torch.Tensor], + fc2_weight: torch.Tensor, + fc2_bias: Optional[torch.Tensor], + fwd_args: LayerNormMLPFwdArgs, + ) -> Tuple[ + torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor] + ]: + """Forward pass: compute the output and set up the autograd context. + + The tensors are positional so autograd tracks them; they are + re-attached to ``fwd_args`` so every downstream helper takes a single + argument. The weight workspaces are non-differentiable cached tensors + passed in via ``fwd_args`` and the freshly produced workspaces are + returned as separate outputs so the module can refresh its cache. + """ + fwd_args.inp = inp + fwd_args.ln_weight = ln_weight + fwd_args.ln_bias = ln_bias + fwd_args.fc1_weight = fc1_weight + fwd_args.fc1_bias = fc1_bias + fwd_args.fc2_weight = fc2_weight + fwd_args.fc2_bias = fc2_bias + ( + out, + ln_out_return, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_mlp_forward_impl(fwd_args) + if ctx is not None: + bwd_args = LayerNormMLPBwdArgs() + tensors_to_save_from_setup = _layernorm_mlp_setup_ctx( + bwd_args, + fwd_args, + (out, ln_out_return, new_fc1_weight_workspace, new_fc2_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + ctx.backward_objects = bwd_args + if not bwd_args.checkpoint and fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() + + return out, ln_out_return, new_fc1_weight_workspace, new_fc2_weight_workspace + + @staticmethod + def backward( + ctx, + grad_output: torch.Tensor, + grad_ln_out: Optional[torch.Tensor], + _grad_fc1_weight_workspace, + _grad_fc2_weight_workspace, + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward pass: compute gradients and reduce FP8 scaling factors.""" + bwd_args: LayerNormMLPBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.grad_ln_out = grad_ln_out + with get_nvtx_range_context("_LayerNormMLP_backward"): + _layernorm_mlp_recompute(bwd_args, ctx) + ( + dgrad, + dgamma, + dbeta, + fc1_wgrad, + fc1_bias_grad, + fc2_wgrad, + fc2_bias_grad, + ) = _layernorm_mlp_backward_impl(bwd_args) + reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, + # main_grad closures) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + dgrad, dgamma, dbeta, fc1_wgrad, - None, # fc1_weight_workspace - fc1_bias_grad if fc1_bias is not None else None, - fc2_wgrad, # pylint: disable=possibly-used-before-assignment - None, # fc2_weight_workspace + fc1_bias_grad, + fc2_wgrad, fc2_bias_grad, - None, + None, # fwd_args ) @@ -2261,6 +2595,7 @@ 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, @@ -2361,13 +2696,6 @@ def forward( if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): self.fast_setattr("bias_gelu_nvfusion", False) - if is_grad_enabled: - fwd_fn = _LayerNormMLP.apply - autograd_ctx = [] - else: - fwd_fn = _LayerNormMLP.forward - autograd_ctx = [None] - cache_name_fc1 = ( None if (is_first_microbatch is None or self.is_fsdp2) else "fc1_weight" ) @@ -2381,71 +2709,153 @@ def forward( self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None ) - non_tensor_args = ( - self.eps, - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - self.fuse_wgrad_accumulation, - fc1_input_quantizer, - fc1_weight_quantizer, - fc1_output_quantizer, - fc1_grad_input_quantizer, - fc1_grad_weight_quantizer, - fc1_grad_output_quantizer, - fc2_input_quantizer, - fc2_weight_quantizer, - fc2_output_quantizer, - fc2_grad_input_quantizer, - fc2_grad_weight_quantizer, - fc2_grad_output_quantizer, - is_cpu_offload_enabled(), - self.tp_group, - self.tp_size, - self.sequence_parallel, - self.tp_size > 1, - self.activation_dtype, - self.return_layernorm_output, - self.return_layernorm_output_gathered, - self.bias_gelu_nvfusion and not self.fp8 and not debug, - self.set_parallel_mode, - is_grad_enabled, - self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, - self.bwd_ln_sm_margin, - self.zero_centered_gamma, - self.activation, - self.activation_params, - self.normalization, - self.ub_overlap_ag, - self.ub_overlap_rs, - self.ub_overlap_rs_dgrad, - self.ub_bulk_dgrad, - self.ub_bulk_wgrad, - self.gemm_gelu_fusion and not debug, - self.fsdp_group, - self.fp8_meta, - cache_name_fc1 is not None, - skip_fp8_weight_update, - self.symmetric_ar_type, - self.checkpoint, - debug, - self.is_fsdp2, + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD + if self.fp8: + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + 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 + else: + backward_override = None + + if debug: # turn off userbuffers in debug mode + ub_overlap_ag = False + ub_overlap_rs = False + ub_overlap_rs_dgrad = False + ub_bulk_wgrad = False + ub_bulk_dgrad = False + else: + ub_overlap_ag = self.ub_overlap_ag + ub_overlap_rs = self.ub_overlap_rs + ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad + ub_bulk_wgrad = self.ub_bulk_wgrad + ub_bulk_dgrad = self.ub_bulk_dgrad + + fc2_bias_tensor = ( + fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None ) - out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn( - *autograd_ctx, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_weight_workspace, - fc1_bias, - fc2_weight, - fc2_weight_workspace, - fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None, - non_tensor_args, + wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + + fwd_args = LayerNormMLPFwdArgs( + # tensors + inp=inp, + ln_weight=self.layer_norm_weight, + ln_bias=self.layer_norm_bias, + fc1_weight=fc1_weight, + fc1_bias=fc1_bias, + fc2_weight=fc2_weight, + fc2_bias=fc2_bias_tensor, + fc1_weight_workspace=fc1_weight_workspace, + fc2_weight_workspace=fc2_weight_workspace, + # requires_grad flags + input_requires_grad=inp.requires_grad, + ln_weight_requires_grad=self.layer_norm_weight.requires_grad, + ln_bias_requires_grad=( + self.layer_norm_bias.requires_grad + if self.layer_norm_bias is not None + else False + ), + fc1_weight_requires_grad=fc1_weight.requires_grad, + fc1_bias_requires_grad=fc1_bias.requires_grad if fc1_bias is not None else False, + fc2_weight_requires_grad=fc2_weight.requires_grad, + fc2_bias_requires_grad=( + fc2_bias_tensor.requires_grad if fc2_bias_tensor is not None else False + ), + # quantizers + fc1_input_quantizer=fc1_input_quantizer, + fc1_weight_quantizer=fc1_weight_quantizer, + fc1_output_quantizer=fc1_output_quantizer, + fc1_grad_input_quantizer=fc1_grad_input_quantizer, + fc1_grad_weight_quantizer=fc1_grad_weight_quantizer, + fc1_grad_output_quantizer=fc1_grad_output_quantizer, + fc2_input_quantizer=fc2_input_quantizer, + fc2_weight_quantizer=fc2_weight_quantizer, + fc2_output_quantizer=fc2_output_quantizer, + fc2_grad_input_quantizer=fc2_grad_input_quantizer, + fc2_grad_weight_quantizer=fc2_grad_weight_quantizer, + fc2_grad_output_quantizer=fc2_grad_output_quantizer, + # normalization + eps=self.eps, + normalization=self.normalization, + zero_centered_gamma=self.zero_centered_gamma, + fwd_ln_sm_margin=( + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin + ), + bwd_ln_sm_margin=self.bwd_ln_sm_margin, + return_layernorm_output=self.return_layernorm_output, + return_layernorm_output_gathered=self.return_layernorm_output_gathered, + # activation + activation=self.activation, + activation_params=self.activation_params, + bias_gelu_fusion=self.bias_gelu_nvfusion and not self.fp8 and not debug, + gemm_gelu_fusion=self.gemm_gelu_fusion and not debug, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_name_fc1 is not None, + skip_fp8_weight_update=skip_fp8_weight_update, + # tensor / sequence parallelism + set_parallel_mode=self.set_parallel_mode, + tp_group=self.tp_group, + tp_size=self.tp_size, + tensor_parallel=self.tp_size > 1, + sequence_parallel=self.sequence_parallel, + symmetric_ar_type=self.symmetric_ar_type, + # userbuffers + ub_overlap_ag=ub_overlap_ag, + ub_overlap_rs=ub_overlap_rs, + ub_overlap_rs_dgrad=ub_overlap_rs_dgrad, + ub_bulk_dgrad=ub_bulk_dgrad, + ub_bulk_wgrad=ub_bulk_wgrad, + # FSDP + fsdp_group=self.fsdp_group, + is_fsdp2=self.is_fsdp2, + # weight-grad scheduling + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + # activation checkpointing + checkpoint=self.checkpoint, + fp8_meta=self.fp8_meta if self.checkpoint else None, + recompute_for_bwd=False, + # misc + cpu_offloading=is_cpu_offload_enabled(), + 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, + ) + 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, + fwd_args, + ) + if new_fc1_ws is not None and cache_name_fc1 is not None: if isinstance(new_fc1_ws, torch.Tensor): new_fc1_ws = new_fc1_ws.detach() diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 55fc69ef7f..43e252f5f8 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -33,6 +33,7 @@ ) from ._common import ( can_reconstruct_wgrad_input_from_original, + check_fp8_reduce_and_update, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -317,16 +318,6 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: ) # pylint: disable=unbalanced-tuple-unpacking -def _check_fp8_reduce_and_update(): - """Check if this is the first FP8 module (for backward reduce-and-update).""" - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - result = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - return result - - def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: """Output's leading (sequence) dim from the input's: sequence parallelism gathers it (column-parallel) or scatters it (row-parallel).""" @@ -1852,7 +1843,7 @@ def forward( or fwd_args.weight_requires_grad or fwd_args.bias_requires_grad ): - bwd_args.reduce_and_update_bwd_fp8_tensors = _check_fp8_reduce_and_update() + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() if fwd_args.backward_override is not None: bwd_args.reduce_and_update_bwd_fp8_tensors = False