From 54543e87da88494b9337c117758ec73723d48c9c Mon Sep 17 00:00:00 2001 From: Yisong Li Date: Thu, 17 Sep 2026 22:12:51 +0800 Subject: [PATCH 1/3] =?UTF-8?q?perf(moe):=20HybridEP=20capacity=20mode=20?= =?UTF-8?q?=E2=80=94=20non-blocking=20dispatch=20under=20dynamic=20routing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HybridEP's blocking dispatch learns the permuted row count by draining the compute stream (`torch.cuda.current_stream().synchronize()` after metadata preprocessing) on every MoE layer, forward and activation-checkpoint recompute: the CPU cannot run ahead of the GPU and every layer becomes an EP-group barrier. Its non-blocking mode takes `num_permuted_tokens` from the caller and only uses it to size the output buffers; the real counts stay on the device and the handle carries an overflow flag (verified bit-exact on GB200 for capacities of 1x / 1.5x / 2x the actual count). Opt-in `BackendConfig.dispatcher_capacity_factor` (default None = the blocking path): the first dispatch of a layer runs blocking and calibrates capacity = ceil(rows x factor), 4-token aligned, EP-group max; every later dispatch passes that capacity and runs non-blocking, the combine and its backward dispatch take the same host int (no drain there either), and `torch._assert_async` on HybridEP's overflow flag fails loudly instead of training on truncated tokens. GroupedExpertsDeepEP skips its per-microbatch count_nonzero host read in this mode (rows are never empty). Ignored under benchmark_static_routing, which already pins the exact count. Measured on 8 x GB200 (EP16, dynamic routing, Kimi-K3 16-layer twin) together with the equal-token- count pad change of the next commit: 1957.9 -> 2124.0 tok/s/GPU (+8.5%, A/A band 0.08%), loss curves identical; forward output and input gradient bit-identical to the blocking path in a 4-rank end-to-end check. Tests: tests/unit_tests/moe/test_token_dispatcher.py (calibration, capacity passed non-blocking, overflow guard, static routing / no factor unchanged), tests/unit_tests/moe/test_backend_config.py. Signed-off-by: Yisong Li --- .../components/models/common/utils.py | 11 ++++ nemo_automodel/components/moe/experts.py | 13 +++- .../components/moe/megatron/fused_a2a.py | 2 + .../moe/megatron/token_dispatcher.py | 63 +++++++++++++++++++ tests/unit_tests/moe/test_backend_config.py | 7 +++ tests/unit_tests/moe/test_token_dispatcher.py | 60 ++++++++++++++++++ 6 files changed, 155 insertions(+), 1 deletion(-) diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index 5308d6f45b..14dc720ace 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/nemo_automodel/components/models/common/utils.py @@ -442,6 +442,11 @@ class BackendConfig: manager instance across MoE layers. dispatcher_async_dispatch: Whether DeepEP/UCCL-EP dispatch and combine should return asynchronously and allocate their outputs on the communication stream. + dispatcher_capacity_factor: HybridEP only, dynamic routing. Run HybridEP dispatch in its + non-blocking mode with output buffers sized to the first microbatch's permuted row + count times this factor (EP-group max, 4-token aligned) instead of letting every + dispatch drain the compute stream to read the exact count. Overflow trips a + device-side assert. None (default) keeps the blocking reference path. enable_deepep: Removed and ignored. Logs a warning if set; configure "dispatcher" and "experts" explicitly instead. fake_balanced_gate: If True, replace the learned Gate with FakeBalancedGate @@ -501,6 +506,12 @@ class BackendConfig: dispatcher_num_sms: int = 20 dispatcher_share_token_dispatcher: bool = True dispatcher_async_dispatch: bool = False + # HybridEP only, dynamic routing: after one blocking calibration dispatch per MoE layer, size every + # later dispatch's output buffers to ceil(calibrated rows x factor) (EP-group max, aligned) and run + # HybridEP in its non-blocking mode. Removes the per-dispatch compute-stream drain that HybridEP's + # blocking mode needs to learn the permuted row count (and the per-layer barrier it implies); an + # overflow of the capacity trips a device-side assert instead of silently truncating. None = blocking. + dispatcher_capacity_factor: float | None = None mok: MoKBackendConfig = field(default_factory=MoKBackendConfig) enable_deepep: bool | None = None # Removed: ignored with a warning; set dispatcher/experts explicitly fake_balanced_gate: bool = False diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 2ebea56d40..6e6e5644e1 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -959,6 +959,12 @@ def __init__( self.dispatcher_num_sms = dispatcher_num_sms self.dispatcher_share_token_dispatcher = dispatcher_share_token_dispatcher self.dispatcher_async_dispatch = dispatcher_async_dispatch + # HybridEP capacity mode (BackendConfig.dispatcher_capacity_factor): the dispatcher returns + # device-side tokens_per_expert and buffers of a fixed capacity, so the per-microbatch + # count_nonzero host read below is skipped as well (rows are never empty). + self.dispatcher_capacity_factor = ( + getattr(backend, "dispatcher_capacity_factor", None) if backend is not None else None + ) # Allocate projection tensor - size depends on whether activation is gated # Gated (SwiGLU, Quick-GEGLU): [n_experts, dim, 2*inter_dim] @@ -992,6 +998,7 @@ def init_token_dispatcher(self, ep_mesh: DeviceMesh): moe_hybridep_num_sms=self.dispatcher_num_sms, moe_share_token_dispatcher=self.dispatcher_share_token_dispatcher, moe_deepep_async_dispatch=self.dispatcher_async_dispatch, + moe_hybridep_capacity_factor=self.dispatcher_capacity_factor, moe_benchmark_static_routing=self.static_routing, ) @@ -1072,7 +1079,11 @@ def forward( # With static routing (forced balance, no noise) every expert receives tokens by # construction, so the count_nonzero device-to-host read (one per microbatch, and # again per activation-checkpoint recompute) can be skipped. - if self.static_routing or torch.count_nonzero(tokens_per_expert) > 0: + if ( + self.static_routing + or self.dispatcher_capacity_factor is not None + or torch.count_nonzero(tokens_per_expert) > 0 + ): if self.use_torch_mm: tokens_per_expert_gpu = tokens_per_expert.to( device=permuted_local_hidden_states.device, non_blocking=True diff --git a/nemo_automodel/components/moe/megatron/fused_a2a.py b/nemo_automodel/components/moe/megatron/fused_a2a.py index 1eaf5124a8..81e2f6e6a8 100644 --- a/nemo_automodel/components/moe/megatron/fused_a2a.py +++ b/nemo_automodel/components/moe/megatron/fused_a2a.py @@ -686,6 +686,8 @@ def backward(ctx, grad_x): handle=handle, pad_multiple=ctx.pad_multiple, num_permuted_tokens=ctx.num_permuted_tokens, + # Capacity mode hands a host int: the backward dispatch then needs no stream drain either. + non_blocking=isinstance(ctx.num_permuted_tokens, int), ) return dispatched_hidden, None, None, None diff --git a/nemo_automodel/components/moe/megatron/token_dispatcher.py b/nemo_automodel/components/moe/megatron/token_dispatcher.py index a9ff515249..b433d71cd3 100644 --- a/nemo_automodel/components/moe/megatron/token_dispatcher.py +++ b/nemo_automodel/components/moe/megatron/token_dispatcher.py @@ -14,6 +14,7 @@ # limitations under the License. import logging +import math import os from abc import ABC, abstractmethod from dataclasses import dataclass @@ -361,6 +362,17 @@ def forward(self, token_indices: torch.Tensor, token_probs: torch.Tensor) -> tup _STATIC_ROUTING_PAD_PIN = os.environ.get("NEMO_STATIC_ROUTING_PAD_PIN", "1") != "0" +def _assert_no_hybridep_overflow(handle, capacity: int) -> None: + """Device-side guard: HybridEP truncates silently when the capacity is exceeded and sets + ``overflow_flag`` (handle item 10); fail loudly instead of training on dropped tokens.""" + flag = handle[10] if isinstance(handle, (tuple, list)) and len(handle) > 10 else None + if torch.is_tensor(flag): + torch._assert_async( + (flag == 0).reshape(()), + f"HybridEP dispatch overflowed its capacity of {capacity} permuted rows; raise BackendConfig.dispatcher_capacity_factor", + ) + + class _HybridEPManager(_DispatchManager): """ A manager class to handle fused all-to-all communication processes for MoE models using @@ -385,6 +397,7 @@ def __init__( permute_fusion: bool = False, moe_hybridep_num_sms: int = 24, benchmark_static_routing: bool = False, + moe_hybridep_capacity_factor: float | None = None, ): self.group = group self.num_local_experts = num_local_experts @@ -392,6 +405,10 @@ def __init__( self.router_topk = router_topk self.permute_fusion = permute_fusion self.moe_hybridep_num_sms = moe_hybridep_num_sms + # Capacity mode (dynamic routing): after one blocking calibration dispatch, later dispatches + # pass this many rows as num_permuted_tokens and run non-blocking (see dispatch()). + self.hybridep_capacity_factor = moe_hybridep_capacity_factor + self._hybridep_capacity: int | None = None # Benchmark-only (TokenDispatcherConfig.moe_benchmark_static_routing): # persist num_permuted_tokens across dispatches, see dispatch()/reset. self.benchmark_static_routing = benchmark_static_routing @@ -502,6 +519,12 @@ def dispatch( self.routing_map = nn.functional.pad(self.routing_map, (0, 0, 0, pad_tokens)) self.token_probs = nn.functional.pad(self.token_probs, (0, 0, 0, pad_tokens)) + capacity_mode = self.hybridep_capacity_factor is not None and not self.benchmark_static_routing + if capacity_mode and self._hybridep_capacity is not None: + # Non-blocking HybridEP dispatch: the buffers are sized to the calibrated capacity, the + # real counts stay on the device, and no compute-stream drain happens on this call. + self.num_permuted_tokens = self._hybridep_capacity + dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = hybrid_ep_dispatch( x=hidden_states, routing_map=self.routing_map, @@ -515,12 +538,46 @@ def dispatch( ) self.tokens_per_expert = tokens_per_expert + if capacity_mode: + if self._hybridep_capacity is None: + self._calibrate_hybridep_capacity(tokens_per_expert, dispatched_hidden.device) + else: + _assert_no_hybridep_overflow(self.handle, self._hybridep_capacity) + self.num_permuted_tokens = self._hybridep_capacity + return dispatched_hidden self.num_permuted_tokens = self.tokens_per_expert.sum() if self.benchmark_static_routing and getattr(self, "_static_num_permuted_tokens", None) is None: self._static_num_permuted_tokens = self.num_permuted_tokens return dispatched_hidden + def _calibrate_hybridep_capacity(self, tokens_per_expert: torch.Tensor, device: torch.device) -> None: + """Turn the one blocking calibration dispatch into the capacity every later dispatch uses. + + The blocking dispatch already brought ``tokens_per_expert`` to the host, so reading its sum is + free; the capacity is that count times ``hybridep_capacity_factor``, aligned to the HybridEP + token alignment (and ``pad_multiple``), taken as the EP-group maximum once so every rank + allocates the same buffers. + """ + actual = int(tokens_per_expert.sum()) + align = max(int(self.pad_multiple or 0), _HYBRIDEP_TOKEN_ALIGNMENT) + cap = -(-int(math.ceil(actual * self.hybridep_capacity_factor)) // align) * align + if torch.distributed.is_initialized() and torch.distributed.get_world_size(self.group) > 1: + cap_t = torch.tensor(cap, device=device) + torch.distributed.all_reduce(cap_t, op=torch.distributed.ReduceOp.MAX, group=self.group) + cap = int(cap_t) + self._hybridep_capacity = cap + # This dispatch itself was blocking: its combine (and the combine's backward dispatch) can + # use the exact host-side count. + self.num_permuted_tokens = actual + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + logging.getLogger(__name__).info( + "HybridEP capacity mode: calibrated %d permuted rows x %.2f -> capacity %d; later dispatches run non-blocking", + actual, + self.hybridep_capacity_factor, + cap, + ) + def combine( self, hidden_states: torch.Tensor, @@ -583,6 +640,10 @@ class TokenDispatcherConfig: None means no changes for dtype.""" moe_flex_dispatcher_backend: Literal["deepep", "hybridep", "uccl_ep"] = "deepep" + + # HybridEP capacity mode, see BackendConfig.dispatcher_capacity_factor + + moe_hybridep_capacity_factor: float | None = None """Backend for the flex token dispatcher. Options: 'deepep', 'hybridep', or 'uccl_ep'.""" moe_deepep_num_sms: int = 20 @@ -705,6 +766,7 @@ def __init__( permute_fusion=self.config.moe_permute_fusion, moe_hybridep_num_sms=self.config.moe_hybridep_num_sms, benchmark_static_routing=self.config.moe_benchmark_static_routing, + moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor, ) self._comm_manager = MoEFlexTokenDispatcher.shared_hybridep_manager else: @@ -716,6 +778,7 @@ def __init__( permute_fusion=self.config.moe_permute_fusion, moe_hybridep_num_sms=self.config.moe_hybridep_num_sms, benchmark_static_routing=self.config.moe_benchmark_static_routing, + moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor, ) self.hybridep_metadata_processor = _HybridEPMetadataProcessor( num_experts=self.tp_size * self.config.num_moe_experts, diff --git a/tests/unit_tests/moe/test_backend_config.py b/tests/unit_tests/moe/test_backend_config.py index 913ab00813..30278bab91 100644 --- a/tests/unit_tests/moe/test_backend_config.py +++ b/tests/unit_tests/moe/test_backend_config.py @@ -492,3 +492,10 @@ def test_preprocess_requires_router_and_hybridep(self): dispatcher="deepep", cuda_graph=CudaGraphConfig(modules=["moe_router", "moe_preprocess"]), ) + + +def test_dispatcher_capacity_factor_defaults_off(): + from nemo_automodel.components.models.common.utils import BackendConfig + + assert BackendConfig().dispatcher_capacity_factor is None + assert BackendConfig(dispatcher_capacity_factor=1.5).dispatcher_capacity_factor == 1.5 diff --git a/tests/unit_tests/moe/test_token_dispatcher.py b/tests/unit_tests/moe/test_token_dispatcher.py index 0fef97971f..7a8bbbf176 100644 --- a/tests/unit_tests/moe/test_token_dispatcher.py +++ b/tests/unit_tests/moe/test_token_dispatcher.py @@ -258,3 +258,63 @@ def test_pinned_size_smaller_than_a_later_batch_falls_back_to_the_all_reduce(sel assert calls == [4] and sizes == [12, 12] calls2, sizes2 = self._dispatch_n(m, monkeypatch, [16], group_max=16, pin=True) assert calls2 == [16] and sizes2 == [16] and m._static_target_tokens == 16 + + +class TestHybridEPCapacityMode: + """BackendConfig.dispatcher_capacity_factor: one blocking calibration dispatch, then every dispatch passes + the calibrated capacity as num_permuted_tokens (non-blocking) and guards HybridEP's overflow flag.""" + + def _run(self, monkeypatch, factor, num_tokens_seq, tpe_rows, overflow=0, static=False): + import nemo_automodel.components.moe.megatron.token_dispatcher as td + + with patch( + "nemo_automodel.components.moe.megatron.token_dispatcher.hybrid_ep_dispatch", new=lambda *a, **kw: None + ): + m = _HybridEPManager( + group=None, + num_local_experts=2, + num_experts=8, + router_topk=2, + benchmark_static_routing=static, + moe_hybridep_capacity_factor=factor, + ) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + passed, asserts = [], [] + + def fake_dispatch(x, routing_map, probs, num_permuted_tokens=None, **kwargs): + passed.append(num_permuted_tokens) + # blocking mode returns host-side counts; non-blocking returns device-side (cpu stands in) counts + tpe = torch.tensor(tpe_rows, dtype=torch.int64) + handle = tuple([None] * 10 + [torch.tensor(overflow)]) + rows = int(tpe.sum()) if num_permuted_tokens is None else int(num_permuted_tokens) + return x.new_zeros(rows, x.shape[1]), probs, None, tpe, handle + + monkeypatch.setattr(td, "hybrid_ep_dispatch", fake_dispatch) + monkeypatch.setattr(torch, "_assert_async", lambda cond, msg="": asserts.append((bool(cond), msg))) + for n in num_tokens_seq: + m.routing_map = torch.ones(n, 8, dtype=torch.bool) + m.token_probs = torch.full((n, 8), 0.125) + m.dispatch(torch.randn(n, 4)) + return m, passed, asserts + + def test_first_dispatch_calibrates_then_passes_the_capacity(self, monkeypatch): + m, passed, asserts = self._run(monkeypatch, 1.5, [8, 8, 8], tpe_rows=[5, 5]) # 10 rows x 1.5 = 15 -> aligned 16 + assert passed == [None, 16, 16] + assert m._hybridep_capacity == 16 and m.num_permuted_tokens == 16 + assert asserts and all(ok for ok, _ in asserts), "overflow guard checked on every capacity dispatch" + + def test_calibration_dispatch_keeps_the_exact_count_for_its_combine(self, monkeypatch): + m, passed, _ = self._run(monkeypatch, 2.0, [8], tpe_rows=[5, 5]) + assert passed == [None] and m.num_permuted_tokens == 10 and m._hybridep_capacity == 20 + + def test_overflow_flag_trips_the_guard(self, monkeypatch): + _, _, asserts = self._run(monkeypatch, 1.5, [8, 8], tpe_rows=[5, 5], overflow=1) + assert asserts and asserts[-1][0] is False and "capacity" in asserts[-1][1] + + def test_static_routing_ignores_the_factor(self, monkeypatch): + m, passed, asserts = self._run(monkeypatch, 1.5, [8, 8], tpe_rows=[5, 5], static=True) + assert passed[0] is None and m._hybridep_capacity is None and asserts == [] + + def test_no_factor_keeps_the_blocking_path(self, monkeypatch): + m, passed, asserts = self._run(monkeypatch, None, [8, 8], tpe_rows=[5, 5]) + assert passed == [None, None] and asserts == [] From d53033b1599944916478d24a3e6d671a8c298c1d Mon Sep 17 00:00:00 2001 From: Yisong Li Date: Thu, 17 Sep 2026 22:14:34 +0800 Subject: [PATCH 2/3] perf(moe): dispatcher_equal_token_counts skips the HybridEP pad-size all-reduce MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The padding added in #3641 derives every rank's HybridEP dispatch size from an EP-group MAX all-reduce of the local row count followed by an int() host sync, once per MoE layer per forward and per recompute. The operand is a host-known shape: for fixed-shape batches — every batch that is not variable-length or in-batch packed — all ranks hold the same count and the collective only costs a compute-stream drain and a per-layer barrier (static routing already pins it, #3895). Opt-in `BackendConfig.dispatcher_equal_token_counts` (default False) declares equal counts: the pad size becomes the aligned local count with no collective. Keep it False for variable-length inputs, where unequal counts would abort the HybridEP collective. Measured on 8 x GB200 under static routing (the collective isolated): skipping it is +7.4% at that scale, of which 5.8 points are the host sync and 1.6 the barrier; at 64 nodes the same pin measured +0.65%. On the dynamic-routing path it is worth combining with dispatcher_capacity_factor (previous commit), which removes the other per-layer host sync. Tests: tests/unit_tests/moe/test_token_dispatcher.py (no all-reduce and aligned sizes when set; per-dispatch all-reduce when unset), tests/unit_tests/moe/test_backend_config.py. Signed-off-by: Yisong Li --- .../components/models/common/utils.py | 9 ++++ nemo_automodel/components/moe/experts.py | 4 ++ .../moe/megatron/token_dispatcher.py | 13 ++++- tests/unit_tests/moe/test_backend_config.py | 7 +++ tests/unit_tests/moe/test_token_dispatcher.py | 48 +++++++++++++++++++ 5 files changed, 80 insertions(+), 1 deletion(-) diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index 14dc720ace..c8537a84d0 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/nemo_automodel/components/models/common/utils.py @@ -447,6 +447,10 @@ class BackendConfig: count times this factor (EP-group max, 4-token aligned) instead of letting every dispatch drain the compute stream to read the exact count. Overflow trips a device-side assert. None (default) keeps the blocking reference path. + dispatcher_equal_token_counts: HybridEP only. Declare that every EP rank dispatches the + same row count, so the per-dispatch EP-group max all-reduce (and its host sync) + that derives the pad size is skipped and the count is only aligned. Default False; + keep it False for variable-length or in-batch-packed inputs. enable_deepep: Removed and ignored. Logs a warning if set; configure "dispatcher" and "experts" explicitly instead. fake_balanced_gate: If True, replace the learned Gate with FakeBalancedGate @@ -512,6 +516,11 @@ class BackendConfig: # blocking mode needs to learn the permuted row count (and the per-layer barrier it implies); an # overflow of the capacity trips a device-side assert instead of silently truncating. None = blocking. dispatcher_capacity_factor: float | None = None + # HybridEP only: every EP rank dispatches the same number of rows (fixed-shape batches, which is + # every batch that is not variable-length / in-batch packed), so the per-dispatch EP-group MAX + # all-reduce + int() host sync that derives the pad size is skipped and the local count is only + # aligned. Leave False for variable-length inputs: unequal counts abort the HybridEP collective. + dispatcher_equal_token_counts: bool = False mok: MoKBackendConfig = field(default_factory=MoKBackendConfig) enable_deepep: bool | None = None # Removed: ignored with a warning; set dispatcher/experts explicitly fake_balanced_gate: bool = False diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 6e6e5644e1..8009d09638 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -965,6 +965,9 @@ def __init__( self.dispatcher_capacity_factor = ( getattr(backend, "dispatcher_capacity_factor", None) if backend is not None else None ) + self.dispatcher_equal_token_counts = ( + bool(getattr(backend, "dispatcher_equal_token_counts", False)) if backend is not None else False + ) # Allocate projection tensor - size depends on whether activation is gated # Gated (SwiGLU, Quick-GEGLU): [n_experts, dim, 2*inter_dim] @@ -999,6 +1002,7 @@ def init_token_dispatcher(self, ep_mesh: DeviceMesh): moe_share_token_dispatcher=self.dispatcher_share_token_dispatcher, moe_deepep_async_dispatch=self.dispatcher_async_dispatch, moe_hybridep_capacity_factor=self.dispatcher_capacity_factor, + moe_hybridep_equal_token_counts=self.dispatcher_equal_token_counts, moe_benchmark_static_routing=self.static_routing, ) diff --git a/nemo_automodel/components/moe/megatron/token_dispatcher.py b/nemo_automodel/components/moe/megatron/token_dispatcher.py index b433d71cd3..926d2a6ded 100644 --- a/nemo_automodel/components/moe/megatron/token_dispatcher.py +++ b/nemo_automodel/components/moe/megatron/token_dispatcher.py @@ -398,6 +398,7 @@ def __init__( moe_hybridep_num_sms: int = 24, benchmark_static_routing: bool = False, moe_hybridep_capacity_factor: float | None = None, + moe_hybridep_equal_token_counts: bool = False, ): self.group = group self.num_local_experts = num_local_experts @@ -409,6 +410,8 @@ def __init__( # pass this many rows as num_permuted_tokens and run non-blocking (see dispatch()). self.hybridep_capacity_factor = moe_hybridep_capacity_factor self._hybridep_capacity: int | None = None + # Equal row counts across the EP group: the pad size is the aligned local count, no collective. + self.equal_token_counts = moe_hybridep_equal_token_counts # Benchmark-only (TokenDispatcherConfig.moe_benchmark_static_routing): # persist num_permuted_tokens across dispatches, see dispatch()/reset. self.benchmark_static_routing = benchmark_static_routing @@ -498,7 +501,11 @@ def dispatch( if torch.distributed.is_initialized() and torch.distributed.get_world_size(self.group) > 1: num_tokens = hidden_states.shape[0] pin = self.benchmark_static_routing and _STATIC_ROUTING_PAD_PIN - if pin and self._static_target_tokens is not None and self._static_target_tokens >= num_tokens: + if self.equal_token_counts: + # Every rank holds the same row count by construction (fixed-shape batches): the + # group-wide maximum is the local count, so only align it. No collective, no host sync. + target_tokens = -(-num_tokens // _HYBRIDEP_TOKEN_ALIGNMENT) * _HYBRIDEP_TOKEN_ALIGNMENT + elif pin and self._static_target_tokens is not None and self._static_target_tokens >= num_tokens: target_tokens = self._static_target_tokens else: group_max = torch.tensor(num_tokens, device=hidden_states.device) @@ -644,6 +651,8 @@ class TokenDispatcherConfig: # HybridEP capacity mode, see BackendConfig.dispatcher_capacity_factor moe_hybridep_capacity_factor: float | None = None + # HybridEP: skip the per-dispatch pad-size all-reduce, see BackendConfig.dispatcher_equal_token_counts + moe_hybridep_equal_token_counts: bool = False """Backend for the flex token dispatcher. Options: 'deepep', 'hybridep', or 'uccl_ep'.""" moe_deepep_num_sms: int = 20 @@ -767,6 +776,7 @@ def __init__( moe_hybridep_num_sms=self.config.moe_hybridep_num_sms, benchmark_static_routing=self.config.moe_benchmark_static_routing, moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor, + moe_hybridep_equal_token_counts=self.config.moe_hybridep_equal_token_counts, ) self._comm_manager = MoEFlexTokenDispatcher.shared_hybridep_manager else: @@ -779,6 +789,7 @@ def __init__( moe_hybridep_num_sms=self.config.moe_hybridep_num_sms, benchmark_static_routing=self.config.moe_benchmark_static_routing, moe_hybridep_capacity_factor=self.config.moe_hybridep_capacity_factor, + moe_hybridep_equal_token_counts=self.config.moe_hybridep_equal_token_counts, ) self.hybridep_metadata_processor = _HybridEPMetadataProcessor( num_experts=self.tp_size * self.config.num_moe_experts, diff --git a/tests/unit_tests/moe/test_backend_config.py b/tests/unit_tests/moe/test_backend_config.py index 30278bab91..8ce16c6fbf 100644 --- a/tests/unit_tests/moe/test_backend_config.py +++ b/tests/unit_tests/moe/test_backend_config.py @@ -499,3 +499,10 @@ def test_dispatcher_capacity_factor_defaults_off(): assert BackendConfig().dispatcher_capacity_factor is None assert BackendConfig(dispatcher_capacity_factor=1.5).dispatcher_capacity_factor == 1.5 + + +def test_dispatcher_equal_token_counts_defaults_off(): + from nemo_automodel.components.models.common.utils import BackendConfig + + assert BackendConfig().dispatcher_equal_token_counts is False + assert BackendConfig(dispatcher_equal_token_counts=True).dispatcher_equal_token_counts is True diff --git a/tests/unit_tests/moe/test_token_dispatcher.py b/tests/unit_tests/moe/test_token_dispatcher.py index 7a8bbbf176..c9aeb1746f 100644 --- a/tests/unit_tests/moe/test_token_dispatcher.py +++ b/tests/unit_tests/moe/test_token_dispatcher.py @@ -318,3 +318,51 @@ def test_static_routing_ignores_the_factor(self, monkeypatch): def test_no_factor_keeps_the_blocking_path(self, monkeypatch): m, passed, asserts = self._run(monkeypatch, None, [8, 8], tpe_rows=[5, 5]) assert passed == [None, None] and asserts == [] + + +class TestHybridEPEqualTokenCounts: + """BackendConfig.dispatcher_equal_token_counts: the pad size is the aligned local row count, with no + EP-group all-reduce and no host sync; unset keeps the per-dispatch MAX all-reduce.""" + + def _run(self, monkeypatch, equal, num_tokens_seq, group_max): + import nemo_automodel.components.moe.megatron.token_dispatcher as td + + with patch( + "nemo_automodel.components.moe.megatron.token_dispatcher.hybrid_ep_dispatch", new=lambda *a, **kw: None + ): + m = _HybridEPManager( + group=None, + num_local_experts=2, + num_experts=8, + router_topk=2, + benchmark_static_routing=False, + moe_hybridep_equal_token_counts=equal, + ) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + calls, sizes = [], [] + + def fake_all_reduce(tensor, op=None, group=None): + calls.append(int(tensor)) + tensor.fill_(group_max) + + def fake_dispatch(x, routing_map, probs, **kwargs): + sizes.append(x.shape[0]) + return x, probs, None, routing_map.sum(dim=0), "handle" + + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + monkeypatch.setattr(td, "hybrid_ep_dispatch", fake_dispatch) + for n in num_tokens_seq: + m.routing_map = torch.ones(n, 8, dtype=torch.bool) + m.token_probs = torch.full((n, 8), 0.125) + m.dispatch(torch.randn(n, 4)) + return calls, sizes + + def test_equal_counts_skip_the_collective_and_only_align(self, monkeypatch): + calls, sizes = self._run(monkeypatch, True, [6, 8, 9], group_max=99) + assert calls == [] and sizes == [8, 8, 12] + + def test_default_keeps_the_per_dispatch_all_reduce(self, monkeypatch): + calls, sizes = self._run(monkeypatch, False, [6, 6], group_max=6) + assert calls == [6, 6] and sizes == [8, 8] From c0cf6854f7bac1c0e1e04df3488c2bd7aa1b0ff8 Mon Sep 17 00:00:00 2001 From: Yisong Li Date: Fri, 18 Sep 2026 05:22:55 +0800 Subject: [PATCH 3/3] perf(moe): record the HybridEP capacity as the checkpoint-replay extent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #3684 replays the checkpoint-forward HybridEP layout on recompute and sizes the replayed dispatch to `int(tokens_per_expert.sum().item())`. Under capacity mode the forward output is sized to the capacity, not to this dispatch's token count, so the replay must reuse the same integer or the recomputed activation no longer matches the saved one. When the forward already ran with a host-side extent (capacity mode, static-routing pin), record that integer; `finalize` then only reduces the entries of blocking dispatches, so capacity mode keeps its forward free of the per-layer device-to-host copy that the reduction would add back. Tests: tests/unit_tests/moe/test_fused_a2a.py — a recorded host extent is kept without reducing tokens_per_expert; a checkpointed dispatch with a capacity extent replays with that extent. Signed-off-by: Yisong Li --- .../components/moe/megatron/fused_a2a.py | 13 ++++++-- tests/unit_tests/moe/test_fused_a2a.py | 33 +++++++++++++++++-- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/nemo_automodel/components/moe/megatron/fused_a2a.py b/nemo_automodel/components/moe/megatron/fused_a2a.py index 81e2f6e6a8..01cc468af2 100644 --- a/nemo_automodel/components/moe/megatron/fused_a2a.py +++ b/nemo_automodel/components/moe/megatron/fused_a2a.py @@ -140,8 +140,15 @@ def __init__(self) -> None: self._cursor = 0 self.replay_misses = 0 - def record(self, handle, tokens_per_expert) -> None: - self._records.append([handle, tokens_per_expert, None]) + def record(self, handle, tokens_per_expert, num_permuted_tokens=None) -> None: + """Log one checkpoint-forward dispatch. + + When the forward already ran with a host-side extent (capacity mode, static-routing pin), + that integer is the extent its output was sized to and the one the replay must reuse; + recording it also spares ``finalize`` the device-to-host reduction for that dispatch. + """ + extent = num_permuted_tokens if isinstance(num_permuted_tokens, int) else None + self._records.append([handle, tokens_per_expert, extent]) def finalize(self) -> None: """Cache each layout extent after the checkpoint-forward op context exits.""" @@ -643,7 +650,7 @@ def forward( if recorder is not None and _hybridep_dispatch_replay_state.mode == "record": # Keep only the reusable layout and its output extent. Recomputed # activations and probabilities are still redispatched through it. - recorder.record(handle, tokens_per_expert) + recorder.record(handle, tokens_per_expert, num_permuted_tokens) return ( dispatched_hidden, dispatched_probs, diff --git a/tests/unit_tests/moe/test_fused_a2a.py b/tests/unit_tests/moe/test_fused_a2a.py index 43d5a78a17..501762be6b 100644 --- a/tests/unit_tests/moe/test_fused_a2a.py +++ b/tests/unit_tests/moe/test_fused_a2a.py @@ -112,7 +112,7 @@ def combine_with_unpermute(self, *, hidden, probs=None, **kwargs): return combined_hidden, combined_probs -def _run_checkpointed_hybridep(context_fn): +def _run_checkpointed_hybridep(context_fn, num_permuted_tokens=None): x = torch.randn(4, 3, requires_grad=True) routing_map = torch.ones(4, 2, dtype=torch.bool) probs = torch.full((4, 2), 0.5, requires_grad=True) @@ -126,7 +126,7 @@ def block(hidden, token_probs): 1, 24, 24, - None, + num_permuted_tokens, None, ) return dispatched_hidden.sin().sum() + dispatched_probs.square().sum() @@ -181,3 +181,32 @@ def save_replay_sensitive_ops(ctx, func, *args, **kwargs): assert buffer.cached_dispatches == 1 assert buffer.replayed_num_permuted_tokens == 5 assert isinstance(buffer.replayed_num_permuted_tokens, int) + + +def test_hybridep_recorder_keeps_a_host_extent_without_reducing_tokens_per_expert(): + recorder = fused_a2a.HybridEPDispatchReplayRecorder() + tokens_per_expert = mock.MagicMock(spec=torch.Tensor) + # capacity mode / static pin: the forward already ran with a host-side extent + recorder.record("layout", tokens_per_expert, 24) + # blocking dispatch: the extent comes from the reduction, after the checkpoint context exits + recorder.record("layout", torch.tensor([2, 3])) + recorder.finalize() + + assert recorder.take() == ["layout", tokens_per_expert, 24] + assert recorder.take()[2] == 5 + tokens_per_expert.sum.assert_not_called() + + +def test_hybridep_checkpoint_replay_reuses_the_forward_capacity_extent(): + from nemo_automodel.components.moe.parallelizer import _replay_hybridep_dispatch_on_recompute + + buffer = _DriftingHybridEPBuffer() + fused_a2a._hybrid_ep_buffer = buffer + context_fn = _replay_hybridep_dispatch_on_recompute(lambda: (nullcontext(), nullcontext())) + + _run_checkpointed_hybridep(context_fn, num_permuted_tokens=24) + + assert buffer.full_dispatches == 1 + assert buffer.cached_dispatches == 1 + # the recompute output must be sized like the forward's (capacity rows), not to this dispatch's token count + assert buffer.replayed_num_permuted_tokens == 24