EP clip_grad_norm undercounts per-expert MoE LoRA gradients by 1/sqrt(ep_size) when the ep_plan omits the LoRA factors

[qywu]

Summary

With Expert Parallelism, clip_grad_norm under-reports the global gradient norm by exactly 1/sqrt(ep_size) for per-expert MoE LoRA factors when the model's ep_plan does not enumerate them. minimax_m3's plan lists no LoRA patterns at all; nemotron_h's omits gate_proj_lora_A/B. Those factors are then treated as EP-replicated and counted once instead of once per rank, so clipping is correspondingly too weak at the same max_grad_norm.

Reproduction

Needs 2 or more GPUs. Save as repro_ep_replicated.py at the repo root:

"""Repro: EP grad-norm undercount for per-expert MoE LoRA factors.

The ground truth is measured, not assumed: a gradient counts once toward the
global norm iff it is bit-identical on every EP rank (MAX-MIN spread == 0),
otherwise once per rank.  No MoE kernels are involved -- gradients are injected
with a known replication structure, so only the classification + norm
arithmetic is under test.

    torchrun --nproc_per_node=4 repro_ep_replicated.py qwen3_moe    # control: exact
    torchrun --nproc_per_node=4 repro_ep_replicated.py minimax_m3   # 0.5x undercount
"""

import sys

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed._composable.fsdp import fully_shard
from torch.distributed.tensor import DTensor, Shard

from xorl.distributed.fsdp2.clip_grad_norm import clip_grad_norm
from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state
from xorl.distributed.torch_parallelize import _build_ep_param_groups
from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig

MODEL = sys.argv[1] if len(sys.argv) > 1 else "minimax_m3"
E, H, I, R = 8, 64, 128, 8


def main():
    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(rank)
    dev = torch.device("cuda", rank)
    init_parallel_state(dp_size=world, dp_shard_size=world, ep_size=world, dp_mode="fsdp2", device_type="cuda")
    ps = get_parallel_state()
    n_local = E // world

    # per-expert MoE LoRA (hybrid_shared=False): no factor is EP-replicated
    experts = MoEExpertsLoRA(
        num_experts=E, hidden_dim=H, intermediate_size=I, moe_implementation="native",
        lora_config=MoELoRAConfig(r=R, lora_alpha=2 * R, hybrid_shared=False,
                                  target_modules=["gate_proj", "up_proj", "down_proj"]),
    )
    mlp = nn.Module(); mlp.experts = experts
    layer = nn.Module(); layer.mlp = mlp
    inner = nn.Module(); inner.layers = nn.ModuleList([layer])
    root = nn.Module(); root.model = inner
    root = root.to(device=dev, dtype=torch.bfloat16)
    for n, p in root.named_parameters():
        p.requires_grad_(not (n.endswith("gate_up_proj") or n.endswith("down_proj")))

    # expert params already at EP-local shapes, as with skip_weight_loading
    # (torch_parallelize.py:382 -> ParallelPlan.apply(already_local=True))
    with torch.no_grad():
        for name, p in list(experts.named_parameters()):
            if p.shape[0] == E:
                setattr(experts, name, nn.Parameter(p.data[rank * n_local:(rank + 1) * n_local].clone(),
                                                    requires_grad=p.requires_grad))

    plan = __import__(f"xorl.models.transformers.{MODEL}.parallelize", fromlist=["get_ep_plan"]).get_ep_plan()
    root._fqn2spec_info = plan.apply(root, ps.ep_fsdp_device_mesh, already_local=True)
    fully_shard(experts, mesh=ps.ep_fsdp_device_mesh["ep_fsdp"], shard_placement_fn=lambda p: Shard(1))
    _build_ep_param_groups(root)

    # inject grads: every per-expert factor gets rank-unique values
    for i, (n, p) in enumerate(sorted(root.named_parameters())):
        if not p.requires_grad:
            continue
        local = p.to_local() if isinstance(p, DTensor) else p
        g = torch.Generator(device="cpu").manual_seed(5000 + i + 977 * (rank + 1))
        gl = (torch.randn(local.shape, generator=g) * 0.1).to(dev, torch.bfloat16)
        p.grad = (DTensor.from_local(gl, p.device_mesh, p.placements, run_check=False)
                  if isinstance(p, DTensor) else gl)

    # measured ground truth
    sq_rep = torch.zeros((), device=dev, dtype=torch.float64)
    sq_uniq = torch.zeros((), device=dev, dtype=torch.float64)
    for n, p in root.named_parameters():
        if p.grad is None:
            continue
        g = (p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad).detach().float()
        gmax, gmin = g.clone(), g.clone()
        dist.all_reduce(gmax, op=dist.ReduceOp.MAX); dist.all_reduce(gmin, op=dist.ReduceOp.MIN)
        rep = (gmax - gmin).abs().max().item() == 0.0
        (sq_rep if rep else sq_uniq).add_((g.double() ** 2).sum())
    dist.all_reduce(sq_rep); dist.all_reduce(sq_uniq)
    true_norm = (sq_rep / world + sq_uniq).sqrt().item()

    reported = clip_grad_norm(root, max_norm=1e9, norm_type=2.0).item()  # huge max_norm: no clipping
    ids = {id(p) for p in root._ep_param_groups.get("ep_replicated", [])}
    if rank == 0:
        print(f"model={MODEL}  ep_size={world}  per-expert MoE LoRA")
        print(f"  measured-replicated grads : none (every factor is rank-unique)")
        print(f"  classified 'ep_replicated': {[n.split('experts.')[-1] for n, p in root.named_parameters() if id(p) in ids]}")
        print(f"  true global grad norm     : {true_norm:.6f}")
        print(f"  clip_grad_norm reported   : {reported:.6f}")
        print(f"  reported / true           : {reported / true_norm:.6f}")
    dist.barrier(); dist.destroy_process_group()


