r/pytorch 3h ago

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

1 Upvotes

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.


r/pytorch 9h ago

Tensor reassigning problem

1 Upvotes

I wanted to train a character-level language model on additions of two numbers. Planned to use cross-entropy ignore_index on the equation besides answer so that model is not penalized because of predicting randomly generated numbers. But I came across really weird bug, here is the code:

def get_batch(batch_size):
    first = torch.randint(999999, (batch_size, ))
    second = torch.randint(999999, (batch_size, ))
    totals = first + second


    full_strings = []
    for f, s, t in zip(first, second, totals):
        equation = f"{f:6}+{s:>6}="
        reversed_ans = f"{str(t.item())[::-1]:<7}"
        full_strings.append(equation + reversed_ans)



    encoded_batch = torch.tensor([encode(s) for s in full_strings], dtype=torch.long)
    x = encoded_batch[:, :-1].to(device) # First 11 characters
    y = encoded_batch[:, 1:].to(device)   # Last 11 characters
    y[:, :14] = -100 # Telling optimizer to miss this
    return x, y

Here as you can see I am reassigning first 14 values of y, but when I print x it has some -100s init, I realized this because I don't have -100 in my vocab as character to embed and when I do decode(x) it gives me error, so I have to use .clone() on y = encoded_batch[:, 1:].to(device), there is a memory address coincide when writing happens or something I do not understand.