From f67fab6bdd7885d28a065264e60b8218105f8e19 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 11:05:25 +0800 Subject: [PATCH] Neuron MoE: two-level all-reduce (device, then across devices) (#27) Microbenchmark at 64 ranks, [1, 5120] fp32 in a compiled graph: world all-reduce 3.37 ms, 8 consecutive ranks 0.41 ms, 4-core device group then a 16-device group (one core each) 2.01 ms; a strided 8-device group is rejected by the runtime. The replicated-token MoE all-reduce (one a layer) now takes the two-level form (DEEPINFER_NEURON_HIER_ALLREDUCE=0 restores one world all-reduce). Host tiny cp=8 matches golden. Co-Authored-By: Claude Opus 5.5 --- .../model_executor/models/v4_modules/mlp.py | 4 ++ deepinfer/neuron_ops/static_moe.py | 38 +++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/deepinfer/model_executor/models/v4_modules/mlp.py b/deepinfer/model_executor/models/v4_modules/mlp.py index 4b233a0..cee0641 100644 --- a/deepinfer/model_executor/models/v4_modules/mlp.py +++ b/deepinfer/model_executor/models/v4_modules/mlp.py @@ -196,6 +196,10 @@ class DeepseekMoE(nn.Module): self.register_buffer( "_neuron_rank_t", torch.tensor([PG.expert_parallel_group().rank()], dtype=torch.int32), persistent=False ) + from deepinfer.neuron_ops.static_moe import hierarchical_groups + + # Two-level all-reduce for the replicated-token MoE (static_moe): within a device, then across devices. + self._neuron_hier_ar = hierarchical_groups(PG.expert_parallel_group()) self.ep_buffer: EPBuffer | None = None self.sym_buffer: deep_gemm.SymmBuffer | None = None self.ep_group = PG.expert_parallel_group() diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 393f4a0..85de881 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -74,7 +74,37 @@ def tokens_replicated(moe: DeepseekMoE) -> bool: return moe.ep_size > 1 and PG.context_parallel_group().size() == moe.ep_size -def _all_reduce(t: torch.Tensor, group: torch.distributed.ProcessGroup) -> torch.Tensor: +#: Cores of one Trainium2 device (LNC=2): ranks r // 4 share a device. +_DEVICE_CORES = 4 +_HIER: dict = {} + + +def hierarchical_groups(group: torch.distributed.ProcessGroup): # noqa: ANN201 + """(intra-device group, cross-device group) for ``group`` = the world, or None. Created once, on every rank in the + same order (call at module init, not inside a graph). Measured at 64 ranks, [1, 5120] fp32: one world all-reduce + 3.37 ms, 4-core device all-reduce then a 16-device all-reduce (one core each) 2.01 ms (issue #27).""" + import torch.distributed as dist + + from deepinfer.model_executor.parallel_utils.parallel_state import _split_group + + world = dist.get_world_size() + if os.environ.get("DEEPINFER_NEURON_HIER_ALLREDUCE", "1") != "1" or group.size() != world or world % _DEVICE_CORES or world <= _DEVICE_CORES: + return None + if world not in _HIER: + from deepinfer import PLATFORM + + backend = PLATFORM.distributed_backend + local = _split_group([list(range(r, r + _DEVICE_CORES)) for r in range(0, world, _DEVICE_CORES)], backend=backend) + cross = _split_group([list(range(c, world, _DEVICE_CORES)) for c in range(_DEVICE_CORES)], backend=backend) + _HIER[world] = (local, cross) + return _HIER[world] + + +def _all_reduce(t: torch.Tensor, group: torch.distributed.ProcessGroup, hier=None) -> torch.Tensor: # noqa: ANN001 + if hier is not None: + local, cross = hier + t = funcol.wait_tensor(funcol.all_reduce(t, "sum", local)) + return funcol.wait_tensor(funcol.all_reduce(t, "sum", cross)) return funcol.wait_tensor(funcol.all_reduce(t, "sum", group)) @@ -143,7 +173,8 @@ def static_moe_forward( _expert_gemm(down_in, down_w, down, e, moe.use_fp8_act, recipe_a, recipe_b) acc = acc + down.float() - out = _all_reduce(acc, group) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] + hier = getattr(moe, "_neuron_hier_ar", None) + out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] shared = moe.shared_forward(hidden_states) if shared is not None: out = out + shared.float() @@ -235,7 +266,8 @@ def _nki_moe_forward( ) acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float() regions.split() - out = _all_reduce(acc, group) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] + hier = getattr(moe, "_neuron_hier_ar", None) + out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] regions.split() if moe.has_shared: shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)