From 9149068b5cb68b2d6d697d2bff343929529c7443 Mon Sep 17 00:00:00 2001 From: mystri Date: Thu, 15 Oct 2026 10:15:43 +0800 Subject: [PATCH 1/4] [MoE] Absorb router scores before W2 Apply router scores to the expert activations before the down (W2) projection instead of multiplying the expert output in combine, so the combine reduces to a plain sum: - Token dispatchers expose absorb_router_scores. When set, dispatch returns routed_scores_R in its metadata and combine skips its score multiply. Covered backends: standard (local + all-to-all), TorchAO (EP=1 padded and EP>1), and DeepEP/HybridEP (opt-in). - GroupedLinear.forward takes routed_scores_R and scales the projection input in float32, restoring the input dtype before the grouped GEMM. RoutedExperts passes it to W2 only, and only when a score is actually absorbed, so W2 modules that do not know about scores are unaffected. - GptOssGroupedLinear scales its post-W2 bias rows by the same scores, so the result stays score * (h @ W2 + b2). - DeepEP/HybridEP dispatchers take the per-row scores the backend already tracks and clear them from the dispatch state, which is what makes the backend combine skip its fused score multiply. - make_routed_experts_config leaves the standard backend on auto and the DeepEP family off. Auto resolves when RoutedExperts is built: the score is absorbed unless output_postprocess owns the expert outputs, because absorbing scales W2's input while the postprocess runs after W2, and a scale-invariant postprocess (e.g. post-expert RMSNorm) would drop the score. Asking for both explicitly raises a ValueError. - swap_token_dispatcher propagates the flag through MoE quantization. Replay: with router scores absorbed, combine no longer scales expert outputs, so the DeepEP compact combine would read the W2 output from a bare _unpermute_tokens on the replay path. Fold the score multiply, unpermute and combine kernel into one never-replayed region, which keeps the W2 and combine outputs off the saved-tensor set. On the standard all-to-all path the absorbed scores are an output of the saved dispatch region, so mark them before W2's bare scaling; that marker is a no-op on the DeepEP/HybridEP paths and without an active checkpoint. Squashed from the router-score-pre-w2-review series and rebased onto current upstream main. Absorbed and post-combine paths agree to bf16 rounding; loss convergence on a representative dataset is still pending. --- torchtitan/distributed/deepep/deepep.py | 57 +++---- torchtitan/models/common/config_utils.py | 8 + torchtitan/models/common/linear/linear.py | 21 ++- torchtitan/models/common/moe/moe.py | 29 +++- .../models/common/moe/token_dispatcher.py | 144 +++++++++++++++--- torchtitan/models/gpt_oss/moe.py | 19 ++- torchtitan/quantization/utils.py | 2 + 7 files changed, 228 insertions(+), 52 deletions(-) diff --git a/torchtitan/distributed/deepep/deepep.py b/torchtitan/distributed/deepep/deepep.py index f337ba5328f..f9a1ef56ca9 100644 --- a/torchtitan/distributed/deepep/deepep.py +++ b/torchtitan/distributed/deepep/deepep.py @@ -614,7 +614,6 @@ def combine_tokens( state: DispatchState, *, remat_region_name: str, - recompute: bool, ) -> torch.Tensor: """Combine expert outputs back to tokens via DeepEP v2. @@ -631,8 +630,6 @@ def combine_tokens( hidden_states: Raw (unweighted) expert outputs [num_recv, hidden]. state: Dispatch state from ``dispatch_tokens``. remat_region_name: Name for the combine communication region. - recompute: Whether activation remat replays the combine communication during - backward. Must match the ``recompute`` passed to ``dispatch_tokens``. Returns: Combined tokens [num_tokens, hidden_dim]. @@ -641,29 +638,33 @@ def combine_tokens( # backward) or leaves it for combine-backward. Evaluated here, before the op. will_backward = torch.is_grad_enabled() - if not state.cuda_graph_compatible: - assert state.permuted_indices is not None - if state.permuted_scores is not None: - hidden_states = hidden_states * state.permuted_scores.to( - hidden_states.dtype - ).reshape(-1, 1) - hidden_states = _unpermute_tokens( - hidden_states, state.permuted_indices, state.num_recv_tokens - ) - elif state.recv_scores is not None: - # One routing score per received row (each row is one token->expert assignment). - # Collapse the trailing dim with sum so this is correct whether recv_scores is - # [num_recv], [num_recv, 1], or [num_recv, topk] with a single valid entry/row. - per_row_score = state.recv_scores.reshape(hidden_states.shape[0], -1).sum( - dim=-1, keepdim=True - ) - hidden_states = hidden_states * per_row_score.to(hidden_states.dtype) - - combined = remat.region( - torch.ops.deepep.combine, + def _scored_combine(hidden_states: torch.Tensor) -> torch.Tensor: + """Apply scores/unpermute and launch the reduction op. + + Runs as one never-replayed region, so the expert outputs are read only + by this saved body; nothing on the replay path touches them. + """ + if not state.cuda_graph_compatible: + assert state.permuted_indices is not None + if state.permuted_scores is not None: + hidden_states = hidden_states * state.permuted_scores.to( + hidden_states.dtype + ).reshape(-1, 1) + hidden_states = _unpermute_tokens( + hidden_states, state.permuted_indices, state.num_recv_tokens + ) + elif state.recv_scores is not None: + # One routing score per received row (each row is one token->expert assignment). + # Collapse the trailing dim with sum so this is correct whether recv_scores is + # [num_recv], [num_recv, 1], or [num_recv, topk] with a single valid entry/row. + per_row_score = state.recv_scores.reshape(hidden_states.shape[0], -1).sum( + dim=-1, keepdim=True + ) + hidden_states = hidden_states * per_row_score.to(hidden_states.dtype) + return torch.ops.deepep.combine(hidden_states, state.handle_id, will_backward) + + return remat.region( + _scored_combine, remat_region_name, - recompute=recompute, - )(hidden_states, state.handle_id, will_backward) - # The caller consumes this output outside another remat region. - remat.recompute_needs_tensor(combined) - return combined + recompute=False, # never replayed; the W2 output is read only by this body + )(hidden_states) diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py index 4c6c4fa595f..43d5f210d40 100644 --- a/torchtitan/models/common/config_utils.py +++ b/torchtitan/models/common/config_utils.py @@ -469,11 +469,18 @@ def make_routed_experts_config( top_k: int, param_init: dict[str, Callable], enable_ep: bool, + absorb_router_scores: bool | None = None, ) -> RoutedExperts.Config: """Build sharded routed experts with gate/up and down projections. Expert weights use sparse ``{EP}`` placements when ``enable_ep`` and dense replicated placements otherwise. + + ``absorb_router_scores`` is forwarded to the standard dispatcher. ``None`` + (auto) absorbs the router scores before W2 unless RoutedExperts carries an + output_postprocess, which must see unscaled expert outputs. DeepEP-family + dispatchers are installed later by TokenDispatcherTransform; their fused + combines apply the scores themselves, so they keep their own opt-in default. """ missing = {"w1_EFD", "w2_EDF", "w3_EFD"} - param_init.keys() if param_init and missing: @@ -500,6 +507,7 @@ def make_routed_experts_config( token_dispatcher=AllToAllTokenDispatcher.Config( num_experts=num_experts, top_k=top_k, + absorb_router_scores=absorb_router_scores, ), sharding_config=routed_experts_sharding_config(), ) diff --git a/torchtitan/models/common/linear/linear.py b/torchtitan/models/common/linear/linear.py index 562002960c3..abb3ef40034 100644 --- a/torchtitan/models/common/linear/linear.py +++ b/torchtitan/models/common/linear/linear.py @@ -405,17 +405,36 @@ def __init__(self, config: Config): torch.empty(config.group_size, *output_shape, config.in_features) ) - def forward(self, input_RI: torch.Tensor, offsets_E: torch.Tensor) -> torch.Tensor: + def forward( + self, + input_RI: torch.Tensor, + offsets_E: torch.Tensor, + *, + routed_scores_R: torch.Tensor | None = None, + ) -> torch.Tensor: """Apply each grouped linear to rows selected by cumulative offsets. Args: input_RI: Input rows grouped by the selected linear. offsets_E: Exclusive cumulative row end for each group. + routed_scores_R: Optional per-row router scores to absorb into the + input, one per input row. Router scores are produced in + float32; rows are scaled in float32 and cast back to the input + dtype so the grouped GEMM sees its usual operand type. Returns: Output rows with an optional ``num_linears`` axis before the output feature axis. """ + if routed_scores_R is not None: + # On the standard all-to-all path the scores are an output of the + # dispatcher's saved dispatch region; this scaling is a bare + # consumer, so replay needs them persisted. No-op when the scores + # are bare-produced (DeepEP/HybridEP) or without a checkpoint. + remat.recompute_needs_tensor(routed_scores_R) + input_RI = (input_RI.float() * routed_scores_R.reshape(-1, 1)).to( + input_RI.dtype + ) output_shape = self.weight.shape[1:-1] weight_EOI = self.weight.flatten(1, -2) output_RO = remat.region( diff --git a/torchtitan/models/common/moe/moe.py b/torchtitan/models/common/moe/moe.py index 9bf5614f54a..e6e2c564e97 100644 --- a/torchtitan/models/common/moe/moe.py +++ b/torchtitan/models/common/moe/moe.py @@ -85,6 +85,22 @@ def __post_init__(self) -> None: raise ValueError("w13 output and w2 input dimensions must match") if self.w2.num_linears != 1: raise ValueError("w2 must contain one down projection") + if ( + self.output_postprocess is not None + and self.token_dispatcher.absorb_router_scores + ): + # Absorbed scores scale W2's input, so the postprocess would see + # score-scaled outputs; the post-combine path instead applies the + # score after the postprocess. Only a linear postprocess keeps + # those equivalent, and a scale-invariant one (e.g. RMSNorm) + # would drop the score entirely. ``absorb_router_scores=None`` + # (auto) is resolved against this postprocess at build time. + raise ValueError( + "output_postprocess and absorb_router_scores cannot be " + "combined: absorbed router scores are applied before W2, " + "while output_postprocess runs after it. Set " + "absorb_router_scores=False or drop output_postprocess." + ) def __init__(self, config: Config): super().__init__() @@ -97,6 +113,11 @@ def __init__(self, config: Config): else None ) self.token_dispatcher = config.token_dispatcher.build() + if config.token_dispatcher.absorb_router_scores is None: + # Auto: resolve here rather than in the config builder, since a + # caller may attach output_postprocess to an already-built config. + # A post-W2 postprocess must see unscaled expert outputs. + self.token_dispatcher.absorb_router_scores = self.output_postprocess is None def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: del buffer_device @@ -139,6 +160,7 @@ def forward( dtype=torch.int32, ) + routed_scores_R = getattr(metadata, "routed_scores_R", None) with maybe_set_sparse_mesh(): # w13 and w2 declare their own remat regions (.grouped_mm). # The bf16 cast reads the dispatched tokens with bare ops. @@ -146,7 +168,12 @@ def forward( gate_up_R2F = self.w13(routed_input_RD.bfloat16(), offsets_E) remat.recompute_needs_tensor(gate_up_R2F) hidden_RF = self.activation_fn(gate_up_R2F, offsets=offsets_E) - routed_output_RD = self.w2(hidden_RF, offsets_E) + # Widen W2's contract only when a score is actually absorbed, so + # W2 modules that do not know about scores keep working unchanged. + w2_kwargs = ( + {} if routed_scores_R is None else {"routed_scores_R": routed_scores_R} + ) + routed_output_RD = self.w2(hidden_RF, offsets_E, **w2_kwargs) # A real dtype cast and the output postprocess read the w2 output with # bare ops, so pin it only then. In the common bf16 case without a # postprocess, type_as is a no-op and the w2 output goes straight to diff --git a/torchtitan/models/common/moe/token_dispatcher.py b/torchtitan/models/common/moe/token_dispatcher.py index 1795caba33e..b21dc1820dd 100644 --- a/torchtitan/models/common/moe/token_dispatcher.py +++ b/torchtitan/models/common/moe/token_dispatcher.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from abc import ABC, abstractmethod from dataclasses import dataclass @@ -17,13 +18,16 @@ from torchtitan.models.base.module import Module from torchtitan.ops.scatter_add import deterministic_scatter_add +logger = logging.getLogger(__name__) + @dataclass(frozen=True, kw_only=True) class LocalDispatchMetadata: """Metadata returned by LocalTokenDispatcher.dispatch() for use in combine().""" token_indices_experts_sorted_N: torch.Tensor # noqa: N815 - topk_scores_experts_sorted_N: torch.Tensor # noqa: N815 + topk_scores_experts_sorted_N: torch.Tensor | None # noqa: N815 + routed_scores_R: torch.Tensor | None = None # noqa: N815 @dataclass(frozen=True, kw_only=True) @@ -47,11 +51,16 @@ class LocalTokenDispatcher(Module): class Config(Module.Config): num_experts: int top_k: int + # None = auto: RoutedExperts decides from its output_postprocess, since + # absorbing the score scales W2's input and a post-W2 postprocess must + # not see scaled expert outputs. + absorb_router_scores: bool | None = False def __init__(self, config: Config): super().__init__() self.num_experts = config.num_experts self.top_k = config.top_k + self.absorb_router_scores = config.absorb_router_scores def init_buffer(self) -> None: """Initialize backend communication buffers, if any.""" @@ -129,7 +138,12 @@ def dispatch( ) = self._local_reorder(x_TD, topk_scores_TK, topk_expert_ids_TK) metadata = LocalDispatchMetadata( token_indices_experts_sorted_N=token_indices_experts_sorted_N, - topk_scores_experts_sorted_N=topk_scores_experts_sorted_N, + topk_scores_experts_sorted_N=( + None if self.absorb_router_scores else topk_scores_experts_sorted_N + ), + routed_scores_R=( + topk_scores_experts_sorted_N if self.absorb_router_scores else None + ), ) return routed_input_RD, num_local_tokens_per_expert_E, metadata @@ -160,17 +174,22 @@ def combine( def _score_and_scatter_add( self, routed_output_ND: torch.Tensor, - topk_scores_N: torch.Tensor, + topk_scores_N: torch.Tensor | None, token_indices_N: torch.Tensor, x_TD: torch.Tensor, ) -> torch.Tensor: - """Weight expert-sorted outputs by their router scores and sum them per token.""" - # Type promotion computes the product in float32 without materializing a - # float32 copy of the routed output, which the multiply would save for - # backward (twice the bytes of the routed output itself). - routed_output_ND = ( - routed_output_ND * topk_scores_N.reshape(-1, 1).to(torch.float32) - ).to(routed_output_ND.dtype) + """Weight expert-sorted outputs by their router scores and sum them per token. + + ``topk_scores_N`` is None when the dispatcher absorbed the router scores + before W2, which leaves the combine a pure reduction. + """ + if topk_scores_N is not None: + # Type promotion computes the product in float32 without materializing a + # float32 copy of the routed output, which the multiply would save for + # backward (twice the bytes of the routed output itself). + routed_output_ND = ( + routed_output_ND * topk_scores_N.reshape(-1, 1).to(torch.float32) + ).to(routed_output_ND.dtype) return deterministic_scatter_add( torch.zeros_like(x_TD), token_indices_N.reshape(-1, 1).expand(-1, x_TD.shape[-1]), @@ -420,6 +439,7 @@ def dispatch( permuted_indices, input_splits, output_splits, + routed_scores_R, ) = remat.region( self._dispatch, self.remat_region_name("dispatch"), @@ -442,13 +462,16 @@ def dispatch( output_splits_list = output_splits.tolist() metadata = AllToAllDispatchMetadata( token_indices_experts_sorted_N=token_indices_experts_sorted_N, - topk_scores_experts_sorted_N=topk_scores_experts_sorted_N, + topk_scores_experts_sorted_N=( + None if self.absorb_router_scores else topk_scores_experts_sorted_N + ), # Pre-permute shape: the all-to-all output, which padding # dispatchers (TorchAO) may grow in _permute. input_shape=torch.Size((sum(output_splits_list), x_TD.shape[-1])), permuted_indices=permuted_indices, input_splits=input_splits_list, output_splits=output_splits_list, + routed_scores_R=routed_scores_R, ) return routed_input_RD, num_global_tokens_per_local_expert_e, metadata @@ -464,7 +487,8 @@ def _dispatch( Runs as the single ``.dispatch`` remat region, so a saved dispatch replays neither the collectives nor the device-to-host split sync, and a recomputed one replays all of them. Regions return only tensors, so - the splits come back as CPU tensors. + the splits come back as CPU tensors. ``routed_scores_R`` is the one + exception: it is None unless the dispatcher absorbs router scores. """ assert self.ep_mesh is not None ep_size = self.ep_mesh.size() @@ -474,6 +498,9 @@ def _dispatch( token_indices_experts_sorted_N, topk_scores_experts_sorted_N, ) = self._local_reorder(x_TD, topk_scores_TK, topk_expert_ids_TK) + routed_scores_N = ( + topk_scores_experts_sorted_N if self.absorb_router_scores else None + ) with maybe_set_sparse_mesh(): pg = "ep" @@ -484,6 +511,10 @@ def _dispatch( routed_input_ND = spmd.reinterpret_mesh( routed_input_ND, spmd.current_mesh() ) + if routed_scores_N is not None: + routed_scores_N = spmd.reinterpret_mesh( + routed_scores_N, spmd.current_mesh() + ) # generate the input splits and output splits for all-to-all with torch.no_grad(): @@ -506,6 +537,16 @@ def _dispatch( output_splits.tolist(), input_splits.tolist(), ) + routed_scores_rank_major_R = None + if routed_scores_N is not None: + # The scores ride the same all-to-all pattern as the tokens so + # their expert-major rows stay aligned with routed_input_RD. + routed_scores_rank_major_R = self._dispatch_token_exchange( + routed_scores_N.reshape(-1, 1), + pg, + output_splits.tolist(), + input_splits.tolist(), + ).reshape(-1) # Reorder from rank-major to expert-major via _permute. # # num_global_tokens_per_local_expert_E layout after all-to-all @@ -521,6 +562,13 @@ def _dispatch( permuted_indices, num_global_tokens_per_local_expert_e, ) = self._permute(routed_input_RD, num_global_tokens_per_local_expert_E) + routed_scores_R = ( + None + if routed_scores_rank_major_R is None + else self._gather_routed_scores( + routed_scores_rank_major_R, permuted_indices + ) + ) return ( routed_input_RD, @@ -530,8 +578,16 @@ def _dispatch( permuted_indices, input_splits, output_splits, + routed_scores_R, ) + def _gather_routed_scores( + self, + routed_scores_rank_major_R: torch.Tensor, + permuted_indices: torch.Tensor, + ) -> torch.Tensor: + return routed_scores_rank_major_R[permuted_indices] + def _permute( self, routed_input_RD, @@ -723,14 +779,23 @@ def dispatch( num_tokens_per_local_expert_padded_e, ) = self._permute(routed_input_ND, num_local_tokens_per_expert_E) + routed_scores_R = None + if self.absorb_router_scores: + routed_scores_R = self._gather_routed_scores( + topk_scores_experts_sorted_N, permuted_indices + ) + metadata = AllToAllDispatchMetadata( token_indices_experts_sorted_N=token_indices_experts_sorted_N, - topk_scores_experts_sorted_N=topk_scores_experts_sorted_N, + topk_scores_experts_sorted_N=( + None if self.absorb_router_scores else topk_scores_experts_sorted_N + ), input_shape=input_shape, permuted_indices=permuted_indices, # Unused in the EP=1 combine path (no all-to-all to reverse). input_splits=[], output_splits=[], + routed_scores_R=routed_scores_R, ) return routed_input_RD, num_tokens_per_local_expert_padded_e, metadata @@ -764,6 +829,21 @@ def combine( x_TD, ) + def _gather_routed_scores( + self, + routed_scores_rank_major_R: torch.Tensor, + permuted_indices: torch.Tensor, + ) -> torch.Tensor: + # Pad positions index the appended zero sentinel row via -1, which the + # scores tensor lacks; appending one zero score makes pad rows read 0 + # instead of wrapping to the last real score. + return torch.cat( + [ + routed_scores_rank_major_R, + routed_scores_rank_major_R.new_zeros(1), + ] + )[permuted_indices] + def _permute( self, routed_input_RD, @@ -815,6 +895,7 @@ class EPDispatchMetadata: """Metadata for DeepEP and HybridEP token dispatch.""" state: object # Backend-specific dispatch state. + routed_scores_R: torch.Tensor | None = None # noqa: N815 class DeepEPTokenDispatcher(BaseEPTokenDispatcher): @@ -925,7 +1006,26 @@ def dispatch( cuda_graph_compatible=self.cuda_graph_compatible, ) - metadata = EPDispatchMetadata(state=state) + routed_scores_R = None + if self.absorb_router_scores: + # Take the scores DeepEP tracked for the dispatched rows, clearing + # them so its combine skips the score multiply. The expand + # (CUDA-graph) layout keeps one score per received row per top-k + # slot; the compact layout keeps expert-sorted per-row scores. + if state.cuda_graph_compatible: + recv_scores = state.recv_scores + state.recv_scores = None + if recv_scores is not None: + routed_scores_R = ( + recv_scores.reshape(hidden_states_RD.shape[0], -1) + .sum(dim=-1) + .reshape(-1) + ) + else: + routed_scores_R = state.permuted_scores + state.permuted_scores = None + + metadata = EPDispatchMetadata(state=state, routed_scores_R=routed_scores_R) return hidden_states_RD, num_global_tokens_per_local_expert_e, metadata # pyrefly: ignore [bad-override] @@ -939,14 +1039,10 @@ def combine( del x_TD from torchtitan.distributed.deepep.deepep import combine_tokens, sync_combine - # combine_tokens applies routing scores with bare ops. - remat.recompute_needs_tensor(routed_output_RD) - combined_TD = combine_tokens( routed_output_RD, metadata.state, # pyrefly: ignore [bad-argument-type] remat_region_name=self.remat_region_name("ep_communication.combine"), - recompute=self._replays_communication(), ) sync_combine() return combined_TD @@ -1057,7 +1153,14 @@ def dispatch( pad_multiple=self.pad_multiple, ) - metadata = EPDispatchMetadata(state=state) + routed_scores_R = None + if self.absorb_router_scores: + # Take the padded per-row scores HybridEP tracked, clearing them so + # its combine skips the score multiply. + routed_scores_R = state.permuted_scores + state.permuted_scores = None + + metadata = EPDispatchMetadata(state=state, routed_scores_R=routed_scores_R) return hidden_states_RD, num_global_tokens_per_local_expert_e, metadata # pyrefly: ignore [bad-override] @@ -1072,7 +1175,8 @@ def combine( from torchtitan.distributed.deepep import hybridep - # combine_tokens applies routing scores and combines with bare ops. + # combine_tokens applies routing scores and the combine kernel reads the + # expert outputs with bare ops (no remat region on this path). remat.recompute_needs_tensor(routed_output_RD) combined_TD = hybridep.combine_tokens( routed_output_RD, diff --git a/torchtitan/models/gpt_oss/moe.py b/torchtitan/models/gpt_oss/moe.py index c6bf6ba276e..04cba50c0a8 100644 --- a/torchtitan/models/gpt_oss/moe.py +++ b/torchtitan/models/gpt_oss/moe.py @@ -27,13 +27,28 @@ def __init__(self, config: Config): super().__init__(config) self.bias = nn.Parameter(torch.empty(self.weight.shape[:-1])) - def forward(self, input_RI: torch.Tensor, offsets_E: torch.Tensor) -> torch.Tensor: - output_RO = super().forward(input_RI, offsets_E) + def forward( + self, + input_RI: torch.Tensor, + offsets_E: torch.Tensor, + *, + routed_scores_R: torch.Tensor | None = None, + ) -> torch.Tensor: + output_RO = super().forward( + input_RI, + offsets_E, + routed_scores_R=routed_scores_R, + ) # The bias add below reads the grouped_mm output with bare ops. remat.recompute_needs_tensor(output_RO) bias_RO = self._expand_grouped_bias( self.bias.flatten(1), offsets_E, output_RO.shape[0] ).reshape_as(output_RO) + if routed_scores_R is not None: + # Absorbing the score into W2's input leaves the post-W2 bias + # unscaled, and the dispatcher's combine no longer applies the + # score, so scale the bias rows here to keep score * (h @ W2 + b2). + bias_RO = bias_RO.float() * routed_scores_R.reshape(-1, 1) return self._add_grouped_bias(output_RO, bias_RO) @staticmethod diff --git a/torchtitan/quantization/utils.py b/torchtitan/quantization/utils.py index 97a5bf8d78a..1240fb44653 100644 --- a/torchtitan/quantization/utils.py +++ b/torchtitan/quantization/utils.py @@ -63,6 +63,7 @@ def swap_token_dispatcher(routed_experts_config, pad_multiple: int) -> None: num_experts=dispatcher.num_experts, top_k=dispatcher.top_k, pad_multiple=pad_multiple, + absorb_router_scores=dispatcher.absorb_router_scores, ) elif isinstance(dispatcher, HybridEPTokenDispatcher.Config): routed_experts_config.token_dispatcher = HybridEPTokenDispatcher.Config( @@ -72,6 +73,7 @@ def swap_token_dispatcher(routed_experts_config, pad_multiple: int) -> None: pad_multiple=pad_multiple, hidden_dim=dispatcher.hidden_dim, num_max_tokens_per_rank=dispatcher.num_max_tokens_per_rank, + absorb_router_scores=dispatcher.absorb_router_scores, ) else: raise ValueError( From c755ef87f4719369f41dbc6ff4cffc24ec4b72d9 Mon Sep 17 00:00:00 2001 From: mystri Date: Mon, 19 Oct 2026 21:42:17 +0800 Subject: [PATCH 2/4] [MoE] Never-replayed combine regions for the standard dispatchers Give the local, all-to-all and TorchAO dispatchers the combine structure DeepEP got with router-score absorption: one never-replayed ".combine" region whose body reads the w2 output, instead of bare ops guarded by recompute_needs_tensor markers. - Local and TorchAO EP=1 ran the scatter-add (and the padded unpermute) as bare ops, so replay needed the w2 output persisted: the marker pinned the full expert-output tensor (R x D) per MoE layer even with absorbed scores, where no backward consumer needs it (the scatter-add saves only indices). - The EP>1 combine region was policy-replayable. Replaying it re-ran the combine all-to-all and required the w2 output real on the replay path (pinning it under policies that save w2), while saving it retains only the small permutation indices and scores, so no policy benefits from replaying it. Dispatch stays policy-driven everywhere, like DeepEP's dispatch. Replay: the combine body runs only on the original forward; replay skips it, so nothing on the replay path reads the w2 or combine outputs, and with absorbed scores the w2 output is freed after the forward. Also fix the stale DeepEP docs and AC test left by the absorption commit: the DeepEP combine region is never replayed, so only dispatch call counts grow under the full policy. Tests: - New tests/unit_tests/cpu/models/common/test_token_dispatcher.py: the combine body runs exactly once under every policy while the w2 projection's replay still follows the policy; with absorbed scores and w2 saved, the w2 output storage is freed after the forward; loss and grads match the un-checkpointed run bitwise. It fails on the previous code. - test_remat_all_to_all: the combine never replays, so the full policy replays 2 collectives (was 3) and saving dispatch alone stops all replay; the w2-retention test now saves only w2.grouped_mm, pinning that the combine is always saved. Docs: REMAT.md and the activation-remat config docstrings now state that the dispatcher combines are always saved and only dispatch follows the policy. Numerics: remat-structure change only. On CPU, loss and gradients are bitwise identical to the un-checkpointed run (and to the previous code) for the local and EP=1 paths and for the EP>1 path through a 1-rank EP mesh, absorbed and not, under the full, default and save-w2-only policies. --- .../models/common/test_token_dispatcher.py | 267 ++++++++++++++++++ .../test_deepep_activation_checkpoint.py | 17 +- .../gpu/distributed/test_remat_all_to_all.py | 23 +- torchtitan/distributed/REMAT.md | 33 ++- torchtitan/distributed/activation_remat.py | 7 +- torchtitan/distributed/deepep/deepep.py | 10 +- .../models/common/moe/token_dispatcher.py | 83 ++++-- 7 files changed, 382 insertions(+), 58 deletions(-) create mode 100644 tests/unit_tests/cpu/models/common/test_token_dispatcher.py diff --git a/tests/unit_tests/cpu/models/common/test_token_dispatcher.py b/tests/unit_tests/cpu/models/common/test_token_dispatcher.py new file mode 100644 index 00000000000..d35f8fcc8d2 --- /dev/null +++ b/tests/unit_tests/cpu/models/common/test_token_dispatcher.py @@ -0,0 +1,267 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""The local combine remat region, shared by the standard dispatchers at EP=1. + +``LocalTokenDispatcher.combine`` -- also the AllToAll dispatcher's EP=1 +fallback -- runs as one never-replayed ``token_dispatcher.combine`` region, +like the DeepEP combine: the expert outputs are read only by that saved body, +so nothing pins them for a replay that never comes. + +These tests pin the three consequences on CPU: + +- the combine body runs exactly once (never replayed) under every policy, + while the policy still decides the w2 projection's replay; +- with absorbed router scores and w2 saved, the w2 output's storage is freed + after the forward, since no persistence marker pins it anymore; +- forward and backward match the un-checkpointed run bitwise. +""" + +import gc +import unittest +from unittest.mock import patch + +import torch +import torch.nn as nn +import torch_remat as remat +from torch.multiprocessing.reductions import StorageWeakRef + +import torchtitan.models.common.moe.token_dispatcher as token_dispatcher_module +from torchtitan.distributed.activation_remat import ( + ActivationRematConfig, + DefaultActivationRematConfig, + FullActivationRematConfig, +) +from torchtitan.models.base.module import Module, ModuleDict +from torchtitan.models.common.moe.token_dispatcher import ( + AllToAllTokenDispatcher, + LocalTokenDispatcher, +) + +_MODEL_DIM = 8 +_NUM_TOKENS = 16 +_NUM_EXPERTS = 4 +_TOP_K = 2 + + +class _W2Region(Module): + """Stand-in for the w2 grouped projection: one policy-controlled region. + + The matmul allocates a fresh output and, like a grouped GEMM, saves its + input and weight (not its output) for backward, so only the remat + machinery could retain the output. ``num_forwards`` counts region-body + executions, so a skipped replay does not count. + """ + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(_MODEL_DIM, _MODEL_DIM)) + self.num_forwards = 0 + + def forward(self, x_RD: torch.Tensor) -> torch.Tensor: + return remat.region( + self._project, + self.remat_region_name("grouped_mm"), + recompute=self.remat_should_recompute("grouped_mm"), + )(x_RD) + + def _project(self, x_RD: torch.Tensor) -> torch.Tensor: + self.num_forwards += 1 + return x_RD @ self.weight + + +class _DispatcherBlock(Module): + """Routed-experts block: dispatch, one projection, combine. + + The router scores depend on the gate, so a wrong score placement or a + misaligned combine changes the gradients, as in training. + """ + + def __init__(self, dispatcher_config: LocalTokenDispatcher.Config): + super().__init__() + self.token_dispatcher = dispatcher_config.build() + self.w2 = _W2Region() + self.gate = nn.Linear(_MODEL_DIM, _NUM_EXPERTS, bias=False) + self.num_experts = _NUM_EXPERTS + self.top_k = dispatcher_config.top_k + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + num_tokens = x_TD.shape[0] + token_ids_T1 = torch.arange(num_tokens, device=x_TD.device).unsqueeze(-1) + offsets_1K = torch.arange(self.top_k, device=x_TD.device).unsqueeze(0) + expert_ids_TK = (token_ids_T1 + offsets_1K) % self.num_experts + scores_TK = torch.sigmoid(self.gate(x_TD)).gather(1, expert_ids_TK) + num_tokens_per_expert_E = torch.bincount( + expert_ids_TK.flatten(), minlength=self.num_experts + ) + routed_input_RD, _num_tokens_E, metadata = self.token_dispatcher.dispatch( + x_TD, scores_TK, expert_ids_TK, num_tokens_per_expert_E + ) + routed_scores_R = getattr(metadata, "routed_scores_R", None) + if routed_scores_R is not None: + # Absorbed scores scale the projection input, as GroupedLinear + # does before its grouped_mm region. + remat.recompute_needs_tensor(routed_scores_R) + routed_input_RD = ( + routed_input_RD.float() * routed_scores_R.reshape(-1, 1) + ).to(routed_input_RD.dtype) + routed_output_RD = self.w2(routed_input_RD) + out_TD = self.token_dispatcher.combine(routed_output_RD, metadata, x_TD) + # The loss is a bare consumer of the combine output. + remat.recompute_needs_tensor(out_TD) + return out_TD.double().sum() + + +class _Model(Module): + def __init__(self, block: Module): + super().__init__() + self.layers = ModuleDict({"0": block}) + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + return self.layers["0"](x_TD) + + +def _forward_backward( + model: Module, x_TD: torch.Tensor +) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Run one forward/backward; return the loss and input/parameter grads.""" + model.zero_grad(set_to_none=True) + input_TD = x_TD.detach().clone().requires_grad_(True) + loss = model(input_TD) + loss.backward() + assert input_TD.grad is not None + grads = [input_TD.grad.detach().clone()] + grads += [parameter.grad.detach().clone() for parameter in model.parameters()] + return loss.detach().clone(), grads + + +class TestLocalCombineRematRegion(unittest.TestCase): + @staticmethod + def _dispatcher_cases(): + # name, dispatcher Config. Without a sparse mesh the AllToAll + # dispatcher falls back to the local dispatch/combine methods. + return ( + ("local", LocalTokenDispatcher.Config), + ("alltoall_ep1", AllToAllTokenDispatcher.Config), + ) + + @staticmethod + def _build(config_cls, absorb_router_scores: bool) -> _Model: + torch.manual_seed(42) + return _Model( + _DispatcherBlock( + config_cls( + num_experts=_NUM_EXPERTS, + top_k=_TOP_K, + absorb_router_scores=absorb_router_scores, + ) + ) + ) + + def test_combine_never_replays_and_matches_eager(self): + for name, config_cls in self._dispatcher_cases(): + for absorb in (False, True): + for policy, expected_w2_forwards in ( + (FullActivationRematConfig(), 2), + (DefaultActivationRematConfig(), 1), + ): + with self.subTest( + dispatcher=name, + absorb=absorb, + policy=type(policy).__qualname__, + ): + baseline = self._build(config_cls, absorb) + remat_model = self._build(config_cls, absorb) + remat_model.load_state_dict(baseline.state_dict()) + policy.apply(remat_model) + block = remat_model.layers["0"] + + num_scatter_adds = 0 + original_scatter_add = ( + token_dispatcher_module.deterministic_scatter_add + ) + + def counted_scatter_add(*args, **kwargs): + nonlocal num_scatter_adds + num_scatter_adds += 1 + return original_scatter_add(*args, **kwargs) + + torch.manual_seed(0) + x_TD = torch.randn(_NUM_TOKENS, _MODEL_DIM) + expected = _forward_backward(baseline, x_TD) + with patch.object( + token_dispatcher_module, + "deterministic_scatter_add", + counted_scatter_add, + ): + actual = _forward_backward(remat_model, x_TD) + + # The combine body runs only on the original forward; + # the policy still decides the w2 projection's replay. + self.assertEqual(num_scatter_adds, 1) + self.assertEqual(block.w2.num_forwards, expected_w2_forwards) + torch.testing.assert_close( + actual[0], expected[0], rtol=0, atol=0 + ) + for actual_grad, expected_grad in zip(actual[1], expected[1]): + torch.testing.assert_close( + actual_grad, expected_grad, rtol=0, atol=0 + ) + + def test_absorbed_combine_does_not_pin_w2_output(self): + # With the router scores absorbed, no backward consumer needs the w2 + # output (the scatter-add saves only indices), and the combine region + # reads it only from its saved body, so saving w2 must not keep the + # w2 output alive. + for name, config_cls in self._dispatcher_cases(): + with self.subTest(dispatcher=name): + baseline = self._build(config_cls, absorb_router_scores=True) + remat_model = self._build(config_cls, absorb_router_scores=True) + remat_model.load_state_dict(baseline.state_dict()) + ActivationRematConfig(save_regions=["w2.grouped_mm"]).apply(remat_model) + w2 = remat_model.layers["0"].w2 + w2_output_storages = [] + original_forward = _W2Region.forward + + def recorded_forward(module, *args, **kwargs): + output = original_forward(module, *args, **kwargs) + # The forward also runs (with its region skipped) during + # replay. + if module is w2 and not remat.is_recomputing(): + # remat retains a detached alias, so track the storage. + w2_output_storages.append( + StorageWeakRef(output.untyped_storage()) + ) + return output + + torch.manual_seed(0) + x_TD = torch.randn(_NUM_TOKENS, _MODEL_DIM) + expected = _forward_backward(baseline, x_TD) + with patch.object( + _W2Region, "forward", autospec=True, side_effect=recorded_forward + ): + input_TD = x_TD.detach().clone().requires_grad_(True) + loss = remat_model(input_TD) + gc.collect() + retained = not w2_output_storages[0].expired() + loss.backward() + + self.assertFalse(retained) + self.assertEqual(len(w2_output_storages), 1) + torch.testing.assert_close(loss.detach(), expected[0], rtol=0, atol=0) + actual_grads = [input_TD.grad.detach().clone()] + actual_grads += [ + parameter.grad.detach().clone() + for parameter in remat_model.parameters() + ] + for actual_grad, expected_grad in zip(actual_grads, expected[1]): + torch.testing.assert_close( + actual_grad, expected_grad, rtol=0, atol=0 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/gpu/distributed/test_deepep_activation_checkpoint.py b/tests/unit_tests/gpu/distributed/test_deepep_activation_checkpoint.py index 9e8614c0b35..8ff715da394 100644 --- a/tests/unit_tests/gpu/distributed/test_deepep_activation_checkpoint.py +++ b/tests/unit_tests/gpu/distributed/test_deepep_activation_checkpoint.py @@ -128,7 +128,8 @@ def test_dispatch_and_combine_are_saved_not_replayed(self): # DeepEP fills receive slots with atomics, so a replayed dispatch can # receive rows in another order than the forward's handle records. # Backward would then pair each token's gradient with another token's - # activations. These policies save dispatch and combine instead. + # activations. These policies save the dispatch communication, and the + # combine region is never replayed regardless. self._check_deepep_activation_checkpointing( ( DefaultActivationRematConfig(), @@ -139,9 +140,13 @@ def test_dispatch_and_combine_are_saved_not_replayed(self): @with_comms def test_region_ac_replay_is_exact(self): - # The full policy and policies without ep_communication saved replay dispatch and - # combine. DeepEP then uses a deterministic buffer, which receives rows in - # the forward's order, so the gradients match the run without AC bitwise. + # The full policy and policies without ep_communication saved replay + # the dispatch communication. DeepEP then uses a deterministic buffer, + # which receives rows in the forward's order, so the gradients match + # the run without AC bitwise. The combine region is never replayed: + # it folds the score multiply, unpermute and combine kernel into one + # saved body, so combine runs only in the forward and the dispatch + # backward. self._check_deepep_activation_checkpointing( ( FullActivationRematConfig(), @@ -230,7 +235,9 @@ def call(*args, **kwargs): self.assertEqual(baseline_calls, {"dispatch": 2, "combine": 2}) self.assertEqual( num_calls, - {"dispatch": 2 + num_replays, "combine": 2 + num_replays}, + # Only the dispatch communication replays; the combine + # region is never replayed. + {"dispatch": 2 + num_replays, "combine": 2}, msg=f"unexpected DeepEP replays: {num_calls}", ) if num_replays: diff --git a/tests/unit_tests/gpu/distributed/test_remat_all_to_all.py b/tests/unit_tests/gpu/distributed/test_remat_all_to_all.py index b95286f6e1e..e0aa8175344 100644 --- a/tests/unit_tests/gpu/distributed/test_remat_all_to_all.py +++ b/tests/unit_tests/gpu/distributed/test_remat_all_to_all.py @@ -156,14 +156,15 @@ def test_collective_replay_follows_region_policy(self): ) # dispatch holds the count exchange, its device-to-host sync, and the - # dispatch all-to-all; combine holds the combine all-to-all. The default - # policy recomputes the routed-expert projections but must retain both - # dispatcher regions under the same routed_experts prefix. + # dispatch all-to-all; combine holds the combine all-to-all. The + # dispatch region follows the policy; the combine region is never + # replayed (like DeepEP's), so only the dispatch collectives replay, + # and a combine-only save pattern changes nothing. dispatch = "routed_experts.token_dispatcher.dispatch" combine = "routed_experts.token_dispatcher.combine" for policy_config, expected_replay_collectives, expected_replay_syncs in ( - (ActivationRematConfig(save_regions=[]), 3, 1), - (ActivationRematConfig(save_regions=[dispatch]), 1, 0), + (ActivationRematConfig(save_regions=[]), 2, 1), + (ActivationRematConfig(save_regions=[dispatch]), 0, 0), (ActivationRematConfig(save_regions=[combine]), 2, 1), (ActivationRematConfig(save_regions=[dispatch, combine]), 0, 0), (DefaultActivationRematConfig(), 0, 0), @@ -223,9 +224,10 @@ def counted_sync(*args, **kwargs): @with_comms def test_saved_w2_output_is_not_retained_under_ep(self): # Unpermute and the combine all-to-all do not need the w2 output in - # backward, so saving w2 and the combine must not keep the w2 output - # alive. bf16 activations, as in mixed-precision training, need no cast - # after w2. + # backward, so saving w2 must not keep the w2 output alive: the + # combine region is always saved and reads it only from its body, so + # it needs no persisting for a replay. bf16 activations, as in + # mixed-precision training, need no cast after w2. mesh = init_device_mesh( self.device_type, (self.world_size,), @@ -249,10 +251,7 @@ def test_saved_w2_output_is_not_retained_under_ep(self): ) remat_model.load_state_dict(baseline.state_dict()) ActivationRematConfig( - save_regions=[ - "routed_experts.w2.grouped_mm", - "routed_experts.token_dispatcher.combine", - ] + save_regions=["routed_experts.w2.grouped_mm"] ).apply(remat_model) w2 = remat_model.layers["0"].routed_experts.w2 w2_output_refs = [] diff --git a/torchtitan/distributed/REMAT.md b/torchtitan/distributed/REMAT.md index de3d5185bdd..868018e5fb6 100644 --- a/torchtitan/distributed/REMAT.md +++ b/torchtitan/distributed/REMAT.md @@ -23,7 +23,8 @@ applied to every transformer block with `config.apply(model)`: those matching `recompute_regions`, and recomputes everything else. - `FullActivationRematConfig` is a fixed policy with `save_regions=[]`. It recomputes every block operation except the regions model code always - retains with `recompute=False`, such as routing decisions and trailing adds. + retains with `recompute=False`, such as routing decisions, trailing adds, + and the token-dispatcher combines. - `DefaultActivationRematConfig` is the trainer's default policy: `save_regions=["*"]` and `recompute_regions=["*routed_experts.w13.*"]`. Its regions may change as measurements improve; use `ActivationRematConfig` @@ -288,16 +289,21 @@ communication regions: - Routed-expert `w13` and `w2` grouped projections (`w13.grouped_mm`, `w2.grouped_mm`). -- Token-dispatcher `dispatch` and `combine`. With the all-to-all dispatcher - (EP > 1), each is one region: `dispatch` covers expert sorting, the - token-count exchange and its device-to-host sync, the dispatch all-to-all - and the expert-major permute; `combine` covers the unpermute, the combine - all-to-all and the score-weighted scatter-add. The DeepEP and HybridEP - dispatchers instead declare `ep_communication.dispatch` and - `ep_communication.combine` around their kernels. DeepEP saves or replays - both together, following `ep_communication`. A replay must receive tokens - in the forward's order, so DeepEP then uses a deterministic buffer, which - adds a sort to every dispatch. `FullActivationRematConfig` replays them. +- Token-dispatcher `dispatch` and `combine`. The combine is always saved + (never replayed), like the DeepEP combine: its region covers the unpermute, + the combine all-to-all (with EP > 1) and the score-weighted scatter-add, + so the `w2` output is read only by that saved body and no bare consumer on + the replay path needs it pinned. Saving it retains only the small + permutation indices and scores, so replaying it would re-run the combine + all-to-all without freeing anything. The `dispatch` region covers expert + sorting, the token-count exchange and its device-to-host sync, the dispatch + all-to-all and the expert-major permute, and follows the remat policy. The + DeepEP and HybridEP dispatchers instead declare + `ep_communication.dispatch` and `ep_communication.combine` around their + kernels; only the dispatch region follows the policy there. A DeepEP + replay must receive tokens in the forward's order, so DeepEP then uses a + deterministic buffer, which adds a sort to every dispatch. + `FullActivationRematConfig` replays the dispatch communication. - Shared-expert linear regions. The shared `w2.tp_reduce` region is the `Partial -> Shard(0)` reduce-scatter when sequence parallelism is enabled; save it together with `w2.linear`. @@ -312,7 +318,6 @@ ActivationRematConfig( "moe.routed_experts.w13.grouped_mm", "moe.routed_experts.w2.grouped_mm", "moe.routed_experts.token_dispatcher.dispatch", - "moe.routed_experts.token_dispatcher.combine", "moe.shared_experts.*", "moe.tp_output_reduction", ] @@ -320,7 +325,9 @@ ActivationRematConfig( ``` Without EP there is no communication to save, so the local expert ordering -and the score-weighted scatter-add are ordinary operations. Operations outside +stays an ordinary recomputed operation; the local score-weighted +scatter-add runs as the always-saved `token_dispatcher.combine` region, so +the `w2` output needs no persistence marker. Operations outside these regions, including token-shard zero-fill and branch addition, are recomputed. Routing decisions are retained separately to keep expert selection identical during replay. diff --git a/torchtitan/distributed/activation_remat.py b/torchtitan/distributed/activation_remat.py index f7687e76b84..c9e8d5bf7d5 100644 --- a/torchtitan/distributed/activation_remat.py +++ b/torchtitan/distributed/activation_remat.py @@ -32,7 +32,8 @@ class ActivationRematConfig: Activation rematerialization, also known as activation checkpointing: ``apply`` checkpoints each transformer block with ``torch_remat``. Regions that model code always retains with ``recompute=False``, such as MoE - routing decisions, are saved under every policy. + routing decisions and the token-dispatcher combines, are saved under + every policy. """ save_regions: list[str] @@ -169,7 +170,9 @@ class FullActivationRematConfig(ActivationRematConfig): Every operation in a transformer block is recomputed except the regions model code always retains with ``recompute=False``, such as MoE routing - decisions, auxiliary-loss accumulation, and trailing residual adds. + decisions, auxiliary-loss accumulation, trailing residual adds, and the + token-dispatcher combines (which read the expert outputs only from their + saved bodies, so replaying them would pin those outputs for nothing). """ save_regions: list[str] = field(default_factory=list) diff --git a/torchtitan/distributed/deepep/deepep.py b/torchtitan/distributed/deepep/deepep.py index f9a1ef56ca9..a94aaad0d9a 100644 --- a/torchtitan/distributed/deepep/deepep.py +++ b/torchtitan/distributed/deepep/deepep.py @@ -109,9 +109,9 @@ def _get_next_handle_id() -> torch.Tensor: # Ordered effects, so graph-level selective checkpointing (the graph trainer's memory # policy) saves both ops instead of replaying them in backward. DeepEP assigns receive # slots with atomics, so a replayed dispatch can receive rows in another order, while -# backward routes the gradients with the forward's handle. torch_remat ignores effects and -# follows the ``recompute`` argument of dispatch_tokens/combine_tokens; dispatch_tokens -# only replays with a deterministic buffer. +# backward routes the gradients with the forward's handle. torch_remat ignores effects +# and follows the ``recompute`` argument of dispatch_tokens (combine_tokens is never +# replayed); dispatch_tokens only replays with a deterministic buffer. _lib._register_effectful_op("deepep::dispatch", torch.library.EffectType.ORDERED) _lib._register_effectful_op("deepep::combine", torch.library.EffectType.ORDERED) @@ -505,8 +505,8 @@ def dispatch_tokens( must not exceed that maximum. remat_region_name: Name for the dispatch communication region. recompute: Whether activation remat replays the dispatch communication during - backward. Requires a deterministic buffer (see below), and must match - the ``recompute`` passed to ``combine_tokens``. + backward. Requires a deterministic buffer (see below). The combine + region is never replayed, so it has no matching argument. cuda_graph_compatible: If True, use the static, no-host-sync expand layout so the forward is CUDA-graph-capturable (inference only -- both prefill and decode -- no backward); note it is forced False whenever grad is enabled. If False, use the compact diff --git a/torchtitan/models/common/moe/token_dispatcher.py b/torchtitan/models/common/moe/token_dispatcher.py index b21dc1820dd..43e85c66aa4 100644 --- a/torchtitan/models/common/moe/token_dispatcher.py +++ b/torchtitan/models/common/moe/token_dispatcher.py @@ -44,7 +44,8 @@ class LocalTokenDispatcher(Module): """Token dispatcher for EP=1. Handles local token reordering only. Dispatchers are parameterless modules so they can own remat region names - and policy state. + and policy state. Dispatch stays bare (cheap to recompute); combine is + one never-replayed ``.combine`` region, like the DeepEP combine. """ @dataclass(kw_only=True, slots=True) @@ -155,6 +156,12 @@ def combine( ) -> torch.Tensor: """Score and scatter_add routed expert outputs. + Runs as the single never-replayed ``.combine`` region, like the + DeepEP combine: the expert outputs are read only by this saved body, + so nothing pins them for a replay that never comes. With absorbed + router scores the scatter-add saves only indices, and the w2 output + is freed after the forward instead of being persisted for it. + Args: routed_output_RD: ``(R, D)`` expert outputs metadata: LocalDispatchMetadata from dispatch() @@ -162,9 +169,11 @@ def combine( Returns: out_TD: ``(T, D)`` combined output. """ - # The scatter-add reads the expert outputs with bare ops. - remat.recompute_needs_tensor(routed_output_RD) - return self._score_and_scatter_add( + return remat.region( + self._score_and_scatter_add, + self.remat_region_name("combine"), + recompute=False, # never replayed; the w2 output is read only by this body + )( routed_output_RD, metadata.topk_scores_experts_sorted_N, metadata.token_indices_experts_sorted_N, @@ -259,7 +268,11 @@ class AllToAllTokenDispatcher(BaseEPTokenDispatcher): With EP, each of dispatch and combine is one remat region, ``.dispatch`` and ``.combine``, covering its local reordering, - collectives, and (for combine) the score-weighted scatter-add. + collectives, and (for combine) the score-weighted scatter-add. Dispatch + follows the remat policy; combine is never replayed, like the DeepEP + combine: a replay would re-run the combine all-to-all and need the w2 + output real on the replay path (pinning it under policies that save w2), + while saving it retains only the small permutation indices and scores. The EP mesh is resolved from the ambient SPMD runtime. """ @@ -676,7 +689,9 @@ def combine( out_TD = remat.region( self._combine, self.remat_region_name("combine"), - recompute=self.remat_should_recompute("combine"), + # Never replayed, like the DeepEP combine: the w2 output is read + # only by this saved body, so nothing pins it for a replay. + recompute=False, )( routed_output_RD, metadata.topk_scores_experts_sorted_N, @@ -702,7 +717,9 @@ def _combine( ) -> torch.Tensor: """Unpermute, all-to-all back to token ranks, then score and scatter-add. - Runs as the single ``.combine`` remat region. + Runs as the single ``.combine`` remat region, which is never + replayed: the w2 output is read only by this saved body, and replay + would re-run the combine all-to-all for nothing. """ with maybe_set_sparse_mesh(): routed_output_RD = self._unpermute( @@ -814,18 +831,39 @@ def combine( # EP=1: strip the padding (via _unpermute) to recover expert-sorted # order, then apply the local score + scatter_add used by the EP=1 - # path. Mirrors LocalTokenDispatcher.combine, plus the unpad. + # path. Mirrors LocalTokenDispatcher.combine, plus the unpad, and + # runs as one never-replayed ``combine`` region like it. assert isinstance(metadata, AllToAllDispatchMetadata) - # The unpermute reads the expert outputs with bare ops. - remat.recompute_needs_tensor(routed_output_RD) - routed_output_RD = self._unpermute( - routed_output_RD, metadata.input_shape, metadata.permuted_indices - ) - - return self._score_and_scatter_add( + return remat.region( + self._combine_padded_local, + self.remat_region_name("combine"), + recompute=False, # never replayed; the w2 output is read only by this body + )( routed_output_RD, metadata.topk_scores_experts_sorted_N, metadata.token_indices_experts_sorted_N, + metadata.permuted_indices, + metadata.input_shape, + x_TD, + ) + + def _combine_padded_local( + self, + routed_output_ND, + topk_scores_N, + token_indices_N, + permuted_indices, + input_shape, + x_TD, + ): + """Strip the padded permute, then score and scatter-add (EP=1 combine).""" + routed_output_ND = self._unpermute( + routed_output_ND, input_shape, permuted_indices + ) + return self._score_and_scatter_add( + routed_output_ND, + topk_scores_N, + token_indices_N, x_TD, ) @@ -906,11 +944,14 @@ class DeepEPTokenDispatcher(BaseEPTokenDispatcher): from its deduplicated output into expert-major order; expand dispatch already returns the static expert-major layout. Combine is synchronized before returning its result. - Dispatch and combine share one remat policy (``ep_communication``) because combine - consumes the handle produced by dispatch. A remat policy replays both unless - ``ep_communication`` is saved (the full policy never saves it); a replay - must receive rows in the forward's order, so ``init_buffer`` then makes the buffer - deterministic (a sort per dispatch). + Dispatch and combine declare ``ep_communication.dispatch`` and + ``ep_communication.combine``. The dispatch region follows the remat + policy: a policy replays it unless ``ep_communication`` is saved (the + full policy never saves it); a replay must receive rows in the forward's + order, so ``init_buffer`` then makes the buffer deterministic (a sort + per dispatch). The combine region is never replayed: it folds the score + multiply, unpermute and combine kernel into one saved body, which reads + the w2 output only there, like the standard dispatchers' combines. """ @dataclass(kw_only=True, slots=True) @@ -947,7 +988,7 @@ def __init__(self, config: Config): from torchtitan.distributed.deepep import deepep # noqa: F401 def _replays_communication(self) -> bool: - """Whether activation remat replays dispatch/combine during backward.""" + """Whether activation remat replays the dispatch communication during backward.""" # Without activation remat the regions have no fqn and do not affect execution # (remat_should_recompute is then True, since nothing is saved). return bool(self._remat_module_fqn) and self.remat_should_recompute( From 1c5177158be0e2ec70bcfd923daac1974b651d58 Mon Sep 17 00:00:00 2001 From: mystri Date: Mon, 19 Oct 2026 22:19:00 +0800 Subject: [PATCH 3/4] [MoE] DeepSeek V3 debug configs for the all-to-all and DeepEP dispatchers Add two named debug recipes so the token-dispatcher backend matrix is runnable on DeepSeek V3, alongside the existing HybridEP config: - deepseek_v3_debugmodel_alltoall: the stock dispatcher at EP=2, so routed tokens actually cross ranks. The plain deepseek_v3_debugmodel (EP=1) already runs the same dispatcher's local fallback, which covers the local backend without a duplicate config. - deepseek_v3_debugmodel_deepep: DeepEP v2 at EP=2 via TokenDispatcherTransform, following qwen3_moe_deepep: the compact, host-synced training path, buffer capacity from the microbatch shape, CUDA graphs off like the stock all-to-all, plus the NVLink-only env notes for local devgpu. Both keep the default activation-remat policy, under which every dispatcher's combine is saved (never replayed) and the standard backend absorbs router scores before W2, so these exercise the paths changed by the never-replayed-combine work. Registered in the debug-config defaults test. Validated on CPU: both configs build, DeepEPTokenDispatcher.Config gets hidden_dim=256 (from w13.in_features) and num_max_tokens_per_rank=16384 (the 8 x 2048 microbatch over CP x TP), and the base EP=1 config keeps CUDA graphs on while both EP=2 configs disable them. --- .../cpu/recipes/test_debug_config_defaults.py | 4 ++ .../tests/models/deepseek_v3.py | 63 ++++++++++++++++++- 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/cpu/recipes/test_debug_config_defaults.py b/tests/unit_tests/cpu/recipes/test_debug_config_defaults.py index 45fe5b91086..727bc11aad5 100644 --- a/tests/unit_tests/cpu/recipes/test_debug_config_defaults.py +++ b/tests/unit_tests/cpu/recipes/test_debug_config_defaults.py @@ -26,6 +26,8 @@ ) from torchtitan_recipes.tests.models.deepseek_v3 import ( deepseek_v3_debugmodel, + deepseek_v3_debugmodel_alltoall, + deepseek_v3_debugmodel_deepep, deepseek_v3_debugmodel_hybridep, deepseek_v3_debugmodel_mtp, deepseek_v3_debugmodel_mxfp8, @@ -90,6 +92,8 @@ _DEBUG_CONFIG_FACTORIES: tuple[DebugConfigFactory, ...] = ( deepseek_v3_debugmodel, + deepseek_v3_debugmodel_alltoall, + deepseek_v3_debugmodel_deepep, deepseek_v3_debugmodel_hybridep, deepseek_v3_debugmodel_mtp, deepseek_v3_debugmodel_mxfp8, diff --git a/torchtitan_recipes/tests/models/deepseek_v3.py b/torchtitan_recipes/tests/models/deepseek_v3.py index 5151969065e..abf8e53fb4a 100644 --- a/torchtitan_recipes/tests/models/deepseek_v3.py +++ b/torchtitan_recipes/tests/models/deepseek_v3.py @@ -30,7 +30,10 @@ decoder_vocab_size, DEFAULT_DEBUG_MODEL_SEQ_LEN, ) -from torchtitan.models.common.moe.token_dispatcher import HybridEPTokenDispatcher +from torchtitan.models.common.moe.token_dispatcher import ( + DeepEPTokenDispatcher, + HybridEPTokenDispatcher, +) from torchtitan.models.deepseek_v3 import build_model_config from torchtitan.models.deepseek_v3.mtp import MTPLoss @@ -154,6 +157,64 @@ def deepseek_v3_debugmodel_mxfp8( return config +def deepseek_v3_debugmodel_alltoall( + seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, +) -> Trainer.Config: + """DeepSeek V3 debug pretraining with the stock all-to-all dispatcher, EP=2. + + EP=2 so routed tokens actually cross ranks; the same dispatcher at EP=1 + (the plain ``deepseek_v3_debugmodel``) runs its local fallback, which + reorders tokens without communication. The token-count exchange syncs + with the CPU, so CUDA graphs stay off, as in the base config at EP>1. + """ + return deepseek_v3_debugmodel( + seq_len=seq_len, + parallelism=ParallelismConfig(expert_parallel_degree=2), + ) + + +def deepseek_v3_debugmodel_deepep( + seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, +) -> Trainer.Config: + """DeepSeek V3 debug pretraining with the DeepEP v2 backend (compact training path), EP=2. + + The MoE expert dispatch uses the DeepEP v2 ElasticBuffer all-to-all; under + autograd it takes the compact, host-synced, backward-able path, so like the + stock all-to-all backend it keeps CUDA graphs off. EP=2 so the dispatch + actually moves tokens across ranks; the training shape determines the fixed + per-rank buffer capacity. Numerics match the standard all-to-all backend + (step-1 bitwise, reduction-order drift thereafter). Needs deep_ep v2 + (ElasticBuffer) in the env. + + Local devgpu (no RDMA NIC) needs these env vars so the ElasticBuffer inits + NVLink-only: + - EP_DISABLE_GIN=1 skip the NCCL GIN / RDMA requirement (no RDMA NIC) + - EP_REUSE_NCCL_COMM=0 avoid the ElasticBuffer null-device-comm segfault + - NVSHMEM_REMOTE_TRANSPORT=none + NVSHMEM_DISABLE_MNNVL=1 intra-node NVLink only + - LD_LIBRARY_PATH must include the deep_ep wheels' nvshmem + nccl lib dirs + Then launch with NGPU=2 ./run_train.sh (none of this is needed on RDMA/RoCE hosts). + """ + config = deepseek_v3_debugmodel( + seq_len=seq_len, parallelism=ParallelismConfig(expert_parallel_degree=2) + ) + return apply_transforms( + config, + [ + TokenDispatcherTransform( + dispatcher=DeepEPTokenDispatcher, + # Each EP rank dispatches one CP x TP shard of a microbatch. + num_max_tokens_per_rank=( + config.training.num_tokens_per_microbatch_per_dp_rank + // ( + config.parallelism.context_parallel_degree + * config.parallelism.tensor_parallel_degree + ) + ), + ) + ], + ) + + def deepseek_v3_debugmodel_hybridep( seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, ) -> Trainer.Config: From 84ba2c882e8a188d5a949370b9cc1f13e1d5bf50 Mon Sep 17 00:00:00 2001 From: mystri Date: Tue, 20 Oct 2026 00:32:49 +0800 Subject: [PATCH 4/4] [MoE] fp64 parity tests for absorbed router scores One fp64 parity test per dispatcher family: with rounding gone, the absorbed and post-combine paths compute the same function, so a dropped, doubled or misaligned router score is an O(1) error. Both tests assert the loss and every input/parameter gradient at the float32 scale W2's input scaling applies. - Standard dispatchers (GPU, world 2): local, all-to-all at EP=1 (local fallback) and EP=2, TorchAO at EP=1 (padded permute, zero-sentinel scores) and EP=2, through the existing _grouped_mm override seam. - DeepEP (CPU, gloo, world 2): stub deep_ep's ElasticBuffer with a pure-PyTorch implementation (all_gather + local selection, deterministic layout, autograd-correct dispatch/combine backward), so the dispatcher, the torchtitan deepep wrapper, its custom ops and their autograd run in float64 without the bf16-only CUDA kernels. The stub module is injected only when deep_ep is absent, and the CPU op impls reuse the wrapper's CUDA impls. Replaces the bf16-tolerance and activation-replay variants and the HybridEP-only file. Verified the DeepEP test fails on double-applied and on misaligned (rolled) scores. --- .../test_deepep_absorb_router_scores.py | 417 ++++++++++++++++++ ...andard_dispatchers_absorb_router_scores.py | 260 +++++++++++ 2 files changed, 677 insertions(+) create mode 100644 tests/unit_tests/cpu/distributed/test_deepep_absorb_router_scores.py create mode 100644 tests/unit_tests/gpu/distributed/test_standard_dispatchers_absorb_router_scores.py diff --git a/tests/unit_tests/cpu/distributed/test_deepep_absorb_router_scores.py b/tests/unit_tests/cpu/distributed/test_deepep_absorb_router_scores.py new file mode 100644 index 00000000000..506f8890403 --- /dev/null +++ b/tests/unit_tests/cpu/distributed/test_deepep_absorb_router_scores.py @@ -0,0 +1,417 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""fp64 parity of absorbed router scores through the DeepEP dispatcher. + +The real deep_ep kernels are bf16-only, so this test replaces the +ElasticBuffer with a pure-PyTorch stub (all_gather + local selection) and +runs the dispatcher, the torchtitan deepep wrapper, its custom ops and their +autograd on CPU in float64. With rounding gone, absorbing the router scores +into W2's input and applying them after W2 compute the same function, so a +dropped, doubled or misaligned score is an O(1) error. +""" + +import importlib +import sys +import types +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch_remat as remat +from torch.distributed.device_mesh import init_device_mesh +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +from torchtitan.distributed.spmd_types import set_current_spmd_mesh, set_spmd_meshes +from torchtitan.models.base.module import Module, ModuleDict +from torchtitan.models.common.activation import SwiGLU +from torchtitan.models.common.linear import GroupedLinear +from torchtitan.models.common.moe import RoutedExperts +from torchtitan.models.common.moe.token_dispatcher import DeepEPTokenDispatcher + +_MODEL_DIM = 8 +_NUM_TOKENS = 64 +_TOP_K = 2 + + +# --------------------------------------------------------------------------- +# deep_ep kernel stub +# --------------------------------------------------------------------------- + + +class _FakeEvent: + """No-op stand-in for deep_ep's async combine event.""" + + def current_stream_wait(self) -> None: + pass + + +class _FakeHandle: + """Routing layout of one dispatch, computed identically on every rank. + + ``per_rank_recv[r]`` is ``(src_rank_R, src_token_R)``: row i of rank r's + received tokens is token ``per_rank_recv[r][1][i]`` of rank + ``per_rank_recv[r][0][i]``. + """ + + def __init__(self, per_rank_recv, num_recv_tokens_per_expert_list, num_tokens): + self.per_rank_recv = per_rank_recv + self.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list + self.num_tokens = num_tokens + self.num_sms = 0 + + +def _all_gather_rows(rows: torch.Tensor, group) -> list[torch.Tensor]: + """all_gather a tensor whose leading dimension varies per rank.""" + num_ranks = dist.get_world_size(group) + count = torch.tensor([rows.shape[0]], device=rows.device) + counts = [torch.empty_like(count) for _ in range(num_ranks)] + dist.all_gather(counts, count, group=group) + max_rows = max(int(c.item()) for c in counts) + padded = rows.new_zeros(max_rows, *rows.shape[1:]) + padded[: rows.shape[0]] = rows + gathered = [torch.empty_like(padded) for _ in range(num_ranks)] + dist.all_gather(gathered, padded, group=group) + return [g[: int(c.item())] for g, c in zip(gathered, counts)] + + +class _FakeElasticBuffer: + """Pure-PyTorch stand-in for deep_ep.ElasticBuffer, compact layout only. + + Dispatch gathers every rank's tokens and keeps the rows with at least one + local expert pick (deduplicated, one row per token). Combine sends rows + back to their owners, which sum the rows of each token. The layout is a + deterministic function of the gathered inputs, so all ranks agree on it + and repeated dispatches receive rows in the same order. + """ + + def __init__( + self, + group, + *, + num_bytes, + num_max_tokens_per_rank, + hidden, + num_topk, + use_fp8_dispatch=False, + deterministic=False, + explicitly_destroy=False, + ): + self.group = group + self.num_bytes = num_bytes + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.deterministic = deterministic + + @staticmethod + def get_buffer_size_hint( + group, num_max_tokens_per_rank, hidden, *, num_topk, use_fp8_dispatch=False + ): + return 1 + + def get_theoretical_num_sms(self, num_experts, num_topk): + return 1 + + def dispatch( + self, + x, + *, + topk_idx=None, + topk_weights=None, + num_experts=None, + num_max_tokens_per_rank=None, + num_sms=0, + handle=None, + do_expand=False, + do_cpu_sync=True, + ): + if handle is not None: + return self._dispatch_from_handle(x, handle) + assert not do_expand, "the stub implements the compact layout only" + return self._dispatch_fresh(x, topk_idx, topk_weights, num_experts) + + def _dispatch_fresh(self, x, topk_idx, topk_weights, num_experts): + num_ranks = dist.get_world_size(self.group) + rank = dist.get_rank(self.group) + num_local_experts = num_experts // num_ranks + xs = _all_gather_rows(x, self.group) + idxs = _all_gather_rows(topk_idx, self.group) + ws = _all_gather_rows(topk_weights, self.group) + + per_rank_recv = [] + recv_x, recv_idx, recv_scores = [], [], [] + for r in range(num_ranks): + lo = r * num_local_experts + src_ranks, src_tokens = [], [] + for s in range(num_ranks): + # -1 picks (no selection) and other ranks' experts fall outside + # [lo, lo + num_local_experts). + is_local = (idxs[s] >= lo) & (idxs[s] < lo + num_local_experts) + keep = is_local.any(dim=1) + src_ranks.append( + torch.full((int(keep.sum()),), s, dtype=torch.long, device=x.device) + ) + src_tokens.append(keep.nonzero(as_tuple=False).squeeze(-1)) + if r == rank: + recv_x.append(xs[s][keep]) + recv_idx.append(torch.where(is_local[keep], idxs[s][keep] - lo, -1)) + recv_scores.append(ws[s][keep]) + per_rank_recv.append((torch.cat(src_ranks), torch.cat(src_tokens))) + recv_idx = torch.cat(recv_idx) + handle = _FakeHandle( + per_rank_recv=per_rank_recv, + num_recv_tokens_per_expert_list=torch.bincount( + recv_idx[recv_idx >= 0], minlength=num_local_experts + ).tolist(), + num_tokens=x.shape[0], + ) + return ( + torch.cat(recv_x), + recv_idx, + torch.cat(recv_scores), + handle, + _FakeEvent(), + ) + + def _dispatch_from_handle(self, grad_TD, handle): + # Combine backward: a received row's gradient is its owner token's row. + num_ranks = dist.get_world_size(self.group) + rank = dist.get_rank(self.group) + grads = _all_gather_rows(grad_TD, self.group) + src_rank_R, src_token_R = handle.per_rank_recv[rank] + grad_rows = grad_TD.new_zeros(src_rank_R.numel(), grad_TD.shape[1]) + for s in range(num_ranks): + from_s = src_rank_R == s + if bool(from_s.any()): + grad_rows[from_s] = grads[s][src_token_R[from_s]] + empty = grad_rows.new_zeros(0) + return grad_rows, empty.long(), empty, handle, _FakeEvent() + + def combine(self, x, *, handle, topk_weights=None, async_with_compute_stream=False): + # Forward combine and dispatch backward share this reduction: each + # owner token sums the rows routed to it across all ranks. + num_ranks = dist.get_world_size(self.group) + rank = dist.get_rank(self.group) + xs = _all_gather_rows(x, self.group) + ws = ( + None if topk_weights is None else _all_gather_rows(topk_weights, self.group) + ) + out = x.new_zeros(handle.num_tokens, x.shape[1]) + out_weights = ( + None + if ws is None + else topk_weights.new_zeros(handle.num_tokens, topk_weights.shape[1]) + ) + for r in range(num_ranks): + src_rank_R, src_token_R = handle.per_rank_recv[r] + mine = src_rank_R == rank + if bool(mine.any()): + out.index_add_(0, src_token_R[mine], xs[r][mine]) + if out_weights is not None: + out_weights.index_add_(0, src_token_R[mine], ws[r][mine]) + return out, out_weights, _FakeEvent() + + +def _import_wrapper_with_stub(): + """Import the deepep wrapper against the stub when deep_ep is absent.""" + try: + import deep_ep # noqa: F401 + except ImportError: + stub = types.ModuleType("deep_ep") + stub.ElasticBuffer = _FakeElasticBuffer + sys.modules["deep_ep"] = stub + injected = True + else: + injected = False + wrapper = importlib.import_module("torchtitan.distributed.deepep.deepep") + if injected: + # Sibling tests' importorskip("deep_ep") must still skip honestly. + del sys.modules["deep_ep"] + return wrapper + + +_deepep_wrapper = _import_wrapper_with_stub() +# The wrapper registers its op impls for the CUDA key only; the stub runs on +# CPU, so point the same impls at the CPU key. +_cpu_impl_library = torch.library.Library("deepep", "FRAGMENT") +_cpu_impl_library.impl("dispatch", _deepep_wrapper._dispatch_op_impl, "CPU") +_cpu_impl_library.impl("combine", _deepep_wrapper._combine_op_impl, "CPU") + + +# --------------------------------------------------------------------------- +# fp64 parity test +# --------------------------------------------------------------------------- + + +class _ExactGroupedLinear(GroupedLinear): + """GroupedLinear whose GEMM runs in the weight dtype. + + ``torch._grouped_mm`` is bf16-only, so the fp64 test swaps in this + per-group matmul through the documented ``_grouped_mm`` override seam. + ``RoutedExperts.forward`` hands w13 a bf16 input; both parity paths see + the same bf16 rounding, so upcasting it keeps the comparison sharp. + """ + + @dataclass(kw_only=True, slots=True) + class Config(GroupedLinear.Config): + pass + + def _grouped_mm(self, *, input_RI, weight_EOI, offsets_E): + input_RI = input_RI.to(weight_EOI.dtype) + starts = [0, *offsets_E[:-1].tolist()] + return torch.cat( + [ + input_RI[start:end] @ weight_EOI[expert].transpose(-2, -1) + for expert, (start, end) in enumerate(zip(starts, offsets_E.tolist())) + ], + dim=0, + ) + + +class _DeepEPBlock(Module): + """One routed-experts block behind a DeepEP dispatcher.""" + + def __init__(self, num_experts: int, *, absorb_router_scores: bool): + super().__init__() + # RoutedExperts.forward runs after EP has selected this rank's local + # expert-weight shard. This test bypasses parallelization, so construct + # that local view directly while the dispatcher retains the global E. + num_local_experts = num_experts // dist.get_world_size() + routed_experts = RoutedExperts.__new__(RoutedExperts) + Module.__init__(routed_experts) + routed_experts.w13 = _ExactGroupedLinear.Config( + group_size=num_local_experts, + in_features=_MODEL_DIM, + out_features=_MODEL_DIM, + num_linears=2, + ).build() + routed_experts.w2 = _ExactGroupedLinear.Config( + group_size=num_local_experts, + in_features=_MODEL_DIM, + out_features=_MODEL_DIM, + ).build() + routed_experts.activation_fn = SwiGLU.Config().build() + routed_experts.output_postprocess = None + routed_experts.token_dispatcher = DeepEPTokenDispatcher.Config( + num_experts=num_experts, + top_k=_TOP_K, + hidden_dim=_MODEL_DIM, + num_max_tokens_per_rank=_NUM_TOKENS, + absorb_router_scores=absorb_router_scores, + ).build() + self.gate = torch.nn.Linear(_MODEL_DIM, num_experts, bias=False) + with torch.no_grad(): + for parameter in routed_experts.parameters(): + parameter.normal_(std=0.5) + self.routed_experts = routed_experts + self.num_experts = num_experts + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + num_tokens = x_TD.shape[0] + token_ids_T1 = torch.arange(num_tokens, device=x_TD.device).unsqueeze(-1) + offsets_1K = torch.arange(_TOP_K, device=x_TD.device).unsqueeze(0) + expert_ids_TK = (token_ids_T1 + offsets_1K) % self.num_experts + # Scores depend on the gate, so a misplaced score moves the gradients. + scores_TK = torch.sigmoid(self.gate(x_TD)).gather(1, expert_ids_TK) + num_tokens_per_expert_E = torch.bincount( + expert_ids_TK.flatten(), minlength=self.num_experts + ) + out_TD = self.routed_experts( + x_TD, + scores_TK, + expert_ids_TK, + num_tokens_per_expert_E, + ) + remat.recompute_needs_tensor(out_TD) + return out_TD.square().sum() + + +class _Model(Module): + def __init__(self, block: Module): + super().__init__() + self.layers = ModuleDict({"0": block}) + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + return self.layers["0"](x_TD) + + +def _forward_backward( + model: Module, + x_TD: torch.Tensor, +) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Run one forward/backward; return the loss and input/parameter grads.""" + model.zero_grad(set_to_none=True) + input_TD = x_TD.detach().clone().requires_grad_(True) + loss = model(input_TD) + loss.backward() + assert input_TD.grad is not None + grads = [input_TD.grad.detach().clone()] + grads += [parameter.grad.detach().clone() for parameter in model.parameters()] + return loss.detach().clone(), grads + + +class TestDeepEPAbsorbRouterScoresParity(DTensorTestBase): + @property + def world_size(self) -> int: + return 2 + + @with_comms + def test_absorbed_equals_post_combine(self): + # Force the stub buffer even when the real deep_ep is importable. + saved = (_deepep_wrapper.ElasticBuffer, _deepep_wrapper._buffer) + _deepep_wrapper.ElasticBuffer = _FakeElasticBuffer + _deepep_wrapper._buffer = None + try: + self._check_absorb_parity() + finally: + _deepep_wrapper.ElasticBuffer, _deepep_wrapper._buffer = saved + + def _check_absorb_parity(self) -> None: + mesh = init_device_mesh( + self.device_type, + (self.world_size,), + mesh_dim_names=("ep",), + ) + set_spmd_meshes( + dense_mesh=mesh, + sparse_mesh=mesh, + dense_sp_enabled=False, + ) + num_experts = 4 * self.world_size + with ( + torch.autograd.set_multithreading_enabled(False), + set_current_spmd_mesh(mesh), + ): + torch.manual_seed(42) + baseline = _Model(_DeepEPBlock(num_experts, absorb_router_scores=False)).to( + self.device_type, torch.float64 + ) + absorbed = _Model(_DeepEPBlock(num_experts, absorb_router_scores=True)).to( + self.device_type, torch.float64 + ) + absorbed.load_state_dict(baseline.state_dict()) + # Both models share one stub buffer. + baseline.layers["0"].routed_experts.token_dispatcher.init_buffer() + absorbed.layers["0"].routed_experts.token_dispatcher.init_buffer() + + torch.manual_seed(self.rank) + x_TD = torch.randn( + _NUM_TOKENS, + _MODEL_DIM, + device=self.device_type, + dtype=torch.float64, + ) + expected_loss, expected = _forward_backward(baseline, x_TD) + actual_loss, actual = _forward_backward(absorbed, x_TD) + + # The absorbed path scales W2's input in float32, so agreement is + # at float32 scale, not float64 epsilon. + torch.testing.assert_close(actual_loss, expected_loss, rtol=1e-5, atol=1e-6) + for actual_grad, expected_grad in zip(actual, expected): + torch.testing.assert_close( + actual_grad, expected_grad, rtol=1e-5, atol=1e-6 + ) diff --git a/tests/unit_tests/gpu/distributed/test_standard_dispatchers_absorb_router_scores.py b/tests/unit_tests/gpu/distributed/test_standard_dispatchers_absorb_router_scores.py new file mode 100644 index 00000000000..b86ed473bb6 --- /dev/null +++ b/tests/unit_tests/gpu/distributed/test_standard_dispatchers_absorb_router_scores.py @@ -0,0 +1,260 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""fp64 parity of absorbed router scores across the standard dispatchers. + +With ``absorb_router_scores`` the dispatcher returns the per-row router +scores for W2 to scale its input, and combine skips its score multiply. In +fp64 both paths compute the same function, so a dropped, doubled or +misaligned score is an O(1) error. Covered: + +- LocalTokenDispatcher (no EP), +- AllToAllTokenDispatcher at EP=1 (falls back to the local methods) and + EP=2 (tokens cross ranks), +- TorchAOTokenDispatcher at EP=1 (local padded permute; exercises the + zero-sentinel for pad rows) and EP=2. + +The DeepEP dispatcher has its own fp64 test with a stubbed deep_ep kernel +under tests/unit_tests/cpu/distributed/. +""" + +import unittest +from dataclasses import dataclass + +import pytest +import torch +import torch_remat as remat +from torch.distributed.device_mesh import init_device_mesh +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +from torchtitan.distributed.spmd_types import set_current_spmd_mesh, set_spmd_meshes +from torchtitan.models.base.module import Module, ModuleDict +from torchtitan.models.common.activation import SwiGLU +from torchtitan.models.common.linear import GroupedLinear +from torchtitan.models.common.moe import RoutedExperts +from torchtitan.models.common.moe.token_dispatcher import ( + AllToAllTokenDispatcher, + LocalTokenDispatcher, + TorchAOTokenDispatcher, +) + +pytest.importorskip("torchao") + +pytestmark = pytest.mark.multi_gpu + +_MODEL_DIM = 8 +_NUM_TOKENS = 64 +_TOP_K = 2 +# TorchAO's fused permute pads rows to this multiple. +_PAD_MULTIPLE = 16 + + +class _ExactGroupedLinear(GroupedLinear): + """GroupedLinear whose GEMM runs in the weight dtype. + + ``torch._grouped_mm`` is bf16-only, so the fp64 test swaps in this + per-group matmul through the documented ``_grouped_mm`` override seam. + ``RoutedExperts.forward`` hands w13 a bf16 input; both parity paths see + the same bf16 rounding, so upcasting it keeps the comparison sharp. + """ + + @dataclass(kw_only=True, slots=True) + class Config(GroupedLinear.Config): + pass + + def _grouped_mm(self, *, input_RI, weight_EOI, offsets_E): + input_RI = input_RI.to(weight_EOI.dtype) + starts = [0, *offsets_E[:-1].tolist()] + return torch.cat( + [ + input_RI[start:end] @ weight_EOI[expert].transpose(-2, -1) + for expert, (start, end) in enumerate(zip(starts, offsets_E.tolist())) + ], + dim=0, + ) + + +class _RoutedExpertsBlock(Module): + """One routed-experts block behind a configurable token dispatcher.""" + + def __init__( + self, + num_experts: int, + num_local_experts: int, + dispatcher_config: LocalTokenDispatcher.Config, + ): + super().__init__() + # RoutedExperts.forward runs after EP has selected this rank's local + # expert-weight shard. This test bypasses parallelization, so construct + # that local view directly while the dispatcher retains the global E. + routed_experts = RoutedExperts.__new__(RoutedExperts) + Module.__init__(routed_experts) + routed_experts.w13 = _ExactGroupedLinear.Config( + group_size=num_local_experts, + in_features=_MODEL_DIM, + out_features=_MODEL_DIM, + num_linears=2, + ).build() + routed_experts.w2 = _ExactGroupedLinear.Config( + group_size=num_local_experts, + in_features=_MODEL_DIM, + out_features=_MODEL_DIM, + ).build() + routed_experts.activation_fn = SwiGLU.Config().build() + routed_experts.output_postprocess = None + routed_experts.token_dispatcher = dispatcher_config.build() + self.gate = torch.nn.Linear(_MODEL_DIM, num_experts, bias=False) + with torch.no_grad(): + for parameter in routed_experts.parameters(): + parameter.normal_(std=0.5) + self.routed_experts = routed_experts + self.num_experts = num_experts + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + num_tokens = x_TD.shape[0] + token_ids_T1 = torch.arange(num_tokens, device=x_TD.device).unsqueeze(-1) + offsets_1K = torch.arange(_TOP_K, device=x_TD.device).unsqueeze(0) + expert_ids_TK = (token_ids_T1 + offsets_1K) % self.num_experts + # Scores depend on the gate, so a misplaced score moves the gradients. + scores_TK = torch.sigmoid(self.gate(x_TD)).gather(1, expert_ids_TK) + num_tokens_per_expert_E = torch.bincount( + expert_ids_TK.flatten(), minlength=self.num_experts + ) + out_TD = self.routed_experts( + x_TD, + scores_TK, + expert_ids_TK, + num_tokens_per_expert_E, + ) + remat.recompute_needs_tensor(out_TD) + return out_TD.square().sum() + + +class _Model(Module): + def __init__(self, block: Module): + super().__init__() + self.layers = ModuleDict({"0": block}) + + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + return self.layers["0"](x_TD) + + +def _forward_backward( + model: Module, + x_TD: torch.Tensor, +) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Run one forward/backward; return the loss and input/parameter grads.""" + model.zero_grad(set_to_none=True) + input_TD = x_TD.detach().clone().requires_grad_(True) + loss = model(input_TD) + loss.backward() + assert input_TD.grad is not None + grads = [input_TD.grad.detach().clone()] + grads += [parameter.grad.detach().clone() for parameter in model.parameters()] + return loss.detach().clone(), grads + + +@unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") +class TestStandardDispatchersAbsorbRouterScoresParity(DTensorTestBase): + @property + def world_size(self) -> int: + return 2 + + @staticmethod + def _cases(): + # name, dispatcher Config, extra kwargs, ep size (1 = no sparse mesh) + return ( + ("local", LocalTokenDispatcher.Config, {}, 1), + ("alltoall_ep1", AllToAllTokenDispatcher.Config, {}, 1), + ("alltoall_ep2", AllToAllTokenDispatcher.Config, {}, 2), + ( + "torchao_ep1", + TorchAOTokenDispatcher.Config, + {"pad_multiple": _PAD_MULTIPLE}, + 1, + ), + ( + "torchao_ep2", + TorchAOTokenDispatcher.Config, + {"pad_multiple": _PAD_MULTIPLE}, + 2, + ), + ) + + @with_comms + def test_absorbed_equals_post_combine(self): + mesh = init_device_mesh( + self.device_type, + (self.world_size,), + mesh_dim_names=("ep",), + ) + num_experts = 4 * self.world_size + for name, config_cls, kwargs, ep_size in self._cases(): + set_spmd_meshes( + dense_mesh=mesh, + sparse_mesh=mesh if ep_size > 1 else None, + dense_sp_enabled=False, + ) + with ( + self.subTest(dispatcher=name), + torch.autograd.set_multithreading_enabled(False), + set_current_spmd_mesh(mesh), + ): + self._check_absorb_parity( + config_cls, + kwargs, + num_experts, + num_experts // ep_size, + ) + + def _check_absorb_parity( + self, + config_cls, + kwargs, + num_experts: int, + num_local_experts: int, + ) -> None: + torch.manual_seed(42) + baseline = _Model( + _RoutedExpertsBlock( + num_experts, + num_local_experts, + config_cls( + num_experts=num_experts, + top_k=_TOP_K, + absorb_router_scores=False, + **kwargs, + ), + ) + ).to(self.device_type, torch.float64) + absorbed = _Model( + _RoutedExpertsBlock( + num_experts, + num_local_experts, + config_cls( + num_experts=num_experts, + top_k=_TOP_K, + absorb_router_scores=True, + **kwargs, + ), + ) + ).to(self.device_type, torch.float64) + absorbed.load_state_dict(baseline.state_dict()) + + torch.manual_seed(self.rank) + x_TD = torch.randn( + _NUM_TOKENS, _MODEL_DIM, device=self.device_type, dtype=torch.float64 + ) + expected_loss, expected = _forward_backward(baseline, x_TD) + actual_loss, actual = _forward_backward(absorbed, x_TD) + + # The absorbed path scales W2's input in float32, so agreement is at + # float32 scale, not float64 epsilon. + torch.testing.assert_close(actual_loss, expected_loss, rtol=1e-5, atol=1e-6) + for actual_grad, expected_grad in zip(actual, expected): + torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-5, atol=1e-6)