clip_grad_norm avoid O(global_params) replication
Author: AmineDiro
Pull request: yes

Followup to the https://github.com/huggingface/transformers/pull/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

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46113). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cc93871](https://github.com/huggingface/transformers/pull/46113/commits/cc93871b82d5bd61e54d6a4151f7bc46cb9cea74)

buckets params based on the mesh and placements. So we can buckets the all_reduce calls 

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

View the CircleCI Test Summary for this PR:

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