Fully address memory management issues: profiling, KV cache lifecycle docs, tunable chunking, tests at scale (#2190) - #2300
Closed
OnePunchMonk wants to merge 13 commits into
Closed
Conversation
…KV cache Previously, growing max_seq_length beyond the current KV mask cache size only printed a warning and continued, so forgetting to call set_kv_cache() afterward led to a cryptic IndexError deep inside attention math (or worse, silently wrong results). Now forward() raises a clear RuntimeError pointing at the fix. Fixes Lightning-AI#2190
OnePunchMonk
requested review from
andyland,
k223kim,
lianakoleva and
t-vi
as code owners
August 16, 2026 15:20
OnePunchMonk
marked this pull request as draft
August 16, 2026 15:24
…-AI#2190) Adds litgpt/scripts/profile_memory.py, which uses torch.profiler to measure peak backward-pass memory of chunked_cross_entropy across a sweep of chunk_size values. Confirms the "workaround hack" comment in utils.py: on CPU, aten::_log_softmax_backward_data is the actual allocation spike (500MB for B=2,T=2048,V=32000), and any chunking (chunk_size>0) cuts peak memory by ~25% vs the unchunked baseline. Committed output under docs/profiling/ for reference in the PR.
set_kv_cache and clear_kv_cache already form the KV cache's init/resize and destroy primitives, but neither was documented as part of a lifecycle. Add docstrings clarifying how they pair together and when to use each. No behavior change.
…AI#2190) The chunked_cross_entropy chunk_size was a hardcoded magic number (128) at every call site. Add TrainArgs.cross_entropy_chunk_size (default 128, preserves current behavior) and thread it through the pretrain and finetune train-step / pretrain-validate loss calls, so users can tune the memory/compute tradeoff via --train.cross_entropy_chunk_size instead of editing source. The finetune scripts' validate() calls keep their existing hardcoded chunk_size=0 (exact loss), which is intentional and unrelated to the training memory hack.
…ghtning-AI#2190) - test_chunked_cross_entropy_equivalence_at_scale: correctness at B=2,T=2048,V=32000 (CPU-safe, runs in CI). - test_chunked_cross_entropy_peak_memory_decreases_with_smaller_chunks: confirms chunking actually lowers CUDA peak memory at B=4,T=4096, V=32000 (gated behind _RunIf(min_cuda_gpus=1), since CI has no GPU runners). - test_kv_cache_full_context_length: forwards a pythia-14m-sized model at its full block_size with batch_size=4, instead of only the artificial block_size=25 configs used elsewhere in this file.
OnePunchMonk
changed the title
OnePunchMonk
marked this pull request as ready for review
August 17, 2026 16:31
Ran the chunked_cross_entropy backward-pass profile on an NVIDIA T4
(via Modal) using torch.profiler's CPU+CUDA activities, confirming the
CPU-only finding with real CUDA memory numbers: aten::_log_softmax and
its backward dominate peak allocation, and chunking cuts peak CUDA
memory from ~2.1GB (unchunked) to ~1.6GB, flat across chunk_size in
[32, 512].
profile_memory.py now supports --device {cpu,cuda} (auto-detected by
default) and can export a memory-timeline PNG per chunk size via
--memory-plot-chunk-sizes.
for more information, see https://pre-commit.ci
- litgpt/chat/base.py: process_prompt's dynamic kv-cache growth path (the interactive REPL loop) called set_kv_cache() without clear_kv_cache() first, unlike the equivalent growth path in LLM.generate (api.py). This held the old, too-small cache and the newly allocated one in memory at the same time during every mid-session cache growth. Now clears first, matching api.py. Regression test added (fails against the old code, passes against the fix). - litgpt/utils.py: added auto_cross_entropy_chunk_size(), which derives chunk_size from a memory_budget_bytes and the model's vocab_size/ dtype instead of a hardcoded constant. The byte-per-chunk-element estimate is fit directly against torch.profiler measurements on a real T4 (docs/profiling/op_table_gpu.md): predicted 32.77MB vs. measured 16.0MB self CUDA mem for chunk_size=128, vocab_size=32000 -- within the function's built-in 2x safety margin. chunked_cross_entropy(..., chunk_size="auto", memory_budget_bytes=...) resolves through it. TrainArgs.cross_entropy_chunk_size now accepts "auto" (default unchanged at 128), with a new cross_entropy_memory_budget_bytes field, threaded through pretrain.py and all finetune scripts. Verified on a T4: chunk_size="auto" gives the same peak-memory reduction as a hand-picked chunk_size.
…formula sweep, end-to-end step 1. torch.compile vs eager (chunked_cross_entropy, T4): compiling cuts peak memory ~2x and wall time ~20-28x for both chunk_size=0 and 128. Notably, compiled unchunked (1048MB) uses about the same memory as compiled chunked (1065MB) -- torch.compile appears to fuse the softmax/nll_loss backward well enough that the chunking hack's memory benefit is largely moot under compile. Eager is still the default in this codebase's training scripts, so the chunking fix stands, but this is worth flagging for anyone using fabric/torch.compile with this loss. 2. auto_cross_entropy_chunk_size swept across vocab_size (8k-152k, GPT-2 to Llama-3/Qwen2.5 scale) at a fixed memory_budget_bytes：the chosen chunk_size scales down as vocab_size grows, and measured per-chunk peak memory stays exactly flat across every vocab_size -- the formula generalizes, not just a fit to one config. It was consistently 1.5x the target budget, though, so recalibrated _CROSS_ENTROPY_BYTES_PER_CHUNK_ELEMENT from 2 to 3 to close that gap (the old factor of 2 was fit from a profiler op table's self-CUDA- mem for one op; the new factor of 3 is fit from the real torch.cuda.max_memory_allocated() peak, which is what a memory budget is actually supposed to bound). 3. End-to-end training step (real pythia-14m forward+backward, not just the isolated loss op): chunking still cuts a full step's peak CUDA memory by ~22.8% (3580MB -> 2765MB), and the op table shows aten::mm (matmuls) actually edges out log_softmax as the single largest CUDA-memory consumer in a real step -- so the cross-entropy spike is real but not the whole story once other ops fill the memory line, which is the actual answer to "does this matter in practice."
for more information, see https://pre-commit.ci
This was referenced Aug 28, 2026
Open
Open
Contributor
Author
|
[split for review] This PR was too big to review in one go, so I've split it into three independent PRs, same code, no changes:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #2190, a report of four separate memory-management concerns. This PR now addresses all four:
Profile memory with PyTorch's profiler, find the actual leak/spike source.
Added
litgpt/scripts/profile_memory.py, which profileschunked_cross_entropy's backward pass across a sweep ofchunk_sizevalues withtorch.profiler(profile_memory=True,record_shapes=True), on both CPU and CUDA (--device, auto-detected). Results are committed underdocs/profiling/for reproducibility.GPU results (NVIDIA T4,
B=2, T=2048, V=32000):Peak CUDA memory allocated drops from 2097MB unchunked → ~1580MB for any
chunk_sizein[32, 512](~25% reduction), flat across that range at this scale — matching the CPU sweep below. The op table (docs/profiling/op_table_gpu.md) pins the spike onaten::_log_softmaxandaten::_log_softmax_backward_data, each holding a 500MB CUDA allocation simultaneously in the unchunked case, confirming thechunked_cross_entropycomment's suspicion about the backward pass.The memory-timeline plots make the mechanism visible directly — unchunked allocates one 500MB→980MB→1950MB block that isn't freed until the whole backward pass finishes, while chunking turns that into a sawtooth of 32 smaller allocate/free cycles that never all coexist:
chunk_size=0)chunk_size=128)The kernel-launch timelines (derived from the same
torch.profilerchrome trace, CPU call-stack depth vs. CUDA kernel launches) make the launch-overhead tradeoff visible too: unchunked is two clean wide blocks (onelog_softmaxop, one CUDA kernel each), while chunked is 32 repeated narrow bursts — a picket fence of CUDA kernel launches instead of two big ones:chunk_size=0)chunk_size=128)GPU profiler op table, chunk_size=0 (unchunked) — aten::_log_softmax + backward each hold 500MB CUDA mem concurrently
Full op tables for every profiled
chunk_size(0, 32, 64, 128, 256, 512) are indocs/profiling/op_table_gpu.md. They also surface the actual compute-for-memory tradeoff (on CPU, chunking measured faster, which hid this):Self CUDA time totalrises from 19.0ms unchunked to 25–28ms chunked — chunking addsaten::cat/aten::split/aten::narrowcalls (extra kernel launches, scaling with the number of chunks) that unchunked doesn't pay for. Socross_entropy_chunk_sizeis a real memory/compute knob on GPU, not a free lunch, which is exactly why item 3 below makes it tunable per training run instead of hardcoding one value.CPU results (kept for the wider
chunk_sizesweep and side-by-side comparison):Same qualitative finding on CPU:
aten::_log_softmax_backward_datais the spike (500MB self CPU mem, unchunked), and any chunking cuts peak profiler-tracked memory by ~25%.CPU profiler op table, chunk_size=0 (unchunked) — single 500MB log_softmax + backward allocation
CPU profiler op table, chunk_size=128 (default) — same ops, 32 smaller calls instead of 1 big one
The total bytes allocated over the run (
CPU Memcolumn) end up similar either way — chunking doesn't reduce total work, it spreads the samelog_softmax/backward allocations across 32 smaller, sequential calls instead of one large one, so each chunk's memory is freed before the next is allocated. That's exactly what caps the peak concurrent memory (the timeline-based numbers in the charts above). This CPU microbenchmark actually measured chunking as faster (446ms → 145ms self CPU time) — but the T4 numbers above show the real story: chunking adds real kernel-launch overhead on GPU, which is whycross_entropy_chunk_sizeis exposed as a tunable rather than hardcoded to a single "best" value.Make the KV cache lifecycle explicit (init/use/clear/destroy).
GPT.set_kv_cache()(init/resize) andGPT.clear_kv_cache()(destroy) already existed as the lifecycle primitives, but neither was documented, andclear_kv_cache()was called in exactly one place in the whole codebase (LLM.generate's dynamic-growth path). Added docstrings to both describing the lifecycle and how they pair together.forward()now raises a clearRuntimeErrorinstead of silently continuing when a stale, too-small mask cache is used — previously this only printed a warning and could produce wrong results or crash with a crypticIndexError.Auditing every
set_kv_cache/clear_kv_cachecall site (grep -rn "set_kv_cache\|clear_kv_cache" litgpt/) turned up a real instance of the bug this lifecycle is meant to prevent:litgpt/chat/base.py'sprocess_prompt, called in the interactive chat REPL'swhile Trueloop, grows the kv cache mid-session viamodel.set_kv_cache(...)whenever the conversation outgrowsmax_seq_length— but never calledclear_kv_cache()first, unlike the equivalent growth path inLLM.generate(api.py, which doesclear_kv_cache()→set_kv_cache(...)). Every cache growth during a long-running chat session was therefore holding the old, too-small cache and the newly allocated one in memory at the same time, instead of a clean destroy-then-reallocate — a small but real instance of exactly the "creeping memory in long-running inference servers" symptom from the issue. Fixed to clear first, matchingapi.py. Addedtest_process_prompt_clears_kv_cache_before_growing(tests/test_chat.py), which fails against the old code and passes against the fix.Add a proper memory budget system instead of the chunking hack.
chunked_cross_entropy'schunk_sizewas a hardcoded magic number (128) baked into every call site — tunable is not the same as budgeted. Addedlitgpt.utils.auto_cross_entropy_chunk_size(vocab_size, dtype, memory_budget_bytes), which deriveschunk_sizefrom an actual memory budget instead of a guess. The byte-per-chunk-element estimate it uses is fit directly against thetorch.profilerT4 measurements above, not assumed: forchunk_size=128, vocab_size=32000, fp32, the formula predicts 32.77MB (with its built-in 2x forward+backward safety margin) against the 16.0MB actually measured — a good enough fit to be a real memory-budget dial rather than a rebranded second magic number.chunked_cross_entropy(..., chunk_size="auto", memory_budget_bytes=...)resolves through it.TrainArgs.cross_entropy_chunk_sizenow acceptsint | Literal["auto"](default unchanged at128, so existing configs are unaffected), plus a newTrainArgs.cross_entropy_memory_budget_bytes(default 32MiB), both threaded throughpretrain.pyand every finetune script's train-step (and pretrain'svalidate()) loss calls. Set--train.cross_entropy_chunk_size=autoto size chunks from the budget instead of picking a number by hand;0still disables chunking entirely (trades peak memory for extra compute, per the profiling above). The finetune scripts'validate()calls intentionally keepchunk_size=0(exact loss), unrelated to this change. Verified on the T4:chunk_size="auto"gives the same peak-memory reduction as a hand-pickedchunk_size=128(test_chunked_cross_entropy_auto_reduces_peak_memory_like_manual_chunking, CUDA-gated).Test with full context lengths and realistic batch sizes.
test_chunked_cross_entropy_equivalence_at_scale(tests/test_utils.py) — correctness atB=2, T=2048, V=32000, CPU-safe, runs in CI.test_chunked_cross_entropy_peak_memory_decreases_with_smaller_chunks(tests/test_utils.py) — confirms chunking actually lowers CUDA peak memory atB=4, T=4096, V=32000; gated@_RunIf(min_cuda_gpus=1)since this repo's CI has no GPU runners. Verified passing on an NVIDIA T4.test_auto_cross_entropy_chunk_size/test_chunked_cross_entropy_auto_matches_manual_chunk_size(tests/test_utils.py) — CPU-safe correctness of the new"auto"chunk-size path.test_chunked_cross_entropy_auto_reduces_peak_memory_like_manual_chunking(tests/test_utils.py) —chunk_size="auto"actually reduces CUDA peak memory vs. unchunked at realistic scale; gated@_RunIf(min_cuda_gpus=1), verified passing on an NVIDIA T4.test_process_prompt_clears_kv_cache_before_growing(tests/test_chat.py) — regression test for the KV cache double-allocation fix in item 2; fails against the pre-fix code.test_kv_cache_full_context_length(tests/test_model.py) — forwards apythia-14m-sized model at its fullblock_sizewithbatch_size=4, instead of only the artificialblock_size=25configs used elsewhere in this file.Follow-up profiling (T4)
Three more experiments to stress-test the claims above, all run on the same NVIDIA T4:
torch.compilevs eager. This codebase's training scripts run eager, but it's worth knowing what changes undertorch.compile:Compiling
chunked_cross_entropycuts peak CUDA memory ~2x and wall time ~20–28x for bothchunk_size=0and128. More interesting: compiled-unchunked (1049MB) and compiled-chunked (1065MB) land at almost the same peak memory — undertorch.compile, the chunking hack's memory benefit is largely gone (the compiler evidently fuses/schedules the softmax+nll_loss backward well enough on its own). This doesn't change anything in this PR (the training scripts run eager andtorch.compilesupport for the full training loop is a separate concern), but it's a real data point for anyone consideringtorch.compilehere: the memory tradeoff this PR profiles is an eager-mode phenomenon.Does the
"auto"budget formula generalize past the one config it was fit on? Sweptauto_cross_entropy_chunk_sizeacrossvocab_sizefrom 8k (small/custom vocabs) to 152k (roughly Qwen2.5/Llama-3 scale), fixedmemory_budget_bytes:The chosen
chunk_sizescales down correctly asvocab_sizegrows (left), and the measured real peak memory (torch.cuda.max_memory_allocated(), not just a profiler op's self-CUDA-mem) stays exactly flat across everyvocab_sizetested (right) — the formula generalizes, it isn't just a fit to one shape. It was consistently 1.5x the target budget, though, so this run directly motivated recalibrating_CROSS_ENTROPY_BYTES_PER_CHUNK_ELEMENTfrom2to3inlitgpt/utils.pyto close that gap — the number in this PR's code is the corrected one.Does the cross-entropy spike actually matter in a real training step, or is it dwarfed by everything else? Ran a full forward+backward step on a real model (
pythia-14m,B=8, its fullblock_size=512,V=50304) instead of just the isolated loss op:Chunking still cuts the whole step's peak CUDA memory by ~22.8% (3580MB → 2765MB) — the effect survives in context, not just in the isolated microbenchmark. The step-level op table (
docs/profiling/end_to_end_step_op_table.md) adds honest context though: in a real step,aten::mm(matmuls) is the single largest CUDA-memory consumer (909MB), edging outaten::_log_softmax(786MB) — so the cross-entropy spike is real and worth fixing, but it's not the only thing filling the memory line once the rest of the model is in the picture.Test plan
pytest litgpt/ tests/ --timeout=180 -q— full suite (mirrors CI'scpu-tests.ymlinvocation).pytest tests/test_model.py -k "kv_cache or stale"— confirms the pre-existingtest_kv_cache[23]xfail (expects rawIndexError) is unaffected by the earlier stale-cacheRuntimeErrorfix (it hits the unwrappedcos/sinindexing, not the wrappedmask_cacheindexing).python -m litgpt.scripts.profile_memory— regenerates the profiling artifacts underdocs/profiling/.ruff checkon all changed files.