TGI container OOM crash loading 30B model on 16GiB pod

Problem Description

In a staging environment a tgi:latest container is deployed in a Kubernetes pod with 16 GiB of RAM. During startup the container crashes with an out‑of‑memory (OOM) termination:

kubectl describe pod tgi-staging-abc123
...
State:          Waiting
  Reason:       OOMKilled
...
Events:
  Type     Reason     Age   From               Message
  ----     ------     ----  ----               -------
  Warning  OOMKill    2m    kubelet, node-01   Container tgi was killed due to memory pressure

The TGI logs show the exact allocation failure:

[2024-08-26 10:12:34] INFO  - Loading model "bigscience/30b" (30B parameters)
[2024-08-26 10:12:38] ERROR - Failed to allocate memory for tensor of size 12000 MB
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _apply
    param = apply(param)
  File "/usr/local/lib/python3.9/site-packages/torch/_utils.py", line 442, in apply
    return fn(t)
RuntimeError: CUDA out of memory. Tried to allocate 12.00 GiB (GPU 0; 0.00 GiB total capacity; 0.00 GiB reserved)

Key observations:

• Pod memory limit is 16 GiB, no GPU is attached.
• Model checkpoint is a 30‑billion‑parameter checkpoint (~30 GiB of FP16 weights).
• The container is killed by the kernel OOM killer before any request is processed.

Root Cause Analysis

The TGI container follows the official model loading guide. When a model is loaded, TGI uses torch to map the entire weight tensor into RAM (or VRAM). For a 30B parameter model the unquantized FP16 representation requires roughly 30 GiB of memory (30 B × 2 bytes ≈ 60 GiB for FP32, halved for FP16). Even with torch.backends.cudnn.benchmark optimizations, the loader attempts to allocate a contiguous block for each weight shard. On a CPU‑only node with only 16 GiB, the allocation request of ~12 GiB (one shard) exceeds the available memory, triggering PyTorch’s RuntimeError and ultimately the Kubernetes OOM killer.

Relevant documentation points:

• The TGI Docker README states that --max-model-size or quantization flags must be used for models larger than the host memory (Docker README).
• PyTorch’s CUDA memory management notes explain that on CPU fallback the same allocation logic applies, and torch.cuda.empty_cache cannot free CPU RAM (PyTorch docs).
• GitHub issue #1245 confirms that a 30B model requires >30 GiB RAM unless --load-in-8bit or --max-model-size is set.

Investigation and Debugging

Step‑by‑step diagnostics that reproduced the failure:

1. Inspect pod resources

kubectl get pod tgi-staging-abc123 -o yaml | grep -A3 resources

Output confirms limits:

resources:
  limits:
    memory: "16Gi"
  requests:
    memory: "16Gi"

2. Check container logs for allocation errors

kubectl logs tgi-staging-abc123 -c tgi

Shows the Failed to allocate memory for tensor of size 12000 MB error (see Problem Description).

3. Run a lightweight model to verify environment

kubectl exec -it tgi-staging-abc123 -- bash -c "\
python - <<'PY'\n\
import torch, transformers\n\
model = transformers.AutoModelForCausalLM.from_pretrained('gpt2')\n\
print('Loaded small model successfully')\n\
PY"

Success indicates the container image and PyTorch runtime are functional.

4. Query PyTorch memory limits

kubectl exec -it tgi-staging-abc123 -- bash -c "python - <<'PY'\n\
import torch, os\n\
print('Available RAM:', os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024**3), 'GiB')\n\
PY"

Confirms the node reports ~16 GiB total.

5. Attempt loading with quantization flags

kubectl exec -it tgi-staging-abc123 -- bash -c "\
text-generation-launcher --model-id bigscience/30b \\
  --device cpu \\
  --load-in-8bit"

Fails with Memory allocation failed: cannot allocate tensor of size ..., matching the community issue #1123 where 8‑bit loading still exceeds available RAM for a 30B model on 16 GiB.

Resolution

The fix consists of three coordinated changes:

1. Reduce the effective model size with 4‑bit quantization

Switch to bitsandbytes 4‑bit loading, which reduces memory footprint to ~6‑7 GiB for a 30B model.

# Before (Docker run)
docker run -e MODEL_ID=bigscience/30b \
           -e DEVICE=cpu \
           -p 80:80 \
           huggingface/text-generation-inference:latest

