C.W.K.
Stream
Lesson 01 of 04 · published

The 75 GB Surprise

~13 min · inference-mode, autograd, memory, oom, war-story

Level 0Tool Renter
0 XP0/33 lessons0/12 achievements
0/100 XP to next level100 XP to go0% complete
"A forward-only path that should have used 2 GB used 75. The model wasn't too big. The framework was quietly saving its homework."

The Crash That Didn't Make Sense

A modest image model, generating a single image, took the machine down — twice. The numbers were absurd: a path that should fit in about 2 GB of memory had ballooned to roughly 75 GB before the display server itself fell over. Nothing about the model's size explained it. The model was small. The image was small. Something invisible was eating memory by the tens of gigabytes.

What Autograd Does When You Aren't Looking

The culprit was the automatic differentiation engine. By default, a deep-learning framework assumes you might want to train — so on every forward pass, it quietly retains the intermediate activations of every layer, because it'll need them to compute gradients during the backward pass. For training, that retention is essential. For inference, where you never call backward, it is pure dead weight — and on a deep diffusion model with many timesteps, that dead weight stacks up into tens of gigabytes of activations nobody will ever use.

A framework's default assumes its most demanding use case. Deep-learning frameworks default to 'you might train this,' because training needs the most bookkeeping. On an inference-only path, that default is invisible waste. Know what your tools assume by default, because the default is rarely tuned for your specific path.

The One-Line Fix

The fix is to tell the framework, explicitly, that this path will never need gradients: wrap the forward pass in an inference/no-grad context. With that, the engine stops retaining activations, and the 75 GB collapses back to the ~2 GB the work actually requires. One context manager, an order-of-magnitude memory difference. The bug wasn't in the model or the math — it was a missing declaration of intent.

The biggest wins are often declarations, not optimizations. No algorithm changed. No model shrank. A single statement told the framework what the code was actually doing, and the framework stopped doing expensive work for a case that didn't apply. Telling your tools your intent is sometimes worth more than any clever optimization.

Why It Took Down the Whole Machine

On a unified-memory Apple Silicon system, GPU and system memory share one pool. So a runaway inference allocation doesn't just crash the Python process — it starves the operating system, and the display server (the thing drawing your screen) goes down with it. That's why the failure was so dramatic: there's no separate VRAM to absorb the mistake. Unified memory is wonderful for big models and unforgiving when something leaks.

Shared resources mean local bugs become global crashes. When the GPU and the OS draw from the same memory pool, an inference leak isn't contained — it takes down everything. On systems with shared resources, a bug that would be a contained crash elsewhere becomes a full-system failure. Budget the shared resource as if the whole machine depends on it, because it does.

Pippa's Confession

I'd read about inference_mode a hundred times and treated it as a nice-to-have, a small speedup. Watching it turn a machine-killing 75 GB back into a calm 2 GB rewired that for me. It's not an optimization — it's the difference between an engine that runs and one that periodically murders the whole computer. Some one-liners aren't polish; they're load-bearing. I stopped skimming past the 'boring' context managers after this.

Code

One decorator, an order of magnitude·python
import torch

# WITHOUT the declaration: autograd retains every activation 'just in case'
# you call backward(). On a deep diffusion forward pass over many timesteps,
# that retention stacks into tens of GB. ~2 GB of real work -> ~75 GB used.
def infer_leaky(model, latents, conditioning, steps):
    x = latents
    for t in range(steps):
        x = model.denoise(x, t, conditioning)   # activations retained each step
    return x

# WITH the declaration: tell the framework this path never needs gradients.
# Activations are not retained. 75 GB collapses back to ~2 GB.
@torch.inference_mode()                          # the load-bearing one-liner
def infer_clean(model, latents, conditioning, steps):
    x = latents
    for t in range(steps):
        x = model.denoise(x, t, conditioning)
    return x
Guard where it can't be forgotten·python
# Put the guard at the ADAPTER ENTRY POINT, not per-module.
# A new model family that forgets its own no_grad is still covered.
class LocalAdapter(Adapter):
    @torch.inference_mode()              # covers EVERY family, present + future
    def _infer_sync(self, request):
        backbone = self.load(request.model)
        return self.sampler.sample(backbone, request)
    # per-module no_grad stays too -> second-of-two net (defense in depth)

External links

Exercise

Find an inference path in any ML code (yours or an example) and check whether it declares no-grad / inference mode. If it does, remove it mentally and reason about what would accumulate. If it doesn't, that's a latent memory bomb. Either way, write one sentence explaining what the declaration is buying.
Hint
The question to answer is always 'does this path ever call backward?' If the answer is no, gradient bookkeeping is pure waste and the declaration that disables it is mandatory, not optional.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.