TensorFlow OOM — Fix ResourceExhaustedError

Your training run dies 40 minutes in with a wall of text ending in ResourceExhaustedError: OOM when allocating tensor with shape[32,512,128,128] and type float on /job:localhost/replica:0/task:0/device:GPU:0. No stack trace points at your code. The shapes look innocent. Restarting just fails again at the same step.

This error means one GPU allocation didn't fit — that's all. But which allocation, and why now, hides in the details: the bracketed shape names the exact tensor that broke the budget, and the allocator log above it shows everything already resident. Reading those two pieces turns a panic into arithmetic: shape times dtype bytes versus free VRAM.

The usual suspects form a short lineup: a batch too big for the card, TensorFlow's default grab of nearly all GPU memory starving everything else, an input pipeline buffering gigabytes in host RAM, float32 where mixed precision would halve the cost, or simply a model bigger than the hardware. Each has a different fix, and the wrong one — shrinking a model when the batch was the problem — wastes days.

This guide teaches the reading order: decode the shape in the message, cut batch size intelligently, tame pre-allocation with growth and caps, switch on mixed precision safely, right-size the model, and stream the input pipeline. You'll turn OOM from a feared crash into a budgeted resource.

Read the Shape in the Message

The error message is a measurement, not a complaint: ResourceExhaustedError: OOM when allocating tensor with shape[32,512,128,128] and type float on device GPU:0. The shape is the failed allocation's dimensions in elements — here 32 times 512 times 128 times 128, about 268 million numbers. The type sets bytes each: 4 for float32, 2 for float16, 1 for int8. Multiply: this tensor alone wanted roughly 1GB, and the allocator couldn't find it contiguous.

Above the error sits the allocator log — lines showing total reserved bytes, bytes in use, and the largest free chunk. Reserved near 100% with modest in-use bytes means fragmentation or pre-allocation overhead, not a truly full card. A largest-free-chunk smaller than your tensor means the memory exists but not contiguously, which points at variable shapes rather than sheer size.

The device tag matters too. GPU:0 tags blame device memory; CPU tags blame host RAM and exonerate the card entirely. Teams that skip this line tune batch sizes for a host-RAM problem and wonder why nothing changes. Read device first, shape second, allocator third — that order never wastes time.

Do the arithmetic on paper before changing code. Elements times bytes gives the failed tensor; parameter counts times optimizer multipliers give the resident floor. Most OOMs resolve the moment someone multiplies — the fix becomes obvious when the numbers are visible.

import tensorflow as tf

# Must run before ANY tensor op: report growth + log the real budget.
gpus = tf.config.list_physical_devices("GPU")
if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)

print("logical GPUs:", tf.config.list_logical_devices("GPU"))
# After a failure, read the peak to size the next attempt:
# info = tf.config.experimental.get_memory_info("GPU:0")
# print("peak bytes:", info["peak"], "current:", info["current"])

# Size a suspect tensor on paper, in code:
import math
shape = (32, 512, 128, 128)
print("float32 GB:", math.prod(shape) * 4 / 1e9)  # ~1.07

Batch Size: The Fastest Lever

Batch size scales activation memory almost linearly, which makes it the highest-leverage knob: halving the batch roughly halves the activation footprint of every layer at once. When the OOM names an activation tensor — shapes with the batch dimension first, like [32,512,128,128] — cut the batch first and ask questions later. Drop from 64 to 32, rerun, and watch the allocator peak fall.

But batch size also steers optimization: smaller batches mean noisier gradients and more steps per epoch, which shifts the learning-rate sweet spot. Naive halving can clear the OOM while silently degrading the model. The professional fix pairs the cut with gradient accumulation — run N micro-batches, sum their gradients, apply one update — restoring the original effective batch exactly while paying only the micro-batch's memory.

Sequence length and image resolution deserve the same scrutiny because they multiply the same activations. Doubling side length quadruples every spatial activation; doubling context length quadruples attention matrices. A resolution bump from 224 to 384 nearly triples activation memory with zero new parameters — the classic invisible OOM behind preprocessing pull requests.

Settle batch size from measured headroom, not habit. Read the allocator peak at a candidate batch, keep 10-20% reserve for allocator overhead and eval passes, and record the chosen size beside the card model in the run config.

import tensorflow as tf

# Halve the per-step batch, accumulate 4 micro-steps -> same effective batch.
base_batch, accum_steps = 64, 4
micro_batch = base_batch // accum_steps  # 16 fits the card

dataset = dataset.batch(micro_batch).prefetch(tf.data.AUTOTUNE)
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

@tf.function
def micro_grads(xb, yb):
    with tf.GradientTape() as tape:
        loss = loss_fn(yb, model(xb, training=True))
    return loss, tape.gradient(loss, model.trainable_variables)

accum = [tf.zeros_like(v) for v in model.trainable_variables]
for step, (xb, yb) in enumerate(dataset):
    loss, g = micro_grads(xb, yb)
    accum = [a + x / accum_steps for a, x in zip(accum, g)]
    if (step + 1) % accum_steps == 0:
        optimizer.apply_gradients(zip(accum, model.trainable_variables))
        accum = [tf.zeros_like(v) for v in model.trainable_variables]

Growth vs Pre-allocation: Share the Card

TensorFlow's default GPU behavior shocks newcomers: the first process reserves nearly all device memory at startup, whether it needs it or not. This upfront grab reduces fragmentation and speeds allocation, but on shared cards it means job one starves job two — a debugger, a second experiment, or an inference sidecar — which then OOMs on an apparently idle GPU. nvidia-smi shows the truth: one process holding gigabytes it never touches.

