grad clipping computes a different "global" norm on every EP group with TE experts

what happens

with model.backend.experts: te and ep_size > 1, grad clipping does not compute one global norm for the whole model. each EP group computes its own norm and clips with it.

the reason is that TE expert params are plain tensors (or DTensors sharded only on ep_shard), so the clip code never adds up their part of the norm across the EP axis. it only reduces over the dims it can see on each param's own mesh.

why it matters

dense params (attention, embeddings, etc) are shared across EP groups through FSDP. each group clips its shard of the same dense param by a different factor, so the shared weights slowly drift apart between groups. clipping is on by default (max_norm 1.0), so this is active in normal runs. nothing in the logs shows it, every rank prints a grad norm that looks fine.

repro

cpu, 2 gloo ranks acting as 2 EP groups, real clip_grad_norm. each rank has different expert grads plus one dense param sharded across both ranks:

correct global grad norm (same on every rank): 40.201493
rank 0: total_norm=40.199502  dense_grad_after_clip=2.487593e-02
rank 1: total_norm=4.019950   dense_grad_after_clip=2.487592e-01
total_norm ratio rank0/rank1: 10.000x

same dense param, clipped 10x differently depending on the rank. with ep_shard > 1 (experts DTensor-wrapped on the ep_shard mesh only) the norms agree inside an EP group but differ 7.1x across groups.

who is affected

every shipped yaml with experts: te and ep_size > 1 that uses the shared MoE experts, 20 configs on main today: qwen3 moe 30b/235b, deepseek_v3, gpt_oss 20b/120b, and moonlight. models that leave experts at its default (torch_mm) are not affected.

related

while fixing this we found the ep grad scaling regex only matches mlp.experts.*. models that attach the moe block as moe (gemma4_moe, diffusion_gemma, step3p5) miss the existing ep grad scaling too.

fix

pr with the fix and multi-rank tests: #3817