nvidia-smi를 찍어보면 이상한 숫자가 보인다. TP=4로 Llama-3 70B를 올렸는데 GPU 0의 메모리 사용량이 GPU 13보다 12GB 높다. 배치를 늘릴수록 격차는 커지고, OOM은 언제나 GPU 0에서 먼저 터진다.
처음 보면 드라이버 버그나 초기화 이슈라고 의심하게 된다. 실제로 vLLM GitHub Issues에는 'OOM on rank 0' 관련 리포트가 꾸준히 올라온다. 그런데 이건 버그가 아니라 구조적 설계 결과다. 임베딩, logit 계산, NCCL 버퍼, KV 캐시 스케줄링이 각각의 이유로 rank 0에 더 많은 메모리를 요구한다.
임베딩 레이어와 LM Head가 rank 0에 모이는 이유
Megatron-LM 논문(Shoeybi et al., 2019)이 정립한 Tensor Parallel 구조에서 Attention과 FFN 가중치는 TP 크기에 맞게 균등하게 분할된다. TP=4라면 각 rank가 전체 파라미터의 1/4을 담당하고, 가중치 메모리는 이론상 정확히 4분의 1이다.
문제는 vocab 임베딩 레이어와 LM Head에서 생긴다.
vLLM의 VocabParallelEmbedding은 어휘 테이블을 Column Parallel 방식으로 분할한다. Llama-3 70B의 vocab_size는 128,256, hidden_size는 8,192다. TP=4 환경에서 각 rank는 32,064개 토큰 행을 담당한다. 가중치 자체는 균등하게 나뉜다.
forward pass에서 비대칭이 발생한다. 각 rank는 자신이 담당하는 토큰에 대한 partial logit만 계산한다. 최종 샘플링을 위해서는 전체 vocab에 걸친 logit이 필요하고, vLLM은 driver worker(rank 0)에서 full logit gather를 수행해 샘플링을 진행한다. 이 과정에서 rank 0는 배치 전체의 full vocab logit을 담는 추가 버퍼를 잡는다.
배치 크기 32 기준:
- Rank 0의 gather 버퍼: 128,256 × 32 × 4 bytes ≈ 16.4MB
- Rank 1~3의 partial buffer: 32,064 × 32 × 4 bytes ≈ 4.1MB
배치 크기가 128로 늘면 rank 0는 65.5MB, 나머지는 16.4MB다. 49MB 차이가 transient 텐서로 매 decode step마다 존재하고, profile_run이 이걸 peak activation으로 측정하면서 rank 0의 KV 캐시 가용 메모리를 그만큼 깎는다.
gather 버퍼 크기 계산식:
rank 0 추가 버퍼 = vocab_size × (1 - 1/tp_size) × batch_size × dtype_bytes
TP=4, BF16 logit이라면: 128,256 × (1 - 0.25) × batch_size × 2 bytes
NCCL 통신 버퍼의 rank 0 편중
Attention과 FFN 레이어 끝에는 각각 all-reduce가 들어간다. Llama-3 70B의 80개 레이어는 총 160번의 all-reduce를 수행한다. NCCL은 이 연산을 위한 workspace 버퍼를 GPU 메모리에 사전 할당한다.
vLLM의 init_device는 NCCL 초기화를 메모리 스냅샷 이전에 의도적으로 완료한다. 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)
NCCL 버퍼가 스냅샷에 이미 포함되어 있으니 KV 캐시 계산에서 자동으로 차감된다. 여기까지는 설계 의도대로다. 문제는 NCCL workspace 크기 자체가 rank별로 균일하지 않다는 점이다. ring/tree 알고리즘에서 루트 역할을 하는 rank 0는 추가 통신 조율 버퍼를 잡는다. TP degree가 올라갈수록 이 편중도 커진다.
all-reduce workspace의 이론적 추정식:
workspace ≈ 2 × hidden_size × tp_size × dtype_bytes
Llama-3 70B (hidden=8,192), TP=8, BF16 기준으로 공식상 약 256KB지만, 실제 NCCL workspace는 알고리즘 선택(ring vs. double binary tree)과 버퍼 풀 설정에 따라 수십~수백 MB까지 늘어날 수 있다. Rank 0의 추가 오버헤드 절대 크기는 환경마다 달라지지만 방향은 항상 rank 0 쪽이다.
KV 캐시 블록 계산이 비대칭을 확정시키는 지점
각 worker는 determine_available_memory를 독립적으로 실행한다:
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는 모든 rank에서 동일하다. A100 80GB, gpu_memory_utilization=0.90이면 72GB다.
non_kv_cache_memory에는 가중치 + profile_run에서 측정한 peak activation이 포함된다. Rank 0는 logit gather 버퍼가 크기 때문에 peak activation이 더 높게 찍히고, available_kv_cache_memory_bytes가 낮게 나온다.
vLLM engine core는 모든 worker의 보고값 중 최솟값으로 KV 캐시 블록 수를 결정한다. 최솟값은 언제나 rank 0에서 나온다. Rank 0의 가용 메모리가 전체 클러스터의 KV 캐시 상한선을 결정하는 병목이 된다.
이게 OOM으로 이어지는 경로다. Engine core가 정한 블록 수는 rank 0의 여유 메모리를 거의 꽉 채운다. 런타임에서 배치 처리 중 logit 계산 순간 transient 텐서가 올라오면 rank 0는 공간이 없어 OOM이 난다. rank 1~3은 같은 순간에 여유가 있어 멀쩡하다.
TP degree별 logit gather 버퍼 비대칭
Llama-3 70B (vocab_size=128,256), float32 logit 기준 이론 계산값이다.
| TP | rank 0 추가 버퍼 (batch=32) | rank 0 추가 버퍼 (batch=128) |
|---|---|---|
| TP=1 | 0 MB | 0 MB |
| TP=2 | ~7.8 MB | ~31.3 MB |
| TP=4 | ~11.7 MB | ~47.0 MB |
| TP=8 | ~13.7 MB | ~54.8 MB |
TP=8에서 batch=128이면 rank 0는 다른 rank보다 54.8MB를 더 써야 한다. 이 차이만으로 KV 캐시 블록이 수백 개 줄어든다. Llama-3 70B의 KV 블록 하나는 토큰 16개 분량의 K·V를 저장하므로, 54.8MB는 컨텍스트 용량으로 환산하면 꽤 큰 손실이다.
배치가 커질수록 격차가 선형으로 비례한다는 점도 문제다. 높은 throughput을 위해 배치를 키울수록 rank 0의 압박이 커진다. 두 목표가 서로 충돌한다.
완화 전략과 그 비용
gpu_memory_utilization을 낮추기
0.90에서 0.85로 낮추면 requested_memory가 A100 80GB 기준 72GB에서 68GB로 4GB 줄어든다. Rank 0의 여유 공간이 생겨 OOM 임계 배치 크기가 올라간다. 단점은 KV 캐시 블록 수가 줄어들면서 처리 가능한 동시 시퀀스가 줄고 throughput이 감소한다는 것이다. utilization 5%p 하락에 따른 throughput 손실은 배치 패턴에 따라 수 퍼센트 범위다.
CUDA Graph 유지 (enforce_eager=False)
반직관적이지만, CUDA Graph를 켜두면 forward pass 중 transient 텐서 할당이 graph 캡처 시점에 고정된다. Rank 0의 logit gather 버퍼도 캡처 크기로 확정되어 런타임 동적 할당이 억제된다. enforce_eager=True로 graph를 끄면 transient 할당이 자유로워져 오히려 OOM 위험이 높아지는 경우가 있다.
--kv-cache-memory 직접 지정
kv_cache_memory_bytes를 직접 지정하면 profile_run 추정 없이 고정 크기를 KV 캐시에 배정한다. 서빙 전에 아래 코드로 rank 0의 실제 여유를 측정하고 그보다 낮게 잡으면 런타임 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")
대신 모델 업데이트나 배치 설정 변경 시 수동 재조정이 필요하다.
분산 sampling으로 gather 제거
근본 원인인 logit gather를 없애는 방향이다. 각 rank에서 partial logit으로 직접 top-k/top-p를 계산하는 distributed sampling은 gather 버퍼 비대칭 자체를 제거한다. vLLM의 VocabParallelEmbedding 구현은 이 방향으로 지속적으로 개선 중이지만, 빔 서치 등 일부 샘플링 방법에서는 full logit gather가 여전히 필요하다.
위 방법들은 모두 rank 0의 압박을 줄이거나 OOM 임계점을 뒤로 미루는 것이지, 비대칭 구조 자체를 제거하지는 못한다.
OOM이 없어도 throughput이 줄어드는 경로
OOM까지 가지 않더라도 rank 0의 메모리 압박은 조용한 throughput 손실을 만든다.
vLLM 스케줄러는 가용 KV 캐시 블록 수를 기준으로 신규 요청 수락 여부를 결정한다. 블록 수가 rank 0를 기준으로 결정되어 있으니, rank 1~7이 여유가 있어도 스케줄러는 rank 0의 한도 내에서만 배치를 잡는다. 전체 GPU 가용 메모리의 합산이 아니라 병목 rank의 상한이 throughput 한계가 된다.
배치 크기를 보수적으로 운영하거나 max_num_seqs를 낮게 설정한 환경에서도 이 현상은 발생한다. 표면적으로는 "GPU 가속 성능" 한계처럼 보이지만, 실제 병목이 rank 0의 메모리 비대칭에서 비롯된 스케줄링 제약인 경우가 있다.
서빙 중에 torch.cuda.memory_stats(device=0)과 torch.cuda.memory_stats(device=1)의 reserved_bytes.all.current를 비교해보면 비대칭 규모를 빠르게 확인할 수 있다. 두 값의 차이가 수백 MB 이상이라면 OOM이 없더라도 rank 0가 throughput 병목으로 작동하고 있을 가능성이 높다.