Why GPU 0 Always OOMs in Tensor Parallel LLM Serving

nvidia-smi shows a puzzling number. With TP=4 serving Llama-3 70B, GPU 0's memory usage is 1–2 GB higher than GPUs 1–3. The gap widens as batch size increases, and OOM always hits GPU 0 first.

At first glance this looks like a driver bug or an initialization issue. Indeed, 'OOM on rank 0' reports surface regularly in the vLLM GitHub Issues. But this isn't a bug — it's a consequence of structural design decisions. The embedding layer, logit computation, NCCL buffers, and KV cache scheduling each have their own reasons for demanding more memory on rank 0.

Why the Embedding Layer and LM Head Concentrate on Rank 0

In the Tensor Parallel architecture established by the Megatron-LM paper (Shoeybi et al., 2019), Attention and FFN weights are split evenly across the TP degree. With TP=4, each rank holds 1/4 of the total parameters, and weight memory is theoretically split exactly four ways.

The asymmetry originates in the vocab embedding layer and the LM Head.

vLLM's VocabParallelEmbedding partitions the vocabulary table in a Column Parallel fashion. Llama-3 70B has vocab_size=128,256 and hidden_size=8,192. With TP=4, each rank is responsible for 32,064 token rows. The weights themselves are split evenly.

The asymmetry appears in the forward pass. Each rank computes only the partial logits for the tokens it owns. Final sampling requires logits over the full vocabulary, so vLLM performs a full logit gather on the driver worker (rank 0) before sampling. This forces rank 0 to allocate an extra buffer large enough to hold the full-vocab logits for the entire batch.

At batch size 32:

• Rank 0 gather buffer: 128,256 × 32 × 4 bytes ≈ 16.4 MB
• Ranks 1–3 partial buffer: 32,064 × 32 × 4 bytes ≈ 4.1 MB

At batch size 128, rank 0 needs 65.5 MB while the others need 16.4 MB — a 49 MB gap. This transient tensor exists on every decode step, and when profile_run measures it as the peak activation, it reduces the KV cache memory available to rank 0 by exactly that amount.

Gather buffer size formula:

rank 0 extra buffer = vocab_size × (1 - 1/tp_size) × batch_size × dtype_bytes

For TP=4 with BF16 logits: 128,256 × (1 - 0.25) × batch_size × 2 bytes

NCCL Communication Buffer Skew Toward Rank 0

Each Attention and FFN layer ends with an all-reduce. Llama-3 70B's 80 layers perform 160 all-reduces in total. NCCL pre-allocates workspace buffers in GPU memory for these operations.

vLLM's init_device intentionally completes NCCL initialization before taking a memory snapshot. Looking at gpu_worker.py:

# Initialize the distributed environment BEFORE taking memory snapshot
# This ensures NCCL buffers are allocated before we measure available memory
init_worker_distributed_environment(...)

# Now take memory snapshot after NCCL is initialized
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)

Since NCCL buffers are already included in the snapshot, they are automatically subtracted from the KV cache calculation — that part is working as designed. The problem is that the NCCL workspace size is not uniform across ranks. In ring/tree algorithms, rank 0 acts as the root and allocates additional coordination buffers. This skew grows with TP degree.

Theoretical estimate for all-reduce workspace:

workspace ≈ 2 × hidden_size × tp_size × dtype_bytes

For Llama-3 70B (hidden=8,192), TP=8, BF16, the formula gives roughly 256 KB, but in practice the NCCL workspace can reach tens to hundreds of MB depending on the algorithm selected (ring vs. double binary tree) and buffer pool configuration. The absolute size of rank 0's extra overhead varies by environment, but the direction is always toward rank 0.

Where KV Cache Block Calculation Locks In the Asymmetry

Each worker runs determine_available_memory independently:

self.available_kv_cache_memory_bytes = (
    self.requested_memory
    - profile_result.non_kv_cache_memory
    - cudagraph_memory_estimate_applied
)

requested_memory = gpu_memory_utilization × total_GPU_memory is identical on every rank. With an A100 80 GB and gpu_memory_utilization=0.90, that's 72 GB.

