Fix hellaswag memory leak and extract cleanup config#519
Open
haltingstate wants to merge 3 commits intokarpathy:masterfrom
Open
Fix hellaswag memory leak and extract cleanup config#519haltingstate wants to merge 3 commits intokarpathy:masterfrom
haltingstate wants to merge 3 commits intokarpathy:masterfrom
Conversation
Add is_mps_device() and should_use_torch_compile() to nanochat/common.py Disable torch.compile on macOS MPS devices (prevents indefinite hanging) Add conditional torch.compile in base_train.py and chat_sft.py Add memory monitoring with 32GB inference / 96GB training limits Reference: Task-20, Task-18, Task-19, Task-28, Task-39
ROOT CAUSE:
GPU tensors (outputs, losses, predictions, input_ids) not explicitly freed
after use, causing memory fragmentation and progressive slowdown. Each
forward pass creates ~411MB output logits tensor that lingers in memory
until Python GC triggers. Over 10,000+ HellaSwag examples, accumulates
4.4GB tensors, exhausts available headroom on 32GB unified memory systems.
SYMPTOMS:
- Progressive slowdown: Example 0: 2.5s → Example 300: 6.2s (+148%)
- Unbounded memory growth: 20-50MB per 100 examples
- Mac Studio (32GB) crashes with OOM after 8000-9000 examples
- HellaSwag-specific (10,000 examples vs MMLU: 100-1000)
MECHANISM:
1. PyTorch caching allocator fragments memory over time
2. Allocator performance degrades (O(1) → O(N) search for free blocks)
3. Python GC lazy, doesn't free promptly
4. No explicit cleanup: no torch.cuda.empty_cache(), no gc.collect()
5. Memory fragmentation + accumulated tensors = progressive slowdown
FIXES IMPLEMENTED:
1. forward_model (lines 166-168): Explicit tensor cleanup
- Added: del outputs, del target_ids
- Impact: Frees ~411MB output logits + 16KB target_ids per call
- outputs tensor: batch_size × seq_len × vocab_size float32
= 4 choices × 512 tokens × 50,257 vocab × 4 bytes = 411MB
2. evaluate_example (lines 246-247): Cleanup after result extraction
- Added: del losses, predictions, input_ids
- Impact: Frees tensors immediately after .item() extracts scalar
- Prevents retention until function returns
3. evaluate_task (lines 262-283): Periodic cache cleanup
- Added: gc.collect() + torch.cuda.empty_cache() every 100 examples
- Impact: Resets allocator state, prevents fragmentation accumulation
- Small cost: ~10-50ms per 100 examples
- Final cleanup after task completes (line 287-289)
EXPECTED IMPROVEMENT:
- Memory growth: <100MB total (vs unbounded before)
- Slowdown: <5% variation (vs 400%+ before)
- Completion: HellaSwag completes in ~7-8 hours without OOM
- Timing: Constant 2.5-2.6s per example throughout evaluation
TESTING:
Before deploying to production, verify:
- MMLU accuracy unchanged (within 0.5% of baseline)
- Memory growth <100MB over 1000 examples
- Time per example: last 100 within 10% of first 100
- HellaSwag completes without OOM crash
WHY HELLASWAG AFFECTED:
- 10,000+ examples (vs MMLU: 100-1000, GSM8K: 1319, HumanEval: 164)
- 4 forward passes per example (multiple choice)
- Runs 8.3 hours (vs MMLU: 40 min)
- More time for fragmentation to accumulate
- MMLU completes before memory pressure becomes severe
TECHNICAL DETAILS:
- @torch.no_grad() prevents gradient graphs, not tensor allocation
- del only removes Python references, GC frees actual memory
- torch.cuda.empty_cache() releases cached memory back to GPU
- gc.collect() forces immediate garbage collection (slow but thorough)
Fixes: Issue karpathy#427 (hellaswag memory leak and progressive slowdown)
Related: kcg-llm task-47.fix-hellaswag-memory-leak-progressive-slowdown.pending
Analysis: kcg-llm/b1.tasks/task-47*/task-47.10-memory-leak-analysis.txt
Extract hardcoded memory cleanup interval (100 → 256) and enable flags to eval_config.py for better maintainability and tuning flexibility. Changes: 1. Created nanochat/eval_config.py: - CACHE_CLEANUP_INTERVAL = 256 (changed from hardcoded 100) - ENABLE_PERIODIC_CLEANUP = True (allows disabling cleanup) - ENABLE_FINAL_CLEANUP = True (allows skipping final cleanup) - Documented rationale for 256: balances overhead vs fragmentation 2. Updated nanochat/core_eval.py: - Import eval_config module - Use eval_config.CACHE_CLEANUP_INTERVAL instead of hardcoded 100 - Check eval_config.ENABLE_PERIODIC_CLEANUP flag before cleanup - Check eval_config.ENABLE_FINAL_CLEANUP flag for final cleanup Rationale for 256 vs 100: - Power of 2 (efficient modulo operation) - Lower overhead: HellaSwag 10,000 examples: 39 cleanups (~2s) vs 100 cleanups (~5s) - Still frequent enough to prevent fragmentation - For MMLU (100-1000 examples): 0-4 cleanups (negligible impact) Benefits: - Centralizes tuning parameters in one location - Allows easy experimentation with cleanup intervals - Can disable cleanup for debugging/profiling - Documents tradeoffs in config comments - No magic numbers in evaluation code Related: Previous commit a7066b8 (hellaswag memory leak fix)
svlandeg
reviewed
Feb 9, 2026
Collaborator
There was a problem hiding this comment.
Hi! I'm somewhat confused by the diff - half of this PR is about memory cleanup, and the other half is an alternative to #319, is that correct? Was #319 not sufficient for you to fix the issues on MPS? It would probably be easier if you'd keep both issues in separate PRs + commented on earlier/other approaches to help us understand what works best.
svlandeg
reviewed
Feb 9, 2026
| @@ -0,0 +1,31 @@ | |||
| """ | |||
Collaborator
There was a problem hiding this comment.
As stated in the readme, nanochat shouldn't become a configuration monster, so we probably don't want to merge a whole new Python module of 31 lines just to set 3 variables 🙃
I think we can probably do without ENABLE_PERIODIC_CLEANUP and set ENABLE_FINAL_CLEANUP always to True?
CACHE_CLEANUP_INTERVAL can just be defined in core_eval.py.
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
2 participants
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 Issue #427: HellaSwag memory leak causing progressive slowdown and OOM crashes.
Changes
Commit 1: Fix hellaswag memory leak (a7066b8)
forward_model,evaluate_example,evaluate_taskCommit 2: Extract cleanup settings to config (c4a183d)
nanochat/eval_config.pywith tunable parametersRoot Cause
GPU tensors (outputs, losses, predictions, input_ids) not freed explicitly, causing:
Impact
Before:
After:
Testing
Before merging to production, verify:
Configuration
New settings in
nanochat/eval_config.py:CACHE_CLEANUP_INTERVAL = 256(cleanup frequency)ENABLE_PERIODIC_CLEANUP = True(enable/disable)ENABLE_FINAL_CLEANUP = True(final cleanup toggle)Files Changed
nanochat/core_eval.py(+35 lines): Memory cleanup logicnanochat/eval_config.py(+34 lines): Configuration parametersRelated