GPT-4 inference OOM during concurrent A/B test variants in production

Problem – OOM During Concurrent A/B‑Test Variants of GPT‑4 Inference

In production we run an A/B testing framework that sends several prompt variants to the OpenAI GPT‑4 endpoint in parallel. Under load the inference workers start throwing Memory allocation failed or CUDA out of memory exceptions, causing request failures and downstream SLA breaches.

Root Cause – Uncontrolled Model Instantiation and Memory Fragmentation

The failure is not a transient GPU pressure spike; it is deterministic when the number of parallel variants exceeds the memory budget of a single GPU node. The primary causes are:

• Per‑request model loading: The client library (e.g., openai Python SDK wrapped by torch for local caching) creates a new GPT4Model object for every variant. Each object allocates a full copy of the 80 GB FP16 model weights, quickly exhausting the 96 GB GPU.
• Process‑level duplication: The A/B test harness spawns a separate multiprocessing.Process per variant. CUDA contexts are not shared across processes, so each process repeats the full model load.
• Missing torch.no_grad() and torch.cuda.empty_cache(): Inference code retains intermediate activation tensors after each request, leading to fragmentation that prevents the allocator from re‑using freed memory.

Debug – Investigation Steps and Evidence

1. Log Inspection

2024-10-24 12:03:41,872 ERROR worker-7 CUDA out of memory. Tried to allocate 12.00 MiB.
2024-10-24 12:03:41,873 INFO  worker-7 GPU memory usage: 93.2/96.0 GiB
2024-10-24 12:03:41,874 DEBUG A/B variant "control" completed in 124ms
2024-10-24 12:03:41,875 ERROR worker-8 Memory allocation failed: CUDA error: out of memory

2. Runtime Metrics

Running nvidia-smi while the test is active shows a linear increase in memory per variant:

+-----------------------------------------------------------------------------+
| GPU  Name        Persistence-M| Bus-Id | Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap| Memory-Usage | GPU-Util  Compute M. |
|===============================+===============================+============|
|   0  NVIDIA A100  Off  | 00000000:3B:00.0 Off | 0%   30C    P0   250W / 400W |
|                               |  0MiB /  81920MiB |      0%      Default |
+-----------------------------------------------------------------------------+
...
|   0  NVIDIA A100  Off  | 00000000:3B:00.0 Off | 0%   30C    P0   250W / 400W |
|                               |  81920MiB / 81920MiB |   99%      Default |
+-----------------------------------------------------------------------------+

3. Process Inspection

$ ps -ef | grep inference_worker
root   12458 12430  0 12:03 ?  00:00:00 python inference_worker.py --variant control
root   12459 12430  0 12:03 ?  00:00:00 python inference_worker.py --variant treatment

Each worker loads the model independently, confirmed by tracing torch.cuda.memory_allocated() at startup:

>>> import torch, gc
>>> torch.cuda.memory_allocated()
0
>>> model = GPT4Model()
>>> torch.cuda.memory_allocated()
82208358400   # ~76.5 GiB per process

4. Code Review – Before Change

def run_variant(variant_id, prompt):
    # A new process per variant
    from multiprocessing import Process

    def worker():
        client = OpenAI(api_key=API_KEY)   # triggers local cache load
        response = client.completions.create(
            model="gpt-4",
            prompt=prompt,
            temperature=0.7,
        )
        print(f"{variant_id} result: {response}")

    p = Process(target=worker)
    p.start()
    p.join()

Solution – Shared Model Pool, Concurrency Limiting, and Memory Hygiene

1. Centralised Model Loader

Instantiate the model once per GPU node and share it across all variant workers using a thread‑safe pool. This eliminates duplicate CUDA contexts.

# model_pool.py
import torch
from threading import Lock

class GPT4Pool:
    _instance = None
    _lock = Lock()

    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance.model = GPT4Model().half().to('cuda')
                cls._instance.model.eval()
            return cls._instance

    def get(self):
        return self.model

2. Async Queue with Bounded Concurrency

Replace multiprocessing.Process with an asyncio.Semaphore that caps concurrent inferences to MAX_CONCURRENT=4 (adjusted to fit GPU memory).

# inference_worker_async.py
import asyncio
from model_pool import GPT4Pool
import openai

