Classify EP replicas by declared EP_SUM reduction, not the singleton axis

[qywu]

Follow-up to #72. _build_ep_param_groups classified any Replicate()-stamped parameter as an EP replica, but ParallelPlan.apply stamps Replicate() for several unrelated reasons: a genuine shared factor, the fallback for every parameter no ep_plan pattern matched, and rank-unique per-expert slices whose expert axis happens to be singleton (ep_size == num_experts). Grad-norm clipping divides EP-replica norms by ep_world, so any misclassified rank-unique gradient under-reported the global norm by ~√(ep_size); worse, a misclassified stamp also misleads checkpoint/export consumers, and a genuinely-shared plan-gap factor was "replicated" for the norm while never actually EP-summed — silently diverging across ranks every step.

The fix: one authoritative signal

A replica is only coherent across EP — and its norm may only be averaged across EP — when its gradient is EP-summed at the optimizer boundary. That is exactly what the owner modules already declare (_ep_gradient_reduction_by_parameter / _ep_gradient_reduction_domain, populated by MoEExpertsLoRA, the QLoRA expert module, and the GLM-5.2 exact experts). So:

• _build_ep_param_groups: is_ep_replicated = Replicate placement and gradient_reduction is EP_SUM. Shape plays no role. Frozen replicas stay out of the sync queue. A trainable EP_SUM replica that lands off the EP mesh now fails closed with an actionable error instead of the previous fail-open empty sync list (it would have silently missed the optimizer-boundary reduction and diverged).
• ParallelPlan.apply singleton branch: replicates only declared-EP_SUM factors; an already_local singleton is annotated as a per-expert Shard; anything else raises rather than guessing between a replica and a rank-unique expert slice.
• ParallelPlan.apply fallback: stamps the owner's declared reduction domain, so plan-gap shared factors (minimax_m3 lists no LoRA patterns; nemotron_h omits gate_proj_lora_A/B) keep their EP_SUM sync contract instead of losing it to the dataclass default.
• MoEExpertsLoRA: a post-EP construction declares its per-expert factors already-local (and exposes num_local_experts), so EP application annotates instead of slicing them a second time — or misreading the num_local_experts == 1 singleton as shared.
• inject_lora_into_experts: carries _xorl_ep_load_presliced onto the replacement module (translating gate_proj → gate_up_proj), so the checkpoint pre-shrink record survives injection and base expert weights are never EP-sliced twice.

Tests

New tests pin the semantics the old fixtures couldn't distinguish (both used leading-dim-1 shapes): declared-reduction classification is shape-agnostic (a [1, 4] NONE param is excluded, a [2, 4] EP_SUM param is included), frozen replicas skip sync, an off-mesh EP_SUM replica raises, undeclared singletons raise at the stamp site, already_local singletons annotate Shard, fallback specs carry the declared domain, post-EP LoRA construction declares its locals, and injection preserves the pre-slice record.

Verified locally: test_ep_clip_grad_norm + test_parallel_plan_meta_slice + test_ep_lora_weight_slicing + test_ep_gradient_reduction_contract + test_moe_expert_weight_init (38 tests), plus the blast radius on the same tree: QLoRA/GLM-52 expert-contract and construction suites (61), and canonical-MoE / adapter-gradient / tests/lora (102).

[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.