Decouple expert parallelism from context parallelism with MoE parallel folding

[haok1402]

Why

A single (pp, dp, cp, ep) mesh makes world_size = pp * dp * cp * ep, so CP and EP multiply and compete for the same ranks. Data parallelism is whatever survives that product, which collapses exactly when both long context and wide expert sharding are wanted at once. On 32 ranks with pp=1, asking for cp=4 and ep=8 leaves dp = 32 / (1 * 4 * 8) = 1, so a mesh that looks reasonable on paper trains with no data parallelism at all.

The coupling also reaches the data map. Reading a distinct chunk per rank was keyed to dp * ep, which makes expert sharding decide which tokens a rank owns. Those are separate concerns: EP names where experts live, and the tokens routed to an expert arrive there through the all-to-all rather than by being read there.

Folding factors the same rank space twice instead of once. The cp=4, ep=8 case above becomes dp=8 with expt_dp=4, keeping full data parallelism while the attention and expert sides each get the degree they need.

What changes

Attention and the experts now get independent parallelism over the same ranks, following MoE parallel folding section 3.2.

setup_device_mesh builds two device meshes instead of one. attn_mesh is (pp, dp, cp) and expt_mesh is (pp, dp, ep), two factorizations of the same world_size // pp ranks of one pipeline stage, so dp * cp == ep * expt_dp == world_size // pp. PP is the one axis both views share and it stays outermost in each, which is what makes the two factorizations describe the same rank space.

Three properties follow. cp and ep each need only divide the stage size rather than each other, so EP may span CP or sit inside it. Both the CP group and the EP group remain contiguous rank blocks, so each stays on NVLink independently. And dp_rank alone decides which data a rank loads, while ep_rank names only the experts a rank hosts.

FSDP shards the routed expert weights over expt_mesh["dp"] and everything else over the flattened attn_mesh["dp", "cp"], with a plain-sum gradient reduction on both, so the single divide by the global token count in the training step is the only normalization. DistributedCfg.hsdp_replica above 1 splits both groups into that many replicas, sharding within one and all-reducing across them.

Correctness validation

Four meshes across two models, every one a PASS. The method is the validate-correctness skill: run the base branch twice and the feature branch once, then read the base-vs-feature difference as a ratio against the base-vs-base run-to-run floor. Thresholds are 3x warn and 5x fail. Every run is 32 steps with real routing, benchmark = False, warmup from 1e-6 to 1e-5.

model	mesh	metric	envelope	signal	ratio	verdict
qwen3-30b-a3b	pp4-dp8-cp1-ep1-seq2048-bf16	cross-entropy-loss	0.0148	0.0074	0.50	PASS
		load-balance-loss	0.0355	0.0208	0.59	PASS
		step-1 forward	0.0000	0.0000
qwen3-30b-a3b	pp1-dp8-cp4-ep8-seq2048-bf16	cross-entropy-loss	0.0062	0.0056	0.90	PASS
		load-balance-loss	0.0127	0.0137	1.08	PASS
		step-1 forward	0.0001	0.0000
deepseek-v2-lite	pp2-dp8-cp2-ep8-seq4096-bf16	cross-entropy-loss	0.0568	0.0581	1.02	PASS
		load-balance-loss	0.0063	0.0058	0.92	PASS
		step-1 forward	0.0035	0.0034
deepseek-v2-lite	pp2-dp4-cp4-ep2-seq4096-bf16	cross-entropy-loss	0.0074	0.0045	0.60	PASS
		load-balance-loss	0.0043	0.0038	0.89	PASS
		step-1 forward	0.0001	0.0000

What each mesh covers

pp4-cp1-ep1 is the arm where the branch is mathematically identical to the base, so its PASS isolates the mesh relabeling from every behavioral change. Its step-1 forward delta of exactly 0.0000 also establishes that the forward is bit-reproducible while the backward is not, which is the premise the whole envelope method rests on.

pp1-cp4-ep8 is the first folded mesh, with EP spanning CP. The gradient-scale correction changes expert gradients here by a factor of ep_size and still does not move the loss, because the fixture optimizes with AdamW, whose update g / (sqrt(v) + eps) is invariant to a uniform per-parameter rescaling of the gradient. Where the correction does bite is the global gradient norm feeding clip_grad_norm_.

