fix(moe): require a singleton expert axis to call an EP param replicated

[qywu]

Fixes #70.

_build_ep_param_groups treated any EP-group parameter with a Replicate placement as an EP replica. ParallelPlan.apply stamps Replicate() for two different reasons, though: a genuine shared factor (singleton expert axis, parallel_plan.py:224) and the fallback for every parameter no ep_plan pattern matched (parallel_plan.py:301-303). The fallback is not evidence of replication, so a plan gap made rank-unique per-expert gradients count once instead of once per rank, scaling the reported grad norm by 1/sqrt(ep_size) and clipping correspondingly too weakly at the same max_grad_norm.

The fix requires the singleton expert axis that both the plan and the EP LoRA backends already use as the shared-factor marker. Every ep_plan shards the expert axis on dim 0, and every trainable parameter that is genuinely EP-replicated today carries a leading dim of 1 — the hybrid-shared MoEExpertsLoRA / QLoRAMoeExperts factors, and the GLM-5 exact routed-experts QLoRA factors (gate_proj_lora_A, up_proj_lora_A, down_proj_lora_B, all (1, ...), with the per-expert ones in _ep_force_shard_parameter_names).

Verification

Using the reproduction from #70 (measured ground truth: a gradient counts once iff it is bit-identical on every EP rank), 4×H100, ep_size=4, per-expert MoE LoRA:

	before	after
minimax_m3 (plan lists no LoRA patterns)	classifies all 6 factors replicated, reported/true = 0.500000	classifies none, reported/true = 1.000000
qwen3_moe control (plan lists them)	1.000000	1.000000

The new unit test is mutation-checked: it fails without the production change and passes with it. tests/distributed/test_ep_clip_grad_norm.py and tests/distributed/test_ep_gradient_reduction_contract.py otherwise pass unchanged, including test_shared_ep_replica_is_recorded_separately_for_clipping, whose shared factor keeps its leading dim of 1.

One caveat on my local run: test_real_two_rank_ep_clip_and_nonfinite_gate fails for me with ncclInvalidUsage in its spawned two-rank subprocess, but it fails identically on a pristine origin/main worktree in the same environment, and I ran against torch 2.10 rather than the pinned 2.12.1+cu132 — so I read it as environmental and unrelated to this change. Worth confirming in CI under the pinned torch.

Not addressed here

minimax_m3's ep_plan lists no LoRA patterns and nemotron_h's omits gate_proj_lora_A/B. Adding those patterns is a separate change with its own slicing implications, and this fix makes the norm correct either way. If you would rather surface such gaps loudly instead of inferring from shape, the fallback SpecInfo could carry a marker and _build_ep_param_groups could warn (or fail closed) when an EP-group parameter arrives with one — happy to follow up with that.

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.

 

[kiddyboots216]

I'm not entirely sure this is the right fix -it's more like, minimax m3 and nemotron h should be changed?

[qywu]

Superseded by #76: the shape-heuristic still misclassifies rank-unique per-expert factors when ep_size == num_experts (they are also [1, ...]-shaped), and it leaves plan-gap shared factors unsynced. #76 classifies by the owner's declared EP_SUM gradient-reduction domain instead, fails closed on the ambiguous cases, and closes #70.