Repository navigation
feat(wideep): multi-node wide expert parallelism on top of #1603 - #1942
Draft
JiaoliangYu wants to merge 12 commits into
Draft
JiaoliangYu wants to merge 12 commits into
JiaoliangYu wants to merge 12 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
JiaoliangYu
marked this pull request as ready for review
August 18, 2026 08:19
added 11 commits
August 25, 2026 13:54
PR ROCm#1603 parameterises multi-node DP as data_parallel_size_local plus a data_parallel_rank offset, not nnodes/node_rank. Add from_parallel_config so the two do not become competing CLI surfaces for one concept, and stop storing raw -tp/-dp: CoreManager's fold makes them unrecoverable, and a second copy of an authoritative field is how the two drift apart. Fix ep_size and gpu_per_node outside DP-attention. FusedMoEParallelConfig only flattens every rank into one EP group under DP-attention; otherwise ep_size stays at tp_size and there are dp_size independent groups. The view reported global_dp_size unconditionally, so -tp 8 -dp 2 read as ep=2/gpu_per_node=2 instead of 8/16. The old test missed it by asserting the multinode identity, which only checks the view against itself. Reject ragged node splits: EPLB's hierarchical placement asserts num_gpus % num_nodes == 0, so a split that cannot be expressed as equal slices has to fail at construction rather than inside a rebalance.
`rebalance_experts(enable_hierarchical=True, num_nodes>1)` has no coverage because EPLBManager._nnodes is hardcoded to 1, so every rebalance so far took the flat branch. WideEP is about to depend on it. Shaped like the real target: DSv4-Pro 384 logical experts in 8 groups, EP16 over 2 nodes, num_redundant = ep_size. Asserts what the node-aware branch is for -- a group's experts stay on one node -- plus the invariants any placement owes its consumers: every logical expert has a slot, and logcnt agrees with p2l (routing reads p2l while the dispatch map is built from logcnt, so a mismatch drops replicas silently). test_flat_branch_does_split_groups pins the discriminating power: measured, the flat branch splits 16/16 group-layer pairs across nodes and the hierarchical branch splits 0/16. Without it a hierarchical branch that degraded into the flat one would still pass the co-location test. Verified on g64 (rocm/atom-dev:latest, torch 2.10, CPU only): 15 passed. The placement code turned out to be correct as written -- no fix needed.
EPLBManager._nnodes was hardcoded to 1, so hierarchical placement never ran and expert groups were free to straddle nodes. PR ROCm#1603 made the derivation possible: data_parallel_size_local now defaults to data_parallel_size instead of 1, so the ratio reads 1 on a single node rather than dp_size. Wired outside the best-effort eplb_config block on purpose. A wrong node count does not fail — it spreads groups across nodes and makes every rebalance pay inter-node traffic for the rest of the run, which is only visible as unexplained latency. Cross-check ep_size against the topology for the same reason. --fake-eplb is the one case the DP fields cannot answer: it shrinks data_parallel_size_local to the visible device count while data_parallel_size keeps the simulated width, which is byte-identical to node 0 of a real two-node slice. from_parallel_config now refuses that input instead of reporting a node count that does not exist, and the caller handles the simulated case explicitly (nnodes=1, today's behaviour, so no change to fake-eplb benchmarks).
A rearrange that looks cheap on one node is not cheap across two: every expert the plan moves off-node is a weight tensor over RDMA, and it lands on the critical path because migration shares the forward stream. count_migration_transfers reports what a plan would actually move, split at the node boundary. It is exact, not an estimate -- it applies the same recv rule and source selection as the planner -- but derives it once from the global maps rather than replaying per-rank planning world_size times. test_migration_cost_matches_planner pins the two together, since the shortcut is only safe while it agrees. Every rank sees the same maps, so every rank reaches the same verdict with no collective. That is what makes the skip safe: half the ranks migrating while the other half skip would leave placement inconsistent. The cap defaults to off. Single-node moves zero cross-node bytes, so a default would only ever bind multi-node runs, and no link measurement exists yet to set it from -- a guessed threshold that silently stops EPLB rebalancing is worse than no cap. Counting and logging run either way, with cumulative skipped rounds and bytes, so a guard that is suppressing every rebalance reads as suppression rather than as an idle EPLB. Tests are written but UNRUN: no machine was reachable (g64 needs a slurm job, mia1's login node drops sessions after auth). Both eplb test modules now import through the aiter stub so a CPU box can run them for real rather than skipping green.
Multi-node narrows the configuration surface rather than adding an orthogonal dimension: most parallelism knobs reshape the rank space, and none of these combinations have run across a node boundary. Deny by default, widen as each is verified. What this avoids is not a crash but a run that starts, serves, and is quietly wrong -- or hangs with nothing in the log. Rejected: no DP-attention, no expert parallelism, PP > 1, PCP/DCP > 1, --moe-backend mega (FlyDSL raises NotImplementedError for internode), --enable-rapidserve (it selects DisaggCoreManager, a second rank layout), and the plugin path (its own multi-node world, plus plugin/config.py assigns data_parallel_size_local the global size). TBO warns instead: it should work once kernel selection is right, but is not verified. Each message carries why and what to do. A validator that only reports "invalid" moves the work to whoever reads the traceback, and for the hang-shaped failures here that is the expensive path. Two rendezvous rules ROCm#1603 left open: - An unset --data-parallel-base-port is now rejected. Its default is get_open_port(), evaluated per process, so nodes silently pick different ports and wait on addresses nobody bound. - Loopback master IP warns rather than rejects. Gate 1 runs both logical nodes on one box, where loopback is correct; rejecting it would block our own single-machine verification path. Also adds the divisibility check ROCm#1603 is missing: data_parallel_size % data_parallel_size_local. EPLB's hierarchical placement asserts num_gpus % num_nodes == 0, so a ragged split passes startup and then fails inside a rebalance thousands of steps later. Tests UNRUN -- no machine reachable this session.
gpu_per_node reaches MoRI as the node width, and InterNodeV1 computes its RDMA/XGMI peer targets from it. A wrong value does not raise -- it addresses the wrong ranks. Before ROCm#1603 it was 1 on every run and nothing noticed, because the intra-node kernel never reads it. So check it where it is built, against the topology: ep_size has to partition into exactly nnodes groups of gpu_per_node. Only under DP-attention, which is the case where every rank joins one EP group; otherwise EP spans a TP group and the identity does not hold. MoRI reads its heap settings when the heap is created, inside the first collective, so mori_env applies them before init_dist_env. It sets MORI_SHMEM_HEAP_SIZE (InterNodeV1 stages through five buffers where intra-node uses one) and MORI_EP_LAUNCH_CONFIG_MODE, and never overrides what an operator exported -- they chose against a specific fabric and this knows less than they do. MORI_SOCKET_IFNAME=lo is rejected across nodes. A single-node run needs that value, which is exactly how it survives into a multi-node launcher, and wrong it hangs the first collective with nothing in the log. node_count() is now the single place that answers "how many nodes", including the one case the DP fields cannot: --fake-eplb runs one box pretending to be wider, which reads identically to node 0 of a real split. EPLB now goes through it too. 41 CPU tests pass (topology + mori_env, neither needs torch). The torch-dependent EPLB tests remain unrun -- no machine reachable.
Diagnosing a multi-node misconfiguration means comparing ranks across nodes, which today means reading the same fields out of several logs in different places. This puts them on one line behind a fixed prefix, so the question "is this rank where it thinks it is" is one grep. Emitted only when nnodes > 1: single-node adds nothing and the existing ModelRunner line already covers it. Everything on it is known at worker init; nothing is computed for the message. The other M-OBS self-checks already landed with the modules they belong to -- the loopback interface check in mori_env, the ep_size == gpu_per_node * nnodes assertion in FusedMoEParallelConfig.make, and EPLB's node count in its init log.
Three barriers wait forever today. On one node that is right -- the peers are sibling processes and their death is visible other ways. Across nodes the common failure is that the other node never started or was pointed at a different master, and all three present the same way: a wait that never ends and never logs. That is the hardest shape of failure to attribute, so each gets a bound and something to read. 1. The gloo DP group. stateless_init_torch_distributed_process_group takes a timeout and threads it through rendezvous, the store and the group handshake. None keeps the backend default, so single-node is byte-identical. 2. The NCCL world. aiter calls init_process_group with no timeout, and we do not own aiter -- but it skips creation when a world exists, so ModelRunner builds one first with the same rank arithmetic. Opt-in behind multi-node AND an explicit --dp-sync-timeout: a run that did not ask for a bound should not silently take an init path no single-node test covers. It also drops aiter's HIP_VISIBLE_DEVICES fallback, which sets range(world_size) -- wrong past one node, where a 16-rank world is not 16 local devices. Warns when unset. 3. CoreManager waiting for READY. Now bounded, and the timeout names the engines that never reported grouped by node. Which node is missing is the whole diagnosis, and the coordinator is the only process that can see it: the engines that did come up are blocked on the same barrier. get_next_dp_init_port advanced by data_parallel_rank, so rank 0 stayed on the base port while every other rank moved a different distance -- the second group was formed on a different port per rank and could never complete. All ranks now walk the same sequence. M-KV: the block-count MIN all-reduce now also covers multi-node. Under DP-attention unequal counts corrupt nothing -- the tighter rank preempts and recomputes sooner while the rest idle -- but nodes rarely have byte-equal free memory and ATOM logs no preemption, so the only trace is throughput that is lower than it should be for no visible reason. Logs local vs group either way. Tests UNRUN except the 44 that need no torch.
weight_iterator splits the prefetch by node because page cache is per node, and it reads LOCAL_RANK/LOCAL_WORLD_SIZE to find its place. Nobody sets them. Its fallback triggers as soon as the world is wider than the visible devices -- which is every multi-node run -- and hands back (0, 1): every rank prefetches the whole checkpoint. Safe, since the page cache absorbs the duplicates, but it gives up the split entirely on the runs that need it most. Derived from the device rank and gpu_per_node, so it stays right whether or not DP-attention folded TP away. setdefault, not assignment: a launcher that already exported them knows more about the layout. Multi-node only. On one node the existing fallback already resolves correctly through world <= visible.
JiaoliangYu
changed the title
JiaoliangYu
marked this pull request as draft
August 25, 2026 09:46
JiaoliangYu
force-pushed
the
wideep
branch
2 times, most recently
from
September 10, 2026 07:49
23d63e3 to
b8f9341
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Multi-node wide EP on top of #1603, which landed the process/rank/socket layer.
Everything here is additive: no change to the coordinator architecture, and every
behavioural change is gated on
nnodes > 1.model_engine/topology.py, new) — derived view overParallelConfig:nnodes,gpu_per_node,ep_size, rank/device arithmetic.node_count()is thesingle place that answers "how many nodes", including
--fake-eplb, whichpost-fold reads identically to node 0 of a real two-node split.
(no DP-attention, no EP, PP/PCP/DCP > 1, mega backend, rapidserve, plugin path),
each with a reason and a fix. Plus
data_parallel_size % data_parallel_size_local,which EPLB's hierarchical placement requires and multi-node dp support #1603 does not check.
gpu_per_nodeis now asserted where it is built. A wrong value doesnot raise, it addresses the wrong ranks; it was 1 on every run before multi-node dp support #1603 and
nothing noticed.
mori_envsets the heap/launch vars beforeinit_dist_envandrejects
MORI_SOCKET_IFNAME=loacross nodes._nnodeswas hardcoded to 1, so hierarchical placement never ran.Adds an exact cross-node migration count (pinned against the planner by test) and
an optional stall cap, off by default until the link is measured.
waited forever. Across nodes that is how "the other node never started" presents.
The READY message names the missing engines by node. The NCCL one is opt-in behind
an explicit
--dp-sync-timeout.LOCAL_RANK/LOCAL_WORLD_SIZEforweight_iterator's node-local prefetch (nothing set them, so every rank prefetchedthe whole checkpoint); one
[wideep]summary line per worker.One change is not gated on
nnodes > 1:get_next_dp_init_portadvanced bydata_parallel_rank, so rank 0 stayed on the base port while every other rank moveda different distance — a second group could never form. Inert if the function is
called once, which is what single-node does today.
Draft: 4 of 6 test suites are unrun (no machine was reachable). Passing:
test_wideep_topology.py(31),test_mori_env.py(13). Unrun: the EPLB migration,config-validation and timeout suites;
test_eplb_hierarchical.pypassed on a GPU boxbefore the rebase. Please don't merge on the strength of the test count.
ruffcomparedper file against the merge base: zero new findings.
Not here: M0-A (the measurement this turns on — whether the freed expert weights beat
~2.7x all-to-all latency, against two independent EP8 instances as the baseline),
Gate 1/Gate 2, and DP idle-loop backoff.