From 43a6cea3eebc5ddbb89cb092910c9f44b1417f75 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sun, 24 May 2026 23:58:23 -0700 Subject: [PATCH 01/15] Generalized Tensor Parallelism (GTP) init commit Co-authored-by: Jieming Zhang Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/distributed.py | 142 ++++++++++++------ transformer_engine/pytorch/module/base.py | 76 +++++++++- .../pytorch/module/grouped_linear.py | 115 +++++++++++--- .../pytorch/module/layernorm_linear.py | 60 +++++++- transformer_engine/pytorch/module/linear.py | 72 ++++++++- 5 files changed, 387 insertions(+), 78 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index a0d4ac3530..acfb6ddff8 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable -from contextlib import contextmanager, AbstractContextManager, ContextDecorator +from contextlib import contextmanager, AbstractContextManager, ContextDecorator, nullcontext from functools import lru_cache from dataclasses import dataclass import math @@ -918,7 +918,10 @@ def fork(self, name: str = "model-parallel-rng"): def reduce_scatter_along_first_dim( - inp: torch.Tensor, tp_group: dist_group_type, async_op: bool = False + inp: torch.Tensor, + tp_group: dist_group_type, + async_op: bool = False, + output: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: """Reduce-scatter the input tensor across model parallel group.""" world_size = get_distributed_world_size(tp_group) @@ -936,7 +939,8 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + if output is None: + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1281,11 +1285,13 @@ def _post_process_nvfp4_gather( handle = None # Fix the interleaved transposed data from gathering along first dim. - out._columnwise_scale_inv = _swap_first_dims(columnwise_scale_inv_interleaved, world_size) - out._columnwise_data = _swap_first_dims(columnwise_data_interleaved, world_size) + # In-place .copy_() (not `=` rebind) to keep the storage address stable + # for CUDA graph capture — replays see the same pointer they captured. + out._columnwise_scale_inv.copy_(_swap_first_dims(columnwise_scale_inv_interleaved, world_size)) + out._columnwise_data.copy_(_swap_first_dims(columnwise_data_interleaved, world_size)) - # Optionally pad the scaling inverse if needed. - out._columnwise_scale_inv = pad_columnwise_scale_inv(out._columnwise_scale_inv) + # Optionally pad the scaling inverse if needed (same in-place pattern). + out._columnwise_scale_inv.copy_(pad_columnwise_scale_inv(out._columnwise_scale_inv)) @dataclass @@ -1299,17 +1305,25 @@ class _NVFP4AllGatherAsyncHandle: async_handle: torch.distributed.Work _synchronized: bool = False - def wait(self) -> None: - """Wait for the async operation to complete and post-process the tensor.""" - if self._synchronized: - return - self.async_handle.wait() + def post_process_nvfp4_gather(self) -> None: + """Fix interleaved transposed data + pad scale_inv after the async AG completes. + + Idempotent: gated by ``_synchronized`` in :meth:`wait`. + """ _post_process_nvfp4_gather( self.output, self.columnwise_data_interleaved, self.columnwise_scale_inv_interleaved, self.world_size, ) + + def wait(self) -> None: + """Wait for the async operation to complete and post-process the tensor.""" + if self._synchronized: + return + if self.async_handle is not None: + self.async_handle.wait() + self.post_process_nvfp4_gather() self._synchronized = True @@ -1320,6 +1334,8 @@ def _all_gather_nvfp4( async_op: bool = False, quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, + output_tensor=None, + grouped=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1383,6 +1399,12 @@ def _all_gather_nvfp4( out = quantizer(out) return out, None + # Construct NVFP4 output tensor + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + # Cast input tensor to NVFP4 with required data if not isinstance(inp, NVFP4TensorStorage): inp = quantizer(inp) @@ -1395,17 +1417,19 @@ def _all_gather_nvfp4( ) inp = quantizer(inp.dequantize(dtype=dtype)) - # Construct NVFP4 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - - # Coalesce NCCL collectives for gathering data and scale inverses. - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as gather_coalescing_manager: + if not grouped: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage + out_columnwise_data = None if quantizer.rowwise_usage: # Remove padding from NVFP4 scale-inverses @@ -1433,8 +1457,9 @@ def _all_gather_nvfp4( group=process_group, ) - # Transfer amax to output. - out._amax_rowwise = inp._amax_rowwise + # Transfer amax to output via in-place .copy_() so the storage + # address stays stable for CUDA graph capture. + out._amax_rowwise.copy_(inp._amax_rowwise) # Gather the transposed NVFP4 data along first dimension. Fix format later. if quantizer.columnwise_usage: @@ -1483,17 +1508,24 @@ def _all_gather_nvfp4( ) # Transfer amax to output. - out._amax_columnwise = inp._amax_columnwise + out._amax_columnwise.copy_(inp._amax_columnwise) - handle = gather_coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if async_op and quantizer.columnwise_usage: - handle = _NVFP4AllGatherAsyncHandle( - out, out_columnwise_data, out_scale_inv, world_size, handle - ) - elif quantizer.columnwise_usage: - _post_process_nvfp4_gather(out, out_columnwise_data, out_scale_inv, world_size, handle) + if quantizer.columnwise_usage: + if async_op or grouped: + # Defer post-processing: either the async op hasn't completed yet, or an + # external coalescing manager owns the NCCL ops and hasn't flushed them. + inner_handle = handle if async_op else None + handle = _NVFP4AllGatherAsyncHandle( + out, out_columnwise_data, out_scale_inv, world_size, inner_handle + ) + else: + _post_process_nvfp4_gather(out, out_columnwise_data, out_scale_inv, world_size, handle) + else: + if handle is not None: + handle.output = out return out, handle @@ -1505,6 +1537,8 @@ def _all_gather_mxfp8( async_op: bool = False, quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, + output_tensor: torch.Tensor = None, + grouped: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1570,15 +1604,22 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Coalesce NCCL collectives - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as coalescing_manager: + if not grouped: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather MXFP8 data for row-wise usage if quantizer.rowwise_usage: @@ -1625,7 +1666,7 @@ def _all_gather_mxfp8( group=process_group, ) - handle = coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None return out, handle @@ -1634,6 +1675,8 @@ def gather_along_first_dim( process_group: dist_group_type, async_op: bool = False, quantizer: Optional[Quantizer] = None, + output_tensor: torch.Tensor = None, + grouped: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. @@ -1724,6 +1767,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + grouped=grouped, ) # NVFP4 case @@ -1738,6 +1783,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + grouped=grouped, ) # High-precision communication for quantized tensors @@ -1767,19 +1814,20 @@ def gather_along_first_dim( inp = inp.dequantize() # Communication for plain PyTorch tensors - out = torch.empty( - out_shape, - dtype=inp.dtype, - device=inp.device, - memory_format=torch.contiguous_format, - ) + if output_tensor is None: + output_tensor = torch.empty( + out_shape, + dtype=inp.dtype, + device=inp.device, + memory_format=torch.contiguous_format, + ) handle = torch.distributed.all_gather_into_tensor( - out, + output_tensor, inp.contiguous(), group=process_group, async_op=async_op, ) - return out, handle + return output_tensor, handle # Global cache to store symmetric memory tensors diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 746177ec78..64853c54c4 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -61,7 +61,13 @@ from ...debug.pytorch.debug_quantization import DebugQuantizer, DebugQuantizedTensor from ...debug.pytorch.utils import next_iter_when_debug_should_be_run, any_feature_enabled -__all__ = ["initialize_ub", "destroy_ub", "UserBufferQuantizationMode"] +__all__ = [ + "initialize_ub", + "destroy_ub", + "UserBufferQuantizationMode", + "register_gtp_hooks", + "maybe_wrap_gtp", +] _2X_ACC_FPROP = False _2X_ACC_DGRAD = True @@ -72,6 +78,47 @@ layers_atomic_ring_exchange = [] +# GTP hook slots. An external integrator (currently ``megatron.experimental.gtp``) +# populates these via ``register_gtp_hooks`` at its own import time. When the +# slots stay ``None``, the ``gtp_group=`` codepath in TE modules is a no-op +# and TE has no ``from megatron...`` dependency. +_gtp_slice_fn = None +_gtp_finalize_fn = None +_gtp_wrap_fn = None + + +def register_gtp_hooks(*, slice_fn=None, finalize_fn=None, wrap_fn=None): + """Register GTP integration hooks. Hooks left as ``None`` are unchanged. + + slice_fn(module, name, param, *, expert_idx) -> GTPShardedParam | None + Fires per weight during ``reset_parameters``, before FP8 quantize. + finalize_fn(module, weight_names) -> None + Fires after the per-weight loop in ``reset_parameters``. + wrap_fn(module, weight_names, gtp_group, is_grouped=False) -> None + Fires at the end of a module's ``__init__`` to finalize GTP wiring. + """ + global _gtp_slice_fn, _gtp_finalize_fn, _gtp_wrap_fn + if slice_fn is not None: + _gtp_slice_fn = slice_fn + if finalize_fn is not None: + _gtp_finalize_fn = finalize_fn + if wrap_fn is not None: + _gtp_wrap_fn = wrap_fn + + +def maybe_wrap_gtp(module, weight_names, gtp_group, is_grouped=False): + """Finalize GTP wiring on a module if a wrap hook is registered. + + No-op when ``gtp_group`` is None or no GTP integrator has called + ``register_gtp_hooks``. Called from each TE module's ``__init__`` after + ``reset_parameters`` finishes; the per-weight slice already happened + inside ``reset_parameters`` via ``_gtp_slice_fn``. + """ + if gtp_group is None or _gtp_wrap_fn is None: + return + _gtp_wrap_fn(module, weight_names, gtp_group, is_grouped=is_grouped) + + class UserBufferQuantizationMode(Enum): """ UserBufferQuantizationMode is an enum that represents the quantization mode of the UserBuffer. @@ -1604,7 +1651,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: if defer_init: return - for name, param in self.named_parameters(recurse=False): + # Names of GTP-sharded weights, for GroupedLinear's post-loop finalize. + _gtp_sharded_weight_names = [] + + for idx, (name, param) in enumerate(self.named_parameters(recurse=False)): # Check if parameter is a DTensor (FSDP2) or regular tensor is_dtensor = isinstance(param, DTensor) dtensor_param = param if is_dtensor else None @@ -1626,10 +1676,23 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: with get_rng_state_tracker().fork(): init_fn(param) + # GTP slice: shard the freshly-init weight into a GTPShardedParam; + # the FP8 quantize block below is skipped for it. + gtp_sharded = None + if ( + not is_dtensor + and getattr(self, "_gtp_group", None) is not None + and _gtp_slice_fn is not None + ): + gtp_sharded = _gtp_slice_fn(self, name, param, expert_idx=idx) + if gtp_sharded is not None: + param = gtp_sharded + _gtp_sharded_weight_names.append(name) + # Wrap parameters in QuantizedTensor if needed fp8_meta_index = self.param_init_meta[name].fp8_meta_index high_precision_init_val = None - if self.primary_weights_in_fp8 and fp8_meta_index is not None: + if self.primary_weights_in_fp8 and fp8_meta_index is not None and gtp_sharded is None: # Keep high-precision values on CPU if needed if self.preserve_high_precision_init_val: @@ -1657,6 +1720,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: # NOTE: Currently this can only be broken when primary weights are in Fp8 but # re-applying the nn.Parameter() wrap is a no-op when the input is already # a parameter so we always re-apply it just for extra safety. + # Skip the wrap for GTPShardedParam (Parameter.__new__ would drop attrs). if is_dtensor: # recreate the DTensor from the parameter. dtensor_param = DTensor.from_local( @@ -1667,7 +1731,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: stride=dtensor_param.stride(), ) dtensor_param = torch.nn.Parameter(dtensor_param) - else: + elif gtp_sharded is None: param = torch.nn.Parameter(param) # Keep high-precision values on CPU if needed @@ -1705,6 +1769,10 @@ def clear(self): else: self.module_setattr(name, dtensor_param) + # GroupedLinear post-loop finalize hook (no-op outside GroupedLinear). + if _gtp_sharded_weight_names and _gtp_finalize_fn is not None: + _gtp_finalize_fn(self, _gtp_sharded_weight_names) + @abstractmethod def forward(self): """Needs override.""" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 627144345c..5f85f9c56e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -24,6 +24,7 @@ _2X_ACC_WGRAD, ) from ._common import WeightGradStore +from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, @@ -100,6 +101,7 @@ def forward( skip_fp8_weight_update, save_original_input, debug, + gtp_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -114,6 +116,14 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + weights_gtp_sharded = weights + if gtp_size > 1: + weights = weights[0].batched_all_gather_and_prefetch( + fwd=True, + skip_weight_cast=is_first_microbatch is False, + cast_noop_flag=skip_fp8_weight_update, + ) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -276,12 +286,19 @@ def forward( else: inputmats = [None] * num_gemms - tensors_to_save, tensor_objects = prepare_for_saving( - *inputmats, - *weights_fp8, - *weights, - *biases, - ) + if gtp_size == 1: + tensors_to_save, tensor_objects = prepare_for_saving( + *inputmats, + *weights_fp8, + *weights, + *biases, + ) + else: + tensors_to_save, tensor_objects = prepare_for_saving( + *inputmats, + *weights_gtp_sharded, + *biases, + ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects @@ -303,6 +320,10 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif gtp_size > 1: + ctx.main_grad_funcs = [ + weights_gtp_sharded[i].get_wgrad_tensor for i in range(num_gemms) + ] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -332,6 +353,7 @@ def forward( ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers + ctx.gtp_size = gtp_size # backward overrides if backward_override is not None: @@ -357,17 +379,30 @@ def backward( with get_nvtx_range_context("_GroupedLinear_backward"): saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms - inputmats = saved_tensors[:N] - weights = saved_tensors[N : 2 * N] - saved_weights = saved_tensors[2 * N : 3 * N] - biases = saved_tensors[3 * N : 4 * N] + if ctx.gtp_size == 1: + inputmats = saved_tensors[:N] + weights = saved_tensors[N : 2 * N] + saved_weights = saved_tensors[2 * N : 3 * N] + biases = saved_tensors[3 * N : 4 * N] + gtp_origin_weights = None + else: + inputmats = saved_tensors[:N] + gtp_origin_weights = saved_tensors[N : 2 * N] + biases = saved_tensors[2 * N : 3 * N] + weights = None # Restore from weakrefs to get original weight python objects # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + if ctx.gtp_size > 1: + # GTP: origin_weights come from saved tensors; main_grads are + # get_wgrad_tensor scratch (do not assign to param.main_grad). + origin_weights = gtp_origin_weights + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] @@ -428,13 +463,18 @@ def backward( ctx.m_splits, ) - if ctx.is_first_microbatch is not None: + if ctx.gtp_size > 1: + 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 ) else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + if ctx.gtp_size > 1: + weights = origin_weights[0].batched_all_gather_and_prefetch_bwd() + if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD if ctx.fp8 or ctx.debug: @@ -485,6 +525,14 @@ def backward( use_split_accumulator=dgrad_gemm_use_split_accumulator, ) + # Gathered weights are no longer needed after dgrad GEMM. + # For nvfp4, the NVFP4TensorStorage and its sub-tensors (scale_inv etc.) + # would otherwise survive until function return via this local ref. + w_shape = None + if ctx.gtp_size > 1: + w_shape = list(weights[0].size()) + del weights + if ctx.weights_requires_grad: wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD if ctx.fp8: @@ -496,7 +544,7 @@ def backward( if ctx.fuse_wgrad_accumulation: wgrad_list = main_grads else: - weight_shape = list(weights[0].size()) + weight_shape = w_shape if ctx.gtp_size > 1 else list(weights[0].size()) wgrad_list = tex.bulk_allocate( [weight_shape] * ctx.num_gemms, [ctx.activation_dtype] * ctx.num_gemms, @@ -553,7 +601,8 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + if ctx.gtp_size == 1 + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), ) @@ -595,10 +644,19 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] + if ctx.gtp_size > 1: + wgrad_list = origin_weights[0].batched_wgrad_reduce_scatter(wgrad_list) + # Drop Python refs to wgrad input buffers. The async RS on rs_stream + # still holds C++ refs (via NCCL Work); those are released when + # _wait_reduce_scatter calls handle.wait() + self.handle = None. + # Without this del, main_grads keeps the tensors alive until function + # return, wasting memory during graph capture warmup. + del main_grads + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] else: wgrad_list = [None] * ctx.num_gemms @@ -716,6 +774,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, + gtp_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -771,6 +830,11 @@ def __init__( "Because the TP communication is handled outside of this module." ) + if gtp_group is None: + self.gtp_size = 1 + else: + self.gtp_size = get_distributed_world_size(gtp_group) + self.parallel_mode = parallel_mode if self.parallel_mode not in GemmParallelModes: raise ValueError( @@ -822,9 +886,18 @@ def __init__( if self.primary_weights_in_fp8: self.init_fp8_metadata(num_gemms=self.num_gemms) + self.weight_names = [f"weight{idx}" for idx in range(self.num_gemms)] is_meta = torch.device(device).type == "meta" + if gtp_group is not None: + # Stashed before reset_parameters so the slice hook can see it; + # _gtp_is_grouped routes through the GroupedLinear finalize path. + self._gtp_group = gtp_group + self._gtp_is_grouped = True + self.reset_parameters(defer_init=is_meta) + maybe_wrap_gtp(self, self.weight_names, gtp_group, is_grouped=True) + if self.wgrad_store.delay_wgrad_compute(): for name, param in self.named_parameters(): if name in ("weight", "bias"): @@ -1148,6 +1221,11 @@ def forward( weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() + if self.gtp_size > 1: + weight_tensors[0].setup( + weight_quantizer=self._get_weight_quantizers(), + ) + quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() if debug: @@ -1202,6 +1280,7 @@ def forward( None, # skip_fp8_weight_update self.save_original_input, debug, + self.gtp_size, ) out, new_workspaces = linear_fn( *autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 8c88f3ee82..8e000e3764 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -28,6 +28,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) +from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( assert_dim_for_fp8_exec, @@ -143,6 +144,7 @@ def forward( symmetric_ar_type, debug, is_fsdp2, + gtp_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -297,6 +299,16 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + + weight_gtp_sharded = weight + if gtp_size > 1: + weight = weight.all_gather_and_prefetch( + fwd=True, + skip_weight_cast=is_first_microbatch is False, + cast_noop_flag=skip_fp8_weight_update, + ) + out_features = weight.shape[0] + new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -484,8 +496,9 @@ def forward( wt_save = None tensors_to_save, tensor_objects = prepare_for_saving( inputmat, - wt_save, - weight, + # GTP: save the sharded reference only; backward re-gathers it. + wt_save if gtp_size == 1 else None, + weight if gtp_size == 1 else weight_gtp_sharded, bias, ln_weight, ln_out_to_save, @@ -512,6 +525,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad + elif gtp_size > 1: + ctx.main_grad_func = weight_gtp_sharded.get_wgrad_tensor else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -554,6 +569,7 @@ def forward( qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug + ctx.gtp_size = gtp_size # backward overrides if backward_override is not None: @@ -605,6 +621,9 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) + if ctx.gtp_size > 1: + weight = saved_weight.all_gather_and_prefetch_bwd() + # 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. @@ -622,7 +641,7 @@ def backward( ), "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: + if main_grad is not None and ctx.gtp_size == 1: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -929,7 +948,10 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: + if ctx.gtp_size > 1: + # GTP: accumulation happens downstream in wgrad_reduce_scatter. + 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 ) @@ -1001,6 +1023,9 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) + if ctx.gtp_size > 1: + wgrad = saved_weight.wgrad_reduce_scatter(wgrad) + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1080,7 +1105,10 @@ def wgrad_gemm( if ctx.requires_wgrad: # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): + if ctx.gtp_size > 1: + # GTP: skip — wgrad RS already produced the correct shard. + pass + elif 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( @@ -1247,6 +1275,7 @@ def __init__( delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, name: Optional[str] = None, + gtp_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1277,6 +1306,11 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() + if gtp_group is None: + self.gtp_size = 1 + else: + self.gtp_size = get_distributed_world_size(gtp_group) + self.parallel_mode = parallel_mode assert ( self.parallel_mode in GemmParallelModes @@ -1471,8 +1505,18 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() + if gtp_group is not None: + # Stashed before reset_parameters so the slice hook can see it. + self._gtp_group = gtp_group + self._gtp_is_grouped = False + self.reset_parameters(defer_init=device == "meta") + maybe_wrap_gtp(self, self.weight_names, gtp_group) + if gtp_group is not None: + # Free the full-size backing buffer; GTP replaced it with a sharded param. + del weight_tensor + # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM if self.parallel_mode == "row" and self.apply_bias: @@ -1635,6 +1679,11 @@ def forward( # Get concatenated weight and bias tensors weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + if self.gtp_size > 1: + weight_tensor.setup( + weight_quantizer=self._get_weight_quantizers(), + ) + quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug @@ -1705,6 +1754,7 @@ def forward( self.symmetric_ar_type, debug, self.is_fsdp2, + self.gtp_size, ) out, ln_out, new_weight_workspace = fwd_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index dcbb9eaf93..7b4e18f066 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -28,6 +28,7 @@ _2X_ACC_WGRAD, ) from ._common import noop_cat, WeightGradStore +from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, @@ -153,6 +154,9 @@ class LinearFwdArgs: cpu_offloading: bool is_grad_enabled: bool + # --- Generalized tensor parallelism --- + gtp_size: int = 1 + @dataclass(slots=True) class LinearBwdArgs: @@ -222,6 +226,9 @@ class LinearBwdArgs: cpu_offloading: bool = False owns_input: bool = False + # --- Generalized tensor parallelism --- + gtp_size: int = 1 + # --- Per-backward scratch state (populated inside _linear_backward) --- ub_obj_gradout: Optional[Any] = None @@ -399,6 +406,15 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + # GTP: rebind `weight` to the all-gathered tensor; `args.weight` keeps + # the GTPShardedParam reference for backward re-gather / wgrad RS. + if args.gtp_size > 1: + weight = weight.all_gather_and_prefetch( + fwd=True, + skip_weight_cast=is_first_microbatch is False, + cast_noop_flag=args.skip_fp8_weight_update, + ) + new_weight_workspace = None weightmat = weight if fp8 or debug: @@ -576,6 +592,9 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # GTP: don't save the workspace; backward re-gathers it. + if args.gtp_size > 1: + wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. @@ -680,11 +699,14 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad + elif fwd_args.gtp_size > 1: + bwd_args.main_grad_func = weight.get_wgrad_tensor else: bwd_args.main_grad_func = lambda: weight.main_grad # Misc bwd_args.cpu_offloading = fwd_args.cpu_offloading + bwd_args.gtp_size = fwd_args.gtp_size if backward_override is not None: bwd_args.fp8 = False @@ -751,7 +773,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - origin_weight_python_object.main_grad = main_grad + if bwd_args.gtp_size == 1: + origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already @@ -921,6 +944,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None + + # GTP: re-gather the sharded weight; runs even when requires_dgrad=False + # so the prev_w prefetch is issued for the next layer's bwd. + if bwd_args.gtp_size > 1: + weight_fp8 = saved_weight.all_gather_and_prefetch_bwd() + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when @@ -1102,7 +1131,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.is_first_microbatch is not None: + if bwd_args.gtp_size > 1: + # GTP: accumulation happens downstream in wgrad_reduce_scatter. + accumulate_wgrad_into_param_main_grad = False + elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) @@ -1178,6 +1210,11 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + # GTP: reduce-scatter the freshly computed wgrad (async; overlap + # with the next layer's bwd via the cascade). + if bwd_args.gtp_size > 1: + wgrad = saved_weight.wgrad_reduce_scatter(wgrad) + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1223,15 +1260,19 @@ def wgrad_gemm( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True + # Use the param's local shape (sharded under GTP) so the dummy wgrad + # matches the saved weight shape; main_grad_func() under GTP returns + # an unsharded scratch and would otherwise mismatch. + wgrad_shape = list(origin_weight_python_object.shape) if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, ) elif bwd_args.fuse_wgrad_accumulation: @@ -1447,6 +1488,7 @@ def __init__( symmetric_ar_type: Optional[str] = None, save_original_input: bool = False, name: Optional[str] = None, + gtp_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1475,6 +1517,11 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() + if gtp_group is None: + self.gtp_size = 1 + else: + self.gtp_size = get_distributed_world_size(gtp_group) + self.parallel_mode = parallel_mode assert ( self.parallel_mode in GemmParallelModes @@ -1644,8 +1691,18 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() + if gtp_group is not None: + # Stashed before reset_parameters so the slice hook can see it. + self._gtp_group = gtp_group + self._gtp_is_grouped = False + self.reset_parameters(defer_init=device == "meta") + maybe_wrap_gtp(self, self.weight_names, gtp_group) + if gtp_group is not None: + # Free the full-size backing buffer; GTP replaced it with a sharded param. + del weight_tensor + # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM if self.parallel_mode == "row" and self.apply_bias: @@ -1776,6 +1833,11 @@ def forward( try: weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + if self.gtp_size > 1: + weight_tensor.setup( + weight_quantizer=self._get_weight_quantizers(), + ) + quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug @@ -1894,6 +1956,8 @@ def forward( # misc cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, + # generalized tensor parallelism + gtp_size=self.gtp_size, ) out, new_weight_workspace = linear_fn( *autograd_ctx, From a532120009c35f397e9c34d52fb415da05f69000 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 1 Jun 2026 05:33:35 -0700 Subject: [PATCH 02/15] GTP + gmm fusion Signed-off-by: Shiqing Fan --- .../pytorch/ops/fused/backward_grouped_mlp.py | 31 +++- .../pytorch/ops/fused/forward_grouped_mlp.py | 142 ++++++++++++------ 2 files changed, 125 insertions(+), 48 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 25ccad1377..be8b741fde 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -161,6 +161,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() + gtp_size = getattr(ctx, "gtp_size", 1) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -193,7 +194,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - if fc_op._accumulate_into_main_grad: + # EGTP: the GEMM produces full-sized wgrads but main_grad is sharded, so use a + # full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and gtp_size == 1: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -209,6 +212,12 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() + if gtp_size > 1 and delay_wgrad: + raise RuntimeError( + "EGTP + cuteDSL fused grouped-MLP does not support delay_wgrad / " + "overlap_dispatch_backward_with_experts_wgrad yet; set " + "delay_wgrad_compute=False." + ) if cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( @@ -232,9 +241,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - - # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - if fc_op._accumulate_into_main_grad: + # EGTP: reduce-scatter the full per-rank wgrads into each sharded main_grad + # (also fires the Megatron grad-accum hook). + if gtp_size > 1: + weights[0].batched_wgrad_reduce_scatter(w_list) + + # Return dummy wgrads when grad is already in main_grad (wgrad fusion, or the EGTP + # reduce-scatter above) so Megatron-LM's wgrad fusion doesn't double-add. + if fc_op._accumulate_into_main_grad or gtp_size > 1: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -505,6 +519,11 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) + # EGTP: forward gathered rowwise; col-AG the columnwise layout the dgrad needs, + # right before use (one weight live at a time). + if getattr(fc2_ctx, "gtp_size", 1) > 1: + grouped_fc2_weight = fc2_op.weight0.batched_all_gather_and_prefetch_bwd() + if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -696,6 +715,10 @@ def fuser_backward( "use_dynamic_sched": True, } + # EGTP: col-AG the columnwise layout for the FC1 dgrad (only when dgrad runs). + if getattr(fc1_ctx, "gtp_size", 1) > 1: + grouped_fc1_weight = fc1_op.weight0.batched_all_gather_and_prefetch_bwd() + if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc1_weight_for_gemm = grouped_fc1_weight.copy() diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index a0c5f766c5..9d5e7e445b 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -165,6 +165,20 @@ def fuser_forward( num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + + # EGTP: expert weights are sharded 1/N along out_features; the fused kernels read + # the full shape, so all-gather the full weight first. gtp_size==1 is a no-op. + fc1_gtp_size = getattr(fc1_weight_param, "gtp_size", 1) + fc2_gtp_size = getattr(fc2_weight_param, "gtp_size", 1) + assert fc1_gtp_size == fc2_gtp_size, "FC1/FC2 must share one EGTP group." + if fc1_gtp_size > 1: + assert not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight, ( + "EGTP + cuteDSL fused grouped-MLP only supports the discrete " + "(single_grouped_weight=False) expert-weight layout." + ) + assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None + assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None + device = fc1_weight_param.device if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") @@ -231,7 +245,22 @@ def fuser_forward( None, ) else: - fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] + if fc1_gtp_size > 1: + # Init per-shard quantizers once (quantize-then-gather). + if fc1_op.weight0._quantizer is None: + fc1_op.weight0.setup( + weight_quantizer=[ + fc1_op.get_quantizer("forward", 2 * idx + 1) + for idx in range(num_groups) + ] + ) + # All-gather the full per-expert weights (returns a list of N full tensors). + # TODO: pass in is_first_microbatch flag to skip redundant quantization after the first microbatch in each training step. + fc1_weights = fc1_op.weight0.batched_all_gather_and_prefetch( + fwd=True, skip_weight_cast=False, cast_noop_flag=None + ) + else: + fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] quantized_fc1_weights = [] for idx, weight in enumerate(fc1_weights): quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) @@ -242,48 +271,11 @@ def fuser_forward( quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights - # Prepare FC2 grouped weight tensor for fused kernels. - if fc2_op.single_grouped_weight: - if not isinstance(fc2_op.weight, GroupedTensor): - raise RuntimeError( - "FC2 expected GroupedTensor weight with single_grouped_weight=True." - ) - if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - fc2_op.weight.quantizer = fc2_weight_quantizer - grouped_fc2_weight = fc2_op.weight - else: - if fc2_op.weight.rowwise_data is None: - raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - grouped_fc2_weight = tex.group_quantize( - fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), - fc2_weight_quantizer, - num_groups, - None, - ) - else: - fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc2_weights = [] - for idx, weight in enumerate(fc2_weights): - quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc2_weights.append(quantizer(weight)) - else: - quantized_fc2_weights.append(weight) - grouped_fc2_weight = quantized_fc2_weights - # Some wrapper-copy paths may drop grouped storage metadata; enforce defaults. if getattr(grouped_fc1_weight, "_with_gemm_swizzled_scales", None) is None and isinstance( grouped_fc1_weight, GroupedTensor ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if getattr(grouped_fc2_weight, "_with_gemm_swizzled_scales", None) is None and isinstance( - grouped_fc2_weight, GroupedTensor - ): - grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) @@ -398,6 +390,59 @@ def fuser_forward( fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + # Prepare FC2 grouped weight, deferred to after the FC1 GEMM launch so the EGTP + # FC2 all-gather prefetch overlaps it (perf-neutral for non-EGTP). + if fc2_op.single_grouped_weight: + if not isinstance(fc2_op.weight, GroupedTensor): + raise RuntimeError( + "FC2 expected GroupedTensor weight with single_grouped_weight=True." + ) + if fc2_op.weight.quantizer is not None: + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc2_op.weight.quantizer = fc2_weight_quantizer + grouped_fc2_weight = fc2_op.weight + else: + if fc2_op.weight.rowwise_data is None: + raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + grouped_fc2_weight = tex.group_quantize( + fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), + fc2_weight_quantizer, + num_groups, + None, + ) + else: + if fc2_gtp_size > 1: + # Init per-shard quantizers once (quantize-then-gather); gated so the + # quantizer list isn't rebuilt every forward. + if fc2_op.weight0._quantizer is None: + fc2_op.weight0.setup( + weight_quantizer=[ + fc2_op.get_quantizer("forward", 2 * idx + 1) + for idx in range(num_groups) + ] + ) + # All-gather the full per-expert weights (returns a list of N full tensors). + fc2_weights = fc2_op.weight0.batched_all_gather_and_prefetch( + fwd=True, skip_weight_cast=False, cast_noop_flag=None + ) + else: + fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] + quantized_fc2_weights = [] + for idx, weight in enumerate(fc2_weights): + quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + if not is_quantized_tensor(weight): + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + quantized_fc2_weights.append(quantizer(weight)) + else: + quantized_fc2_weights.append(weight) + grouped_fc2_weight = quantized_fc2_weights + if getattr(grouped_fc2_weight, "_with_gemm_swizzled_scales", None) is None and isinstance( + grouped_fc2_weight, GroupedTensor + ): + grouped_fc2_weight._with_gemm_swizzled_scales = False + # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous # logical dims. @@ -526,9 +571,13 @@ def fuser_forward( # FC1 saved-tensor layout. # [split_sizes, base_split_offsets, split_points, # grouped_fc1_x, *fc1_weight_tensors] - fc1_weight_tensors = ( - [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight - ) + # EGTP: save the small sharded params; backward col-AGs the layout it needs. + if fc1_gtp_size > 1: + fc1_weight_tensors = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] + else: + fc1_weight_tensors = ( + [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight + ) fc1_ctx.save_for_backward( split_sizes, base_split_offsets, @@ -545,6 +594,7 @@ def fuser_forward( fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad + fc1_ctx.gtp_size = fc1_gtp_size # Activation activation_ctx.save_for_backward(activation_in, scales) @@ -559,9 +609,12 @@ def fuser_forward( # [split_sizes, base_split_offsets, split_points, # (fc2_scales if _scale_bias), # grouped_fc2_x, *fc2_weight_tensors] - fc2_weight_tensors = ( - [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight - ) + if fc2_gtp_size > 1: + fc2_weight_tensors = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] + else: + fc2_weight_tensors = ( + [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight + ) fc2_saved: list[Optional[torch.Tensor]] = [ split_sizes, base_split_offsets, @@ -581,6 +634,7 @@ def fuser_forward( fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad + fc2_ctx.gtp_size = fc2_gtp_size fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x return fc2_out, [(), (), ()] From 8bb26f047f690bb8a7e5884a1f9bf419a34fe002 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sun, 14 Jun 2026 07:32:44 -0700 Subject: [PATCH 03/15] [fix] Respect per-op activation-offload markers in fused grouped MLP Signed-off-by: Shiqing Fan --- .../pytorch/ops/fused/grouped_mlp.py | 29 +++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 034fb0766e..b40e5b225e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,7 +17,12 @@ import transformer_engine_torch as tex from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload +from ...cpu_offload import ( + is_cpu_offload_enabled, + mark_activation_offload, + mark_not_offload, + start_offload, +) from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe @@ -1509,9 +1514,29 @@ def fuser_forward( grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None + # Per-op fine-grained offload markers. + offload_fc1_x = bool(getattr(fc1_op, "fine_grained_activation_offloading", False)) + offload_act = bool(getattr(activation_op, "fine_grained_activation_offloading", False)) + fine_grained_offload = offload_fc1_x or offload_act + saved_activations = ( + (grouped_fc1_x, offload_fc1_x), + (activation_in, offload_act), + (saved_grouped_fc2_x, offload_act), + ) + + # The hook-based offloader is opt-out, so explicitly keep the + # non-selected tensors resident (mark_not_offload sets _TE_do_not_offload). + if fine_grained_offload: + keep = [t for t, sel in saved_activations if t is not None and not sel] + if keep: + mark_not_offload(*keep) + if cpu_offloading: + # TE-native path; with no markers, offload everything saved (legacy). activation_tensors = [ - t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None + t + for t, sel in saved_activations + if t is not None and (sel or not fine_grained_offload) ] start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) From 8aaa5a16887d303b6c00bc3fedc4210e7a3ff3a6 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sat, 27 Jun 2026 08:27:07 -0700 Subject: [PATCH 04/15] Code clean: rename GTP weight-sharding axis to gtp_remat Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/module/base.py | 12 +++--- .../pytorch/module/grouped_linear.py | 42 +++++++++---------- .../pytorch/module/layernorm_linear.py | 42 +++++++++---------- transformer_engine/pytorch/module/linear.py | 40 +++++++++--------- .../pytorch/ops/fused/grouped_mlp.py | 24 +++++------ 5 files changed, 80 insertions(+), 80 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f4e8ad430d..f672796317 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -85,7 +85,7 @@ # GTP hook slots. An external integrator (currently ``megatron.experimental.gtp``) # populates these via ``register_gtp_hooks`` at its own import time. When the -# slots stay ``None``, the ``gtp_group=`` codepath in TE modules is a no-op +# slots stay ``None``, the ``gtp_remat_group=`` codepath in TE modules is a no-op # and TE has no ``from megatron...`` dependency. _gtp_slice_fn = None _gtp_finalize_fn = None @@ -99,7 +99,7 @@ def register_gtp_hooks(*, slice_fn=None, finalize_fn=None, wrap_fn=None): Fires per weight during ``reset_parameters``, before FP8 quantize. finalize_fn(module, weight_names) -> None Fires after the per-weight loop in ``reset_parameters``. - wrap_fn(module, weight_names, gtp_group, is_grouped=False) -> None + wrap_fn(module, weight_names, gtp_remat_group, is_grouped=False) -> None Fires at the end of a module's ``__init__`` to finalize GTP wiring. """ global _gtp_slice_fn, _gtp_finalize_fn, _gtp_wrap_fn @@ -111,17 +111,17 @@ def register_gtp_hooks(*, slice_fn=None, finalize_fn=None, wrap_fn=None): _gtp_wrap_fn = wrap_fn -def maybe_wrap_gtp(module, weight_names, gtp_group, is_grouped=False): +def maybe_wrap_gtp(module, weight_names, gtp_remat_group, is_grouped=False): """Finalize GTP wiring on a module if a wrap hook is registered. - No-op when ``gtp_group`` is None or no GTP integrator has called + No-op when ``gtp_remat_group`` is None or no GTP integrator has called ``register_gtp_hooks``. Called from each TE module's ``__init__`` after ``reset_parameters`` finishes; the per-weight slice already happened inside ``reset_parameters`` via ``_gtp_slice_fn``. """ - if gtp_group is None or _gtp_wrap_fn is None: + if gtp_remat_group is None or _gtp_wrap_fn is None: return - _gtp_wrap_fn(module, weight_names, gtp_group, is_grouped=is_grouped) + _gtp_wrap_fn(module, weight_names, gtp_remat_group, is_grouped=is_grouped) def is_ub_initialized() -> bool: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 37347df25d..fa08462a7c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -425,7 +425,7 @@ def forward( skip_fp8_weight_update, save_original_input, debug, - gtp_size, + gtp_remat_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -441,7 +441,7 @@ def forward( weight_requires_grad = weights[0].requires_grad weights_gtp_sharded = weights - if gtp_size > 1: + if gtp_remat_size > 1: weights = weights[0].batched_all_gather_and_prefetch( fwd=True, skip_weight_cast=is_first_microbatch is False, @@ -651,12 +651,12 @@ def forward( # Python parameter attributes without keeping the parameter alive here. saved_weights = ( weights - if backward_override == "high_precision" and inp.requires_grad and gtp_size == 1 + if backward_override == "high_precision" and inp.requires_grad and gtp_remat_size == 1 else [None] * num_gemms ) tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, - *weights_fp8 if gtp_size == 1 else weights_gtp_sharded, + *weights_fp8 if gtp_remat_size == 1 else weights_gtp_sharded, *saved_weights, *biases, ) @@ -681,7 +681,7 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] - elif gtp_size > 1: + elif gtp_remat_size > 1: ctx.main_grad_funcs = [ weights_gtp_sharded[i].get_wgrad_tensor for i in range(num_gemms) ] @@ -714,7 +714,7 @@ def forward( ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers - ctx.gtp_size = gtp_size + ctx.gtp_remat_size = gtp_remat_size # backward overrides if backward_override is not None: @@ -939,7 +939,7 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: # GTP: origin_weights come from saved tensors; main_grads are # get_wgrad_tensor scratch (do not assign to param.main_grad). origin_weights = weights @@ -1006,7 +1006,7 @@ def backward( ctx.m_splits, ) - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: accumulate_wgrad_into_param_main_grad = False elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1015,7 +1015,7 @@ def backward( else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: weights = origin_weights[0].batched_all_gather_and_prefetch_bwd() if ctx.requires_dgrad: @@ -1086,7 +1086,7 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: # Gathered weights are no longer needed after dgrad GEMM. del weights @@ -1139,7 +1139,7 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if ctx.gtp_size == 1 + if ctx.gtp_remat_size == 1 and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), @@ -1182,7 +1182,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: wgrad_list = origin_weights[0].batched_wgrad_reduce_scatter(wgrad_list) # Drop Python refs to wgrad input buffers. The async RS on rs_stream # still holds C++ refs (via NCCL Work); those are released when @@ -1313,7 +1313,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, - gtp_group: Optional[dist_group_type] = None, + gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1369,10 +1369,10 @@ def __init__( "Because the TP communication is handled outside of this module." ) - if gtp_group is None: - self.gtp_size = 1 + if gtp_remat_group is None: + self.gtp_remat_size = 1 else: - self.gtp_size = get_distributed_world_size(gtp_group) + self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) self.parallel_mode = parallel_mode if self.parallel_mode not in GemmParallelModes: @@ -1427,15 +1427,15 @@ def __init__( self.weight_names = [f"weight{idx}" for idx in range(self.num_gemms)] is_meta = torch.device(device).type == "meta" - if gtp_group is not None: + if gtp_remat_group is not None: # Stashed before reset_parameters so the slice hook can see it; # _gtp_is_grouped routes through the GroupedLinear finalize path. - self._gtp_group = gtp_group + self._gtp_group = gtp_remat_group self._gtp_is_grouped = True self.reset_parameters(defer_init=is_meta) - maybe_wrap_gtp(self, self.weight_names, gtp_group, is_grouped=True) + maybe_wrap_gtp(self, self.weight_names, gtp_remat_group, is_grouped=True) if self.wgrad_store.delay_wgrad_compute(): for name, param in self.named_parameters(): @@ -1786,7 +1786,7 @@ def forward( weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() - if self.gtp_size > 1: + if self.gtp_remat_size > 1: weight_tensors[0].setup( weight_quantizer=self._get_weight_quantizers(), ) @@ -1843,7 +1843,7 @@ def forward( skip_fp8_weight_update, self.save_original_input, debug, - self.gtp_size, + self.gtp_remat_size, ) out, new_workspaces = linear_fn( *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 7e74282e4a..374297d7d5 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -146,7 +146,7 @@ def forward( symmetric_ar_type, debug, is_fsdp2, - gtp_size, + gtp_remat_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -303,7 +303,7 @@ def forward( # ------------------------------------------------------ weight_gtp_sharded = weight - if gtp_size > 1: + if gtp_remat_size > 1: weight = weight.all_gather_and_prefetch( fwd=True, skip_weight_cast=is_first_microbatch is False, @@ -504,8 +504,8 @@ def forward( tensors_to_save, tensor_objects = prepare_for_saving( inputmat, # GTP: save the sharded reference only; backward re-gathers it. - wt_save if gtp_size == 1 else None, - weight if gtp_size == 1 else weight_gtp_sharded, + wt_save if gtp_remat_size == 1 else None, + weight if gtp_remat_size == 1 else weight_gtp_sharded, bias, ln_weight, ln_out_to_save, @@ -532,7 +532,7 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad - elif gtp_size > 1: + elif gtp_remat_size > 1: ctx.main_grad_func = weight_gtp_sharded.get_wgrad_tensor else: ctx.main_grad_func = lambda: weight.main_grad @@ -576,7 +576,7 @@ def forward( qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug - ctx.gtp_size = gtp_size + ctx.gtp_remat_size = gtp_remat_size # backward overrides if backward_override is not None: @@ -628,7 +628,7 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: weight = saved_weight.all_gather_and_prefetch_bwd() # Restore from weakref to get original weight python object @@ -648,7 +648,7 @@ def backward( ), "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 ctx.gtp_size == 1: + if main_grad is not None and ctx.gtp_remat_size == 1: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -983,7 +983,7 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: # GTP: accumulation happens downstream in wgrad_reduce_scatter. accumulate_wgrad_into_param_main_grad = False elif ctx.is_first_microbatch is not None: @@ -1058,7 +1058,7 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: wgrad = saved_weight.wgrad_reduce_scatter(wgrad) # Update grad bias if needed @@ -1140,7 +1140,7 @@ def wgrad_gemm( if ctx.requires_wgrad: # Handle custom DDP from mcore. - if ctx.gtp_size > 1: + if ctx.gtp_remat_size > 1: # GTP: skip — wgrad RS already produced the correct shard. pass elif ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): @@ -1310,7 +1310,7 @@ def __init__( delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, name: Optional[str] = None, - gtp_group: Optional[dist_group_type] = None, + gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1341,10 +1341,10 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - if gtp_group is None: - self.gtp_size = 1 + if gtp_remat_group is None: + self.gtp_remat_size = 1 else: - self.gtp_size = get_distributed_world_size(gtp_group) + self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) self.parallel_mode = parallel_mode assert ( @@ -1556,15 +1556,15 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() - if gtp_group is not None: + if gtp_remat_group is not None: # Stashed before reset_parameters so the slice hook can see it. - self._gtp_group = gtp_group + self._gtp_group = gtp_remat_group self._gtp_is_grouped = False self.reset_parameters(defer_init=device == "meta") - maybe_wrap_gtp(self, self.weight_names, gtp_group) - if gtp_group is not None: + maybe_wrap_gtp(self, self.weight_names, gtp_remat_group) + if gtp_remat_group is not None: # Free the full-size backing buffer; GTP replaced it with a sharded param. del weight_tensor @@ -1730,7 +1730,7 @@ def forward( # Get concatenated weight and bias tensors weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() - if self.gtp_size > 1: + if self.gtp_remat_size > 1: weight_tensor.setup( weight_quantizer=self._get_weight_quantizers(), ) @@ -1805,7 +1805,7 @@ def forward( self.symmetric_ar_type, debug, self.is_fsdp2, - self.gtp_size, + self.gtp_remat_size, ) out, ln_out, new_weight_workspace = fwd_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 22923d1328..6afb49aa06 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -157,7 +157,7 @@ class LinearFwdArgs: is_grad_enabled: bool # --- Generalized tensor parallelism --- - gtp_size: int = 1 + gtp_remat_size: int = 1 @dataclass(slots=True) @@ -229,7 +229,7 @@ class LinearBwdArgs: owns_input: bool = False # --- Generalized tensor parallelism --- - gtp_size: int = 1 + gtp_remat_size: int = 1 # --- Per-backward scratch state (populated inside _linear_backward) --- ub_obj_gradout: Optional[Any] = None @@ -410,7 +410,7 @@ def _linear_forward_impl( # ------------------------------------------------------ # GTP: rebind `weight` to the all-gathered tensor; `args.weight` keeps # the GTPShardedParam reference for backward re-gather / wgrad RS. - if args.gtp_size > 1: + if args.gtp_remat_size > 1: weight = weight.all_gather_and_prefetch( fwd=True, skip_weight_cast=is_first_microbatch is False, @@ -606,7 +606,7 @@ def _linear_forward_impl( if is_fsdp2 and weightmat is not weight: wt_save = None # GTP: don't save the workspace; backward re-gathers it. - if args.gtp_size > 1: + if args.gtp_remat_size > 1: wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` @@ -712,14 +712,14 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad - elif fwd_args.gtp_size > 1: + elif fwd_args.gtp_remat_size > 1: bwd_args.main_grad_func = weight.get_wgrad_tensor else: bwd_args.main_grad_func = lambda: weight.main_grad # Misc bwd_args.cpu_offloading = fwd_args.cpu_offloading - bwd_args.gtp_size = fwd_args.gtp_size + bwd_args.gtp_remat_size = fwd_args.gtp_remat_size if backward_override is not None: bwd_args.fp8 = False @@ -786,7 +786,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - if bwd_args.gtp_size == 1: + if bwd_args.gtp_remat_size == 1: origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -960,7 +960,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # GTP: re-gather the sharded weight; runs even when requires_dgrad=False # so the prev_w prefetch is issued for the next layer's bwd. - if bwd_args.gtp_size > 1: + if bwd_args.gtp_remat_size > 1: weight_fp8 = saved_weight.all_gather_and_prefetch_bwd() if bwd_args.requires_dgrad: @@ -1171,7 +1171,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.gtp_size > 1: + if bwd_args.gtp_remat_size > 1: # GTP: accumulation happens downstream in wgrad_reduce_scatter. accumulate_wgrad_into_param_main_grad = False elif bwd_args.is_first_microbatch is not None: @@ -1252,7 +1252,7 @@ def wgrad_gemm( # GTP: reduce-scatter the freshly computed wgrad (async; overlap # with the next layer's bwd via the cascade). - if bwd_args.gtp_size > 1: + if bwd_args.gtp_remat_size > 1: wgrad = saved_weight.wgrad_reduce_scatter(wgrad) # Update grad bias if needed @@ -1528,7 +1528,7 @@ def __init__( symmetric_ar_type: Optional[str] = None, save_original_input: bool = False, name: Optional[str] = None, - gtp_group: Optional[dist_group_type] = None, + gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1557,10 +1557,10 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - if gtp_group is None: - self.gtp_size = 1 + if gtp_remat_group is None: + self.gtp_remat_size = 1 else: - self.gtp_size = get_distributed_world_size(gtp_group) + self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) self.parallel_mode = parallel_mode assert ( @@ -1747,15 +1747,15 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() - if gtp_group is not None: + if gtp_remat_group is not None: # Stashed before reset_parameters so the slice hook can see it. - self._gtp_group = gtp_group + self._gtp_group = gtp_remat_group self._gtp_is_grouped = False self.reset_parameters(defer_init=device == "meta") - maybe_wrap_gtp(self, self.weight_names, gtp_group) - if gtp_group is not None: + maybe_wrap_gtp(self, self.weight_names, gtp_remat_group) + if gtp_remat_group is not None: # Free the full-size backing buffer; GTP replaced it with a sharded param. del weight_tensor @@ -1889,7 +1889,7 @@ def forward( try: weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() - if self.gtp_size > 1: + if self.gtp_remat_size > 1: weight_tensor.setup( weight_quantizer=self._get_weight_quantizers(), ) @@ -2013,7 +2013,7 @@ def forward( cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, # generalized tensor parallelism - gtp_size=self.gtp_size, + gtp_remat_size=self.gtp_remat_size, ) out, new_weight_workspace = linear_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index b40e5b225e..10c7afbc70 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -539,7 +539,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() - gtp_size = getattr(ctx, "gtp_size", 1) + gtp_remat_size = getattr(ctx, "gtp_remat_size", 1) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -574,7 +574,7 @@ def _compute_grad_params( if ctx.weight_requires_grad: # EGTP: the GEMM produces full-sized wgrads but main_grad is sharded, so use a # full-sized scratch buffer (the reduce-scatter below lands it in main_grad). - if fc_op._accumulate_into_main_grad and gtp_size == 1: + if fc_op._accumulate_into_main_grad and gtp_remat_size == 1: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -590,7 +590,7 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() - if gtp_size > 1 and delay_wgrad: + if gtp_remat_size > 1 and delay_wgrad: raise RuntimeError( "EGTP + cuteDSL fused grouped-MLP does not support delay_wgrad / " "overlap_dispatch_backward_with_experts_wgrad yet; set " @@ -636,12 +636,12 @@ def _compute_grad_params( gemm_fn(grouped_x, grouped_dy, wgrad_output) # EGTP: reduce-scatter the full per-rank wgrads into each sharded main_grad # (also fires the Megatron grad-accum hook). - if gtp_size > 1: + if gtp_remat_size > 1: weights[0].batched_wgrad_reduce_scatter(w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added # (wgrad fusion, or the EGTP reduce-scatter above) so it doesn't double-add. - if fc_op._accumulate_into_main_grad or gtp_size > 1: + if fc_op._accumulate_into_main_grad or gtp_remat_size > 1: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -906,9 +906,9 @@ def fuser_forward( fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 # EGTP: expert weights are sharded 1/N along out_features; the fused kernels read - # the full shape, so all-gather the full weight first. gtp_size==1 is a no-op. - fc1_gtp_size = getattr(fc1_weight_param, "gtp_size", 1) - fc2_gtp_size = getattr(fc2_weight_param, "gtp_size", 1) + # the full shape, so all-gather the full weight first. gtp_remat_size==1 is a no-op. + fc1_gtp_size = getattr(fc1_weight_param, "gtp_remat_size", 1) + fc2_gtp_size = getattr(fc2_weight_param, "gtp_remat_size", 1) assert fc1_gtp_size == fc2_gtp_size, "FC1/FC2 must share one EGTP group." if fc1_gtp_size > 1: assert not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight, ( @@ -1573,7 +1573,7 @@ def fuser_forward( fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad - fc1_ctx.gtp_size = fc1_gtp_size + fc1_ctx.gtp_remat_size = fc1_gtp_size fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] @@ -1581,7 +1581,7 @@ def fuser_forward( fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x - fc2_ctx.gtp_size = fc2_gtp_size + fc2_ctx.gtp_remat_size = fc2_gtp_size return fc2_out, [(), (), ()] @@ -1823,7 +1823,7 @@ def fuser_backward( # EGTP: forward gathered rowwise; col-AG the columnwise layout the dgrad needs, # right before use (one weight live at a time). - if getattr(fc2_ctx, "gtp_size", 1) > 1: + if getattr(fc2_ctx, "gtp_remat_size", 1) > 1: grouped_fc2_weight = fc2_op.weight0.batched_all_gather_and_prefetch_bwd() if fc2_op.single_grouped_weight: @@ -2105,7 +2105,7 @@ def fuser_backward( } # EGTP: col-AG the columnwise layout for the FC1 dgrad (only when dgrad runs). - if getattr(fc1_ctx, "gtp_size", 1) > 1: + if getattr(fc1_ctx, "gtp_remat_size", 1) > 1: grouped_fc1_weight = fc1_op.weight0.batched_all_gather_and_prefetch_bwd() if fc1_op.single_grouped_weight: From 48d541306a704eec68adf3d7de32186aecd24d83 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 27 Jun 2026 15:31:55 +0000 Subject: [PATCH 05/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/module/grouped_linear.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index fa08462a7c..c0e8a7ba39 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -651,7 +651,9 @@ def forward( # Python parameter attributes without keeping the parameter alive here. saved_weights = ( weights - if backward_override == "high_precision" and inp.requires_grad and gtp_remat_size == 1 + if backward_override == "high_precision" + and inp.requires_grad + and gtp_remat_size == 1 else [None] * num_gemms ) tensors_to_save, tensor_objects = prepare_for_saving( From 5f1cd62df83895012fa39a240a6bfc4d4d7e9783 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Wed, 1 Jul 2026 20:00:51 -0700 Subject: [PATCH 06/15] Revert "[fix] Respect per-op activation-offload markers in fused grouped MLP" This reverts commit 8bb26f047f690bb8a7e5884a1f9bf419a34fe002. Signed-off-by: Shiqing Fan --- .../pytorch/ops/fused/grouped_mlp.py | 29 ++----------------- 1 file changed, 2 insertions(+), 27 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 10c7afbc70..8787a2e10e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,12 +17,7 @@ import transformer_engine_torch as tex from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE -from ...cpu_offload import ( - is_cpu_offload_enabled, - mark_activation_offload, - mark_not_offload, - start_offload, -) +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe @@ -1514,29 +1509,9 @@ def fuser_forward( grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None - # Per-op fine-grained offload markers. - offload_fc1_x = bool(getattr(fc1_op, "fine_grained_activation_offloading", False)) - offload_act = bool(getattr(activation_op, "fine_grained_activation_offloading", False)) - fine_grained_offload = offload_fc1_x or offload_act - saved_activations = ( - (grouped_fc1_x, offload_fc1_x), - (activation_in, offload_act), - (saved_grouped_fc2_x, offload_act), - ) - - # The hook-based offloader is opt-out, so explicitly keep the - # non-selected tensors resident (mark_not_offload sets _TE_do_not_offload). - if fine_grained_offload: - keep = [t for t, sel in saved_activations if t is not None and not sel] - if keep: - mark_not_offload(*keep) - if cpu_offloading: - # TE-native path; with no markers, offload everything saved (legacy). activation_tensors = [ - t - for t, sel in saved_activations - if t is not None and (sel or not fine_grained_offload) + t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None ] start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) From 96ef07149f7635274080fad9d4929a84099c359e Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 6 Jul 2026 00:35:21 -0700 Subject: [PATCH 07/15] Make TE GTP-agnostic at construction Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/module/base.py | 70 +------------------ .../pytorch/module/grouped_linear.py | 21 +----- .../pytorch/module/layernorm_linear.py | 24 +------ transformer_engine/pytorch/module/linear.py | 24 +------ .../pytorch/ops/fused/grouped_mlp.py | 19 ----- 5 files changed, 12 insertions(+), 146 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f672796317..6c7ba8a8ab 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -66,8 +66,6 @@ "initialize_ub", "destroy_ub", "is_ub_initialized", - "maybe_wrap_gtp", - "register_gtp_hooks", "using_cublasmp_backend", "UserBufferQuantizationMode", ] @@ -83,47 +81,6 @@ layers_atomic_ring_exchange = [] -# GTP hook slots. An external integrator (currently ``megatron.experimental.gtp``) -# populates these via ``register_gtp_hooks`` at its own import time. When the -# slots stay ``None``, the ``gtp_remat_group=`` codepath in TE modules is a no-op -# and TE has no ``from megatron...`` dependency. -_gtp_slice_fn = None -_gtp_finalize_fn = None -_gtp_wrap_fn = None - - -def register_gtp_hooks(*, slice_fn=None, finalize_fn=None, wrap_fn=None): - """Register GTP integration hooks. Hooks left as ``None`` are unchanged. - - slice_fn(module, name, param, *, expert_idx) -> GTPShardedParam | None - Fires per weight during ``reset_parameters``, before FP8 quantize. - finalize_fn(module, weight_names) -> None - Fires after the per-weight loop in ``reset_parameters``. - wrap_fn(module, weight_names, gtp_remat_group, is_grouped=False) -> None - Fires at the end of a module's ``__init__`` to finalize GTP wiring. - """ - global _gtp_slice_fn, _gtp_finalize_fn, _gtp_wrap_fn - if slice_fn is not None: - _gtp_slice_fn = slice_fn - if finalize_fn is not None: - _gtp_finalize_fn = finalize_fn - if wrap_fn is not None: - _gtp_wrap_fn = wrap_fn - - -def maybe_wrap_gtp(module, weight_names, gtp_remat_group, is_grouped=False): - """Finalize GTP wiring on a module if a wrap hook is registered. - - No-op when ``gtp_remat_group`` is None or no GTP integrator has called - ``register_gtp_hooks``. Called from each TE module's ``__init__`` after - ``reset_parameters`` finishes; the per-weight slice already happened - inside ``reset_parameters`` via ``_gtp_slice_fn``. - """ - if gtp_remat_group is None or _gtp_wrap_fn is None: - return - _gtp_wrap_fn(module, weight_names, gtp_remat_group, is_grouped=is_grouped) - - def is_ub_initialized() -> bool: """Whether the Userbuffers communicators have been initialized.""" return _ub_initialized @@ -1723,10 +1680,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: if defer_init: return - # Names of GTP-sharded weights, for GroupedLinear's post-loop finalize. - _gtp_sharded_weight_names = [] - - for idx, (name, param) in enumerate(self.named_parameters(recurse=False)): + for name, param in self.named_parameters(recurse=False): # Check if parameter is a DTensor (FSDP2) or regular tensor is_dtensor = isinstance(param, DTensor) dtensor_param = param if is_dtensor else None @@ -1748,23 +1702,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: with get_rng_state_tracker().fork(): init_fn(param) - # GTP slice: shard the freshly-init weight into a GTPShardedParam; - # the FP8 quantize block below is skipped for it. - gtp_sharded = None - if ( - not is_dtensor - and getattr(self, "_gtp_group", None) is not None - and _gtp_slice_fn is not None - ): - gtp_sharded = _gtp_slice_fn(self, name, param, expert_idx=idx) - if gtp_sharded is not None: - param = gtp_sharded - _gtp_sharded_weight_names.append(name) - # Wrap parameters in QuantizedTensor if needed fp8_meta_index = self.param_init_meta[name].fp8_meta_index high_precision_init_val = None - if self.primary_weights_in_fp8 and fp8_meta_index is not None and gtp_sharded is None: + if self.primary_weights_in_fp8 and fp8_meta_index is not None: # Keep high-precision values on CPU if needed if self.preserve_high_precision_init_val: @@ -1794,7 +1735,6 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: # NOTE: Currently this can only be broken when primary weights are in Fp8 but # re-applying the nn.Parameter() wrap is a no-op when the input is already # a parameter so we always re-apply it just for extra safety. - # Skip the wrap for GTPShardedParam (Parameter.__new__ would drop attrs). if is_dtensor: # recreate the DTensor from the parameter. dtensor_param = DTensor.from_local( @@ -1805,7 +1745,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: stride=dtensor_param.stride(), ) dtensor_param = torch.nn.Parameter(dtensor_param) - elif gtp_sharded is None: + else: param = torch.nn.Parameter(param) # Keep high-precision values on CPU if needed @@ -1843,10 +1783,6 @@ def clear(self): else: self.module_setattr(name, dtensor_param) - # GroupedLinear post-loop finalize hook (no-op outside GroupedLinear). - if _gtp_sharded_weight_names and _gtp_finalize_fn is not None: - _gtp_finalize_fn(self, _gtp_sharded_weight_names) - @abstractmethod def forward(self): """Needs override.""" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c0e8a7ba39..eddd0cd59c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -29,7 +29,6 @@ _2X_ACC_WGRAD, ) from ._common import WeightGradStore -from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, @@ -1315,7 +1314,6 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, - gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1371,10 +1369,9 @@ def __init__( "Because the TP communication is handled outside of this module." ) - if gtp_remat_group is None: - self.gtp_remat_size = 1 - else: - self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) + # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites + # this with the real shard count after __init__ for GTP-sharded weights. + self.gtp_remat_size = 1 self.parallel_mode = parallel_mode if self.parallel_mode not in GemmParallelModes: @@ -1429,16 +1426,9 @@ def __init__( self.weight_names = [f"weight{idx}" for idx in range(self.num_gemms)] is_meta = torch.device(device).type == "meta" - if gtp_remat_group is not None: - # Stashed before reset_parameters so the slice hook can see it; - # _gtp_is_grouped routes through the GroupedLinear finalize path. - self._gtp_group = gtp_remat_group - self._gtp_is_grouped = True self.reset_parameters(defer_init=is_meta) - maybe_wrap_gtp(self, self.weight_names, gtp_remat_group, is_grouped=True) - if self.wgrad_store.delay_wgrad_compute(): for name, param in self.named_parameters(): if name in ("weight", "bias"): @@ -1788,11 +1778,6 @@ def forward( weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() - if self.gtp_remat_size > 1: - weight_tensors[0].setup( - weight_quantizer=self._get_weight_quantizers(), - ) - quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() if debug: diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 374297d7d5..caf1861ad6 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -30,7 +30,6 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( assert_dim_for_fp8_exec, @@ -1310,7 +1309,6 @@ def __init__( delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, name: Optional[str] = None, - gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1341,10 +1339,9 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - if gtp_remat_group is None: - self.gtp_remat_size = 1 - else: - self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) + # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites + # this with the real shard count after __init__ for GTP-sharded weights. + self.gtp_remat_size = 1 self.parallel_mode = parallel_mode assert ( @@ -1556,18 +1553,8 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() - if gtp_remat_group is not None: - # Stashed before reset_parameters so the slice hook can see it. - self._gtp_group = gtp_remat_group - self._gtp_is_grouped = False - self.reset_parameters(defer_init=device == "meta") - maybe_wrap_gtp(self, self.weight_names, gtp_remat_group) - if gtp_remat_group is not None: - # Free the full-size backing buffer; GTP replaced it with a sharded param. - del weight_tensor - # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM if self.parallel_mode == "row" and self.apply_bias: @@ -1730,11 +1717,6 @@ def forward( # Get concatenated weight and bias tensors weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() - if self.gtp_remat_size > 1: - weight_tensor.setup( - weight_quantizer=self._get_weight_quantizers(), - ) - quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 6afb49aa06..d8f2795113 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -30,7 +30,6 @@ _2X_ACC_WGRAD, ) from ._common import noop_cat, WeightGradStore -from .base import maybe_wrap_gtp from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, @@ -1528,7 +1527,6 @@ def __init__( symmetric_ar_type: Optional[str] = None, save_original_input: bool = False, name: Optional[str] = None, - gtp_remat_group: Optional[dist_group_type] = None, ) -> None: super().__init__(name) @@ -1557,10 +1555,9 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - if gtp_remat_group is None: - self.gtp_remat_size = 1 - else: - self.gtp_remat_size = get_distributed_world_size(gtp_remat_group) + # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites + # this with the real shard count after __init__ for GTP-sharded weights. + self.gtp_remat_size = 1 self.parallel_mode = parallel_mode assert ( @@ -1747,18 +1744,8 @@ def __init__( if with_fp8_params: self.init_fp8_metadata() - if gtp_remat_group is not None: - # Stashed before reset_parameters so the slice hook can see it. - self._gtp_group = gtp_remat_group - self._gtp_is_grouped = False - self.reset_parameters(defer_init=device == "meta") - maybe_wrap_gtp(self, self.weight_names, gtp_remat_group) - if gtp_remat_group is not None: - # Free the full-size backing buffer; GTP replaced it with a sharded param. - del weight_tensor - # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM if self.parallel_mode == "row" and self.apply_bias: @@ -1889,11 +1876,6 @@ def forward( try: weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() - if self.gtp_remat_size > 1: - weight_tensor.setup( - weight_quantizer=self._get_weight_quantizers(), - ) - quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 8787a2e10e..9beb6c4ab4 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -992,17 +992,7 @@ def fuser_forward( ) else: if fc1_gtp_size > 1: - # Init per-shard quantizers once (quantize-then-gather). - if fc1_op.weight0._quantizer is None: - fc1_op.weight0.setup( - weight_quantizer=[ - fc1_op.get_quantizer("forward", 2 * idx + 1) - for idx in range(num_groups) - ] - ) # All-gather the full per-expert weights (returns a list of N full tensors). - # TODO: pass in is_first_microbatch flag to skip redundant quantization after - # the first microbatch in each training step. fc1_weights = fc1_op.weight0.batched_all_gather_and_prefetch( fwd=True, skip_weight_cast=False, cast_noop_flag=None ) @@ -1269,15 +1259,6 @@ def fuser_forward( ) else: if fc2_gtp_size > 1: - # Init per-shard quantizers once (quantize-then-gather); gated so the - # quantizer list isn't rebuilt every forward. - if fc2_op.weight0._quantizer is None: - fc2_op.weight0.setup( - weight_quantizer=[ - fc2_op.get_quantizer("forward", 2 * idx + 1) - for idx in range(num_groups) - ] - ) # All-gather the full per-expert weights (returns a list of N full tensors). fc2_weights = fc2_op.weight0.batched_all_gather_and_prefetch( fwd=True, skip_weight_cast=False, cast_noop_flag=None From 5252b0e1ca68fead394c35d2fd93e021ea34c117 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sun, 12 Jul 2026 18:57:58 -0700 Subject: [PATCH 08/15] GTP+nvfp4: fix GTP backward GEMM scaling-mode mismatch for bf16-gathered weights Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/module/linear.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index d8f2795113..075edfaf5b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -976,6 +976,16 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif ( + bwd_args.gtp_remat_size > 1 + and bwd_args.fp8 + and bwd_args.weight_quantizer is not None + and not isinstance(weight_fp8, QuantizedTensorStorage) + ): + # GTP re-gathered a BF16 weight (mxfp8-recipe layer kept BF16 for the AG): quantize + # with the layer quantizer so the dgrad operand isn't cast by the delayed recipe. + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(weight_fp8) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): From e59a0548a8b07a34075995482eb66c15cd7469c6 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 13 Jul 2026 00:39:47 -0700 Subject: [PATCH 09/15] Make TE runtime GTP-agnostic via a DistributedWeight protocol Signed-off-by: Shiqing Fan --- tests/pytorch/test_distributed_weight.py | 124 ++++++++++++++++++ .../pytorch/distributed_weight.py | 87 ++++++++++++ .../pytorch/module/grouped_linear.py | 61 ++++----- .../pytorch/module/layernorm_linear.py | 54 ++++---- transformer_engine/pytorch/module/linear.py | 68 ++++------ .../pytorch/ops/fused/grouped_mlp.py | 81 ++++++------ 6 files changed, 328 insertions(+), 147 deletions(-) create mode 100644 tests/pytorch/test_distributed_weight.py create mode 100644 transformer_engine/pytorch/distributed_weight.py diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py new file mode 100644 index 0000000000..b11f9d7f8c --- /dev/null +++ b/tests/pytorch/test_distributed_weight.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. + +These tests exercise TE's weight-parallelism extension point in isolation, with +a tiny in-repo ``FakeDistributedWeight`` stub standing in for a real implementer +(e.g. Megatron's GTPShardedParam). No GPU, process group, or Megatron import is +required: the whole contract lives in TE and is verified against the fake. +""" + +import torch + +from transformer_engine.pytorch.distributed_weight import ( + DistributedWeight, + is_distributed_weight, + materialize_weights_for_forward, + materialize_weights_for_backward, + finalize_weight_grads, +) + + +class FakeDistributedWeight(torch.Tensor): + """Minimal DistributedWeight implementer for dispatcher tests. + + Records how it was called and returns marker tensors so the dispatcher's + behavior (delegation, list normalization, no-op fallback) is observable. + """ + + is_distributed_weight = True + + def __new__(cls, group_size=1): + t = torch.zeros(1).as_subclass(cls) + t.group_size = group_size + t.calls = [] + return t + + def materialize_group_for_forward(self): + self.calls.append("fwd") + out = [torch.full((2, 2), float(i)) for i in range(self.group_size)] + # Match the real GTP contract: single weight returns a bare tensor. + return out if self.group_size > 1 else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls.append(("bwd", kwargs)) + out = [torch.full((2, 2), float(10 + i)) for i in range(self.group_size)] + return out if self.group_size > 1 else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls.append(("finalize", wgrads)) + wl = wgrads if isinstance(wgrads, (list, tuple)) else [wgrads] + out = [w + 100 for w in wl] + return out if self.group_size > 1 else out[0] + + def grad_buffer(self): + return torch.full((2, 2), -1.0) + + +def test_protocol_runtime_checkable(): + """A conforming object passes isinstance; a plain tensor does not.""" + assert isinstance(FakeDistributedWeight(), DistributedWeight) + assert not isinstance(torch.zeros(2), DistributedWeight) + + +def test_is_distributed_weight(): + assert is_distributed_weight(FakeDistributedWeight()) + assert not is_distributed_weight(torch.zeros(2)) + assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) + + +def test_forward_noop_on_plain_tensor(): + """Plain weights pass through unchanged — the critical non-regression.""" + w = torch.nn.Parameter(torch.randn(4, 4)) + out = materialize_weights_for_forward([w]) + assert out == [w] + assert out[0] is w + + +def test_forward_dispatches_single_weight(): + """Linear (N=1): dispatcher delegates and returns a length-1 list.""" + w = FakeDistributedWeight(group_size=1) + out = materialize_weights_for_forward([w]) + assert isinstance(out, list) and len(out) == 1 + assert torch.equal(out[0], torch.zeros(2, 2)) + # The forward dispatcher owns the forward semantic; the leader is delegated to once. + assert w.calls == ["fwd"] + + +def test_forward_dispatches_grouped_weights(): + """GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=3) + out = materialize_weights_for_forward([w, w, w]) + assert isinstance(out, list) and len(out) == 3 + # Leader was called exactly once (coalesced), not once per weight. + assert len(w.calls) == 1 + + +def test_backward_dispatch_and_noop(): + w = FakeDistributedWeight(group_size=2) + out = materialize_weights_for_backward([w, w]) + assert len(out) == 2 and torch.equal(out[0], torch.full((2, 2), 10.0)) + + plain = [torch.zeros(2), torch.ones(2)] + assert materialize_weights_for_backward(plain) == plain + + +def test_finalize_grads_dispatch_and_noop(): + w = FakeDistributedWeight(group_size=1) + wgrads = [torch.zeros(2, 2)] + out = finalize_weight_grads([w], wgrads) + assert len(out) == 1 and torch.equal(out[0], torch.full((2, 2), 100.0)) + + # No-op path leaves the grads untouched. + plain_w = [torch.nn.Parameter(torch.zeros(2))] + g = [torch.ones(2)] + assert finalize_weight_grads(plain_w, g) == g + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + fn() + print(f"{name}: PASS") diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py new file mode 100644 index 0000000000..ac44cd8d59 --- /dev/null +++ b/transformer_engine/pytorch/distributed_weight.py @@ -0,0 +1,87 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP-agnostic weight-parallelism extension point. + +TE owns this contract but ships no implementation; the caller (e.g. Megatron GTP) implements the +protocol on the weight and injects it at construction. Dispatchers are list-shaped (Linear -> 1, +GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. +""" + +from typing import Any, List, Protocol, runtime_checkable + +import torch + +__all__ = [ + "DistributedWeight", + "is_distributed_weight", + "materialize_weights_for_forward", + "materialize_weights_for_backward", + "finalize_weight_grads", +] + + +@runtime_checkable +class DistributedWeight(Protocol): + """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the + grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not + subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + """ + + # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. + is_distributed_weight: bool + + def materialize_group_for_forward(self) -> Any: + """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" + ... + + def materialize_group_for_backward(self) -> Any: + """Re-materialize the full weight(s) for the backward GEMMs.""" + ... + + def finalize_group_grads(self, wgrads: Any) -> Any: + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter).""" + ... + + def grad_buffer(self) -> torch.Tensor: + """The gradient accumulation buffer for this weight.""" + ... + + +def is_distributed_weight(weight: Any) -> bool: + """True if ``weight`` participates in custom weight parallelism (False on plain tensors).""" + return bool(getattr(weight, "is_distributed_weight", False)) + + +def materialize_weights_for_forward(weights: List[Any]) -> List[Any]: + """Materialize a weight group for the forward GEMM. Delegates once to the leader (it may + coalesce) if distributed, else returns the weights unchanged. Always returns a list. + """ + lead = weights[0] + if is_distributed_weight(lead): + out = lead.materialize_group_for_forward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def materialize_weights_for_backward(weights: List[Any]) -> List[Any]: + """Backward counterpart of :func:`materialize_weights_for_forward`.""" + lead = weights[0] + if is_distributed_weight(lead): + out = lead.materialize_group_for_backward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def finalize_weight_grads(weights: List[Any], wgrads: List[Any]) -> List[Any]: + """Post-process the weight grads of a homogeneous group. + + Delegates to the group leader when distributed (e.g. reduce-scatter); + otherwise returns ``wgrads`` unchanged. + """ + lead = weights[0] + if is_distributed_weight(lead): + out = lead.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(wgrads) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index eddd0cd59c..28d882e9ef 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -46,6 +46,12 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weights_for_forward, + materialize_weights_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -424,7 +430,6 @@ def forward( skip_fp8_weight_update, save_original_input, debug, - gtp_remat_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -439,13 +444,9 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad - weights_gtp_sharded = weights - if gtp_remat_size > 1: - weights = weights[0].batched_all_gather_and_prefetch( - fwd=True, - skip_weight_cast=is_first_microbatch is False, - cast_noop_flag=skip_fp8_weight_update, - ) + weight_params = weights + if is_distributed_weight(weights[0]): + weights = materialize_weights_for_forward(weights) # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -648,16 +649,17 @@ def forward( # Original weights are only needed by high_precision dgrad. The weakrefs # used for fused wgrad accumulation serve a different purpose: restoring # Python parameter attributes without keeping the parameter alive here. + _is_dist_weight = is_distributed_weight(weight_params[0]) saved_weights = ( weights if backward_override == "high_precision" and inp.requires_grad - and gtp_remat_size == 1 + and not _is_dist_weight else [None] * num_gemms ) tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, - *weights_fp8 if gtp_remat_size == 1 else weights_gtp_sharded, + *(weight_params if _is_dist_weight else weights_fp8), *saved_weights, *biases, ) @@ -682,10 +684,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] - elif gtp_remat_size > 1: - ctx.main_grad_funcs = [ - weights_gtp_sharded[i].get_wgrad_tensor for i in range(num_gemms) - ] + elif is_distributed_weight(weight_params[0]): + ctx.main_grad_funcs = [weight_params[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -715,7 +715,6 @@ def forward( ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers - ctx.gtp_remat_size = gtp_remat_size # backward overrides if backward_override is not None: @@ -940,9 +939,9 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.gtp_remat_size > 1: - # GTP: origin_weights come from saved tensors; main_grads are - # get_wgrad_tensor scratch (do not assign to param.main_grad). + if is_distributed_weight(weights[0]): + # Distributed weight (e.g. GTP): origin_weights come from saved tensors; + # main_grads are grad_buffer scratch (do not assign to param.main_grad). origin_weights = weights if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] @@ -1007,7 +1006,7 @@ def backward( ctx.m_splits, ) - if ctx.gtp_remat_size > 1: + if is_distributed_weight(origin_weights[0]): accumulate_wgrad_into_param_main_grad = False elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1016,8 +1015,8 @@ def backward( else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - if ctx.gtp_remat_size > 1: - weights = origin_weights[0].batched_all_gather_and_prefetch_bwd() + if is_distributed_weight(origin_weights[0]): + weights = materialize_weights_for_backward(origin_weights) if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD @@ -1087,7 +1086,7 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] - if ctx.gtp_remat_size > 1: + if is_distributed_weight(origin_weights[0]): # Gathered weights are no longer needed after dgrad GEMM. del weights @@ -1140,7 +1139,7 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if ctx.gtp_remat_size == 1 + if not is_distributed_weight(origin_weights[0]) and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), @@ -1183,14 +1182,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - if ctx.gtp_remat_size > 1: - wgrad_list = origin_weights[0].batched_wgrad_reduce_scatter(wgrad_list) - # Drop Python refs to wgrad input buffers. The async RS on rs_stream - # still holds C++ refs (via NCCL Work); those are released when - # _wait_reduce_scatter calls handle.wait() + self.handle = None. - # Without this del, main_grads keeps the tensors alive until function - # return, wasting memory during graph capture warmup. - del main_grads + if is_distributed_weight(origin_weights[0]): + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) else: wgrad_list = [ handle_custom_ddp_from_mcore(weight, main_grad, wgrad) @@ -1369,10 +1362,6 @@ def __init__( "Because the TP communication is handled outside of this module." ) - # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites - # this with the real shard count after __init__ for GTP-sharded weights. - self.gtp_remat_size = 1 - self.parallel_mode = parallel_mode if self.parallel_mode not in GemmParallelModes: raise ValueError( @@ -1426,7 +1415,6 @@ def __init__( self.weight_names = [f"weight{idx}" for idx in range(self.num_gemms)] is_meta = torch.device(device).type == "meta" - self.reset_parameters(defer_init=is_meta) if self.wgrad_store.delay_wgrad_compute(): @@ -1830,7 +1818,6 @@ def forward( skip_fp8_weight_update, self.save_original_input, debug, - self.gtp_remat_size, ) out, new_workspaces = linear_fn( *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index caf1861ad6..748e1d125e 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -55,6 +55,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weights_for_forward, + materialize_weights_for_backward, + finalize_weight_grads, +) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -145,7 +151,6 @@ def forward( symmetric_ar_type, debug, is_fsdp2, - gtp_remat_size, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -301,13 +306,9 @@ def forward( # Prepare weight tensor # ------------------------------------------------------ - weight_gtp_sharded = weight - if gtp_remat_size > 1: - weight = weight.all_gather_and_prefetch( - fwd=True, - skip_weight_cast=is_first_microbatch is False, - cast_noop_flag=skip_fp8_weight_update, - ) + weight_param = weight + if is_distributed_weight(weight): + weight = materialize_weights_for_forward([weight])[0] out_features = weight.shape[0] new_weight_workspace = None @@ -500,11 +501,13 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + _is_dist_weight = is_distributed_weight(weight_param) tensors_to_save, tensor_objects = prepare_for_saving( inputmat, - # GTP: save the sharded reference only; backward re-gathers it. - wt_save if gtp_remat_size == 1 else None, - weight if gtp_remat_size == 1 else weight_gtp_sharded, + # Distributed weight (e.g. GTP): save the sharded reference only; + # backward re-gathers it. + None if _is_dist_weight else wt_save, + weight_param if _is_dist_weight else weight, bias, ln_weight, ln_out_to_save, @@ -531,8 +534,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad - elif gtp_remat_size > 1: - ctx.main_grad_func = weight_gtp_sharded.get_wgrad_tensor + elif is_distributed_weight(weight_param): + ctx.main_grad_func = weight_param.grad_buffer else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -575,7 +578,6 @@ def forward( qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug - ctx.gtp_remat_size = gtp_remat_size # backward overrides if backward_override is not None: @@ -627,9 +629,8 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) - if ctx.gtp_remat_size > 1: - weight = saved_weight.all_gather_and_prefetch_bwd() - + if is_distributed_weight(saved_weight): + weight = materialize_weights_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. @@ -647,7 +648,7 @@ def backward( ), "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 ctx.gtp_remat_size == 1: + if main_grad is not None and not is_distributed_weight(saved_weight): origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -982,8 +983,8 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.gtp_remat_size > 1: - # GTP: accumulation happens downstream in wgrad_reduce_scatter. + if is_distributed_weight(saved_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 = ( @@ -1057,8 +1058,8 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - if ctx.gtp_remat_size > 1: - wgrad = saved_weight.wgrad_reduce_scatter(wgrad) + if is_distributed_weight(saved_weight): + wgrad = finalize_weight_grads([saved_weight], [wgrad])[0] # Update grad bias if needed if grad_bias is None: @@ -1139,8 +1140,8 @@ def wgrad_gemm( if ctx.requires_wgrad: # Handle custom DDP from mcore. - if ctx.gtp_remat_size > 1: - # GTP: skip — wgrad RS already produced the correct shard. + if is_distributed_weight(saved_weight): + # Distributed weight (e.g. GTP): skip — grad finalize already produced the shard. pass elif ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): origin_weight.grad_added_to_main_grad = True @@ -1339,10 +1340,6 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites - # this with the real shard count after __init__ for GTP-sharded weights. - self.gtp_remat_size = 1 - self.parallel_mode = parallel_mode assert ( self.parallel_mode in GemmParallelModes @@ -1787,7 +1784,6 @@ def forward( self.symmetric_ar_type, debug, self.is_fsdp2, - self.gtp_remat_size, ) out, ln_out, new_weight_workspace = fwd_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 075edfaf5b..ac863959d0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -54,6 +54,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weights_for_forward, + materialize_weights_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_gemm, ) @@ -155,9 +161,6 @@ class LinearFwdArgs: cpu_offloading: bool is_grad_enabled: bool - # --- Generalized tensor parallelism --- - gtp_remat_size: int = 1 - @dataclass(slots=True) class LinearBwdArgs: @@ -227,9 +230,6 @@ class LinearBwdArgs: cpu_offloading: bool = False owns_input: bool = False - # --- Generalized tensor parallelism --- - gtp_remat_size: int = 1 - # --- Per-backward scratch state (populated inside _linear_backward) --- ub_obj_gradout: Optional[Any] = None @@ -407,14 +407,11 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ - # GTP: rebind `weight` to the all-gathered tensor; `args.weight` keeps - # the GTPShardedParam reference for backward re-gather / wgrad RS. - if args.gtp_remat_size > 1: - weight = weight.all_gather_and_prefetch( - fwd=True, - skip_weight_cast=is_first_microbatch is False, - cast_noop_flag=args.skip_fp8_weight_update, - ) + # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; + # `args.weight` keeps the sharded-param reference for backward re-gather / grad + # finalize. No-op for a plain weight. + if is_distributed_weight(args.weight): + weight = materialize_weights_for_forward([args.weight])[0] # Refresh out_features from the gathered weight (captured sharded above, pre-gather). out_features = weight.shape[0] @@ -604,8 +601,8 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None - # GTP: don't save the workspace; backward re-gathers it. - if args.gtp_remat_size > 1: + # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. + if is_distributed_weight(args.weight): wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` @@ -711,14 +708,13 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad - elif fwd_args.gtp_remat_size > 1: - bwd_args.main_grad_func = weight.get_wgrad_tensor + 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 - bwd_args.gtp_remat_size = fwd_args.gtp_remat_size if backward_override is not None: bwd_args.fp8 = False @@ -785,7 +781,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - if bwd_args.gtp_remat_size == 1: + if not is_distributed_weight(bwd_args.saved_weight): origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -957,10 +953,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None - # GTP: re-gather the sharded weight; runs even when requires_dgrad=False - # so the prev_w prefetch is issued for the next layer's bwd. - if bwd_args.gtp_remat_size > 1: - weight_fp8 = saved_weight.all_gather_and_prefetch_bwd() + # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when + # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. + if is_distributed_weight(saved_weight): + weight_fp8 = materialize_weights_for_backward([saved_weight])[0] if bwd_args.requires_dgrad: @@ -977,13 +973,13 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) elif ( - bwd_args.gtp_remat_size > 1 + is_distributed_weight(saved_weight) and bwd_args.fp8 and bwd_args.weight_quantizer is not None and not isinstance(weight_fp8, QuantizedTensorStorage) ): - # GTP re-gathered a BF16 weight (mxfp8-recipe layer kept BF16 for the AG): quantize - # with the layer quantizer so the dgrad operand isn't cast by the delayed recipe. + # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer + # so the dgrad operand isn't cast by the delayed recipe. bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(weight_fp8) @@ -1180,8 +1176,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.gtp_remat_size > 1: - # GTP: accumulation happens downstream in wgrad_reduce_scatter. + if is_distributed_weight(bwd_args.saved_weight): + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. accumulate_wgrad_into_param_main_grad = False elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1259,10 +1255,10 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) - # GTP: reduce-scatter the freshly computed wgrad (async; overlap - # with the next layer's bwd via the cascade). - if bwd_args.gtp_remat_size > 1: - wgrad = saved_weight.wgrad_reduce_scatter(wgrad) + # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad + # (async; overlap with the next layer's bwd via the cascade). + if is_distributed_weight(saved_weight): + wgrad = finalize_weight_grads([saved_weight], [wgrad])[0] # Update grad bias if needed if grad_bias is None: @@ -1565,10 +1561,6 @@ def __init__( self.set_tensor_parallel_group(tp_group) self.set_nccl_overlap_warning_if_tp() - # GTP is attached post-init by Megatron (GTP-agnostic construction); Megatron overwrites - # this with the real shard count after __init__ for GTP-sharded weights. - self.gtp_remat_size = 1 - self.parallel_mode = parallel_mode assert ( self.parallel_mode in GemmParallelModes @@ -2004,8 +1996,6 @@ def forward( # misc cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, - # generalized tensor parallelism - gtp_remat_size=self.gtp_remat_size, ) out, new_weight_workspace = linear_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 9beb6c4ab4..2ca0864117 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -19,6 +19,12 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...distributed_weight import ( + is_distributed_weight, + materialize_weights_for_forward, + materialize_weights_for_backward, + finalize_weight_grads, +) from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer @@ -534,7 +540,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() - gtp_remat_size = getattr(ctx, "gtp_remat_size", 1) + is_dist_weight = is_distributed_weight(weights[0]) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -567,9 +573,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - # EGTP: the GEMM produces full-sized wgrads but main_grad is sharded, so use a - # full-sized scratch buffer (the reduce-scatter below lands it in main_grad). - if fc_op._accumulate_into_main_grad and gtp_remat_size == 1: + # Distributed weight: the GEMM produces full-sized wgrads but main_grad is sharded, + # so use a full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and not is_dist_weight: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -585,11 +591,9 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() - if gtp_remat_size > 1 and delay_wgrad: + if is_dist_weight and delay_wgrad: raise RuntimeError( - "EGTP + cuteDSL fused grouped-MLP does not support delay_wgrad / " - "overlap_dispatch_backward_with_experts_wgrad yet; set " - "delay_wgrad_compute=False." + "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." ) if cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) @@ -629,14 +633,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - # EGTP: reduce-scatter the full per-rank wgrads into each sharded main_grad - # (also fires the Megatron grad-accum hook). - if gtp_remat_size > 1: - weights[0].batched_wgrad_reduce_scatter(w_list) + # Distributed weight: reduce-scatter the full per-rank wgrads into each sharded + # main_grad (also fires the Megatron grad-accum hook). + if is_dist_weight: + finalize_weight_grads([weights[0]], w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - # (wgrad fusion, or the EGTP reduce-scatter above) so it doesn't double-add. - if fc_op._accumulate_into_main_grad or gtp_remat_size > 1: + # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. + if fc_op._accumulate_into_main_grad or is_dist_weight: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -900,15 +904,14 @@ def fuser_forward( fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 - # EGTP: expert weights are sharded 1/N along out_features; the fused kernels read - # the full shape, so all-gather the full weight first. gtp_remat_size==1 is a no-op. - fc1_gtp_size = getattr(fc1_weight_param, "gtp_remat_size", 1) - fc2_gtp_size = getattr(fc2_weight_param, "gtp_remat_size", 1) - assert fc1_gtp_size == fc2_gtp_size, "FC1/FC2 must share one EGTP group." - if fc1_gtp_size > 1: + # Distributed weight: expert weights are sharded 1/N along out_features; the fused kernels + # read the full shape, so all-gather the full weight first. A plain weight is a no-op. + fc1_is_dist = is_distributed_weight(fc1_weight_param) + fc2_is_dist = is_distributed_weight(fc2_weight_param) + assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." + if fc1_is_dist: assert not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight, ( - "EGTP + cuteDSL fused grouped-MLP only supports the discrete " - "(single_grouped_weight=False) expert-weight layout." + "distributed-weight fused grouped-MLP requires single_grouped_weight=False." ) assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None @@ -991,11 +994,9 @@ def fuser_forward( None, ) else: - if fc1_gtp_size > 1: + if fc1_is_dist: # All-gather the full per-expert weights (returns a list of N full tensors). - fc1_weights = fc1_op.weight0.batched_all_gather_and_prefetch( - fwd=True, skip_weight_cast=False, cast_noop_flag=None - ) + fc1_weights = materialize_weights_for_forward([fc1_op.weight0]) else: fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] quantized_fc1_weights = [] @@ -1258,11 +1259,9 @@ def fuser_forward( None, ) else: - if fc2_gtp_size > 1: + if fc2_is_dist: # All-gather the full per-expert weights (returns a list of N full tensors). - fc2_weights = fc2_op.weight0.batched_all_gather_and_prefetch( - fwd=True, skip_weight_cast=False, cast_noop_flag=None - ) + fc2_weights = materialize_weights_for_forward([fc2_op.weight0]) else: fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] quantized_fc2_weights = [] @@ -1499,14 +1498,14 @@ def fuser_forward( # Save an internal layout for this joint fused op. The saved state is # intentionally not compatible with the basic GroupedLinear backward. - # EGTP: save the small sharded params; backward col-AGs the layout it needs. - if fc1_gtp_size > 1: + # Distributed weight: save the small sharded params; backward col-AGs the layout. + if fc1_is_dist: fc1_weight_tensors = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] else: fc1_weight_tensors = ( [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight ) - if fc2_gtp_size > 1: + if fc2_is_dist: fc2_weight_tensors = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] else: fc2_weight_tensors = ( @@ -1529,7 +1528,6 @@ def fuser_forward( fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad - fc1_ctx.gtp_remat_size = fc1_gtp_size fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] @@ -1537,7 +1535,6 @@ def fuser_forward( fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x - fc2_ctx.gtp_remat_size = fc2_gtp_size return fc2_out, [(), (), ()] @@ -1777,10 +1774,10 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) - # EGTP: forward gathered rowwise; col-AG the columnwise layout the dgrad needs, - # right before use (one weight live at a time). - if getattr(fc2_ctx, "gtp_remat_size", 1) > 1: - grouped_fc2_weight = fc2_op.weight0.batched_all_gather_and_prefetch_bwd() + # Distributed weight: forward gathered rowwise; col-AG the columnwise layout the dgrad + # needs, right before use (one weight live at a time). + if is_distributed_weight(fc2_op.weight0): + grouped_fc2_weight = materialize_weights_for_backward([fc2_op.weight0]) if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2060,9 +2057,9 @@ def fuser_backward( "use_dynamic_sched": True, } - # EGTP: col-AG the columnwise layout for the FC1 dgrad (only when dgrad runs). - if getattr(fc1_ctx, "gtp_remat_size", 1) > 1: - grouped_fc1_weight = fc1_op.weight0.batched_all_gather_and_prefetch_bwd() + # Distributed weight: col-AG the columnwise FC1 dgrad layout (only when dgrad runs). + if is_distributed_weight(fc1_op.weight0): + grouped_fc1_weight = materialize_weights_for_backward([fc1_op.weight0]) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM From de9dc8122e5ddcade1ceba5c5bcd033f56d1da34 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 13 Jul 2026 21:22:28 -0700 Subject: [PATCH 10/15] Code clean - Take a single leader weight in the DistributedWeight dispatchers - Gather the FC2 grouped weight late in the fused grouped MLP Signed-off-by: Shiqing Fan --- tests/pytorch/test_distributed_weight.py | 20 ++-- .../pytorch/distributed_weight.py | 41 ++++---- .../pytorch/module/grouped_linear.py | 10 +- .../pytorch/module/layernorm_linear.py | 10 +- transformer_engine/pytorch/module/linear.py | 10 +- .../pytorch/ops/fused/grouped_mlp.py | 93 +++++++++---------- 6 files changed, 89 insertions(+), 95 deletions(-) diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index b11f9d7f8c..6757154692 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/tests/pytorch/test_distributed_weight.py @@ -15,8 +15,8 @@ from transformer_engine.pytorch.distributed_weight import ( DistributedWeight, is_distributed_weight, - materialize_weights_for_forward, - materialize_weights_for_backward, + materialize_weight_for_forward, + materialize_weight_for_backward, finalize_weight_grads, ) @@ -72,7 +72,7 @@ def test_is_distributed_weight(): def test_forward_noop_on_plain_tensor(): """Plain weights pass through unchanged — the critical non-regression.""" w = torch.nn.Parameter(torch.randn(4, 4)) - out = materialize_weights_for_forward([w]) + out = materialize_weight_for_forward(w) assert out == [w] assert out[0] is w @@ -80,7 +80,7 @@ def test_forward_noop_on_plain_tensor(): def test_forward_dispatches_single_weight(): """Linear (N=1): dispatcher delegates and returns a length-1 list.""" w = FakeDistributedWeight(group_size=1) - out = materialize_weights_for_forward([w]) + out = materialize_weight_for_forward(w) assert isinstance(out, list) and len(out) == 1 assert torch.equal(out[0], torch.zeros(2, 2)) # The forward dispatcher owns the forward semantic; the leader is delegated to once. @@ -90,7 +90,7 @@ def test_forward_dispatches_single_weight(): def test_forward_dispatches_grouped_weights(): """GroupedLinear (N=k): one coalesced call, full list returned.""" w = FakeDistributedWeight(group_size=3) - out = materialize_weights_for_forward([w, w, w]) + out = materialize_weight_for_forward(w) assert isinstance(out, list) and len(out) == 3 # Leader was called exactly once (coalesced), not once per weight. assert len(w.calls) == 1 @@ -98,21 +98,21 @@ def test_forward_dispatches_grouped_weights(): def test_backward_dispatch_and_noop(): w = FakeDistributedWeight(group_size=2) - out = materialize_weights_for_backward([w, w]) + out = materialize_weight_for_backward(w) assert len(out) == 2 and torch.equal(out[0], torch.full((2, 2), 10.0)) - plain = [torch.zeros(2), torch.ones(2)] - assert materialize_weights_for_backward(plain) == plain + plain = torch.zeros(2) + assert materialize_weight_for_backward(plain) == [plain] def test_finalize_grads_dispatch_and_noop(): w = FakeDistributedWeight(group_size=1) wgrads = [torch.zeros(2, 2)] - out = finalize_weight_grads([w], wgrads) + out = finalize_weight_grads(w, wgrads) assert len(out) == 1 and torch.equal(out[0], torch.full((2, 2), 100.0)) # No-op path leaves the grads untouched. - plain_w = [torch.nn.Parameter(torch.zeros(2))] + plain_w = torch.nn.Parameter(torch.zeros(2)) g = [torch.ones(2)] assert finalize_weight_grads(plain_w, g) == g diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index ac44cd8d59..6c350847e1 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/transformer_engine/pytorch/distributed_weight.py @@ -16,8 +16,8 @@ __all__ = [ "DistributedWeight", "is_distributed_weight", - "materialize_weights_for_forward", - "materialize_weights_for_backward", + "materialize_weight_for_forward", + "materialize_weight_for_backward", "finalize_weight_grads", ] @@ -54,34 +54,31 @@ def is_distributed_weight(weight: Any) -> bool: return bool(getattr(weight, "is_distributed_weight", False)) -def materialize_weights_for_forward(weights: List[Any]) -> List[Any]: - """Materialize a weight group for the forward GEMM. Delegates once to the leader (it may - coalesce) if distributed, else returns the weights unchanged. Always returns a list. +def materialize_weight_for_forward(weight: Any) -> List[Any]: + """Materialize a (leader) weight for the forward GEMM. If distributed, delegate to it (it may + coalesce the whole group and return several tensors); else return it unchanged. Always a list. """ - lead = weights[0] - if is_distributed_weight(lead): - out = lead.materialize_group_for_forward() + if is_distributed_weight(weight): + out = weight.materialize_group_for_forward() return list(out) if isinstance(out, (list, tuple)) else [out] - return list(weights) + return [weight] -def materialize_weights_for_backward(weights: List[Any]) -> List[Any]: - """Backward counterpart of :func:`materialize_weights_for_forward`.""" - lead = weights[0] - if is_distributed_weight(lead): - out = lead.materialize_group_for_backward() +def materialize_weight_for_backward(weight: Any) -> List[Any]: + """Backward counterpart of :func:`materialize_weight_for_forward`.""" + if is_distributed_weight(weight): + out = weight.materialize_group_for_backward() return list(out) if isinstance(out, (list, tuple)) else [out] - return list(weights) + return [weight] -def finalize_weight_grads(weights: List[Any], wgrads: List[Any]) -> List[Any]: - """Post-process the weight grads of a homogeneous group. +def finalize_weight_grads(weight: Any, wgrads: List[Any]) -> List[Any]: + """Post-process the weight grad(s) of a (leader) weight's group. - Delegates to the group leader when distributed (e.g. reduce-scatter); - otherwise returns ``wgrads`` unchanged. + Delegates to the weight when distributed (e.g. reduce-scatter); otherwise returns ``wgrads`` + unchanged. """ - lead = weights[0] - if is_distributed_weight(lead): - out = lead.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + if is_distributed_weight(weight): + out = weight.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) return list(out) if isinstance(out, (list, tuple)) else [out] return list(wgrads) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 28d882e9ef..d52bf697e0 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -48,8 +48,8 @@ ) from ..distributed_weight import ( is_distributed_weight, - materialize_weights_for_forward, - materialize_weights_for_backward, + materialize_weight_for_forward, + materialize_weight_for_backward, finalize_weight_grads, ) from ..cpp_extensions import ( @@ -446,7 +446,7 @@ def forward( weight_params = weights if is_distributed_weight(weights[0]): - weights = materialize_weights_for_forward(weights) + weights = materialize_weight_for_forward(weights[0]) # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -1016,7 +1016,7 @@ def backward( accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation if is_distributed_weight(origin_weights[0]): - weights = materialize_weights_for_backward(origin_weights) + weights = materialize_weight_for_backward(origin_weights[0]) if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD @@ -1183,7 +1183,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): return wgrad if is_distributed_weight(origin_weights[0]): - wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + wgrad_list = finalize_weight_grads(origin_weights[0], wgrad_list) else: wgrad_list = [ handle_custom_ddp_from_mcore(weight, main_grad, wgrad) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 748e1d125e..a9ed210219 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -57,8 +57,8 @@ ) from ..distributed_weight import ( is_distributed_weight, - materialize_weights_for_forward, - materialize_weights_for_backward, + materialize_weight_for_forward, + materialize_weight_for_backward, finalize_weight_grads, ) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type @@ -308,7 +308,7 @@ def forward( weight_param = weight if is_distributed_weight(weight): - weight = materialize_weights_for_forward([weight])[0] + weight = materialize_weight_for_forward(weight)[0] out_features = weight.shape[0] new_weight_workspace = None @@ -630,7 +630,7 @@ def backward( ) = restore_from_func_ctx(ctx) if is_distributed_weight(saved_weight): - weight = materialize_weights_for_backward([saved_weight])[0] + 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. @@ -1059,7 +1059,7 @@ def wgrad_gemm( wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) if is_distributed_weight(saved_weight): - wgrad = finalize_weight_grads([saved_weight], [wgrad])[0] + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] # Update grad bias if needed if grad_bias is None: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ac863959d0..cec1162d36 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -56,8 +56,8 @@ ) from ..distributed_weight import ( is_distributed_weight, - materialize_weights_for_forward, - materialize_weights_for_backward, + materialize_weight_for_forward, + materialize_weight_for_backward, finalize_weight_grads, ) from ..cpp_extensions import ( @@ -411,7 +411,7 @@ def _linear_forward_impl( # `args.weight` keeps the sharded-param reference for backward re-gather / grad # finalize. No-op for a plain weight. if is_distributed_weight(args.weight): - weight = materialize_weights_for_forward([args.weight])[0] + weight = materialize_weight_for_forward(args.weight)[0] # Refresh out_features from the gathered weight (captured sharded above, pre-gather). out_features = weight.shape[0] @@ -956,7 +956,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. if is_distributed_weight(saved_weight): - weight_fp8 = materialize_weights_for_backward([saved_weight])[0] + weight_fp8 = materialize_weight_for_backward(saved_weight)[0] if bwd_args.requires_dgrad: @@ -1258,7 +1258,7 @@ def wgrad_gemm( # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad # (async; overlap with the next layer's bwd via the cascade). if is_distributed_weight(saved_weight): - wgrad = finalize_weight_grads([saved_weight], [wgrad])[0] + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] # Update grad bias if needed if grad_bias is None: diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 2ca0864117..7e2ddcb974 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -21,8 +21,8 @@ from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( is_distributed_weight, - materialize_weights_for_forward, - materialize_weights_for_backward, + materialize_weight_for_forward, + materialize_weight_for_backward, finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD @@ -636,7 +636,7 @@ def _compute_grad_params( # Distributed weight: reduce-scatter the full per-rank wgrads into each sharded # main_grad (also fires the Megatron grad-accum hook). if is_dist_weight: - finalize_weight_grads([weights[0]], w_list) + finalize_weight_grads(weights[0], w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. @@ -910,9 +910,9 @@ def fuser_forward( fc2_is_dist = is_distributed_weight(fc2_weight_param) assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." if fc1_is_dist: - assert not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight, ( - "distributed-weight fused grouped-MLP requires single_grouped_weight=False." - ) + assert ( + not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight + ), "distributed-weight fused grouped-MLP requires single_grouped_weight=False." assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None @@ -994,11 +994,7 @@ def fuser_forward( None, ) else: - if fc1_is_dist: - # All-gather the full per-expert weights (returns a list of N full tensors). - fc1_weights = materialize_weights_for_forward([fc1_op.weight0]) - else: - fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] + fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] quantized_fc1_weights = [] for idx, weight in enumerate(fc1_weights): quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) @@ -1008,6 +1004,40 @@ def fuser_forward( else: quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights + if fc1_is_dist: + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight[0]) + + # Prepare FC2 grouped weight tensor for fused kernels. + if fc2_op.single_grouped_weight: + if not isinstance(fc2_op.weight, GroupedTensor): + raise RuntimeError( + "FC2 expected GroupedTensor weight with single_grouped_weight=True." + ) + if fc2_op.weight.quantizer is not None: + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc2_op.weight.quantizer = fc2_weight_quantizer + grouped_fc2_weight = fc2_op.weight + else: + if fc2_op.weight.rowwise_data is None: + raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + grouped_fc2_weight = _group_quantize_for_grouped_mlp( + fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), + fc2_weight_quantizer, + num_groups, + None, + ) + else: + fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] + quantized_fc2_weights = [] + for idx, weight in enumerate(fc2_weights): + quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) + if not is_quantized_tensor(weight): + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + quantized_fc2_weights.append(quantizer(weight)) + else: + quantized_fc2_weights.append(weight) + grouped_fc2_weight = quantized_fc2_weights # Some wrapper-copy paths may drop grouped storage metadata; enforce defaults. if isinstance(grouped_fc1_weight, GroupedTensor) and not hasattr( @@ -1238,41 +1268,8 @@ def fuser_forward( else: fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) - # Prepare FC2 grouped weight tensor for fused kernels. - if fc2_op.single_grouped_weight: - if not isinstance(fc2_op.weight, GroupedTensor): - raise RuntimeError( - "FC2 expected GroupedTensor weight with single_grouped_weight=True." - ) - if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - fc2_op.weight.quantizer = fc2_weight_quantizer - grouped_fc2_weight = fc2_op.weight - else: - if fc2_op.weight.rowwise_data is None: - raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - grouped_fc2_weight = _group_quantize_for_grouped_mlp( - fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), - fc2_weight_quantizer, - num_groups, - None, - ) - else: - if fc2_is_dist: - # All-gather the full per-expert weights (returns a list of N full tensors). - fc2_weights = materialize_weights_for_forward([fc2_op.weight0]) - else: - fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] - quantized_fc2_weights = [] - for idx, weight in enumerate(fc2_weights): - quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) - if not is_quantized_tensor(weight): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc2_weights.append(quantizer(weight)) - else: - quantized_fc2_weights.append(weight) - grouped_fc2_weight = quantized_fc2_weights + if fc2_is_dist: + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight[0]) if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( grouped_fc2_weight, "_with_gemm_swizzled_scales" ): @@ -1777,7 +1774,7 @@ def fuser_backward( # Distributed weight: forward gathered rowwise; col-AG the columnwise layout the dgrad # needs, right before use (one weight live at a time). if is_distributed_weight(fc2_op.weight0): - grouped_fc2_weight = materialize_weights_for_backward([fc2_op.weight0]) + grouped_fc2_weight = materialize_weight_for_backward(fc2_op.weight0) if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2059,7 +2056,7 @@ def fuser_backward( # Distributed weight: col-AG the columnwise FC1 dgrad layout (only when dgrad runs). if is_distributed_weight(fc1_op.weight0): - grouped_fc1_weight = materialize_weights_for_backward([fc1_op.weight0]) + grouped_fc1_weight = materialize_weight_for_backward(fc1_op.weight0) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM From 0f7226ec5d0d3ca77927dade4f95d0597168cb5a Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Tue, 14 Jul 2026 01:04:35 -0700 Subject: [PATCH 11/15] Simplify the NVFP4 gather post-process; Materialize the EGTP FC1 weight before the NVFP4 dgrad dispatch Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/distributed.py | 36 +++++++------------ .../pytorch/distributed_weight.py | 4 --- .../pytorch/ops/fused/grouped_mlp.py | 31 +++++++--------- 3 files changed, 25 insertions(+), 46 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 6a2f7fee1d..61e0c32168 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1293,13 +1293,11 @@ def _post_process_nvfp4_gather( handle = None # Fix the interleaved transposed data from gathering along first dim. - # In-place .copy_() (not `=` rebind) to keep the storage address stable - # for CUDA graph capture — replays see the same pointer they captured. - out._columnwise_scale_inv.copy_(_swap_first_dims(columnwise_scale_inv_interleaved, world_size)) - out._columnwise_data.copy_(_swap_first_dims(columnwise_data_interleaved, world_size)) + out._columnwise_scale_inv = _swap_first_dims(columnwise_scale_inv_interleaved, world_size) + out._columnwise_data = _swap_first_dims(columnwise_data_interleaved, world_size) - # Optionally pad the scaling inverse if needed (same in-place pattern). - out._columnwise_scale_inv.copy_(pad_columnwise_scale_inv(out._columnwise_scale_inv)) + # Optionally pad the scaling inverse if needed. + out._columnwise_scale_inv = pad_columnwise_scale_inv(out._columnwise_scale_inv) @dataclass @@ -1313,25 +1311,18 @@ class _NVFP4AllGatherAsyncHandle: async_handle: torch.distributed.Work _synchronized: bool = False - def post_process_nvfp4_gather(self) -> None: - """Fix interleaved transposed data + pad scale_inv after the async AG completes. - - Idempotent: gated by ``_synchronized`` in :meth:`wait`. - """ - _post_process_nvfp4_gather( - self.output, - self.columnwise_data_interleaved, - self.columnwise_scale_inv_interleaved, - self.world_size, - ) - def wait(self) -> None: """Wait for the async operation to complete and post-process the tensor.""" if self._synchronized: return if self.async_handle is not None: self.async_handle.wait() - self.post_process_nvfp4_gather() + _post_process_nvfp4_gather( + self.output, + self.columnwise_data_interleaved, + self.columnwise_scale_inv_interleaved, + self.world_size, + ) self._synchronized = True @@ -1465,9 +1456,8 @@ def _all_gather_nvfp4( group=process_group, ) - # Transfer amax to output via in-place .copy_() so the storage - # address stays stable for CUDA graph capture. - out._amax_rowwise.copy_(inp._amax_rowwise) + # Transfer amax to output. + out._amax_rowwise = inp._amax_rowwise # Gather the transposed NVFP4 data along first dimension. Fix format later. if quantizer.columnwise_usage: @@ -1516,7 +1506,7 @@ def _all_gather_nvfp4( ) # Transfer amax to output. - out._amax_columnwise.copy_(inp._amax_columnwise) + out._amax_columnwise = inp._amax_columnwise handle = coalesced_handle if async_op else None diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index 6c350847e1..d9cc30bec8 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/transformer_engine/pytorch/distributed_weight.py @@ -34,19 +34,15 @@ class DistributedWeight(Protocol): def materialize_group_for_forward(self) -> Any: """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" - ... def materialize_group_for_backward(self) -> Any: """Re-materialize the full weight(s) for the backward GEMMs.""" - ... def finalize_group_grads(self, wgrads: Any) -> Any: """Post-process freshly computed weight grad(s) (e.g. reduce-scatter).""" - ... def grad_buffer(self) -> torch.Tensor: """The gradient accumulation buffer for this weight.""" - ... def is_distributed_weight(weight: Any) -> bool: diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 7e2ddcb974..7464faf828 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1493,21 +1493,17 @@ def fuser_forward( start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) - # Save an internal layout for this joint fused op. The saved state is - # intentionally not compatible with the basic GroupedLinear backward. - # Distributed weight: save the small sharded params; backward col-AGs the layout. + fc1_weight_tensors = ( + [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight + ) + fc2_weight_tensors = ( + [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight + ) + # Save the joint op's internal layout; distributed weights save the small shards. if fc1_is_dist: - fc1_weight_tensors = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] - else: - fc1_weight_tensors = ( - [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight - ) + fc1_weight_tensors = fc1_weights if fc2_is_dist: - fc2_weight_tensors = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] - else: - fc2_weight_tensors = ( - [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight - ) + fc2_weight_tensors = fc2_weights fc1_ctx.save_for_backward( split_sizes, base_split_offsets, @@ -1771,8 +1767,6 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) - # Distributed weight: forward gathered rowwise; col-AG the columnwise layout the dgrad - # needs, right before use (one weight live at a time). if is_distributed_weight(fc2_op.weight0): grouped_fc2_weight = materialize_weight_for_backward(fc2_op.weight0) @@ -2000,6 +1994,9 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] + if is_distributed_weight(fc1_op.weight0): + grouped_fc1_weight = materialize_weight_for_backward(fc1_op.weight0) + if use_nvfp4: grad_input = torch.empty(in_shape, dtype=dtype, device=device) if num_groups == 1: @@ -2054,10 +2051,6 @@ def fuser_backward( "use_dynamic_sched": True, } - # Distributed weight: col-AG the columnwise FC1 dgrad layout (only when dgrad runs). - if is_distributed_weight(fc1_op.weight0): - grouped_fc1_weight = materialize_weight_for_backward(fc1_op.weight0) - if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc1_weight_for_gemm = grouped_fc1_weight.copy() From c256929424376abaf30166112526899a4dc2b969 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Tue, 14 Jul 2026 19:06:41 -0700 Subject: [PATCH 12/15] Code clean - Rename gather coalescing flag grouped -> external_coalescing; - Clean up DistributedWeight wiring in TE modules - Restructure _all_gather_nvfp4 Signed-off-by: Shiqing Fan --- transformer_engine/pytorch/distributed.py | 56 +++++++++---------- .../pytorch/module/grouped_linear.py | 37 ++++++------ .../pytorch/module/layernorm_linear.py | 23 ++++---- transformer_engine/pytorch/module/linear.py | 16 +++--- .../pytorch/ops/fused/grouped_mlp.py | 2 + 5 files changed, 70 insertions(+), 64 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 61e0c32168..13e150de55 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1334,7 +1334,7 @@ def _all_gather_nvfp4( quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, output_tensor=None, - grouped=False, + external_coalescing=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1398,12 +1398,6 @@ def _all_gather_nvfp4( out = quantizer(out) return out, None - # Construct NVFP4 output tensor - if output_tensor is not None: - out = output_tensor - else: - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Cast input tensor to NVFP4 with required data if not isinstance(inp, NVFP4TensorStorage): inp = quantizer(inp) @@ -1416,19 +1410,25 @@ def _all_gather_nvfp4( ) inp = quantizer(inp.dequantize(dtype=dtype)) - if not grouped: - # Coalesce NCCL collectives for gathering data and scale inverses. + # Construct NVFP4 output tensor + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + + # Coalesce NCCL collectives for gathering data and scale inverses. + if not external_coalescing: gather_coalescing_manager = torch.distributed._coalescing_manager( group=process_group, device=device, async_ops=async_op, ) else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. gather_coalescing_manager = nullcontext() with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage - out_columnwise_data = None if quantizer.rowwise_usage: # Remove padding from NVFP4 scale-inverses @@ -1511,19 +1511,12 @@ def _all_gather_nvfp4( handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if quantizer.columnwise_usage: - if async_op or grouped: - # Defer post-processing: either the async op hasn't completed yet, or an - # external coalescing manager owns the NCCL ops and hasn't flushed them. - inner_handle = handle if async_op else None - handle = _NVFP4AllGatherAsyncHandle( - out, out_columnwise_data, out_scale_inv, world_size, inner_handle - ) - else: - _post_process_nvfp4_gather(out, out_columnwise_data, out_scale_inv, world_size, handle) - else: - if handle is not None: - handle.output = out + if (async_op or external_coalescing) and quantizer.columnwise_usage: + handle = _NVFP4AllGatherAsyncHandle( + out, out_columnwise_data, out_scale_inv, world_size, handle + ) + elif quantizer.columnwise_usage: + _post_process_nvfp4_gather(out, out_columnwise_data, out_scale_inv, world_size, handle) return out, handle @@ -1536,7 +1529,7 @@ def _all_gather_mxfp8( quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, output_tensor: torch.Tensor = None, - grouped: bool = False, + external_coalescing: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1607,7 +1600,7 @@ def _all_gather_mxfp8( else: out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - if not grouped: + if not external_coalescing: # Coalesce NCCL collectives for gathering data and scale inverses. gather_coalescing_manager = torch.distributed._coalescing_manager( group=process_group, @@ -1615,6 +1608,7 @@ def _all_gather_mxfp8( async_ops=async_op, ) else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. gather_coalescing_manager = nullcontext() with gather_coalescing_manager as coalesced_handle: @@ -1674,10 +1668,16 @@ def gather_along_first_dim( async_op: bool = False, quantizer: Optional[Quantizer] = None, output_tensor: torch.Tensor = None, - grouped: bool = False, + external_coalescing: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. + + ``external_coalescing``: composability flag for callers that batch several gathers into + one outer ``torch.distributed._coalescing_manager``. Coalescing managers cannot nest, so + when set this call skips opening its own manager and defers any post-gather fixup + (e.g. NVFP4 columnwise de-interleave) into the returned handle; ``handle.wait()`` completes + it once the outer manager has closed. Leave ``False`` for standalone gathers. """ # Return immediately if no communication is required @@ -1766,7 +1766,7 @@ def gather_along_first_dim( quantizer=quantizer, out_shape=out_shape, output_tensor=output_tensor, - grouped=grouped, + external_coalescing=external_coalescing, ) # NVFP4 case @@ -1782,7 +1782,7 @@ def gather_along_first_dim( quantizer=quantizer, out_shape=out_shape, output_tensor=output_tensor, - grouped=grouped, + external_coalescing=external_coalescing, ) # High-precision communication for quantized tensors diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d52bf697e0..68b691cc65 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -444,8 +444,9 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad - weight_params = weights - if is_distributed_weight(weights[0]): + origin_weights = weights + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: weights = materialize_weight_for_forward(weights[0]) # Configure quantizers @@ -649,17 +650,18 @@ def forward( # Original weights are only needed by high_precision dgrad. The weakrefs # used for fused wgrad accumulation serve a different purpose: restoring # Python parameter attributes without keeping the parameter alive here. - _is_dist_weight = is_distributed_weight(weight_params[0]) saved_weights = ( weights - if backward_override == "high_precision" - and inp.requires_grad - and not _is_dist_weight + if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + weights_fp8 = [None] * num_gemms + saved_weights = origin_weights tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, - *(weight_params if _is_dist_weight else weights_fp8), + *weights_fp8, *saved_weights, *biases, ) @@ -684,8 +686,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] - elif is_distributed_weight(weight_params[0]): - ctx.main_grad_funcs = [weight_params[i].grad_buffer for i in range(num_gemms)] + elif is_dist_weight: + ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -939,10 +941,9 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if is_distributed_weight(weights[0]): - # Distributed weight (e.g. GTP): origin_weights come from saved tensors; - # main_grads are grad_buffer scratch (do not assign to param.main_grad). - origin_weights = weights + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: @@ -1006,7 +1007,7 @@ def backward( ctx.m_splits, ) - if is_distributed_weight(origin_weights[0]): + if is_dist_weight: accumulate_wgrad_into_param_main_grad = False elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1015,7 +1016,7 @@ def backward( else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - if is_distributed_weight(origin_weights[0]): + if is_dist_weight: weights = materialize_weight_for_backward(origin_weights[0]) if ctx.requires_dgrad: @@ -1086,7 +1087,7 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] - if is_distributed_weight(origin_weights[0]): + if is_dist_weight: # Gathered weights are no longer needed after dgrad GEMM. del weights @@ -1139,7 +1140,7 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not is_distributed_weight(origin_weights[0]) + if not is_dist_weight and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), @@ -1182,7 +1183,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - if is_distributed_weight(origin_weights[0]): + if is_dist_weight: wgrad_list = finalize_weight_grads(origin_weights[0], wgrad_list) else: wgrad_list = [ diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a9ed210219..70b72fbe42 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -305,12 +305,11 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ - - weight_param = weight - if is_distributed_weight(weight): + 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 @@ -501,13 +500,15 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None - _is_dist_weight = is_distributed_weight(weight_param) + # 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 + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, - # Distributed weight (e.g. GTP): save the sharded reference only; - # backward re-gathers it. - None if _is_dist_weight else wt_save, - weight_param if _is_dist_weight else weight, + wt_save, + origin_weight, bias, ln_weight, ln_out_to_save, @@ -534,8 +535,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad - elif is_distributed_weight(weight_param): - ctx.main_grad_func = weight_param.grad_buffer + 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 diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index cec1162d36..02d11dc4f7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -267,6 +267,7 @@ def _linear_forward_impl( """ weight = args.weight + is_dist_weight = is_distributed_weight(args.weight) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -410,7 +411,7 @@ def _linear_forward_impl( # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; # `args.weight` keeps the sharded-param reference for backward re-gather / grad # finalize. No-op for a plain weight. - if is_distributed_weight(args.weight): + if is_dist_weight: weight = materialize_weight_for_forward(args.weight)[0] # Refresh out_features from the gathered weight (captured sharded above, pre-gather). out_features = weight.shape[0] @@ -602,7 +603,7 @@ def _linear_forward_impl( if is_fsdp2 and weightmat is not weight: wt_save = None # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. - if is_distributed_weight(args.weight): + if is_dist_weight: wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` @@ -755,6 +756,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer @@ -781,7 +783,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - if not is_distributed_weight(bwd_args.saved_weight): + if not is_dist_weight: origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -955,7 +957,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. - if is_distributed_weight(saved_weight): + if is_dist_weight: weight_fp8 = materialize_weight_for_backward(saved_weight)[0] if bwd_args.requires_dgrad: @@ -973,7 +975,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) elif ( - is_distributed_weight(saved_weight) + is_dist_weight and bwd_args.fp8 and bwd_args.weight_quantizer is not None and not isinstance(weight_fp8, QuantizedTensorStorage) @@ -1176,7 +1178,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if is_distributed_weight(bwd_args.saved_weight): + if is_dist_weight: # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. accumulate_wgrad_into_param_main_grad = False elif bwd_args.is_first_microbatch is not None: @@ -1257,7 +1259,7 @@ def wgrad_gemm( # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad # (async; overlap with the next layer's bwd via the cascade). - if is_distributed_weight(saved_weight): + if is_dist_weight: wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] # Update grad bias if needed diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 7464faf828..bfff1542b9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1493,6 +1493,8 @@ def fuser_forward( start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) + # Save an internal layout for this joint fused op. The saved state is + # intentionally not compatible with the basic GroupedLinear backward. fc1_weight_tensors = ( [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight ) From ae56984d4e498b1bdad7c73952af57aa2a55f2b6 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Fri, 17 Jul 2026 02:17:18 -0700 Subject: [PATCH 13/15] fix comments Signed-off-by: Shiqing Fan --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_distributed_weight.py | 96 +++++++++++++------ .../pytorch/distributed_weight.py | 56 ++++++++--- .../pytorch/module/grouped_linear.py | 5 +- .../pytorch/ops/fused/grouped_mlp.py | 4 +- 5 files changed, 116 insertions(+), 46 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 91d3be63d0..cff3357ddc 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -45,6 +45,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index 6757154692..19ed7ac73b 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/tests/pytorch/test_distributed_weight.py @@ -10,6 +10,7 @@ required: the whole contract lives in TE and is verified against the fake. """ +import pytest import torch from transformer_engine.pytorch.distributed_weight import ( @@ -57,6 +58,24 @@ def grad_buffer(self): return torch.full((2, 2), -1.0) +class FakeNonTensorWeight: + """DistributedWeight-shaped object that is NOT a torch.Tensor (contract violation).""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + return torch.zeros(2, 2) + + def materialize_group_for_backward(self, **kwargs): + return torch.zeros(2, 2) + + def finalize_group_grads(self, wgrads, **kwargs): + return wgrads + + def grad_buffer(self): + return torch.zeros(2, 2) + + def test_protocol_runtime_checkable(): """A conforming object passes isinstance; a plain tensor does not.""" assert isinstance(FakeDistributedWeight(), DistributedWeight) @@ -69,6 +88,12 @@ def test_is_distributed_weight(): assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) +def test_non_tensor_implementer_rejected(): + """Implementers must be torch.Tensor subclasses; a non-Tensor fails loudly.""" + with pytest.raises(TypeError, match="torch.Tensor subclass"): + is_distributed_weight(FakeNonTensorWeight()) + + def test_forward_noop_on_plain_tensor(): """Plain weights pass through unchanged — the critical non-regression.""" w = torch.nn.Parameter(torch.randn(4, 4)) @@ -77,48 +102,65 @@ def test_forward_noop_on_plain_tensor(): assert out[0] is w -def test_forward_dispatches_single_weight(): - """Linear (N=1): dispatcher delegates and returns a length-1 list.""" - w = FakeDistributedWeight(group_size=1) +@pytest.mark.parametrize("group_size", [1, 3]) +def test_forward_dispatches(group_size): + """Linear (N=1) and GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=group_size) out = materialize_weight_for_forward(w) - assert isinstance(out, list) and len(out) == 1 - assert torch.equal(out[0], torch.zeros(2, 2)) - # The forward dispatcher owns the forward semantic; the leader is delegated to once. + assert isinstance(out, list) and len(out) == group_size + # Leader is delegated to exactly once (coalesced), not once per weight. assert w.calls == ["fwd"] -def test_forward_dispatches_grouped_weights(): - """GroupedLinear (N=k): one coalesced call, full list returned.""" - w = FakeDistributedWeight(group_size=3) - out = materialize_weight_for_forward(w) +def test_forward_accepts_weight_list(): + """The dispatcher accepts the full per-expert list; the leader (index 0) coalesces it.""" + leader = FakeDistributedWeight(group_size=3) + followers = [torch.zeros(2, 2), torch.zeros(2, 2)] + out = materialize_weight_for_forward([leader, *followers]) assert isinstance(out, list) and len(out) == 3 - # Leader was called exactly once (coalesced), not once per weight. - assert len(w.calls) == 1 + # Leader delegated exactly once; the follower entries are not materialized separately. + assert leader.calls == ["fwd"] -def test_backward_dispatch_and_noop(): - w = FakeDistributedWeight(group_size=2) +def test_forward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + out = materialize_weight_for_forward(ws) + assert out == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_backward_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) out = materialize_weight_for_backward(w) - assert len(out) == 2 and torch.equal(out[0], torch.full((2, 2), 10.0)) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + +def test_backward_accepts_weight_list(): + """Backward dispatcher also accepts the full per-expert list; the leader coalesces it.""" + leader = FakeDistributedWeight(group_size=2) + out = materialize_weight_for_backward([leader, torch.zeros(2, 2)]) + assert isinstance(out, list) and len(out) == 2 + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_noop_on_plain_tensor(): plain = torch.zeros(2) assert materialize_weight_for_backward(plain) == [plain] -def test_finalize_grads_dispatch_and_noop(): - w = FakeDistributedWeight(group_size=1) - wgrads = [torch.zeros(2, 2)] +@pytest.mark.parametrize("group_size", [1, 2]) +def test_finalize_grads_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + wgrads = [torch.zeros(2, 2) for _ in range(group_size)] out = finalize_weight_grads(w, wgrads) - assert len(out) == 1 and torch.equal(out[0], torch.full((2, 2), 100.0)) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 100.0)) + - # No-op path leaves the grads untouched. +def test_finalize_grads_noop_on_plain_tensor(): + """No-op path leaves the grads untouched.""" plain_w = torch.nn.Parameter(torch.zeros(2)) g = [torch.ones(2)] assert finalize_weight_grads(plain_w, g) == g - - -if __name__ == "__main__": - for name, fn in sorted(globals().items()): - if name.startswith("test_") and callable(fn): - fn() - print(f"{name}: PASS") diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index d9cc30bec8..453946c693 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/transformer_engine/pytorch/distributed_weight.py @@ -27,6 +27,9 @@ class DistributedWeight(Protocol): """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + + Implementers MUST be ``torch.Tensor`` subclasses (needed by ``ctx.save_for_backward``, DDP + backward hooks, and ``torch.compile``); enforced at runtime by :func:`is_distributed_weight`. """ # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. @@ -46,26 +49,51 @@ def grad_buffer(self) -> torch.Tensor: def is_distributed_weight(weight: Any) -> bool: - """True if ``weight`` participates in custom weight parallelism (False on plain tensors).""" - return bool(getattr(weight, "is_distributed_weight", False)) - + """True if ``weight`` participates in custom weight parallelism (False on plain tensors). -def materialize_weight_for_forward(weight: Any) -> List[Any]: - """Materialize a (leader) weight for the forward GEMM. If distributed, delegate to it (it may - coalesce the whole group and return several tensors); else return it unchanged. Always a list. + Enforces the :class:`DistributedWeight` requirement that an implementer be a ``torch.Tensor`` + subclass, failing loudly here rather than silently breaking autograd downstream. """ - if is_distributed_weight(weight): - out = weight.materialize_group_for_forward() + flag = bool(getattr(weight, "is_distributed_weight", False)) + if flag and not isinstance(weight, torch.Tensor): + raise TypeError( + "DistributedWeight implementers must be torch.Tensor subclasses; got " + f"{type(weight).__name__}." + ) + return flag + + +def materialize_weight_for_forward(weights: Any) -> List[Any]: + """Prepare the weight(s) fed to the forward GEMM, always returned as a list. + + Args: + weights: the module's weight(s) -- a single weight (Linear) or the full per-expert list + (GroupedLinear). A bare weight is treated as a one-element list. + + Returns: + - Distributed group: the leader ``weights[0]`` all-gathers/coalesces the whole group and + returns all N materialized weights; the follower entries ``weights[1:]`` are ignored + here (the leader already holds references to its group). + - Otherwise: the input weights, unchanged. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_forward() return list(out) if isinstance(out, (list, tuple)) else [out] - return [weight] + return list(weights) -def materialize_weight_for_backward(weight: Any) -> List[Any]: - """Backward counterpart of :func:`materialize_weight_for_forward`.""" - if is_distributed_weight(weight): - out = weight.materialize_group_for_backward() +def materialize_weight_for_backward(weights: Any) -> List[Any]: + """Backward-GEMM mirror of :func:`materialize_weight_for_forward` (same contract).""" + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_backward() return list(out) if isinstance(out, (list, tuple)) else [out] - return [weight] + return list(weights) def finalize_weight_grads(weight: Any, wgrads: List[Any]) -> List[Any]: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 439693841e..3a75bad56e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -466,7 +466,7 @@ def forward( origin_weights = weights is_dist_weight = is_distributed_weight(weights[0]) if is_dist_weight: - weights = materialize_weight_for_forward(weights[0]) + weights = materialize_weight_for_forward(weights) # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -1036,7 +1036,7 @@ def backward( accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation if is_dist_weight: - weights = materialize_weight_for_backward(origin_weights[0]) + weights = materialize_weight_for_backward(origin_weights) if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD @@ -1433,7 +1433,6 @@ def __init__( if self.primary_weights_in_fp8: self.init_fp8_metadata(num_gemms=self.num_gemms) - self.weight_names = [f"weight{idx}" for idx in range(self.num_gemms)] is_meta = torch.device(device).type == "meta" self.reset_parameters(defer_init=is_meta) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 2c77b147e2..ce5491f2d1 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1010,7 +1010,7 @@ def fuser_forward( quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights if fc1_is_dist: - grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight[0]) + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1282,7 +1282,7 @@ def fuser_forward( fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) if fc2_is_dist: - grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight[0]) + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( grouped_fc2_weight, "_with_gemm_swizzled_scales" ): From aad67ccaa4220cdc2ad372912ef380af12794b5d Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sun, 19 Jul 2026 19:30:54 -0700 Subject: [PATCH 14/15] Support DistributedWeight in the fusible grouped-linear ops path - Add a self-contained dispatch test with a fake DistributedWeight implementer Signed-off-by: Shiqing Fan --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_distributed_weight.py | 6 + ...t_ops_grouped_linear_distributed_weight.py | 106 ++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 62 ++++++++-- .../pytorch/ops/fused/grouped_mlp.py | 10 +- 5 files changed, 169 insertions(+), 16 deletions(-) create mode 100644 tests/pytorch/test_ops_grouped_linear_distributed_weight.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index cff3357ddc..5d767ba4d1 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -67,6 +67,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" if [ "$RET" -ne 0 ]; then diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index 19ed7ac73b..69b5917d90 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/tests/pytorch/test_distributed_weight.py @@ -150,6 +150,12 @@ def test_backward_noop_on_plain_tensor(): assert materialize_weight_for_backward(plain) == [plain] +def test_backward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + assert materialize_weight_for_backward(ws) == ws + + @pytest.mark.parametrize("group_size", [1, 2]) def test_finalize_grads_dispatches(group_size): w = FakeDistributedWeight(group_size=group_size) diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py new file mode 100644 index 0000000000..77bdfecda6 --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -0,0 +1,106 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). + +TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each hook +so every plumbing point is independently checked -- a wiring bug fails a specific assertion: + + * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use + the materialized weights, so ``out == FWD_SCALE * plain_out``. + * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the + re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. + * ``finalize_group_grads`` scales the wgrads by ``FINALIZE_SCALE`` -> the returned param grad + must be finalize's output, so ``weight.grad == FINALIZE_SCALE * plain_wgrad``. + +Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake +is single-process and only proves the ops.GroupedLinear integration routes weights/grads through +the hooks (and does not bypass them). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te + +FWD_SCALE = 2.0 +BWD_SCALE = 3.0 +FINALIZE_SCALE = 5.0 + + +class _ScaledDistWeight(torch.nn.Parameter): + """Non-identity DistributedWeight leader: each hook applies its own scale (see module docstring), + coalescing its sibling group without any real gather.""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [w * FWD_SCALE for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + return [g * FINALIZE_SCALE for g in wl] + + def grad_buffer(self): + return self.data + + +def _make_scaled_dist_leader(op, num_gemms): + """Replace ``op.weight0`` with a scaled distributed leader referencing the whole group.""" + w0 = op.weight0 + leader = _ScaledDistWeight(w0.data) + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)] + op.weight0 = leader + return leader + + +@pytest.mark.parametrize("num_gemms", [2, 4]) +def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): + """Each DistributedWeight hook's output must be routed through the GEMM flow (not bypassed).""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.manual_seed(0) + in_f, out_f, total_tokens = 32, 64, num_gemms * 8 + dtype, device = torch.bfloat16, "cuda" + + op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference.load_state_dict(op.state_dict()) + + leader = _make_scaled_dist_leader(op, num_gemms) + + m_splits = [total_tokens // num_gemms] * num_gemms + m_splits[-1] += total_tokens - sum(m_splits) + split_sizes = torch.tensor(m_splits, dtype=torch.int64, device=device) + + x = torch.randn(total_tokens, in_f, dtype=dtype, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + out = op(x, split_sizes) + out.sum().backward() + ref_out = reference(ref_x, split_sizes) + ref_out.sum().backward() + + # The dispatch hooks actually fired. + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + tols = dict(rtol=2e-2, atol=2e-2) + # fwd used the materialized (FWD_SCALE) weights. + torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) + # dgrad used the re-materialized (BWD_SCALE) weights. + torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) + # returned param grads are finalize's (FINALIZE_SCALE) output (wgrad GEMM itself is scale-free). + for i in range(num_gemms): + got = getattr(op, f"weight{i}").grad + exp = getattr(reference, f"weight{i}").grad + assert got is not None, f"weight{i} received no grad" + torch.testing.assert_close(got, FINALIZE_SCALE * exp, **tols) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..b1f29df31e 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -51,6 +51,12 @@ view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext +from ...distributed_weight import ( + finalize_weight_grads, + is_distributed_weight, + materialize_weight_for_backward, + materialize_weight_for_forward, +) from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, @@ -898,6 +904,26 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: return [self.weight] return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + def _forward_weight_list(self) -> list[torch.Tensor]: + """Per-expert forward weights, materialized (all-gathered) when distributed.""" + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + if is_distributed_weight(weights[0]): + weights = materialize_weight_for_forward(weights) + return weights + + def _backward_weight_setup(self): + """Return ``(origin_weights, is_dist_weight, dgrad_weights)``; dgrad weights are the + re-materialized (all-gathered) weights when distributed, else ``None``.""" + origin_weights = self._get_weight_tensors() + is_dist_weight = is_distributed_weight(origin_weights[0]) + dgrad_weights = materialize_weight_for_backward(origin_weights) if is_dist_weight else None + return origin_weights, is_dist_weight, dgrad_weights + + def _is_distributed_weight(self) -> bool: + """Whether this op's weights are distributed (materialized per fwd/bwd, not saved).""" + leader = self.weight if self.single_grouped_weight else self.weight0 + return is_distributed_weight(leader) + def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, @@ -1143,7 +1169,7 @@ def _fuser_forward_split_quantize( if weights is None: weights = self.weight.split_into_quantized_tensors() else: - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: bs = self._get_bias_tensors(dtype) @@ -1198,7 +1224,8 @@ def _fuser_forward_split_quantize( out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) # Prepare weight tensors for backward pass - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so we never save the gathered weight + if not input_requires_grad or self._is_distributed_weight(): ws = [None] * num_groups elif with_quantized_compute: for w, weight_param in zip(ws, weights): @@ -1290,7 +1317,7 @@ def _fuser_forward_grouped_tensor( else: # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( - [getattr(self, f"weight{idx}") for idx in range(num_groups)], + self._forward_weight_list(), weight_quantizers, columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, @@ -1333,7 +1360,8 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so never save the gathered weight. + if not input_requires_grad or self._is_distributed_weight(): grouped_weights = None if self.single_grouped_weight else [None] * num_groups if not weight_requires_grad: @@ -1393,7 +1421,7 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device # Saved tensors from forward pass. Layout: @@ -1445,7 +1473,8 @@ def _fuser_backward_split_quantize( grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers. - accumulate_into_main_grad = self._accumulate_into_main_grad + # Distributed weights reduce their own grads (finalize); never accumulate into main_grad. + accumulate_into_main_grad = self._accumulate_into_main_grad and not is_dist_weight grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1495,7 +1524,7 @@ def _fuser_backward_split_quantize( device=device, ) general_grouped_gemm( - ws, + dist_dgrad_weights if is_dist_weight else ws, dys, [grad_input], [None] * num_groups, # quantization_params @@ -1540,11 +1569,15 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + final_weight_grads = finalize_weight_grads(weights[0], grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + elif ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1578,7 +1611,7 @@ def _fuser_backward_grouped_tensor( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype @@ -1679,7 +1712,7 @@ def _fuser_backward_grouped_tensor( tensor_offsets=base_split_offsets * self.in_features, ) general_grouped_gemm_for_grouped_tensor( - ws, + dist_dgrad_weights if is_dist_weight else ws, grouped_dy, grouped_grad_input, layout="NN", @@ -1725,7 +1758,8 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if self._accumulate_into_main_grad: + # Distributed weights finalize wgrads (below); never accumulate into main_grad. + if self._accumulate_into_main_grad and not is_dist_weight: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] @@ -1755,11 +1789,15 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + final_weight_grads = finalize_weight_grads(weights[0], final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + elif ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index ce5491f2d1..e74a5984e4 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1782,8 +1782,9 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) - if is_distributed_weight(fc2_op.weight0): - grouped_fc2_weight = materialize_weight_for_backward(fc2_op.weight0) + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2009,8 +2010,9 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] - if is_distributed_weight(fc1_op.weight0): - grouped_fc1_weight = materialize_weight_for_backward(fc1_op.weight0) + fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + if is_distributed_weight(fc1_leader): + grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) if use_nvfp4: grad_input = torch.empty(in_shape, dtype=dtype, device=device) From 4d29de3e7556922968031b6aae88034df6056f5c Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Mon, 20 Jul 2026 01:36:46 -0700 Subject: [PATCH 15/15] Unify distributed-weight wgrad finalize to return a graph-safe dummy - `finalize_weight_grads` now accepts a weight list or a bare leader, mirroring materialize_weight_for_backward; - Centralize the in-place / dummy / async-None finalize contract in DistributedWeight.finalize_group_grads and delegate the dispatcher docstring to it. Signed-off-by: Shiqing Fan --- ...t_ops_grouped_linear_distributed_weight.py | 71 +++++++++++++------ .../pytorch/distributed_weight.py | 23 ++++-- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_linear.py | 14 ++-- .../pytorch/ops/basic/grouped_linear.py | 10 +-- .../pytorch/ops/fused/grouped_mlp.py | 6 +- 6 files changed, 80 insertions(+), 46 deletions(-) diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py index 77bdfecda6..700be4138e 100644 --- a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -5,15 +5,17 @@ """DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the -*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each hook -so every plumbing point is independently checked -- a wiring bug fails a specific assertion: +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each +materialize hook so every plumbing point is independently checked -- a wiring bug fails a specific +assertion: * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use the materialized weights, so ``out == FWD_SCALE * plain_out``. * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. - * ``finalize_group_grads`` scales the wgrads by ``FINALIZE_SCALE`` -> the returned param grad - must be finalize's output, so ``weight.grad == FINALIZE_SCALE * plain_wgrad``. + * ``finalize_group_grads`` reduce-scatters the wgrads into ``main_grad`` in-place and returns a + dummy -> the real wgrad lands in ``main_grad`` and the ops path returns a throwaway ``.grad`` + (it discards finalize's return; see DistributedWeight.finalize_group_grads). Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake is single-process and only proves the ops.GroupedLinear integration routes weights/grads through @@ -25,14 +27,20 @@ import transformer_engine.pytorch as te +# Distinct powers of two so each scale commutes exactly through the (possibly TF32) GEMM rounding, +# making the linear scale relations bit-exact under a tight tolerance. FWD_SCALE = 2.0 -BWD_SCALE = 3.0 -FINALIZE_SCALE = 5.0 +BWD_SCALE = 4.0 -class _ScaledDistWeight(torch.nn.Parameter): - """Non-identity DistributedWeight leader: each hook applies its own scale (see module docstring), - coalescing its sibling group without any real gather.""" +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight leader for dispatch testing. + + Each materialize hook applies its own scale (see module docstring) so fwd/bwd routing is + observable; ``finalize_group_grads`` models the real main-grad contract -- reduce-scatter (here + an identity accumulate) the wgrads into each shard's ``main_grad`` in-place, flag + ``grad_added_to_main_grad``, and return a dummy that the ops path discards. + """ is_distributed_weight = True @@ -47,16 +55,19 @@ def materialize_group_for_backward(self, **kwargs): def finalize_group_grads(self, wgrads, **kwargs): self.calls["finalize"] += 1 wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] - return [g * FINALIZE_SCALE for g in wl] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # in-place accumulate into main_grad + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # dummy grads (real value is now in main_grad) def grad_buffer(self): return self.data -def _make_scaled_dist_leader(op, num_gemms): - """Replace ``op.weight0`` with a scaled distributed leader referencing the whole group.""" +def _make_fake_dist_leader(op, num_gemms): + """Replace ``op.weight0`` with a fake distributed leader referencing the whole group.""" w0 = op.weight0 - leader = _ScaledDistWeight(w0.data) + leader = _FakeDistWeight(w0.data) leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)] op.weight0 = leader @@ -65,18 +76,28 @@ def _make_scaled_dist_leader(op, num_gemms): @pytest.mark.parametrize("num_gemms", [2, 4]) def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): - """Each DistributedWeight hook's output must be routed through the GEMM flow (not bypassed).""" + """Every DistributedWeight hook must be routed through the GEMM flow (and not bypassed). + + fwd/bwd use the scaled materialized weights; finalize reduce-scatters the wgrad into + ``main_grad`` in-place, and the ops path returns a throwaway dummy ``.grad``. + """ if not torch.cuda.is_available(): pytest.skip("requires CUDA") torch.manual_seed(0) + # fp32 with power-of-two scales: each scale commutes exactly through the GEMM rounding (even + # TF32), so the linear scale relations below are bit-exact under a tight tolerance. in_f, out_f, total_tokens = 32, 64, num_gemms * 8 - dtype, device = torch.bfloat16, "cuda" + dtype, device = torch.float32, "cuda" op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) reference.load_state_dict(op.state_dict()) - leader = _make_scaled_dist_leader(op, num_gemms) + leader = _make_fake_dist_leader(op, num_gemms) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + w.main_grad = torch.zeros((out_f, in_f), dtype=torch.float32, device=device) + w.grad_added_to_main_grad = False # DDP initializes this on every param m_splits = [total_tokens // num_gemms] * num_gemms m_splits[-1] += total_tokens - sum(m_splits) @@ -90,17 +111,21 @@ def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): ref_out = reference(ref_x, split_sizes) ref_out.sum().backward() - # The dispatch hooks actually fired. + # All three dispatch hooks actually fired. assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 - tols = dict(rtol=2e-2, atol=2e-2) + tols = dict(rtol=1e-5, atol=1e-5) # fwd used the materialized (FWD_SCALE) weights. torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) # dgrad used the re-materialized (BWD_SCALE) weights. torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) - # returned param grads are finalize's (FINALIZE_SCALE) output (wgrad GEMM itself is scale-free). for i in range(num_gemms): - got = getattr(op, f"weight{i}").grad - exp = getattr(reference, f"weight{i}").grad - assert got is not None, f"weight{i} received no grad" - torch.testing.assert_close(got, FINALIZE_SCALE * exp, **tols) + w = getattr(op, f"weight{i}") + ref_w = getattr(reference, f"weight{i}") + # The distributed op's grad of record is main_grad (finalize reduce-scattered the identity + # wgrad there); compare it to the plain reference module's ordinary autograd .grad. + torch.testing.assert_close(w.main_grad.to(dtype), ref_w.grad, **tols) + # main_grad is flagged, and op's own .grad is a discarded dummy (not the real grad, which + # lives in main_grad). Nobody writes the real grad into op.weight.grad by design. + assert w.grad_added_to_main_grad is True + assert w.grad is not None, f"weight{i} should receive a dummy .grad" diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index 453946c693..0d491430e5 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/transformer_engine/pytorch/distributed_weight.py @@ -42,7 +42,12 @@ def materialize_group_for_backward(self) -> Any: """Re-materialize the full weight(s) for the backward GEMMs.""" def finalize_group_grads(self, wgrads: Any) -> Any: - """Post-process freshly computed weight grad(s) (e.g. reduce-scatter).""" + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter). + + May consume ``wgrads`` in-place -- reduce-scatter into ``main_grad`` and set + ``grad_added_to_main_grad`` -- returning a dummy grad (or ``None`` for an async collective) + that callers use as the parameter grad(s) or discard. + """ def grad_buffer(self) -> torch.Tensor: """The gradient accumulation buffer for this weight.""" @@ -96,13 +101,17 @@ def materialize_weight_for_backward(weights: Any) -> List[Any]: return list(weights) -def finalize_weight_grads(weight: Any, wgrads: List[Any]) -> List[Any]: - """Post-process the weight grad(s) of a (leader) weight's group. +def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]: + """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`. - Delegates to the weight when distributed (e.g. reduce-scatter); otherwise returns ``wgrads`` - unchanged. + Delegates to the leader's :meth:`DistributedWeight.finalize_group_grads` (which defines the + in-place / dummy / async-``None`` return contract); returns ``wgrads`` unchanged when not + distributed. """ - if is_distributed_weight(weight): - out = weight.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) return list(out) if isinstance(out, (list, tuple)) else [out] return list(wgrads) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 3a75bad56e..3f264900c1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1203,7 +1203,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): return wgrad if is_dist_weight: - wgrad_list = finalize_weight_grads(origin_weights[0], wgrad_list) + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) else: wgrad_list = [ handle_custom_ddp_from_mcore(weight, main_grad, wgrad) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 95fde3f152..cffc5ba1e6 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -641,7 +641,8 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) - if is_distributed_weight(saved_weight): + 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.) @@ -660,7 +661,7 @@ def backward( ), "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_distributed_weight(saved_weight): + if main_grad is not None and not is_dist_weight: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -1004,7 +1005,7 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if is_distributed_weight(saved_weight): + 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: @@ -1079,7 +1080,7 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - if is_distributed_weight(saved_weight): + if is_dist_weight: wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] # Update grad bias if needed @@ -1163,10 +1164,7 @@ def wgrad_gemm( if ctx.requires_wgrad: # Handle custom DDP from mcore. - if is_distributed_weight(saved_weight): - # Distributed weight (e.g. GTP): skip — grad finalize already produced the shard. - pass - elif ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): + 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( diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b1f29df31e..d01bfbc35c 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1570,14 +1570,15 @@ def _fuser_backward_split_quantize( clear_tensor_data(*xs) # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. if ctx.weight_requires_grad and is_dist_weight: assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." - final_weight_grads = finalize_weight_grads(weights[0], grad_weights) + finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - elif ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1790,14 +1791,15 @@ def _fuser_backward_grouped_tensor( wgrad_gemm(grouped_x, grouped_dy, wgrad_output) # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. if ctx.weight_requires_grad and is_dist_weight: assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." - final_weight_grads = finalize_weight_grads(weights[0], final_weight_grads) + finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - elif ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e74a5984e4..c91f3d46bc 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -638,10 +638,10 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - # Distributed weight: reduce-scatter the full per-rank wgrads into each sharded - # main_grad (also fires the Megatron grad-accum hook). + # Distributed weight: reduce-scatter the wgrads into main_grad. + # Return discarded (see finalize_weight_grads); dummy wgrads returned below. if is_dist_weight: - finalize_weight_grads(weights[0], w_list) + finalize_weight_grads(weights, w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add.