"Peak memory usage" reports 0.00MiB on CPU and MPS

base_train.py and chat_sft.py both bind:

get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0

and print it at the end of a run:

print0(f"Peak memory usage: {get_max_memory() / 1024 / 1024:.2f}MiB")

So on any non-CUDA device the run ends with:

Peak memory usage: 0.00MiB
Total training time: 0.22m

That is from an unmodified master on an M4 Max, --depth=6 --max-seq-len=512 --device-batch-size=32. On a machine small enough to care about memory, this is the number that tells you whether the next --depth will fit.

Why the obvious fix does not work

torch has no peak-memory API for MPS. As of 2.9 the full set is:

torch.mps.current_allocated_memory()
torch.mps.driver_allocated_memory()
torch.mps.recommended_max_memory()

driver_allocated_memory() looks like the drop-in replacement, but it reports the Metal caching allocator's reserved pool. That is the analogue of torch.cuda.max_memory_reserved(), not max_memory_allocated(), so it does not mean the same thing as the CUDA number printed under the same label. Same depth-6 run as above:

actual peak allocated      :  1499.78 MiB
driver_allocated_memory()  : 11322.84 MiB

It also does not track spikes well, because a spike that fits inside the pool already reserved does not grow the pool. Controlled bench, 6-layer MLP, batch 32 with a single 8x spike at step 4:

 step    allocated     driver
    3      289.4 MB   1050.7 MB
    4      305.1 MB   1050.8 MB   <- 8x batch
    5      289.4 MB   1050.8 MB

Allocated moves 5.4%. Driver moves 0.01%.

What does work

Sample current_allocated_memory() once per training step and keep the running max. That gives a number in the same units as CUDA's max_memory_allocated(), and it picks up the spike above (305.1 vs 289.4).

One limitation: sampling per step gives a lower bound, not the exact peak. The real high-water mark happens during backward, when activations and partial gradients are both live, and only the allocator can see that. There is no way to get the exact figure on MPS without a change on the torch side.

Related

scripts/infer_bench.py:109:

assert device_type == "cuda", "infer_bench currently assumes a CUDA GPU (for timing and VRAM measurement)"

VRAM measurement is part of why that is CUDA-only, so the same gap affects two places.

Environment

M4 Max, macOS 26.4.1, torch 2.9.1, uv sync --extra cpu.