non_kv_cache_memory includes weights plus the peak activation measured during profile_run. Because rank 0 has a larger logit gather buffer, its peak activation is higher, so its available_kv_cache_memory_bytes comes out lower.

The vLLM engine core determines the number of KV cache blocks using the minimum value reported across all workers. That minimum always comes from rank 0. Rank 0's available memory becomes the bottleneck that caps the KV cache for the entire cluster.

This is the path to OOM. The block count the engine core sets nearly fills rank 0's available memory. When a transient tensor spikes during logit computation mid-batch at runtime, rank 0 has no room left and OOMs. Ranks 1–3 have headroom at that exact moment and are fine.

Logit Gather Buffer Asymmetry by TP Degree

Theoretical values for Llama-3 70B (vocab_size=128,256) with float32 logits.

At TP=8 with batch=128, rank 0 must use 54.8 MB more than any other rank. That difference alone reduces the KV cache block count by hundreds of blocks. Each KV cache block in Llama-3 70B stores the K·V for 16 tokens, so 54.8 MB translates to a significant loss in effective context capacity.

The gap also scales linearly with batch size, which creates a direct conflict: pushing batch size higher for better throughput increases the pressure on rank 0. The two goals work against each other.

Mitigation Strategies and Their Costs

Lower gpu_memory_utilization

Dropping from 0.90 to 0.85 reduces requested_memory from 72 GB to 68 GB on an A100 80 GB — a 4 GB reduction. This gives rank 0 more headroom and raises the OOM-triggering batch size threshold. The downside is fewer KV cache blocks, which reduces the number of concurrent sequences the system can handle and lowers throughput. The throughput loss from a 5 percentage-point drop in utilization is typically a few percent, depending on the request pattern.

Keep CUDA Graph enabled (enforce_eager=False)

Counterintuitively, leaving CUDA Graph on freezes transient tensor allocation during the forward pass at graph capture time. Rank 0's logit gather buffer size is fixed at capture and dynamic allocation is suppressed at runtime. Disabling the graph with enforce_eager=True frees up transient allocation, which can actually increase OOM risk in some cases.

Specify --kv-cache-memory directly

Setting kv_cache_memory_bytes explicitly allocates a fixed amount to the KV cache, bypassing profile_run estimation. Measure rank 0's actual free memory before serving using the code below, then set the value below that figure to prevent runtime OOM.

import torch
stats = torch.cuda.memory_stats(device=0)
free = stats["reserved_bytes.all.current"] - stats["allocated_bytes.all.current"]
print(f"rank 0 실제 여유: {free / 1024**3:.2f} GiB")

The tradeoff is that manual re-tuning is required whenever the model or batch configuration changes.

Distributed sampling to eliminate the gather

This addresses the root cause by removing the logit gather entirely. Distributed sampling computes top-k/top-p directly from partial logits on each rank, eliminating the gather buffer asymmetry. vLLM's VocabParallelEmbedding implementation is being continuously improved in this direction, but some sampling methods such as beam search still require a full logit gather.

All of the above approaches reduce pressure on rank 0 or push the OOM threshold further out; none of them eliminate the underlying asymmetry.

How Throughput Degrades Even Without OOM

Even without hitting OOM, memory pressure on rank 0 causes silent throughput loss.

The vLLM scheduler decides whether to admit new requests based on the number of available KV cache blocks. Since that block count is determined by rank 0, the scheduler caps batch size to rank 0's limit even when ranks 1–7 have headroom. The throughput ceiling is set by the bottleneck rank, not by the sum of available memory across all GPUs.

This happens even when batch size is run conservatively or max_num_seqs is set low. On the surface it looks like hitting a "GPU compute" limit, but the actual bottleneck is a scheduling constraint caused by rank 0's memory asymmetry.

Comparing reserved_bytes.all.current from torch.cuda.memory_stats(device=0) and torch.cuda.memory_stats(device=1) during serving gives a quick read on the asymmetry. If the difference is several hundred MB or more, rank 0 is likely acting as a throughput bottleneck even in the absence of OOM.