PyTorch CUDA out of memory in an inference service: autograd, batch size and fragmentation

We run a GPU service that classifies support tickets into 140 intents with a fine-tuned transformer encoder. It sat behind a small dynamic batcher that groups requests arriving within 20 ms. After we added a bulk re-classification endpoint for a large customer, pods started dying every three to six hours with:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.88 GiB. GPU 0 has a total capacity of 22.05 GiB of which 1.06 GiB is free. Including non-PyTorch memory, this process has 20.97 GiB memory in use. Of the allocated memory 16.42 GiB is allocated by PyTorch, and 4.12 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Kubernetes restarted the pod, in-flight requests failed with 500s, and the cycle repeated. The model itself is about 1.7 GB in fp32. Something else was using 16 GB.

Reproducing it

We replayed a day of production traffic against a single pod on a staging GPU and logged torch.cuda.memory_allocated() after every batch. Two patterns appeared. First, allocated memory never returned to the baseline after a batch; it crept upward by tens of megabytes per minute. Second, the spikes that killed the process were always batches of 64 long tickets, around 512 tokens each, from the bulk endpoint.

What we ruled out

• Another process on the GPU. nvidia-smi showed only our Python process, and the error's "this process has 20.97 GiB in use" matched.
• A leak in the tokenizer or CPU-side code. Host memory was flat; only device memory grew.
• The model being too large for the card. A single 512-token request needed about 2.3 GB of peak memory including weights. The problem was what accumulated around it.

Root cause

We found three contributing causes, in order of impact:

1. Autograd graphs kept alive. The handler called the model without torch.no_grad() or torch.inference_mode(). Normally that only wastes memory during the forward pass, because the graph is freed when the outputs go out of scope. But for drift monitoring we appended each batch's logits to a deque(maxlen=5000) that a background thread summarised every 10 minutes. Each stored logits tensor had a grad_fn, which kept references to the saved activations from every layer of its forward pass. The deque was holding gigabytes of activations.
2. Batches sized by count, not by tokens. The batcher capped batches at 64 requests. With short tickets (median 60 tokens) that is small; with 64 bulk tickets padded to 512 tokens, attention and feed-forward activations are roughly an order of magnitude larger.
3. Fragmentation. Variable sequence lengths produce differently sized allocations every batch. The caching allocator had 4.12 GiB reserved but split into blocks too small for the 1.88 GiB request, as the error message itself pointed out.

What is happening under the hood

When a tensor that requires gradients flows through a model, PyTorch records the operations and saves the intermediate tensors needed for the backward pass. The output's grad_fn is the root of that graph. As long as anything references the output, the saved activations stay in GPU memory. torch.no_grad() turns off recording; torch.inference_mode() does the same and also skips version counter and view tracking, which makes it slightly faster. Tensors created in inference mode cannot later be used in autograd, which is exactly what a serving path wants.

PyTorch's CUDA caching allocator keeps freed blocks reserved for reuse instead of returning them to the driver. That makes allocation fast but means that with many different allocation sizes, free memory can be scattered across blocks that do not fit a new large request. Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True lets the allocator grow segments in place, which reduces this kind of fragmentation for workloads with changing shapes.

Environment

Minimal reproduction

# Excerpt: the original handler.
recent_logits = deque(maxlen=5000)

def predict(batch_texts: list[str]):
    enc = tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to("cuda")
    out = model(**enc)                      # autograd recording is on
    recent_logits.append(out.logits)        # keeps grad_fn and every saved activation alive
    return out.logits.softmax(-1).argmax(-1).tolist()

Running this in a loop with random lengths and printing torch.cuda.memory_allocated() shows steady growth until the deque fills, at a level far above what the model needs.

Who is affected

Any PyTorch serving code that runs models outside inference_mode or no_grad and holds outputs anywhere: caches, metrics buffers, debug lists, closures. Services with variable-length inputs and count-based batching are especially exposed to peak spikes and fragmentation.

The fix

1. Inference mode for the whole serving path, and detached CPU copies for anything stored.

@torch.inference_mode()
def predict(batch_texts: list[str]):
    enc = tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to("cuda")
    with torch.autocast("cuda", dtype=torch.float16):
        logits = model(**enc).logits
    recent_logits.append(logits.float().cpu())   # small, on the host, no graph
    return logits.argmax(-1).tolist()

We also call model.eval() once at startup (it was already there, but we added an assertion), because dropout in training mode would change predictions, not just memory.

