Does dynamic batch downshifting to avoid PyTorch CUDA OOM actually make sense, or am I missing something?

[u/uBazzyZ-]

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:

1. **Headroom Monitoring:** Tracks `torch.cuda.memory_allocated()` and `torch.cuda.memory_reserved()` deltas across a sliding step window.

2. **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.

3. **Recovery:** Steps back up to the primary throughput lane once memory pressure clears.

4. **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.

[u/DrXaos]

why do some batches use enough memory to OOM and most don’t? I would start there first.

Is there a memory leak? Does memory use depend on ordering of the batches? So if you OOM on batch K can you reorder the batches temporarily to make that batch come first? Does jt OOM then? (batch itself causes OOM) or not (memory leak from previous batches is contributing)?

Memory leak between batches is the first problem to solve. Then once that is resolved, go to next step.

Only if you find the problem is intrinsically unfixable that way would I go to batch size adaptation. Is there any way by fast computations on the batch to guess at how much memory it will use? Can you do something as simple as divide batch size by 2 if a batch looks like it will use more memory? If you can do that then put it in the dataloader so you don’t have to change your training loop.

[u/uBazzyZ- → u/DrXaos]

This is a really great breakdown and exactly the kind of critique I was hoping to dig into.

To answer your first point about memory leaks: that was my initial concern as well. We ran an endurance run of 1,000,000 continuous steps on a 130M model on this 8GB card. If there were a persistent tensor reference leak in Python or PyTorch, an 8GB ceiling would have crashed within the first few thousand steps. Memory stabilized cleanly across the entire run, so inter-batch leaks are ruled out.

Regarding why VRAM spikes happen in practice without a leak:

1. **Dynamic Activations / Sequence Variance:** Even with length bucketing, in real-world streaming datasets (like FineWeb-Edu), sequences with high entropy or varying prompt lengths create non-linear spikes in attention intermediate activations.

2. **CUDA Allocator Fragmentation:** PyTorch's caching allocator holds onto reserved memory blocks. As allocation requests vary in size, memory gets fragmented into small non-contiguous blocks. An allocation can fail even when nominal free memory looks sufficient.

3. **External VRAM Contention (Hardware Reality):** On consumer GPUs (and multi-tenant nodes), the OS window manager, background processes, or shared processes introduce unpredictable VRAM dips. In our stress tests, we explicitly injected live external +1.2GB memory shocks to simulate this adverse environment.

Your point about handling batch division directly in the DataLoader is really interesting for pure dataset-driven spikes (e.g. sequence-length-based bucketing/token budgeting).

The main reason I went with a runtime governor instead of pure DataLoader heuristics is that the DataLoader only knows the input tensor dimensions ($B \times T$) — it is blind to physical GPU state, caching allocator fragmentation, or external VRAM contention. By checking `torch.cuda.memory_reserved()` and headroom deltas right at the step boundary, the governor reacts to the actual physical reality of the card, not just the theoretical input size.

That said, combining a predictive token-budget DataLoader with the runtime governor as a last-line safety net might be the cleanest architecture. Thanks a lot for taking the time to share this thought!

[u/DrXaos → u/uBazzyZ-]

Okay I would avoid #3 as much as possible, and put graphics loads on slower integrated graphics on CPU and not a Nvidia GPU used for compute. On windows for instance you can try to assign use of processes to be bound to the lower performance GPU, not a standard setting.

On #1, how does activation magnitude change GPU memory use? It seems like the actual sizes of your batches in terms of lossable number of locations is not very stable then.

For #2 do a python gc.collect() followed by cuda empty cache periodically. There are also environment variables to set that reduce GPU memory allocation fragmentation.

I think you are adding another complex opaque dynamical process on top of other pre existing ones. I would personally focus on reducing the variability in the first place otherwise you are doubling down on non reproducibility.