Memory growth fixes the sharing: tf.config.experimental.set_memory_growth(gpu, True) makes TensorFlow allocate as needed instead of upfront. Every process takes what it uses, and cards become shareable. The call must precede any tensor creation — put it at the very top of the entrypoint, before model building and dataset construction, or it raises and does nothing.

For hard partitioning, virtual device caps slice one card into bounded lanes: set_virtual_device_configuration with memory_limit in MB gives each process a ceiling it cannot cross. This is the right tool for CI runners and shared dev boxes where noisy neighbors must be contained, not just politely sharing. Exceeding your lane OOMs cleanly instead of starving others silently.

Note the trade-off honestly: growth can increase fragmentation under rapidly varying shapes, slightly raising peak versus upfront reservation. For single-job dedicated cards, the default grab is fine and sometimes leaner. Choose per machine: dedicated trainers keep defaults, shared boxes get growth or caps — and record the choice in the run config.

import tensorflow as tf

# Share a card politely: growth MUST come before any tensor op.
gpus = tf.config.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# Or partition strictly (e.g. 8GB lane on a shared CI card):
# tf.config.set_virtual_device_configuration(
#     gpus[0],
#     [tf.config.LogicalDeviceConfiguration(memory_limit=8192)])

print("growth on:", [tf.config.experimental.get_memory_growth(g) for g in gpus])

Mixed Precision: Halve Memory in Two Lines

Mixed precision is the closest thing to free memory in deep learning: compute in float16, keep a float32 master copy of weights, and watch most tensors halve. Activations, gradients, and optimizer moments shrink dramatically while final accuracy typically matches float32 within noise — modern accelerators even run fp16 math faster, so the memory fix arrives with a speed bonus.

Enable it in two lines with a policy plus loss scaling: tf.keras.mixed_precision.set_global_policy('mixed_float16') routes compute to fp16, and LossScaleOptimizer guards the backward pass against fp16 underflow by scaling the loss up before backprop and down afterward. Skipping loss scaling is the classic footgun — tiny gradients flush to zero in fp16, training silently stalls, and the team blames the memory fix for a numerics bug.

Verify the switch worked: model summaries show fp16 compute dtypes beside fp32 variables, allocator peaks drop roughly 30-50%, and loss curves track the fp32 baseline run. If NaNs appear, they're a scaling or epsilon issue (raise the loss scale, check batch-norm placement), not a reason to abandon precision — fp32 master weights already protect the update step.

Treat mixed precision as the default for new runs on capable hardware, not an emergency lever. Enabling it before the first OOM means the budget starts halved; enabling it after means re-validating a training recipe mid-crisis. The two-line change belongs in the template, not the incident thread.

import tensorflow as tf
from tensorflow.keras import mixed_precision

# Two-line switch: fp16 compute, fp32 master weights, scaled loss.
mixed_precision.set_global_policy("mixed_float16")
optimizer = mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(1e-3))
model.compile(optimizer=optimizer, loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

print("compute dtype:", model.layers[0].compute_dtype)  # float16
print("variable dtype:", model.layers[0].variable_dtype)  # float32

Right-Size the Model and Optimizer

Sometimes the model genuinely exceeds the card, and honesty beats hope. Budget it: parameters times bytes per element for weights, the same again for gradients, and twice more for Adam's two moment estimates — roughly 4x parameters in fp32 (2x for SGD, less with Adafactor or 8-bit optimizers). Add per-step activations from the allocator log. If that total exceeds VRAM at batch size 1 with fp16 on, no pipeline trick fits it — the architecture must change.

Shrink in order of accuracy cost. First, swap the optimizer: Adam to SGD or Adafactor sheds moment memory with modest tuning cost. Second, trim width before depth — narrower layers cut parameters quadratically with small accuracy impact. Third, freeze and share: frozen backbones need no gradients or moments, and gradient checkpointing recomputes activations instead of storing them (slower steps, far less memory). Depth cuts and tiny embeddings come last, since they change model character most.

Alternatively, shard instead of shrink: model-parallel splits, pipeline stages, or fully-sharded data parallel spread one model across cards. These are engineering projects, not config tweaks — adopt them when the model must stay big, not to avoid a batch-size conversation.

Record the budget beside the run: params, dtype, optimizer multiplier, batch, peak bytes, card model. Next quarter's capacity planning starts from arithmetic, not folklore.

Stream the Input Pipeline

A training job uses two memories, and OOM can strike either. Host RAM feeds the input pipeline: decoded images, tokenized batches, shuffled buffers, and .cache() contents all live there. When the container dies while nvidia-smi shows a half-empty GPU, the input pipeline — not the model — is the killer, and every GPU-side tweak misses.

The usual culprit is .cache() without a filename, which materializes the entire preprocessed dataset in host RAM after the first epoch. On large datasets that means tens of gigabytes appearing exactly when training looks healthiest — epoch two. Fix it by caching to sharded files, caching a smaller pre-batched stage, or dropping the cache for streaming with prefetch. Order matters too: .batch() before .map() materializes giant raw batches; .map() then .batch() keeps elements small until the end.

Build the pipeline to stream: interleave file reads across shards, prefetch with tf.data.AUTOTUNE to overlap host preprocessing with device compute, and vectorize maps so per-element Python overhead doesn't inflate buffers. Monitor host RAM (docker stats, free -m) on the same dashboard as GPU bytes — two lines that together pinpoint the pool in seconds.

Size shuffle buffers deliberately: a 10k shuffle buffer of 4K images holds gigabytes. Match the buffer to the mixing your training actually needs — full-dataset shuffling is rarely worth its RAM, and file-level plus modest buffer shuffling converges the same.