feat: make cross_entropy_chunk_size tunable, with an auto memory budget (#2190) [3/3]

[OnePunchMonk]

Part of #2190.

Part 3 of 3, split out of #2300 so each piece is reviewable on its own. This is the "make the chunking tunable" ask from #2190.

chunk_size was hardcoded to 128 in chunked_cross_entropy, which is a reasonable default for GPT-2 scale vocabs but not for a 152k vocab model where each chunk is over ten times bigger.

Two changes:

cross_entropy_chunk_size is now a TrainArgs knob, threaded through pretrain.py and all five finetune scripts, so you can set it from a config instead of editing the library.

It accepts "auto", which picks a chunk size from a memory budget rather than a fixed row count. auto_cross_entropy_chunk_size divides memory_budget_bytes (default 32MB) by vocab_size * itemsize * 3, so the chunk gets smaller as the vocab gets bigger and peak memory stays roughly flat instead of scaling with vocab.

The factor of 3 is measured, not guessed. It comes from sweeping vocab_size from 8k to 152k on a T4 and reading actual torch.cuda.max_memory_allocated(), not a profiler table's self-CUDA-mem for a single op. An earlier factor of 2 based on the profiler table undershot the real allocator peak by 1.5x, because it missed the nll_loss and allocator overhead around the log_softmax. The sweep is in #2311 as docs/profiling/budget_formula_sweep.png.

Tests

• test_auto_cross_entropy_chunk_size covers the formula directly.
• test_chunked_cross_entropy_auto_matches_manual_chunk_size checks "auto" gives the same loss as passing the equivalent int.
• test_chunked_cross_entropy_equivalence_at_scale checks chunked and unchunked agree at realistic vocab sizes.
• Two GPU-gated tests confirm peak memory actually drops, both for manual chunking and for "auto".

Default behavior is unchanged, chunk_size is still 128 unless you ask for something else.

Depends on nothing in the other two, but the docstring references a plot that lands in #2311. Related: #2310, #2311.

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 #2311 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.