fix(optimizer): clip by global grad norm across FSDP ranks in BF16Optimizer

[jessiewei7]

fix(optimizer): clip by global grad norm across FSDP ranks in BF16Optimizer

Summary

BF16Optimizer.step() clips with plain torch.nn.utils.clip_grad_norm_ over
its fp32 master copies. All callers (train_dflash.py, train_domino.py,
train_eagle3.py, train_peagle.py, the dataflow runtime, the disagg
example) wrap the model in FSDP with SHARD_GRAD_OP, so after backward each
rank's .grads are its reduce-scattered shard — the clip sees a
rank-local norm, roughly 1/sqrt(world_size) of the true norm. Two
consequences:

• gradients are systematically under-clipped (the threshold effectively
  scales with world size), and
• each rank picks a different clip coefficient for its shard, so the
  clipped global gradient no longer points in the original direction.

specforge/runtime/training/backend.py already acknowledges this — its
step() all-reduces the norm — but it does so after optimizer.step()
has clipped, so only the logged norm was corrected, not the clipping
itself.

Fix

• BF16Optimizer._clip_grad_norm(): accumulate the squared norm, all-reduce
  it across ranks (when dist is initialized), and scale every shard by the
  single global clip coefficient. Non-distributed behavior is unchanged
  (verified against torch.nn.utils.clip_grad_norm_ in the test). The
  returned/logged norm is now the true global norm.
• FSDPTrainingBackend.step(): drop the post-hoc all-reduce — with the fix
  it would reduce twice and inflate the logged norm by ~sqrt(world_size).

Why not FSDP.clip_grad_norm_: it clips the module's bf16 shard gradients,
while BF16Optimizer's whole design is clipping its fp32 master copies —
this keeps that numerics path, applying the same cross-rank reduction that
API does internally. An FSDP2 migration would make this class of bug
impossible (DTensor grads reduce cross-rank automatically) — #242 did that
for the old offline script, but the current dataflow backend is back on
FSDP1 — so the FSDP1 paths need the reduction done explicitly, as here.

Note on behavior: max_grad_norm=0.5 was copied from EAGLE's DeepSpeed
ds_config.json, and DeepSpeed clips by the global norm — so this restores
parity with the reference implementation the constant came from. At scale
the true norm is ~sqrt(world_size)× the previously-compared local norm, so
clipping engages more often; existing multi-GPU loss curves will not
reproduce exactly.

Test plan

tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py (CUDA-free):

• single-process clipping matches torch.nn.utils.clip_grad_norm_ exactly
  (norm and clipped gradients);
• 2-process gloo run with different per-rank gradient shards: both ranks
  compute the same global norm and apply the same clip coefficient.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[jiapingW]

Great findings! We have fixed the bugs in branch.