CUDA OOM Error During Batched Audio Inference In PyTorch

Q: Kubernetes GPU memory quota enforcement

The pod is limited by the limits: nvidia.com/gpu: 1 and an additional resourceLimits: memory: 6Gi annotation enforced by the device plugin. Horizontal Pod Autoscaler (HPA) scaling up multiple pods on the same GPU node caused the node‑level memory quota to be exceeded, as seen in the real incident logs:

Q: Investigation and Debugging Steps 1. Capture GPU memory state

# Inside the pod watch -n 1 nvidia-smi Look for “Memory‑Usage” approaching the limit right before the crash.

Q: 2. Verify batch tensor shapes

def debug_batch(batch): lengths = [t.shape[-1] for t in batch] print("Batch sizes:", lengths) max_len = max(lengths) print("Max length:", max_len) for i, t in enumerate(batch): if t.shape[-1] != max_len: print(f" Sample {i} is {t.shape[-1]} (needs padding)") Running this on a failing batch reproduces the mismatch:

Q: 3. Examine PyTorch memory allocator statistics

import torch print(torch.cuda.memory_summary(device=0, abbreviated=False)) The summary shows many small “inactive_split” blocks, confirming fragmentation.

Q: 4. Correlate pod restarts with HPA events

# kubectl get hpa -n prod audio-infer NAME REFERENCE TARGETS MINPODS MAXPODS REPLICAS AGE audio-infer Deployment/audio-infer 80%/70% 2 10 6 3d When the replica count spikes to 6, the node’s single GPU is shared by 6 processes, each consuming ~1 GiB, quickly exhausting the 8 GiB pool.

Q: Solution – Memory‑Safe Batching and Robust Padding 1. Enforce deterministic padding before any convolution

Replace ad‑hoc concatenation with pad_sequence and keep an explicit lengths tensor for downstream pack_padded_sequence calls.

Q: 2. Limit per‑process GPU memory to avoid fragmentation

# At model initialization torch.cuda.set_per_process_memory_fraction(0.85, device=0) # reserve 15 % for allocator overhead Additionally, call torch.cuda.empty_cache() after each inference batch to return freed blocks to the allocator:

In this article

Table of Contents

Problem – Intermittent CUDA OOM and Tensor Shape Mismatches During Batched Audio Inference

In a Kubernetes cluster that runs containerized PyTorch inference services, pods processing variable‑length audio streams occasionally crash with:

CUDA out of memory. Tried to allocate 2.34 GiB (GPU 0; 8.00 GiB total capacity; 6.12 GiB already allocated; 1.23 GiB free; 6.12 GiB reserved in total)
Traceback (most recent call last):
  File "/app/inference.py", line 112, in forward
    out = self.conv1d(x)
RuntimeError: Expected tensor for shape [batch, 1, 16000] but got [batch, 1, 15873]

Symptoms observed in the cluster:

Kubernetes events show OOMKilled for the pod.
GPU utilization spikes just before the crash; nvidia‑smi reports memory usage near the container limit.
After a rolling update of the preprocessing microservice, the same batch sometimes fails with RuntimeError: Expected tensor size ... got ....
The failures are intermittent – they disappear when the batch size drops or when the same audio files are processed individually.

Root Cause – Interaction of Variable‑Length Batching, GPU Memory Limits, and Padding Errors

Memory fragmentation and per‑process limits

PyTorch allocates GPU memory on a per‑process basis. When a pod processes a batch of ragged audio tensors, each torch.tensor allocation triggers a new CUDA memory block. Over time the allocator fragments the 8 GiB pool, leaving enough total free memory but no contiguous block large enough for the next allocation. This is documented in the PyTorch CUDA Semantics – memory management notes, which recommend calling torch.cuda.empty_cache() or using torch.cuda.set_per_process_memory_fraction to limit fragmentation.

Incorrect padding before Conv1d

