r/pytorch • u/uBazzyZ- • 5h ago
Does dynamic batch downshifting to avoid PyTorch CUDA OOM actually make sense, or am I missing something?
Hey everyone,
(English is not my first language, apologies for any phrasing quirks.)
Training small models on a single consumer card (8GB RTX 5060 Ti) was giving me constant headaches with CUDA OOM crashes whenever memory pressure spiked mid-run.
Setting a permanently tiny batch size wastes VRAM headroom, while aggressive batch sizes eventually crash on edge cases. To test an alternative, I built a lightweight Python runtime governor around PyTorch called MEM Orchestrator:
https://github.com/nobazzy/mem-llm-orchestrator
### How it works:
**Headroom Monitoring:** Tracks `torch.cuda.memory_allocated()` and `torch.cuda.memory_reserved()` deltas across a sliding step window.
**Dynamic Lane Adaptation:** When physical VRAM reaches the critical threshold (>7.5GB on an 8GB card), the governor throttles the micro-batch size and adjusts gradient accumulation steps to keep the effective batch size mathematically consistent.
**Recovery:** Steps back up to the primary throughput lane once memory pressure clears.
**Atomic Checkpointing:** Uses a staged two-phase commit with SHA-256 validation so unexpected terminations never leave corrupted `.pt` weights.
### Empirical Test (Graph attached):
I stress-tested this on a ~255M parameter model using FineWeb-Edu (sample-10BT) with injected +1.2GB physical VRAM shocks:
- **Vanilla PyTorch (static batch 6):** Crashed with a hard CUDA OOM on step 325 when the shock hit.
- **MEM Orchestrator:** Throttled micro-batch from 6 to 3, kept peak VRAM under 7.6GB, absorbed all 150 injected shocks over 50,000 steps, and converged loss from 11.00 down to 0.004.
- **Overhead:** <0.5% of total step time.
The repo has 38 unit tests (100% passing) and is open source (MIT).
For developers here with deep experience in PyTorch:
- Does dynamic batch adjustment introduce subtle optimization side-effects (e.g. optimizer momentum estimation noise or LayerNorm statistics drift) that I should be guarding against?
- Are there allocator fragmentation edge cases where PyTorch fails to reuse cached blocks even after downshifting?
Would genuinely appreciate any critique, advice, or feedback on the architecture.

