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 <noreply@anthropic.com>

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()

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)