loader: freeze MPS graph cache after warmup for graph DataLoaders

[anagnorisis2peripeteia]

Summary

Depends on: pytorch/pytorch#182648

Related issues:

• Memory usage and epoch iteration time increases indefinitely on M1 pro MPS pytorch/pytorch#77753 — MPS graph cache: 1 GB → 7.33 GB by epoch 8, 8s → 45s/iter
• MPS Memory Leaks across core MPS files pytorch/pytorch#164299 — "MPS Memory Leaks across core MPS files" (graphCache root cause)
• [MPS] Unbounded process RSS growth with varying-shape inference on macOS pytorch/pytorch#181213 — Unbounded RSS growth with varying-shape inference on MPS
• Trainer class causes massive memory leak when using mps huggingface/transformers#33717 — Llama fine-tuning: continuous memory growth on Apple Silicon

The problem

PyG's Batch.from_data_list() concatenates variable-size graphs: the shape key is (sum_nodes, sum_edges) — a 2-D combinatorial space. With shuffled DataLoaders (required for SGD convergence), batch compositions never repeat and every batch triggers a new MPSGraph compilation cached indefinitely:

P(repeat) ≈ 1/C(N, batch_size) ≈ 0

A 100-epoch training run on ZINC (10k molecules, BS=32) produces ~31,000 unique batch shapes. At ~55 KB/entry that is ~1.7 GB of graph cache that is never reused — a known source of OOM on Apple Silicon (pytorch/pytorch#77753, pytorch/pytorch#164299).

The fix

After a short warmup that populates the cache with the most common shapes, freeze_graph_cache() makes the cache read-only. Hits continue to be served; rare shapes compile on-the-fly and are immediately discarded. No unbounded growth.

Warmup formula (derived from CLT on the 2-D shape distribution):

eff_2d    = sqrt(BS) * std(node_counts) * sqrt(BS) * std(edge_counts) * 2π
freeze_at = min(total_iters // 4, max(5, int(sqrt(eff_2d))))

Computed from actual dataset statistics in DataLoader.__init__. Not hand-tuned.

Usage

Zero user changes required — the freeze happens automatically:

loader = DataLoader(dataset, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in loader:          # freeze_graph_cache() fires at iter freeze_at, first epoch only
        batch = batch.to(device)
        loss = model(batch).loss
        loss.backward()
        optimizer.step()

Benchmark results

MacBook Pro M1 Max, PyTorch 2.13.0a0, ZINC subset (12k molecules, BS=32).
Training loop (forward + backward + optimizer.step).
eff_2d=9,536, freeze_at=50 — computed from the ZINC node/edge distribution.

Strategy               ΔRSS MB   ms/iter   note
────────────────────────────────────────────────────────────────
always (baseline)      +1370.7    43.2ms   → ~1.7 GB / 100 epochs
freeze_after_warmup       +0.0    69.6ms   (1.61×, flat)
clear_per_iter (prev)     +0.0   112.0ms   (2.59×, flat)
never                     +0.0   143.7ms   (3.33×, flat)

freeze_after_warmup is 38% faster than clearing after each iteration and 52% faster than never-cache. The clear_per_iter overhead comes from destroying forward→backward graph reuse every iteration; freeze preserves compiled graphs for warmup-seen shapes through the full forward+backward cycle.

Reproduce with: python3 benchmark/mps_graph_cache.py

Changes

• torch_geometric/loader/dataloader.py: add _mps_freeze_at() helper and override DataLoader.__iter__ to call torch.mps.freeze_graph_cache() at the computed warmup step
• benchmark/mps_graph_cache.py: self-contained ZINC training benchmark reproducing the above results

The feature is a strict no-op when:

• MPS is unavailable (non-Apple hardware)
• torch.mps.freeze_graph_cache is absent (PyTorch < 2.13)
• the dataset does not expose num_nodes / num_edges
• node or edge counts are uniform (std == 0, caching is already optimal)

Test Plan

• ZINC training on MPS: always → +1370 MB RSS over 200 iters; freeze_after_warmup → +0.0 MB RSS (flat)
• Performance: freeze is 1.61× vs always baseline (vs 2.59× for clear_per_iter)
• Benchmark: benchmark/mps_graph_cache.py reproduces these results
• CI: existing DataLoader tests pass (freeze is a no-op until PyTorch 2.13 merges)

cc @rusty1s @mananshah99 @akihironitta