trainer: clear MPS graph cache via torch_empty_cache_steps
Author: anagnorisis2peripeteia
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=45818)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45818)
<!-- ci-dashboard-badge:end -->

## Summary

Closes #33717

MPSGraph on Apple Silicon can retain compiled graphs for shape-varying training workloads. With variable-length batches, the compiled graph cache grows across training. This PR keeps the default behavior unchanged, but when users opt in to periodic device-cache clearing with `torch_empty_cache_steps`, it also clears the MPS graph cache in that same existing `training_step()` cache-clearing path.

This follows review feedback to use the existing `torch_empty_cache_steps` cache-management knob instead of clearing graph cache unconditionally.

## Why this matters

I used local benchmark scripts to compare MPS graph-cache policies on BERT fine-tuning with WikiText-2, batch size 8, max length 128, 200 iterations, and subprocess-isolated `psutil` RSS. The scripts are intentionally not included in the final PR diff, per review feedback, but the benchmark setup and results are kept here for reproducibility/context.

The main benchmark compared the default cache policy with clearing, freezing, and disabling the graph cache:

| Strategy | RSS delta | ms/iter | note |
|---|---:|---:|---|
| always (default cache policy) | +794 MB | 185 ms | unbounded growth |
| `clear_graph_cache()` every optimizer step | +338 MB | 217 ms | bounded, 1.17x |
| `freeze_graph_cache()` after warmup | +443 MB | 192 ms | bounded, 1.04x |
| never cache | +209 MB | 1412 ms | bounded, 7.63x |

A second local sweep varied `max_len` to show why a user-facing clear option is useful for long-context, high-shape-diversity workloads:

| max_len | unique padded lengths | freeze overhead | clear overhead | lower-overhead bounded option |
|---:|---:|---:|---:|---|
| 64 | 61 | 1.00x | 1.68x | freeze |
| 128 | 125 | 1.03x | 1.15x | freeze |
| 256 | 253 | 1.36x | 1.07x | clear |
| 512 | 428 | 5.88x | 1.51x | clear |
| 1024 | 436 | 5.36x | 1.42x | clear |

The crossover is around 200 unique padded lengths. The workloads that OOM in practice, especially longer-context training with more variable padding, are above that crossover. That is why this PR exposes the bounded-memory behavior through the existing opt-in `torch_empty_cache_steps` setting rather than changing default Trainer behavior.

<details>
<summary>Benchmark script outline</summary>

The local benchmark scripts, kept out of the final diff, used:

- `bert-base-uncased` masked-LM fine-tuning on WikiText-2 raw train text.
- `DataCollatorWithPadding`, so each batch is padded to the maximum sequence length in that batch.
- Fresh subprocesses per strategy, because macOS `ru_maxrss` is monotonic and current RSS needs process isolation for a fair comparison.
- `torch.mps.clear_graph_cache()` after each optimizer step for the clear strategy, matching `torch_empty_cache_steps=1` in this PR.
- `torch.mps.freeze_graph_cache()` after warmup for the freeze strategy.
- `torch.mps.set_graph_cache_policy("never")` for the no-cache strategy.

The core strategy loop was:

```python
for i in range(ITERS):
    optimizer.zero_grad()
    loss = model(**inputs_for_iteration(i)).loss
    loss.backward()
    optimizer.step()

    if strategy == "clear_per_iter":
        torch.mps.clear_graph_cache()
    elif strategy == "freeze_after_warmup" and i == freeze_at:
        torch.mps.freeze_graph_cache()

    torch.mps.synchronize()
```

The sweep repeated the same workload for `max_len` values `[64, 128, 256, 512, 1024]` and compared runtime overhead against the default `always` cache policy.

</details>

## Change

- Add a guarded `torch.mps.clear_graph_cache()` call beside the existing `clear_device_cache()` call in `training_step()`.
- Gate it behind:
  - `self.args.torch_empty_cache_steps is not None`
  - the same step interval check used by the existing cache-clearing path
  - `torch.backends.mps.is_available()`
  - `hasattr(torch.mps, "clear_graph_cache")`
- Remove benchmark scripts and MPS-only tests from the final code diff.

With `torch_empty_cache_steps=1`, MPS graph cache is cleared whenever the existing per-step device-cache clear runs. With larger values, it follows that interval. If unset, behavior is unchanged.

The `hasattr` guard keeps this a no-op on PyTorch versions that do not expose `torch.mps.clear_graph_cache()`.

## Validation

- `make style`
- `make check-code-quality`
- `make check-repo`
- `python -m pytest tests/trainer/test_trainer.py::TrainerIntegrationTest::test_double_train_wrap_once tests/trainer/test_trainer.py::TrainerIntegrationTest::test_training_arguments_are_left_untouched tests/trainer/test_training_args.py::TestTrainingArguments::test_torch_empty_cache_steps_requirements -q`
- `utils/tests_fetcher.py` was run; because `trainer.py` is central, it selected a broad CI slice. Local targeted Trainer checks above passed.

AI assistance was used to action reviewer feedback and run validation; the final code diff was kept to reviewer-requested scope.

## References

- pytorch/pytorch#154329 and pytorch/pytorch#164299 - upstream MPS memory-growth reports involving graph/cache behavior
- pytorch/pytorch#181485 - existing PyTorch MPS policy: `torch.mps.empty_cache()` clears `MPSGraphCache`
- #33717 - Trainer MPS memory growth report


Comments:
cc @sunmarc maybe?

Ah, this looks a lot like the bug I experienced in #45517

> Ah, this looks a lot like the bug I experienced in #45517

Yes MPSGraphCache is the source of the memory swelling. Effectively unbounded growth with variable sized inputs.

what are we waiting on :)

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45818)
**Latest run:** [28949964486:2](https://github.com/huggingface/transformers/actions/runs/28949964486)
**Result:** `success` | **Jobs:** 15 | **Tests:** 172,320 | **Failures:** 0 | **Duration:** 7h 54m

<!-- ci-dashboard-recap:end -->

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