pp2-cp2-ep8 is the first mesh with PP, CP and EP all greater than one at the same time, and the first to exercise MLA-aware zigzag ring attention, which no GQA run reaches. The two branches also factor it differently: the base gets dp=1 and reads data over dp * ep, the feature gets dp=8 and reads over dp alone, so equal loss curves across that gap exercise the data-map redesign rather than bypass it.

pp2-cp4-ep2 inverts the nesting the other three share. With cp=4 and ep=2 the EP group sits inside a single CP group and the expert gradient reduction spans across CP boundaries. It also carries expt_dp=8, four times the largest value anywhere else, so a residual error in the plain-sum reduction would scale up here rather than cancel. The signal lands below the envelope and the feature arm ends 0.0005 from base0 while the two base arms sit 0.0020 apart.

Provenance

Every arm below is listed with the git commit that wandb recorded for the process that produced it. 91018ec is the base branch and 3203f94 is this branch. The check matters because the arms are switched with git checkout, so a run whose srun is queued when a checkout lands executes the other branch while still writing the log name it was launched under. Each run was also gated on a log holding steps 1 through 32 exactly once, in order, before it was accepted.

Project PithTrain/pr-93.

Group correctness/qwen3-30b-a3b:

mesh	arm	run	commit
pp4-dp8-cp1-ep1-seq2048-bf16	base0	tsvqavri	91018ec
	base1	ihsfztyz	91018ec
	feat0	yulfprx3	3203f94
pp1-dp8-cp4-ep8-seq2048-bf16	base0	b32acghj	91018ec
	base1	vz0o3f7f	91018ec
	feat0	54ww38s8	3203f94

Group correctness/deepseek-v2-lite:

mesh	arm	run	commit
pp2-dp8-cp2-ep8-seq4096-bf16	base0	mzou5jji	91018ec
	base1	p056deqr	91018ec
	feat0	7uyz24kw	3203f94
pp2-dp4-cp4-ep2-seq4096-bf16	base0	01nians3	91018ec
	base1	dzpyniwb	91018ec
	feat0	j46jrtdw	3203f94

Mesh selection

Both branches have to be able to express a mesh for it to be testable at all. The base branch derives dp = world / (pp * cp * ep), so pp * cp * ep must divide the world size. This branch needs only pp * cp to divide the world size with ep dividing the stage. On 32 ranks that intersection rules out pp2-cp4-ep8, which is why the deepseek meshes are pp2-cp2-ep8 and pp2-cp4-ep2.

One property has no base arm to compare against: hsdp_replica above 1 does not exist on the base branch. Validating it needs a different framing, hsdp_replica=1 against 2 on this branch, which tests self-consistency rather than equivalence to base.

Performance

No measurable throughput difference. Peak memory is identical on every mesh.

Method: validate-performance, force-balanced routing, 8 steps, medians over steps 2 to 8, project PithTrain/pr-93 group performance/<model>. Each mesh has four runs. base and feat compare the branches. aa0 and aa1 are both the base branch, run back to back, and measure what the protocol reports when the true difference is zero.

model	mesh	base	feat	delta	aa0	aa1	A/A	peak GiB
qwen3-30b-a3b	pp4-dp8-cp1-ep1-seq2048-bf16	2.939	2.861	-2.65%	2.889	2.944	+1.90%	86.62 both
qwen3-30b-a3b	pp1-dp8-cp4-ep8-seq2048-bf16	15.481	15.150	-2.14%	15.548	15.716	+1.08%	52.53 both
deepseek-v2-lite	pp2-dp8-cp2-ep8-seq4096-bf16	4.592	4.508	-1.83%	4.558	4.644	+1.89%	25.69 both
deepseek-v2-lite	pp2-dp4-cp4-ep2-seq4096-bf16	11.108	10.847	-2.35%	10.881	10.964	+0.76%	43.49 both

The four feat deltas span -1.83% to -2.35% and are not a speedup.

pp4-dp8-cp1-ep1 sets cp=1 and ep=1. Both branches then build identical rank layouts and identical collectives, so the true step-time difference on that mesh is zero. It reports -2.65%, the largest delta in the table. Rerunning that mesh with the arms reversed, feat first and base second, gives +1.15%. Two measurements with a known true value of zero therefore read -2.65% and -1.14%, and all four feat deltas fall inside that bracket.

