diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index 5308d6f45b..c8537a84d0 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/nemo_automodel/components/models/common/utils.py @@ -442,6 +442,15 @@ 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. + 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 @@ -501,6 +510,17 @@ 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 + # 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 2ebea56d40..8009d09638 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -959,6 +959,15 @@ 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 + ) + 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] @@ -992,6 +1001,8 @@ 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_hybridep_equal_token_counts=self.dispatcher_equal_token_counts, moe_benchmark_static_routing=self.static_routing, ) @@ -1072,7 +1083,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..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, @@ -686,6 +693,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..926d2a6ded 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,8 @@ def __init__( permute_fusion: bool = False, 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 @@ -392,6 +406,12 @@ 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 + # 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 @@ -481,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) @@ -502,6 +526,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 +545,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 +647,12 @@ 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 + # 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 @@ -705,6 +775,8 @@ 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, + moe_hybridep_equal_token_counts=self.config.moe_hybridep_equal_token_counts, ) self._comm_manager = MoEFlexTokenDispatcher.shared_hybridep_manager else: @@ -716,6 +788,8 @@ 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, + 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 913ab00813..8ce16c6fbf 100644 --- a/tests/unit_tests/moe/test_backend_config.py +++ b/tests/unit_tests/moe/test_backend_config.py @@ -492,3 +492,17 @@ 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 + + +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_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 diff --git a/tests/unit_tests/moe/test_token_dispatcher.py b/tests/unit_tests/moe/test_token_dispatcher.py index 0fef97971f..c9aeb1746f 100644 --- a/tests/unit_tests/moe/test_token_dispatcher.py +++ b/tests/unit_tests/moe/test_token_dispatcher.py @@ -258,3 +258,111 @@ 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 == [] + + +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]