Fix slow Trainer path with 4D attention mask
Author: qflen
Pull request: yes

Fixes #32101

`Trainer` with a pretrained tokenizer and a 4D `attention_mask` is dramatically slower than the 2D-mask path because `PreTrainedTokenizerBase.pad` routes the multi-dim per-sample mask through `to_py_obj` and then rebuilds it with `torch.tensor(deeply_nested_list)` in `BatchEncoding.convert_to_tensors`. That round trip is the dominant cost on the 4D path, and `_pad` itself cannot extend rank > 1 per-sample masks, so the conversion is wasted work.

This PR pops a multi-dim per-sample `attention_mask` out before the per-sample padding loop and stacks it back with `torch.stack` / `np.stack` at the end.

End-to-end on the OP's repro (CPU, `batch=32`, `grad_accum=32`, `seq=1024`, 1 optimizer step) the 4D `Trainer` step drops from 225.59 s to 114.62 s, within 13% of the 2D path's 101.52 s; the collator alone goes from ~3.65 s / batch to ~11.8 ms / batch.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case.
      Issue: https://github.com/huggingface/transformers/issues/32101.
      Maintainer invitation to open a PR: https://github.com/huggingface/transformers/issues/32101#issuecomment-4406013253.
- [ ] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

## Who can review?

@SunMarc (trainer) and @ArthurZucker (tokenizers, text models).

---

## Numbers

End-to-end Trainer benchmark (CPU, single optimizer step, `batch=32`, `grad_accum=32`, `seq=1024`) matching the OP's repro (tiny `LlamaForCausalLM`, 6 layers, hidden=384, vocab=10):

| `attention_mask` | before | after |
| --- | ---: | ---: |
| 2D `(seq,)` per sample | 99.53 s | 101.52 s |
| 4D `(1, q, q)` per sample | 225.59 s | 114.62 s |

Before the fix the 4D step was 2.27x slower than the 2D step. After the fix it is within 13%, since the remaining cost is dominated by the model forward / backward (~100 s on this CPU). On hardware where the model is faster (e.g. the OP's, where the 2D step took 16 s and the 4D step took 340 s), the fix recovers a much larger relative speedup at the same parity point, for the same reason: the collator stops being the bottleneck.

### Collator-only decomposition

`DataCollatorWithPadding` timed in isolation across three shapes; mean ms / batch over `n_repeat` calls after a warm-up:

| (batch, seq) | n_repeat | mask | before (ms / batch) | after (ms / batch) | speedup |
| --- | ---: | --- | ---: | ---: | ---: |
| (8, 128) | 200 | 2D | 0.32 | 0.33 | ~1.0x |
| (8, 128) | 200 | 4D `(1, q, q)` | 13.16 | 0.26 | **~50x** |
| (16, 256) | 50 | 2D | 1.20 | 1.16 | ~1.0x |
| (16, 256) | 50 | 4D `(1, q, q)` | 110.13 | 0.92 | **~120x** |
| (32, 1024) | 20 | 2D | 8.96 | 8.89 | ~1.0x |
| (32, 1024) | 20 | 4D `(1, q, q)` (OP shape) | 3650.73 | 11.80 | **~310x** |

The 2D path is unchanged within noise. The 4D speedup grows with shape because the dominant pre-fix cost is `torch.tensor` on a deeply nested Python list of size O(batch * q * kv); after the fix, a single `torch.stack` scales linearly with elements rather than super-linearly with Python overhead.

Numerical equivalence is verified by the new `test_4d_attention_mask_preserved` regression test: per-sample shape, dtype, and contents match the stacked batch slice exactly, for both torch and numpy inputs.

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

Thanks for catching my style issue, merging!