2. Batch by token budget. The batcher now sorts pending requests by length and fills a batch until batch_size * padded_length reaches 16,384 tokens. Short tickets still batch 64 or more at a time; long ones batch 32.

def take_batch(pending: list[Request], budget: int = 16_384) -> list[Request]:
    pending.sort(key=lambda r: r.num_tokens)
    batch: list[Request] = []
    for r in pending:
        longest = max([b.num_tokens for b in batch] + [r.num_tokens])
        if (len(batch) + 1) * longest > budget:
            break
        batch.append(r)
    return batch

3. Half precision. Mixed precision with torch.autocast roughly halves activation memory. We compared predictions on 20,000 labelled tickets between fp32 and fp16: identical top-1 labels on 99.96% of them and no change in macro F1 to two decimals.

4. Allocator configuration. We set the environment variable in the deployment:

env:
  - name: PYTORCH_CUDA_ALLOC_CONF
    value: "expandable_segments:True"

We did not add torch.cuda.empty_cache() calls in the request path. It returns cached memory to the driver, which helps other processes see free memory but does not fix leaks, and calling it per request slows every allocation afterwards.

Alternatives and trade-offs

• max_split_size_mb stops the allocator from splitting large blocks and is an older remedy for fragmentation. We tested it but expandable_segments behaved better with our shapes.
• Bucketed padding (pad to 64, 128, 256, 512) reduces the number of distinct shapes further and helps fragmentation, at the cost of some wasted compute on padding.
• A dedicated serving framework with built-in dynamic batching and memory management. Worth it for larger fleets; for one model on a few GPUs our batcher was simpler to keep.
• Separate pools for bulk and interactive traffic. We did this later so bulk jobs cannot affect interactive latency.

How to test the fix

We used PyTorch's memory snapshot tooling during a replay to see what was holding memory before and after:

torch.cuda.memory._record_memory_history(max_entries=100_000)
run_replay("samples/2024-12-18.jsonl")
torch.cuda.memory._dump_snapshot("snapshot.pickle")

Loading the snapshot in PyTorch's memory visualiser showed thousands of live activation blocks attributed to the encoder layers before the fix, and only weights plus the current batch after it. We also added a test that runs 500 random-length batches and asserts that torch.cuda.memory_allocated() returns to within 50 MB of the post-warmup baseline:

pytest tests/test_gpu_memory.py -m gpu

Confirming in production

• Peak allocated memory per pod, read from torch.cuda.max_memory_allocated() and exported as a metric, dropped from 20.9 GiB to 5.8 GiB under the same replayed traffic.
• Zero OOM restarts in the 30 days after deploy, compared with 4 to 7 per day before.
• Throughput on the bulk endpoint rose from about 210 to about 480 tickets per second per GPU, thanks to half precision and length-sorted batching.

Regression guard

• The GPU memory test runs nightly on a GPU runner and fails if memory does not return to baseline.
• A lint rule flags model calls in the serving/ package that are not inside a function decorated with torch.inference_mode.
• The service exports memory_allocated, memory_reserved and max_memory_allocated per pod, with an alert when reserved minus allocated stays above 30% of the card for 15 minutes, which is our early warning for fragmentation.

Security implications

An OOM crash is a denial-of-service vector: before the fix, a client able to send many maximum-length tickets to the bulk endpoint could reliably restart our pods. Token-budget batching and a per-customer rate limit on the bulk endpoint close that. Storing logits on the host instead of the GPU does not change data exposure; they were already held in memory, and they are aggregated and discarded every ten minutes.

Prevention checklist

• Wrap every serving path in torch.inference_mode() and call model.eval().
• Store only detached, CPU-side copies of model outputs, and keep those buffers small.
• Batch by total tokens, not request count, and sort by length.
• Measure peak memory with max_memory_allocated under realistic traffic before launch.
• Try expandable_segments:True when the OOM message reports large reserved-but-unallocated memory.

References

• PyTorch documentation: CUDA semantics, memory management and environment variables — https://docs.pytorch.org/docs/stable/notes/cuda.html
• PyTorch documentation: torch.inference_mode — https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.inference_mode.html
• PyTorch documentation: Understanding CUDA memory usage — https://docs.pytorch.org/docs/stable/torch_cuda_memory.html
• PyTorch documentation: Automatic Mixed Precision package, torch.amp — https://docs.pytorch.org/docs/stable/amp.html