vLLM-Omni 분산 계층별 오프로드

vLLM-Omni의 분산 계층별 오프로드(Distributed Layerwise Offload, DLO)는 단일 장치의 HBM 용량을 초과하는 비디오 생성 모델(예: 64B Cosmos3-Super)을 최소한의 호스트 메모리 오버헤드로 여러 NPU 또는 GPU에 걸쳐 실행할 수 있게 합니다. 이 시스템은 메타 디바이스 초기화, 가중치 샤딩, 이중 버퍼 프리페치 파이프라인을 결합함으로써 대규모 Diffusion Transformer(DiT) 모델과 관련된 메모리 병목 문제를 해결합니다.

HBM 및 호스트 메모리 병목 문제 해결

대규모 DiT 모델은 종종 장치의 HBM에 맞지 않으며, 전통적인 오프로딩 또는 병렬화 전략은 상당한 트레이드오프를 수반합니다. 데이터 병렬(DP) 구성에서 순수 계층별 오프로딩은 일반적으로 각 랭크가 모델의 전체 복사본을 호스트 메모리에 저장해야 하므로, 호스트 RAM 요구량이 장치 수에 비례해 증가합니다(O(dp_size × model_size)).

분산 계층별 오프로드는 네 가지 주요 기술적 메커니즘을 통해 이러한 제약을 해결합니다.

1. 메타 디바이스 초기화 및 mmap 가중치 로딩

모델 생성 시 막대한 Resident Set Size(RSS) 급증을 방지하기 위해 vLLM-Omni는 메타 디바이스 초기화와 메모리 매핑(mmap) 로딩을 사용합니다.

• 메커니즘: 시스템은 to_empty(device="meta")를 사용해 DiT 모듈을 메타 디바이스로 변환하여 파라미터 저장을 해제하면서도 메타데이터는 유지합니다. 그런 다음 이 파라미터들을 공유된 OS 페이지 캐시를 직접 가리키는 mmap 뷰로 교체합니다.
• 결과: 모든 랭크가 동일한 safetensors 파일을 mmap하기 때문에 OS는 캐시에 파일 페이지를 단 하나만 유지합니다. Cosmos3-Nano DP4의 경우, 냉시작 시 cgroup에서 보이는 최대 메모리는 178 GB에서 47 GB로 감소(73% 감소)했습니다.

2. 가중치 샤딩 및 AllGather 재구성

각 랭크가 핀된 CPU 메모리에 모델 전체 복사본을 보유할 필요를 제거하기 위해 vLLM-Omni는 가중치 샤딩을 구현합니다.

• 메커니즘: 각 랭크는 모델 가중치의 1/dp_size만 저장합니다. 런타임에 전용 통신 스트림에서 all_gather_into_tensor를 사용해 현재 계층의 전체 가중치를 장치에서 재구성합니다.
• 결과: 전체 핀된 호스트 메모리는 dp_size × model_size에서 단순히 model_size로 감소합니다. Cosmos3-Super DP4의 경우, 랭크당 핀된 메모리는 124 GB에서 31 GB로 감소했습니다.

3. 이중 버퍼 프리페치 파이프라인

데이터 이동 중 GPU가 공백 상태가 되는 것을 방지하고, 모델 깊이에 관계없이 HBM 사용량을 일정하게 유지하기 위해 vLLM-Omni는 이중 버퍼 방식을 사용합니다.

• 메커니즘: 시스템은 모델에서 가장 큰 블록 크기와 정확히 같은 크기의 두 개의 장치 버퍼를 유지합니다. 계산 스트림이 슬롯 0에서 계층 N을 처리하는 동안, 백그라운드 스트림은 슬롯 1에 계층 N+1에 대한 H2D(호스트에서 장치로) 전송과 AllGather를 처리합니다.
• 결과: 가중치에 대한 HBM 사용량은 2 × max_block_size로 제한됩니다. Cosmos3-Nano 및 Cosmos3-Super 테스트에서 모델 크기가 3.8배 증가했음에도 불구하고 최대 HBM 사용량은 단지 22% 증가(23.1 GB에서 28.1 GB)했습니다.

4. DP 다중 동시성을 통한 처리량 향상

AllGather는 요청에 독립적(활성화가 아니라 가중치를 수집함)이므로, vLLM-Omni는 서로 다른 DP 랭크가 병렬로 다른 요청을 처리하면서도 가중치 재구성에 대해 동기화된 상태를 유지할 수 있습니다.

• 메커니즘: 스케줄러는 최대 dp_size개의 요청을 배치합니다. 각 DP 랭크는 목록에서 하나의 요청을 선택하여 파이프라인의 단일 요청 전방 경로를 통해 실행합니다.
• 결과: 이로 인해 AllGather 오버헤드가 분산됩니다. 측정 결과, 4개의 동시 요청은 HSDP 단일 요청 베이스라인 대비 3.3배의 처리량을 달성했으며, 이는 이상적인 선형 확장의 약 83%에 해당합니다.

성능 및 검증 결과

플랫폼 독립성

분산 계층별 오프로드는 NVIDIA GPU(CUDA/NCCL)와 Ascend NPU(CANN/HCCL) 모두에서 작동하는 플랫폼 독립적입니다.

NVIDIA B300 GPU에서 DLO+AG DP4는 4개의 동시 요청을 처리했을 때, HSDP+USP4 대비 1.39배의 처리량을 달성했으며, 1024×1024 T2I 작업에서 HBM 사용량은 30%만(12.6 GiB 대 42.0 GiB)을 차지했습니다. 720p 10초 비디오 생성 작업에서는 DLO+AG+USP4가 HSDP의 지연 시간에 비해 단지 2.13% 내외로 차이가 나면서도 HBM 사용량은 47%만 사용했습니다.

토폴로지 인식 정책(MiniMax-H3 연구)

8× NVIDIA B300 GPU를 사용한 MiniMax-H3 모델에 대한 연구 결과, 최적의 DLO 모드는 하드웨어 토폴로지에 따라 달라집니다:

• DP1×SP8: 최저 지연 시간(P50 34.55초)을 위해 AllGather가 선호됩니다.
• DP4×SP2: 처리량 측면에서 균형 잡힌 점(시간당 151.89개 비디오)을 제공합니다.
• DP8×SP1: 최고의 처리량(시간당 183.78개 비디오)과 최저 에너지 소비(비디오당 43.97 Wh)를 위해 랭크 로컬 DLO가 선호됩니다.

Ascend NPU에서의 메모리 계정

Ascend 하드웨어에서는 pin_memory()가 /dev/davinci_manager를 통해 할당되며, 샤드는 CPU 커널 DMA 메모리에 위치합니다. 이 메모리는 cgroup 메모리 컨트롤러에 보이지 않습니다. 따라서 cgroup에서 보이는 메모리는 O(model_size + dp_size × constant)로 증가하지만, 전체 물리적 RAM에는 여전히 이러한 핀된 DMA 샤드가 포함됩니다.

메모리 스케일링 요약(외삽된 결과)

측정된 메모리 모델을 기반으로, vLLM-Omni는 2TB 시스템 RAM 제한 내에서 매우 큰 모델까지 확장 가능합니다:

모델 크기	dp_size	추정 cgroup 최대값	추정 전체 RAM	2TB에 맞는가?
33 GB	4	47 GB	~80 GB	예
124 GB	4	172 GB	~296 GB	예
185 GB	4	~220 GB	~405 GB	예
400 GB	4	~423 GB	~823 GB	예
400 GB	8	~443 GB	~843 GB	예