vLLM Ray Direct Transport를 사용한 대규모 분할 가중치 전송

개요

vLLM는 온라인 강화학습(RL) 환경에서 트레이너와 추론 워커 간 모델 가중치 동기화를 최적화하기 위해 Ray Direct Transport(RDT)를 활용한 내장형 분할 가중치 전송 엔진을 구현했습니다. 이 시스템은 트리리언 파라미터 규모에서 메모리 병목 현상과 동기화 지연이 자주 발생하는 기존의 NCCL 브로드캐스트 방식을 대체하여, 피크 메모리 사용량을 줄이고 전송 속도를 향상시키는 풀 기반의 분할 접근 방식을 사용합니다.

브로드캐스트 기반 동기화의 한계

표준 가중치 동기화는 일반적으로 NCCL 브로드캐스트에 의존합니다. 이 방식에서는 트레이너가 파라미터를 HuggingFace 형식으로 모두 모아서 모든 추론 워커에 브로드캐스트합니다. 이 접근 방식은 대규모 모델에 대해 두 가지 주요 문제를 야기합니다:

메모리 비효율성: Tensor Parallelism 8(TP8)와 같은 구성에서는 워커가 전체 모델을 수신하지만, 실제 보유하는 가중치는 1/8에 불과하며 나머지는 버립니다. 대규모 Mixture-of-Experts(MoE) 모델에서는 이로 인해 막대한 피크 메모리 부담이 발생합니다.
집단적 동기화: NCCL는 모든 랭크가 동기적으로 참여해야 합니다. 느린 랭크나 복제 실패 시 전체 집단이 정지하게 되어, 동적이고 대규모 환경에서는 문제가 됩니다.

기술적 구현: 분할 가중치 전송

기록 텐서 드라이런

다양한 모델 아키텍처(예: Llama-4의 융합 전문가 또는 다양한 모델의 GQA)와의 호환성을 보장하기 위해 vLLM은 초기화 시 "기록 텐서" 드라이런을 사용합니다. vLLM 로더는 모양과 dtype만 보고하고 실제 데이터는 포함하지 않는 텐서 서브클래스를 제공합니다. 모든 변환(뷰, 네어, 전치, 재형성 등)은 연쇄된 작업(즉, "분할 계획")으로 기록됩니다.

이 계획을 통해 트레이너는 초기 레이아웃 작업(융합, 재배치, 분할, 분할)을 수행하고, 각 vLLM 랭크가 필요로 하는 특정 분할된 가중치만 BF16 형식으로 전송할 수 있습니다. 이로써 과정이 구성적으로 정확함을 보장합니다.

Ray Direct Transport(RDT) 및 NIXL

엔진은 Ray Direct Transport(RDT)와 NIXL 백엔드를 사용하여 Ray 액터 간 GPU-to-GPU 통신을 가능하게 합니다. 이 아키텍처는 추론 랭크가 필요한 분할된 텐서만 매핑된 트레이너 랭크로부터 풀하는 풀 기반 시스템을 가능하게 합니다.

초기화 흐름은 다섯 단계로 구성됩니다:

트레이너 랭크들이 소유권 메타데이터(파라미터 이름, dtype, 모양)를 모두 모읍니다.
랭크 0은 이 메타데이터와 트레이너 Ray 액터 이름을 추론 워커에 전송합니다.
각 vLLM 워커는 기록 텐서 드라이런을 통해 자신의 분할 계획을 생성합니다.
워커들은 로드 밸런싱 방식으로 소스 트레이너 랭크에 매핑됩니다.
생산자와 소비자 모두 사전에 RDT 버퍼를 할당하고 등록합니다.

성능 최적화

vLLM은 Qwen3-235B-A22B 모델(TP4/PP2/EP8 트레이너 → DP16/EP16 vLLM 서버)에 대해 엔드투엔드 지연 시간을 최적화하기 위해 세 가지 버전을 반복적으로 개선했습니다:

V1(간단한 반복자): TP, PP, EP 모든 차원을 하나씩 모아서 가중치를 수집했습니다. 이로 인해 수천 개의 작은 집단적 작업과 중복 메모리 사용이 발생했으며, 동기화 시간은 25.02초였습니다.
V2(PP/EP-로컬): PP-로컬 수집(같은 파이프라인 스테이지 내에서만)과 EP-로컬 전송(전문가를 전혀 수집하지 않음; 추론 랭크는 전문가를 보유한 랭크로부터 직접 풀함)을 구현했습니다. 이로 인해 동기화 시간은 5.61초로 감소했습니다.
V3(파이프라인 실행): 모든 가중치 수집, 재실행 작업, RDMA 전송을 겹쳐 수행했습니다. 디코더 블록 그룹 단위로 가중치를 수집하고 백그라운드에서 처리함으로써 동기화 지연은 3.49초로 감소했습니다.

대규모 검증: Kimi K2

48개 노드(8xH100, 총 32개 트레이너 노드, 16개 추론 노드)에서 Kimi K2 모델에 대한 검증 결과는 다음과 같습니다:

메트릭	값
동기화당 전송된 바이트 수	7.9 TB
가중치 동기화 시간	7.53초
집계 대역폭	1,049 GB/s

이 성능은 vLLM의 계층별 재로딩 로직의 제약 조건을 고려할 때, 이 특정 설정에 대해 예상되는 "빛의 속도"(SoL) 전송 시간의 약 1.5배에 해당합니다.

장애 내성 및 통합

NIXL을 브로드캐스트 집단적 작업 대신 사용함으로써 시스템은 고장에 더 강건합니다. 추론 엔진이 실패하면 라우터는 트래픽을 남은 엔진으로 계속 전달하고, 트레이너는 다음 동기화 시 살아있는 엔진들과만 통신합니다. 실패한 복제가 복구되면 다음 동기화 경계에서 다시 참여하여 업데이트된 가중치를 수신하지만, 전체 수렴에는 영향을 주지 않습니다.

프레임워크 통합

이 엔진은 SkyRL에 통합되어 있습니다. 다른 RL 프레임워크는 WeightSource 반복자를 구현하여 파라미터 메타데이터를 제공하고, 실체화된 텐서를 생성하며, 선택적으로 held_names 메서드를 통해 EP/PP-로컬 최적화를 활성화할 수 있습니다.

현재 제한 사항

로더 제약: 로더는 기록 가능한 작업만 사용해야 하며, 로딩 중 실제 값을 검사하는 로더는 실패합니다.
메모리 예산: RDT 목적지 버퍼는 vLLM의 gpu_memory_utilization 예산 외부에 존재합니다.
호환성: 현재 구현은 vLLM의 EPLB와 호환되지 않습니다.
직렬화: 계층별 재로딩 중 OOM을 방지하기 위해 트레이너 PP 그룹 간 전송이 현재 직렬로 이루어집니다.