The A/A columns measure a different pairing and are not the bracket for a feat delta. All four A/A pairs are positive, +0.76% to +1.90%, and all five branch comparisons are negative, -1.14% to -2.65%. The A/A arms keep one branch checked out for both runs; a branch comparison runs git checkout between them. What the A/A arms establish is the size of the noise: a same-code pair moves 1 to 2 percent on its own.

This protocol resolves to roughly plus or minus 2.5% at 8 steps. Every number in the table is inside that.

Peak GPU memory is identical to the reported precision on all four meshes, so folding costs nothing in memory. Memory does not fluctuate between runs, so this row is a direct result rather than a measurement inside a noise floor.

Do not read step time off the correctness runs above. Those use real routing, where expert imbalance dominates the variance: within-run scatter reaches a coefficient of variation of 9.5% with step-to-step spreads up to 38%. The same two deepseek meshes measured that way gave -6.2% and +3.2%, against -1.83% and -2.35% here.

Other testing

test_dualpipev against a single-device reference at pp=2 ep=2 for cp=1 and cp=2, packed, ragged, and packed with ragged, worst case 3.5e-3 on an expert weight against a 1e-2 threshold. The noise-floor guard now exempts routed expert weights, without which the expert gradients on deepseek-v2-lite all sit below the cutoff and the one test meant to catch an expert gradient-scale error cannot see one.

test_muon_fsdp at ep=1, ep=2, and ep=2 with cp=2, exact to 0.00e+00.

The operator suite and tools/tests, 223 pass with 2 failures that reproduce on a clean base branch.

[haok1402]

@claude review

[claude]

Compactness Review

No compactness issues. This is a disciplined refactor: the device_mesh → attn_mesh/expt_mesh split and the dp*ep → dp data-map simplification are direct substitutions, and the new surface all pays its way — split_replicas dedups a subtle (replicate, shard) contract used twice in apply_fsdp; hsdp_replica reshapes the pre-existing sharding_strategy knob while deleting the old _concatenate branch (net simpler); and the memory-estimator's new stage_size/expt_dp_size properties plus __post_init__ are all exercised (the guard is what forced the dp_size=1→8 test updates). No added indirection, no inference-only knobs, no dead defensiveness.

[claude]

Performance Review

No performance regression identified. This is a mesh-relabeling refactor (single (pp,dp,cp,ep) mesh → folded attn_mesh (pp,dp,cp) + expt_mesh (pp,dp,ep)) that preserves every collective's group size: attn params still shard/reduce over the whole stage, expert params still over stage_size/ep, so FSDP all-gather/reduce-scatter costs and peak memory are unchanged — matching the reported identical peak memory. The train-step gradient scale gains two Python-int factors with no new device sync, and the extra setup cost (a second pp communicator build, the set_gradient_divide_factor(1.0) module walk) is one-time, not per-step. The load-balance all-reduce group grows by a factor of cp but only carries per-expert counts, so it is negligible. The committed evidence — throughput within a ±2.5% A/A noise floor across 4 meshes / 2 models including folded cp>1,ep>1 configs, and identical peak memory — is thorough and consistent with the change.

One gap, not a blocker: hsdp_replica > 1 (the new 2-D (replica, shard) FSDP path) has no throughput or peak-memory number and no base to compare against. It is opt-in and off by default, so it does not touch the default hot path, but a single hsdp_replica=1 vs 2 step-time/peak-memory row would close it.

[claude]

Correctness Review

I traced the gradient and data flow of the folding change against the four validated meshes and found the core correct. Both parameter classes now reduce with a plain sum (set_gradient_divide_factor(1.0)) and the step applies the single 1/(num_tokens * dp_size * cp_size) normalization: the attn class sums over the whole dp x cp stage, and the expt class sums over expt_dp copies each of which already received a full EP group's tokens through the all-to-all, so both cover the same global token set exactly once. The two views stay consistent because pp_idx = rank // stage_size in both, and EP-spans-CP (pp1-cp4-ep8) resolves correctly — each expert's expt_dp copies partition the data slices disjointly. The data map keying to dp_rank alone, the load-balance group move to dp x cp, the clip_grad_norm_ /hsdp_replica over-count correction, and the memory-estimator refactor are all internally consistent. Evidence is strong: 4 meshes with real routing, the AdamW-invariance / grad-norm reasoning, and hsdp self-consistency (replica=1 vs 2, which does exercise the clip correction). No in-diff correctness issues.