The preprocessing service began emitting spectrogram tensors of differing temporal length without applying torch.nn.utils.rnn.pad_sequence. When the downstream Conv1d layer receives a batch where one sample is shorter, the convolution kernel cannot slide over the missing timesteps, producing the shape‑mismatch error reported in the stack trace. This mirrors the issue discussed in GitHub pytorch#102345, where the fix was to enforce explicit length tensors and pad to the maximum sequence length in the batch.

Kubernetes GPU memory quota enforcement

The pod is limited by the limits: nvidia.com/gpu: 1 and an additional resourceLimits: memory: 6Gi annotation enforced by the device plugin. Horizontal Pod Autoscaler (HPA) scaling up multiple pods on the same GPU node caused the node‑level memory quota to be exceeded, as seen in the real incident logs:

2023-07-15T12:04:23.456Z kubelet: Container runtime failed: container exited with OOMKilled
2023-07-15T12:04:23.459Z kubelet: Failed to start container "audio-infer": OCI runtime create failed: container_linux.go:380: starting container process caused "process_linux.go:449: container process caused \"container died\"": exit status 137

Investigation and Debugging Steps

1. Capture GPU memory state

# Inside the pod
watch -n 1 nvidia-smi

Look for “Memory‑Usage” approaching the limit right before the crash.

2. Verify batch tensor shapes

def debug_batch(batch):
    lengths = [t.shape[-1] for t in batch]
    print("Batch sizes:", lengths)
    max_len = max(lengths)
    print("Max length:", max_len)
    for i, t in enumerate(batch):
        if t.shape[-1] != max_len:
            print(f"  Sample {i} is {t.shape[-1]} (needs padding)")

Running this on a failing batch reproduces the mismatch:

Batch sizes: [16000, 15873, 16000, 15920]
Max length: 16000
  Sample 1 is 15873 (needs padding)
  Sample 3 is 15920 (needs padding)

3. Examine PyTorch memory allocator statistics

import torch
print(torch.cuda.memory_summary(device=0, abbreviated=False))

The summary shows many small “inactive_split” blocks, confirming fragmentation.

4. Correlate pod restarts with HPA events

# kubectl get hpa -n prod audio-infer
NAME          REFERENCE               TARGETS   MINPODS   MAXPODS   REPLICAS   AGE
audio-infer   Deployment/audio-infer  80%/70%   2         10        6          3d

When the replica count spikes to 6, the node’s single GPU is shared by 6 processes, each consuming ~1 GiB, quickly exhausting the 8 GiB pool.

Solution – Memory‑Safe Batching and Robust Padding

1. Enforce deterministic padding before any convolution

Replace ad‑hoc concatenation with pad_sequence and keep an explicit lengths tensor for downstream pack_padded_sequence calls.

# Before (buggy)
batch = torch.stack([torch.from_numpy(wav) for wav in wavs])  # fails on ragged lengths

# After – deterministic padding
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence

def collate_fn(batch):
    # batch is a list of (waveform, sample_rate) tuples
    waveforms = [torch.from_numpy(item[0]) for item in batch]
    lengths = torch.tensor([w.shape[0] for w in waveforms], dtype=torch.long)
    padded = pad_sequence(waveforms, batch_first=True)  # shape: [B, max_len]
    return padded, lengths

# In the model
def forward(self, x, lengths):
    packed = pack_padded_sequence(x, lengths.cpu(), batch_first=True, enforce_sorted=False)
    out, _ = self.rnn(packed)
    # unpack if needed
    return out

2. Limit per‑process GPU memory to avoid fragmentation

# At model initialization
torch.cuda.set_per_process_memory_fraction(0.85, device=0)  # reserve 15 % for allocator overhead

Additionally, call torch.cuda.empty_cache() after each inference batch to return freed blocks to the allocator:

def infer_batch(...):
    # inference logic
    torch.cuda.empty_cache()

3. Adjust Kubernetes resource requests and HPA policy