# After (add quantization flags)
docker run -e MODEL_ID=bigscience/30b \
           -e DEVICE=cpu \
           -e QUANTIZATION=bitsandbytes \
           -e BITS=4 \
           -p 80:80 \
           huggingface/text-generation-inference:latest

2. Explicitly set MAX_TOTAL_TOKENS and MAX_INPUT_LENGTH

Limiting the token buffer prevents TGI from pre‑allocating large intermediate tensors.

# Environment variables added to pod spec
env:
  - name: MAX_TOTAL_TOKENS
    value: "4096"
  - name: MAX_INPUT_LENGTH
    value: "1024"

3. Adjust Kubernetes resource requests to give the container headroom for the quantized model

Increase the memory limit to 12 GiB (still below the node’s capacity) and add a memorySwap limit to allow limited swapping if needed.

# Before (pod spec)
resources:
  limits:
    memory: "16Gi"
  requests:
    memory: "16Gi"

# After (reduced to realistic requirement)
resources:
  limits:
    memory: "12Gi"
    memorySwap: "14Gi"
  requests:
    memory: "12Gi"

Why it works:

• 4‑bit quantization compresses weight tensors, dropping the required RAM from ~30 GiB to ~6‑7 GiB.
• MAX_TOTAL_TOKENS and MAX_INPUT_LENGTH cap the size of activation buffers, avoiding additional allocations that would push the process over the limit.
• Lowering the pod memory limit aligns the request with the actual consumption, preventing the OOM killer from mis‑interpreting a temporary spike as a leak.

Validation

After applying the changes, perform the following checks:

1. Pod status

kubectl get pod tgi-staging-abc123
# Expected: STATUS = Running

Container logs for successful load

kubectl logs tgi-staging-abc123 -c tgi | grep "Loading model"
# Expected line: "Model loaded successfully (4-bit, 6.8 GiB RAM used)"

Health endpoint

curl -s http://:80/health
# Expected JSON: {"status":"healthy"}

Run an inference request

curl -X POST http://:80/generate \
  -H "Content-Type: application/json" \
  -d '{"inputs":"Hello, world!", "parameters":{"max_new_tokens":32}}'

Response should contain generated text without latency spikes or OOM errors.

Monitor memory usage

kubectl top pod tgi-staging-abc123
# Expected: MEMORY ~ 7GiB / 12GiB

Prevention and Best Practices

• Always size the pod memory limit to be greater than the expected model footprint plus a 20‑30 % safety margin.
• Prefer quantized loading (--load-in-4bit or --load-in-8bit) for models >10 B parameters when running on CPU‑only nodes.
• Set MAX_TOTAL_TOKENS and MAX_INPUT_LENGTH based on your service’s SLA; smaller limits reduce activation memory.
• Enable torch.backends.cudnn.benchmark = False in CPU mode to avoid aggressive workspace allocation. This can be done via the TORCH_CUDNN_BENCHMARK=0 environment variable.
• Use Kubernetes oomScoreAdj to lower the container’s OOM priority if other critical workloads share the node.
• Instrument Prometheus metrics (tgi_memory_usage_bytes) and set alerts when usage exceeds 80 % of the pod limit.

Related Topic Hub: Model Serving Troubleshooting Hub

FAQ

1. Why does the OOM happen even though the pod limit is 16 GiB?
   The model weight tensor alone requires ~30 GiB. PyTorch attempts to allocate a contiguous block (e.g., 12 GiB) which exceeds the 16 GiB limit, causing the kernel OOM killer to terminate the process.
2. Can I run a 30B model on a 16 GiB pod without quantization?
   No. The unquantized FP16 representation exceeds available RAM. Quantization (4‑bit or 8‑bit) or offloading to disk/Swap is required.
3. What is the difference between --load-in-8bit and --load-in-4bit?
   8‑bit reduces memory by ~2×, 4‑bit reduces it by ~4×. 4‑bit is typically needed for >20 B parameter models on < 12 GiB RAM.
4. How do I know the exact RAM requirement of a model before deployment?
   Use the Hugging Face model card or run torchinfo.summary on a sample checkpoint. Multiply parameter count by bytes per element (2 for FP16, 1 for 8‑bit, 0.5 for 4‑bit) and add ~10 % for activation buffers.
5. Is swapping a viable workaround?
   Swapping can prevent immediate OOM kills but leads to severe latency and possible thrashing. It is only a stop‑gap; proper quantization is the recommended solution.