Fix hellaswag memory leak and progressive slowdown (Issue #427)

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

nanochat/core_eval.py

@@ -146,6 +146,8 @@ def forward_model(model, input_ids):
     """
     Take BxT tensor of token ids, return BxT tensor of losses and argmax predictions.
     The last column of losses is set to nan because we don't have autoregressive targets there.
+
+    MEMORY FIX: Explicitly cleanup intermediate tensors to prevent GPU memory accumulation.
     """
     batch_size, seq_len = input_ids.size()
     outputs = model(input_ids)
@@ -161,6 +163,9 @@ def forward_model(model, input_ids):
     losses[:, -1] = float('nan')
     # Get the argmax predictions at each position
     predictions = outputs.argmax(dim=-1)
+    # MEMORY FIX: Explicitly free large intermediate tensors
+    del outputs  # outputs is largest tensor (B×T×V, ~GB for large models)
+    del target_ids  # target_ids is B×T
     return losses, predictions
 
 
@@ -238,6 +243,9 @@ def evaluate_example(idx, model, tokenizer, data, device, task_meta):
     else:
         raise ValueError(f"Unsupported task type: {task_type}")
 
+    # MEMORY FIX: Explicitly free tensors after extracting scalar result
+    del losses, predictions, input_ids
+
     return is_correct
 
 
@@ -245,18 +253,41 @@ def evaluate_task(model, tokenizer, data, device, task_meta):
     """
     This function is responsible for evaluating one task across many examples.
     It also handles dispatch to all processes if the script is run with torchrun.
+
+    MEMORY FIX: Added periodic cache cleanup to prevent memory accumulation.
     """
+    import gc  # For explicit garbage collection
+
     rank = dist.get_rank() if dist.is_initialized() else 0
     world_size = dist.get_world_size() if dist.is_initialized() else 1
     correct = torch.zeros(len(data), dtype=torch.float32, device=device)
+
     # stride the examples to each rank
     for idx in range(rank, len(data), world_size):
         is_correct = evaluate_example(idx, model, tokenizer, data, device, task_meta)
         correct[idx] = float(is_correct)
+
+        # MEMORY FIX: Periodic cache cleanup every 100 examples
+        # This releases cached GPU memory and triggers Python GC
+        # Prevents progressive slowdown from memory fragmentation
+        if idx % 100 == 0 and idx > 0:
+            # Release PyTorch cached memory back to GPU
+            if torch.cuda.is_available() and device.type == 'cuda':
+                torch.cuda.empty_cache()
+            # Force Python garbage collection
+            gc.collect()
+
     # sync results across all the processes if running distributed
     if world_size > 1:
         dist.barrier()
         dist.all_reduce(correct, op=dist.ReduceOp.SUM)
+
     # compute the mean
     mean_correct = correct.mean().item()
+
+    # MEMORY FIX: Final cleanup after task completes
+    del correct
+    if torch.cuda.is_available() and device.type == 'cuda':
+        torch.cuda.empty_cache()
+
     return mean_correct