clip_grad_norm avoid O(global_params) replication

[AmineDiro]

Followup to the #45028 . I've been integrating the new native FSDP+TP to my benchmark.

Current clip_grad_norm currently has two issues that compose:

• it materializes a full replicated copy of every DTensor gradient via _replicate_dtensor(g).to_local() into a Python list, holding all of them live for the duration of the function. Per-rank scratch is O(global_params) bytes — independent of how sharded the model is. On Qwen3-30B-A3B at tp_size=8, fsdp_size=2 this is a ~60 GB transient peak per rank, OOMing 80 GB H100s regardless of context length, batch size, or optimizer.
• (BUG) The function passes grads (raw gradient tensors) to torch.nn.utils.clip_grads_with_norm_, but that API expects parameters and re-extracts .grad internally. Raw tensors have p.grad = None, so the inner list is empty and the function returns silently. The reported total_norm is correct, but no gradient is ever actually scaled: max_norm is probably ignored.

Mem Profile

torch.cuda.memory._record_memory_history(enabled='all', context='all', stacks='all') at job start (rank 0 only) — captures every allocation/free with full call stack (C++ + Python).

torch.cuda.memory.memory_snapshot() + torch.cuda.memory_stats() at 11 named lifecycle stages — pickled to disk per stage so the OOM can't lose the data:

S0_after_pg_init                # baseline
S1_after_from_pretrained        # post weight load (shard-on-read)
S2_after_batch_ready            # post tokenization
S3_step{0,1}_after_forward      # post forward pass
S4_step{0,1}_after_backward     # post backward
S5_step0_after_clip             # post clip_grad_norm (step 0)
S6_step0_after_optim_step       # post optimizer.step (Adam state alloc)
S7_step0_after_zero_grad        # post zero_grad
S99_OOM_at_step{N}              # captured inside the try/except on OOM

stage	allocated	reserved	active blocks
S0_after_pg_init	0.00 GB	0.00 GB	0
S1_after_from_pretrained	5.17 GB	5.82 GB	531
S2_after_batch_ready	5.17 GB	5.82 GB	533
S3_step0_after_forward	8.13 GB	10.37 GB	596
S4_step0_after_backward	18.27 GB	21.55 GB	683
S5_step0_after_clip	18.27 GB	81.57 GB ← reserved spike	684
S6_step0_after_optim_step	28.60 GB	81.77 GB	2277
S7_step0_after_zero_grad	23.43 GB	81.77 GB	2227
S3_step1_after_forward	18.49 GB	81.78 GB	2191
S4_step1_after_backward	28.60 GB	81.78 GB	2277
S99_OOM_at_step1	29.00 GB	82.07 GB	peak_since_reset = 82.37 GB

reserved jumps from 21.55 GB (post-backward step 0) to 81.57 GB (post-clip step 0) — a 60 GB transient that the caching allocator never gives back to the driver. This is exactly the gradient-replicate scratch buffer clip_grad_norm builds, plus its caching residue. The OOM at step 1 occurs because the second invocation tries to allocate the same 60 GB of scratch on top of 28.6 GB of live tensors plus the ~5 GB FSDP staging + 7 GB autograd retained.

Validation

• ✅ Memory profile diff: reserved at S5_step0 drops from 81.57 GB → 21.55 GB (–60 GB) after skipping the call. Stays stable through 5 steps.

Who can review?

• distributed: @3outeille @ArthurZucker

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[3outeille]

if i understand properly, you are doing all_reduce over every parameters right ? It's going to be costly i think, is there a way to bucket them ?

[AmineDiro]

I missed that, you are right ! I'll fix this !

[AmineDiro]

Fixed in cc93871 👍🏼

[AmineDiro]

cc93871

buckets params based on the mesh and placements. So we can buckets the all_reduce calls

[3outeille]

Nice ! A last polishing pass for readability and it should be good

[3outeille]

can we use all_reduce to make it more explicit ?

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46113&sha=7f56bf

[ArthurZucker]

PR was reverted but happy to have cc @3outeille