One regression, outside this PR's diff so no inline thread: benchmarks/pretraining/qwen3-30b-a3b/h100-8n8g/pp4-dp2-cp1-ep8-seq4096-bf16.py:10 still sets distributed.sharding_strategy = "hsdp". That field was renamed to hsdp_replica, and DistributedCfg is slots=True, so this line now raises AttributeError at launch. It needs distributed.hsdp_replica = <N> (and note the semantics changed — the old "hsdp" sharded within cp x ep and replicated across dp, whereas hsdp_replica splits each replica group into N). docs/user-guide.md:75 documents the same removed distributed.sharding_strategy and is likewise stale.

[claude]

Consistency Review

The diff itself is internally consistent: the files this PR rewrote (AGENTS.md, docs/architecture.md, the two validate-* skills, the distributed contexts/module, the memory estimator) all describe the folded two-mesh model correctly, in present tense, with no post-mortem phrasing. No in-diff inconsistencies found, so there are no inline threads.

However, the rename DistributedCfg.sharding_strategy -> hsdp_replica and the decoupling of DP from EP (dp = total_gpus / (pp x cp), not /(pp x cp x ep); the batch splits by dp alone) left stale copies of the old contract in files outside this diff. All of the following describe the pre-folding state and should be updated:

Removed field sharding_strategy (now hsdp_replica: int):

• docs/user-guide.md:75 — config table row lists distributed.sharding_strategy with "fsdp"/"hsdp" values.
• benchmarks/pretraining/qwen3-30b-a3b/h100-8n8g/pp4-dp2-cp1-ep8-seq4096-bf16.py:10 — sets distributed.sharding_strategy = "hsdp"; the slot no longer exists, so this config now raises AttributeError at launch. It should set distributed.hsdp_replica = <N> instead.

Old DP formula and the "pp x cp x ep must divide" rule (now dp = total_gpus / (pp x cp); only pp x cp must divide the world and ep must divide the stage world // pp):

• docs/user-guide.md:87, 93 — state dp = total_gpus / (pp x cp x ep) and "pp x cp x ep must divide ... whatever is left over becomes DP."
• docs/user-guide.md:97-99 — the worked-example dp values are wrong under folding: pp=1, ep=8 on 8 GPUs gives attention dp=8 (expt_dp=1), not dp=1; likewise pp=1, ep=4 -> dp=8, and pp=2, ep=8 on 16 GPUs -> dp=8.
• docs/user-guide.md:144 — troubleshooting quotes the error string world_size not divisible by pp x cp x ep; the code now raises world_size not divisible by pp (and a separate stage-divisibility error for cp/ep).
• .agents/skills/estimate-memory/SKILL.md:26, 47 — both give dp = total_gpus / (pp x ep x cp); the tool now computes dp = (total_gpus // pp) // cp.
• .agents/skills/capture-nsys-profile/SKILL.md:64 — short-allocation message reports "required pp * cp * ep" ranks, which no longer describes the requirement.

Old batch-split / num_chunks formula (dropped the ep_size factor):

• .agents/skills/estimate-memory/SKILL.md:59 — num_chunks = global_batch_size / (micro_batch_size * dp_size * ep_size); the tool now divides by micro_batch_size * dp_size.
• .agents/skills/add-new-model/reference/testing.md:145 — the test_dualpipev.py label example uses ep_size * num_chunks * micro_batch_size; the test now sizes the global batch as dp_size * num_chunks.
• .agents/skills/capture-nsys-profile/scripts/capture.py:42-43 — computes dp_size = WORLD_SIZE // (pp*cp*ep) then global_batch = 32 * dp_size * ep_size. The product is still numerically correct (32 microbatches/stage), but the variable named dp_size no longer matches the framework's dp and reads misleadingly; recompute as WORLD_SIZE // (pp*cp) and drop the * ep_size.

Minor: pithtrain/modules/distributed.py:21-22 (DistributedCfg class docstring, outside the diff hunk) still says "FSDP2 sharding strategy ... DP is inferred from the world size" — now two DP degrees are derived; worth aligning with the folded model while here.