MAX_CONCURRENT = 4
semaphore = asyncio.Semaphore(MAX_CONCURRENT)

async def run_variant(variant_id, prompt):
    async with semaphore:
        model = GPT4Pool().get()
        # No additional model load; use shared instance
        async with openai.AsyncClient() as client:
            resp = await client.completions.create(
                model="gpt-4",
                prompt=prompt,
                temperature=0.7,
                max_tokens=256,
            )
        print(f"{variant_id} result: {resp}")

# Example entry point
async def main():
    prompts = {...}   # dict of variant_id -> prompt
    await asyncio.gather(*(run_variant(v, p) for v, p in prompts.items()))

if __name__ == "__main__":
    asyncio.run(main())

3. Explicit Memory Cleanup

Wrap inference in torch.no_grad() and call torch.cuda.empty_cache() after each batch.

async def run_variant(...):
    async with semaphore:
        model = GPT4Pool().get()
        with torch.no_grad():
            # inference logic …
            pass
        torch.cuda.empty_cache()

Before vs. After Comparison

Aspect	Before	After
Model loading	One full copy per process (≈76 GiB)	Single shared copy per GPU node
Concurrency control	Unbounded, each variant spawns a process	Bounded semaphore (e.g., 4 concurrent)
Memory hygiene	No explicit cleanup, tensors retained	torch.no_grad() + torch.cuda.empty_cache()
Peak GPU usage	≈96 GiB → OOM	≈30 GiB (single model + overhead)

Verify – Confirming the Fix

1. Run a controlled load test with the same number of variants as production (e.g., 8 concurrent variants).
2. Observe nvidia-smi output; GPU memory should stay below the configured threshold (e.g., <90 % of total).
3. Check logs for the absence of CUDA out of memory entries.
4. Validate functional correctness by comparing a sample of control‑variant outputs before and after the change (they should be identical within floating‑point tolerance).

Sample successful log excerpt:

2024-10-24 13:15:02,112 INFO  worker-0 GPU memory usage: 31.4/96.0 GiB
2024-10-24 13:15:02,115 INFO  variant control completed in 112ms
2024-10-24 13:15:02,117 INFO  variant treatment completed in 118ms

Prevent – Operational Guardrails

• Resource quotas: Define requests and limits for GPU memory in the Kubernetes pod spec (e.g., limits: nvidia.com/gpu: 1 and annotations: gpu.memory: "64Gi").
• Back‑pressure: Integrate a queueing layer (e.g., RabbitMQ or Kafka) that throttles new A/B test jobs when GPUUtilization > 80%.
• Monitoring & alerts: Export torch.cuda.memory_allocated and nvidia_smi metrics to Prometheus; fire an alert if usage exceeds 85 % for more than 5 minutes.
• Health checks: Add an endpoint that runs a lightweight inference (e.g., 1‑token prompt) and verifies memory stays under a safe margin.
• CI validation: Include a test that spawns the maximum expected number of concurrent variants and asserts no OOM occurs.

Related Topic Hub: LLM Systems Troubleshooting Hub

FAQ

1. Why does the OOM happen only when A/B testing, not with single‑variant traffic?

   The single‑variant path reuses the same model instance, keeping GPU memory constant. The A/B framework creates a new process per variant, each loading its own copy of the 80 GB model, which exceeds the GPU capacity.

2. Can I keep the multiprocessing approach and still avoid OOM?

   Only if you use torch.multiprocessing.set_start_method('spawn') and share the CUDA context via torch.cuda.ipc_collect(), but this adds complexity and is less reliable than a thread‑based pool.

3. How do I determine the safe value for MAX_CONCURRENT?

   Run torch.cuda.memory_reserved() after loading a single model, add a safety margin (≈20 %), and divide the remaining memory by the per‑request overhead (typically 1–2 GiB for prompt encoding and activations).

4. Will using DeepSpeed or TensorRT inference eliminate the OOM?

   They reduce per‑request memory but do not solve the root cause of multiple full model copies. Combining a shared model pool with an inference optimizer yields the best results.

5. Is it safe to call torch.cuda.empty_cache() after every request?

   Yes, in a high‑throughput service it incurs a negligible latency penalty and prevents fragmentation that can otherwise trigger OOM under bursty loads.