TGI node OOM during large batch inference with autoscaling disabled

Problem – TGI Node Out‑of‑Memory (OOM) During Large Batch Inference with Autoscaling Disabled

In a production Kubernetes deployment of Text Generation Inference (TGI), a batch job that aggregates thousands of text‑generation requests into a single inference call caused the pod to be killed with an OOM error. Autoscaling was deliberately disabled, so the node could not acquire additional memory resources, leading to repeated job failures.

Root Cause – Unbounded Batch Size Combined with Fixed Memory Limits

The core issue is a mismatch between the max_batch_size used by the client and the memory limits configured on the TGI pod:

• Batch aggregation logic concatenates incoming requests until a configurable threshold is reached. The threshold was set to max_batch_size=1024 tokens per request, but the client sent batch_size=2048 requests in a single call.
• Pod memory limits were set to requests: 8Gi, limits: 12Gi. The TGI container’s internal memory allocator (torch.cuda) does not enforce a hard cap; it allocates as needed until the cgroup limit is hit.
• Autoscaling disabled (both HPA and cluster‑autoscaler) prevented the scheduler from provisioning a larger node when the pod exceeded its limit.
• The combination caused the Linux OOM killer to terminate the container, producing the following log snippet:

  
  [2024-09-12T14:32:07.842Z] INFO  tgi.server: Received batch request: size=2048, max_new_tokens=256
  [2024-09-12T14:32:08.001Z] WARN  torch.cuda: Allocating 13.2 GiB on device 0 (available: 12.0 GiB)
  [2024-09-12T14:32:08.015Z] ERROR kernel: Out of memory: Kill process 31245 (python) score 987 or sacrifice child
  [2024-09-12T14:32:08.016Z] FATAL tgi.server: Process killed due to OOM
  

  Debug – Investigation Steps

  1. Verify Pod Resource Configuration

  kubectl get pod tgi-0 -n inference -o yaml | grep -A4 resources

  Output:

  
  resources:
    limits:
      cpu: "8"
      memory: 12Gi
    requests:
      cpu: "4"
      memory: 8Gi
  

  2. Inspect TGI Server Logs for Batch Parameters

  kubectl logs tgi-0 -n inference | grep "Received batch request"

  Shows multiple entries with size=2048, confirming the client is sending a batch larger than the pod can handle.

  3. Capture Runtime Memory Usage

  kubectl exec -it tgi-0 -n inference -- bash -c "cat /sys/fs/cgroup/memory/memory.usage_in_bytes"

  Typical peak usage: 13,215,874,560 (≈12.3 GiB), exceeding the limit.

  4. Validate Autoscaling Configuration

  kubectl get hpa -n inference tgi-hpa -o yaml

  Result: spec.minReplicas: 3, spec.maxReplicas: 3 (effectively disabled).

  5. Reproduce Locally with Controlled Batch Size

  python - <<'PY'
  import torch
  from transformers import AutoModelForCausalLM, AutoTokenizer
  model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto")
  tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
  inputs = tokenizer(["Hello"]*2048, return_tensors="pt", padding=True).to('cuda')
  with torch.no_grad():
      outputs = model.generate(**inputs, max_new_tokens=256)
  print("generated", len(outputs))
  PY

  Local run fails with CUDA out of memory, confirming the batch size exceeds GPU memory.

  Solution – Align Batch Size with Pod Memory Limits and Enable Reactive Scaling

  1. Enforce a Safe max_batch_size in the TGI Client

  Update the client configuration to cap the batch size at a value that fits within the pod’s memory budget. For a 12 GiB limit on an A100 (40 GiB total), a safe batch size is ~1024 tokens per request.

  Before (client code)
  batch_size = 2048

  After
  batch_size = min(requested_batch, 1024)  # enforce limit

  2. Adjust TGI Container Memory Requests/Limits

  Increase the memory limit to provide a buffer for occasional spikes, and add a memory.swap limit to prevent silent swapping.

  Before (deployment snippet)
  
  resources:
    requests:
      cpu: "4"
      memory: 8Gi
    limits:
      cpu: "8"
      memory: 12Gi
  
  After
  
  resources:
    requests:
      cpu: "4"
      memory: 12Gi
    limits:
      cpu: "8"
      memory: 16Gi
    # optional: prevent swapping
    oomScoreAdj: -1000
  
  3. Enable Horizontal Pod Autoscaling (HPA) for Memory‑Based Scaling
  Add an HPA that reacts to memory usage, allowing the system to spin up additional TGI pods when a single pod approaches its limit.
  
  apiVersion: autoscaling/v2
  kind: HorizontalPodAutoscaler
  metadata:
    name: tgi-hpa
    namespace: inference
  spec:
    scaleTargetRef:
      apiVersion: apps/v1
      kind: Deployment
      name: tgi-deployment
    minReplicas: 3
    maxReplicas: 10
    metrics:
    - type: Resource
      resource:
        name: memory
        target:
          type: Utilization
          averageUtilization: 70
  
  4. Add Runtime Guardrails in TGI Server
  Configure the TGI server to reject batches that would exceed a configurable memory budget.
  # tgi_server.yaml
  max_batch_memory_gb: 10   # abort batches that would allocate >10 GiB
  max_batch_size: 1024
  
  5. Deploy Updated Manifests
  kubectl apply -f tgi-deployment.yaml
  kubectl apply -f tgi-hpa.yaml
  
  Verify – Confirm the Fix Prevents OOM
  1. Load Test with Targeted Batch Size
  hey -c 50 -n 5000 -m POST http://tgi-service.inference.svc.cluster.local/generate \
    -d '{"inputs":"Hello world","parameters":{"max_new_tokens":256}}'
  Observe that the maximum observed memory stays below the limit:
  kubectl top pod -n inference | grep tgi
  tgi-0   8Mi   10Gi
  tgi-1   7Mi   9.8Gi
  
  2. Check HPA Scaling Activity
  kubectl get hpa -n inference tgi-hpa
  During peak load, the replica count should increase (e.g., from 3 to 6) and later scale back down.
  3. Verify No OOM Events in System Logs
  kubectl logs tgi-0 -n inference | grep -i "out of memory"
  Output should be empty.
  Prevent – Operational Guardrails
  
  Monitoring: Export container_memory_working_set_bytes and set an alert at 80 % of the limit.
  Admission Control: Use a validating webhook that rejects Deployments where max_batch_size * avg_input_length * 4 bytes exceeds memory.requests.
  Capacity Planning: Periodically run a synthetic batch that gradually increases size until memory usage plateaus; record the safe ceiling.
  Documentation: Encode the safe batch size in the client SDK as a constant derived from the deployment’s resource profile.
  Fail‑Fast Policy: Enable TGI’s max_batch_memory_gb guardrail to return HTTP 429 instead of OOM when a request would exceed the budget.
  
  Related Topic Hub: Model Serving Troubleshooting Hub
  FAQ
  
  Why does the OOM happen only when autoscaling is disabled?
  
     Without autoscaling, the scheduler cannot place additional pods on larger nodes. The single pod must absorb the entire batch, exceeding its cgroup memory limit and triggering the OOM killer.
  Can I keep autoscaling disabled and still avoid OOM?
  
     Yes, by enforcing a batch size that fits within the pod’s memory limits and by adding server‑side guards (e.g., max_batch_memory_gb) to reject oversized batches.
  How do I determine a safe max_batch_size for a given GPU?
  
     Run a local benchmark that increments the batch size until torch.cuda.memory_allocated() approaches 85 % of the GPU’s total memory, then use that value (rounded down) as the production limit.
  What metric should I alert on to catch future memory pressure early?
  
     Alert on container_memory_working_set_bytes / container_spec_memory_limit_bytes > 0.75 for the TGI pods, and on the HPA’s desiredReplicas > maxReplicas condition.
  Is there a way to make TGI automatically split large batches?
  
     TGI itself does not split batches, but you can wrap the client call with a helper that partitions the input list into chunks of max_batch_size and issues multiple parallel requests.