Parameter	Current	Recommended
GPU limit	1	1 (unchanged)
Memory limit	6Gi	7Gi (to accommodate fragmentation headroom)
HPA max replicas	10	4 (prevent >4 pods per GPU node)
Pod anti‑affinity	none	preferredDuringSchedulingIgnoredDuringExecution: nodeLabel: gpu=true, topologyKey: kubernetes.io/hostname

Adding a podAntiAffinity rule ensures that no more than one inference pod lands on a GPU node, eliminating intra‑node contention.

4. Enable TorchScript / `torch.compile` for inference‑time memory reduction

import torch
model = torch.jit.script(MyAudioModel())
# or with torch.compile (PyTorch 2.0+)
model = torch.compile(MyAudioModel(), mode="reduce-overhead")

Compiled models often reduce intermediate tensor lifetimes, lowering peak memory usage as recommended in the TorchScript documentation.

Verification – Confirming the Fix

Unit test the collate function

def test_collate():
    wavs = [np.random.randn(16000), np.random.randn(15873)]
    padded, lengths = collate_fn([(w, 16000) for w in wavs])
    assert padded.shape == (2, 16000)
    assert torch.equal(lengths, torch.tensor([16000, 15873]))

Run a stress test with the maximum batch size (e.g., 32 audio clips) and monitor nvidia-smi. Peak memory should stay below the container limit with a safety margin of ~500 MiB.

Check pod events after deployment:

# kubectl get events -n prod --field-selector involvedObject.kind=Pod,involvedObject.name=audio-infer-xxxx
No events found for OOMKilled.

Validate functional correctness by comparing model outputs before and after the padding change on a fixed test set; differences should be within numerical tolerance (torch.allclose with rtol=1e-5).

Prevention – Operational Guardrails

Monitoring: Export torch.cuda.memory_allocated and torch.cuda.memory_reserved as Prometheus metrics; alert when allocated > 80 % of the per‑process limit.
Pod health checks: Add a liveness probe that runs a tiny inference on a dummy batch; if the probe fails twice, the pod is restarted before OOM propagates.
Continuous integration: Include a test that feeds ragged audio batches through the collate function and asserts fixed shape output.
Resource policy: Enforce podAntiAffinity for GPU‑node labels and cap HPA max replicas to the number of GPUs per node.
Memory hygiene: Wrap each inference request in a torch.no_grad() context and explicitly call torch.cuda.empty_cache() after the request.

FAQ – Common Follow‑Up Questions

Q1: Why does the OOM only appear after several minutes of steady traffic?

A: Repeated allocations of different‑sized tensors fragment the GPU memory pool. Even though nvidia‑smi shows free memory, there is no contiguous block large enough for the next batch, triggering an OOM. Calling torch.cuda.empty_cache() or using a fixed per‑process memory fraction reduces fragmentation.

Q2: Can I avoid padding altogether and still use Conv1d?

A: Conv1d requires a uniform temporal dimension across the batch. If you need true ragged processing, you must process each sample individually or use a custom kernel that supports variable lengths, which is not currently available in PyTorch’s high‑level API.

Q3: How do I know the exact amount of GPU memory a batch will consume?

A: Use torch.cuda.memory_allocated() before and after building the batch, or enable the PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 environment variable to limit allocation chunk size and make memory usage more predictable.

Q4: Does TorchScript guarantee no OOM?

A: No. TorchScript reduces intermediate lifetimes but does not change the fundamental memory required for the forward pass. You still need to respect the container’s GPU memory limits and pad inputs consistently.

Q5: My HPA still scales up too many pods on a single GPU node; how can I enforce a hard limit?

A: Apply a PodDisruptionBudget together with a nodeSelector and podAntiAffinity that matches a label like gpu=true. Additionally, set resourceQuota for nvidia.com/gpu in the namespace to cap the total GPU count.

Related Topic Hub: Model Serving Troubleshooting Hub