if __name__ == "__main__":
    main()

Control — a model whose ep_plan does enumerate the LoRA factors is exact:

$ torchrun --nproc_per_node=4 repro_ep_replicated.py qwen3_moe
model=qwen3_moe  ep_size=4  per-expert MoE LoRA
  measured-replicated grads : none (every factor is rank-unique)
  classified 'ep_replicated': []
  true global grad norm     : 19.227601
  clip_grad_norm reported   : 19.227600
  reported / true           : 1.000000

Same model size, same injected gradients, plan without LoRA patterns:

$ torchrun --nproc_per_node=4 repro_ep_replicated.py minimax_m3
model=minimax_m3  ep_size=4  per-expert MoE LoRA
  measured-replicated grads : none (every factor is rank-unique)
  classified 'ep_replicated': ['gate_proj_lora_A', 'gate_proj_lora_B', 'up_proj_lora_A', 'up_proj_lora_B', 'down_proj_lora_A', 'down_proj_lora_B']
  true global grad norm     : 19.227601
  clip_grad_norm reported   : 9.613800
  reported / true           : 0.500000

The factor tracks 1/sqrt(ep_size) — the same run on 2 GPUs reports 0.707107.

Where it comes from

ParallelPlan.apply assigns Replicate() as the fallback for every parameter no ep_plan pattern matches (src/xorl/distributed/parallel_plan.py:301-303), and _build_ep_param_groups reads that placement as proof of EP replication (src/xorl/distributed/torch_parallelize.py:829), so unmatched per-expert factors land in the ep_replicated group whose norm contribution is divided by ep_world in clip_grad_norm.

Suggested fix

Require the leading dim of 1 that the EP LoRA backends themselves use as the shared-factor marker, so a fallback placement is not sufficient:

# _build_ep_param_groups
is_ep_replicated = (
    Replicate is not None
    and isinstance(getattr(spec_info, "placement", None), Replicate)
    and p.shape[0] == 1
)

With that change the minimax_m3 run above reports 1.000000 and the qwen3_moe control is unaffected. Adding the LoRA patterns to the minimax_m3 / nemotron_h plans fixes those two models, but the fallback would stay silent for the next plan that misses a parameter.

Environment

4x H100, nproc_per_node=4, ep_size=4, ep_fsdp_size=1, hybrid_shared=False.