profile: measure the chunked_cross_entropy memory spike and find its source (#2190) [2/3]

[OnePunchMonk]

Part of #2190.

Part 2 of 3, split out of #2300 so each piece is reviewable on its own. This is the "profile it and find the actual source" ask from #2190.

chunked_cross_entropy has a comment saying the memory spike is suspected to come from the backward pass, but nobody had measured it. This adds litgpt/scripts/profile_memory.py, which sweeps chunk_size and profiles the forward and backward with torch.profiler (profile_memory=True, record_shapes=True) on CPU or CUDA, and commits the results under docs/profiling/.

What the profile shows

On a T4 at B=2, T=2048, V=32000, peak CUDA memory allocated drops from 2097MB unchunked to about 1580MB for any chunk_size in [32, 512], roughly 25 percent, and it is flat across that whole range at this scale. The CPU sweep agrees.

The op table pins the spike on aten::_log_softmax and aten::_log_softmax_backward_data, each holding a 500MB CUDA allocation at the same time in the unchunked case. So the comment's suspicion was right.

The memory timelines make the mechanism visible: unchunked allocates one block that grows 500MB to 980MB to 1950MB and is not freed until the whole backward finishes, while chunking turns that into a sawtooth of smaller allocate/free cycles that never all coexist.

Three follow-up experiments are included too: torch.compile vs eager, an end to end training step (to check how much of a real step this actually accounts for), and a sweep of the budget formula that part 3 uses.

Notes

• The plots and op tables are committed rather than linked, so the numbers stay reproducible against a known commit. Happy to move them to a gist and keep just the script if you would rather not carry the binaries.
• Script only, no library code changes, so nothing here can affect training.

Related: #2310 (KV cache) and #2312 (tunable chunking).

AI Usage Disclaimer

• AI assistance (Claude Code) was used for this change.

[OnePunchMonk]

The failing tests.yaml / Lit Job check is not a lint failure. Looking at the CI log, it's test_serve_with_generate_strategy[tensor_parallel] in tests/test_serve.py. The server does start (NCCL init, model load, and "Application startup complete" all show up in the log), it just finishes right around the 30s mark, and the test's polling loop only waits 30s total before asserting failure. This PR doesn't touch test_serve.py or serving code, and #2312 hits the identical failure on a completely different diff, so this looks like a pre-existing timeout that's too tight for the 2-GPU tensor_parallel case on this runner, not something caused by this change.