본문으로 건너뛰기

[vLLM] Model Weight Offloading: 가중치 CPU 오프로딩

들어가며

GPU 메모리가 부족할 때, 모델 가중치 일부를 CPU 메모리에 저장하고 필요할 때 가져오는 오프로딩 전략이 있다. vLLM은 vllm/model_executor/offloader/ 모듈에서 UVA(Unified Virtual Addressing)와 Prefetch 두 가지 방식을 제공한다.

공식 문서

vLLM 공식 문서: Conserving Memory

핵심 구조/코드 분석

BaseOffloader 추상 클래스

class BaseOffloader(ABC):
    @abstractmethod
    def wrap_modules(self, modules_generator: Generator[nn.Module, None, None]) -> list[nn.Module]:
        pass
    def post_init(self): return
    def sync_prev_onload(self) -> None: pass
    def join_after_forward(self) -> None: pass

모든 오프로더는 wrap_modules로 모델의 레이어들을 감싸고, 필요한 경우 post_init에서 초기화를 완료한다. 팩토리 패턴으로 자동 선택된다:

def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
    if backend == "auto":
        if prefetch.offload_group_size > 0: backend = "prefetch"
        elif uva.cpu_offload_gb > 0: backend = "uva"
        else: return NoopOffloader()

UVA 오프로딩: 제로카피 접근

class UVAOffloader(BaseOffloader):
    def __init__(self, cpu_offload_max_bytes, cpu_offload_params=None):
        self.pin_memory = is_pin_memory_available()
        self.uva_offloading = is_uva_available()

UVA는 pinned CPU 메모리에 가중치를 두고, CUDA UVA를 통해 GPU가 직접 PCIe로 접근하는 방식이다. 명시적인 복사 없이 GPU 코드에서 CPU 메모리를 투명하게 참조할 수 있다. 단, PCIe 대역폭으로 인한 속도 저하가 있다.

Prefetch 오프로딩: 비동기 사전 로딩

class PrefetchOffloader(BaseOffloader):
    def __init__(self, group_size, num_in_group, prefetch_step, offload_params=None, mode="cpu"):
        self.group_size = group_size
        self.num_in_group = num_in_group
        self.prefetch_step = prefetch_step
        self.copy_stream = torch.cuda.Stream()

Prefetch 방식은 레이어를 그룹으로 나누어, 현재 레이어를 실행하는 동안 다음 레이어의 가중치를 비동기로 GPU에 복사한다. copy_stream이라는 별도의 CUDA 스트림을 사용하여 연산과 전송을 오버랩한다.

StaticBufferPool: 고정 GPU 버퍼

class StaticBufferPool:
    def __init__(self, param_infos, slot_capacity, device):
        for key, info in unique_params.items():
            slot_tensors = []
            for _ in range(slot_capacity):
                buf = torch.empty_strided(
                    size=info.shape, stride=info.stride,
                    dtype=info.dtype, device=device,
                )
                slot_tensors.append(buf)
            self._buffers[key] = slot_tensors

고정된 GPU 버퍼를 사전 할당하여 매번 동적 할당하지 않는다. slot_capacity(=prefetch_step)만큼의 슬롯으로 더블/트리플 버퍼링이 가능하다. 레이어 N은 slot (N % slot_capacity)를 사용한다.

Forward Hook과 torch.compile 호환

def _hook_module_forward(self, index, module):
    original_forward = module.forward
    def forward(*args, **kwargs):
        module.forward = original_forward  # 재귀 방지
        input_tensor = args[0] if args else kwargs.get("hidden_states")
        torch.ops.vllm.wait_prefetch(input_tensor, index)  # 프리페치 대기
        output = original_forward(*args, **kwargs)
        next_index = (index + self.prefetch_step) % len(self.module_offloaders)
        torch.ops.vllm.start_prefetch(output[0], next_index)  # 다음 프리페치 시작
        module.forward = forward
        return output
    module.forward = forward

커스텀 op(wait_prefetch, start_prefetch)을 사용하여 torch.compile과 CUDA 그래프에서도 동작하도록 했다. mutates_args를 통해 데이터 의존성을 만들어 컴파일러가 순서를 보장한다.

CUDA 그래프 캡처와 이벤트 동기화

def start_onload_to_static(self):
    self._prefetch_in_capture = torch.cuda.is_current_stream_capturing()
    fork_event = torch.cuda.Event()
    torch.cuda.current_stream().record_event(fork_event)
    self.copy_stream.wait_event(fork_event)
    with torch.cuda.stream(self.copy_stream):
        for name, offloader in self._param_offloaders.items():
            gpu_buffer.copy_(cpu_storage, non_blocking=True)
    self._copy_done_event.record(self.copy_stream)

CUDA 그래프 캡처 중에는 이벤트 기반 동기화를 사용하고, eager 모드에서는 wait_stream 폴백을 제공한다. Pinned 메모리를 반드시 사용해야 non_blocking=True 복사가 스트림 동기화를 깨뜨리지 않는다.

왜 이 설계인가

  1. 이중 전략 제공: UVA는 설정이 간단하고 코드 변경이 최소화되지만, 매 접근마다 PCIe 대역폭을 소모한다. Prefetch는 구현이 복잡하지만 비동기 파이프라이닝으로 전송 지연을 숨길 수 있다. 워크로드 특성에 따라 적합한 전략이 다르다.

  2. 그룹 기반 레이어 선택: group_size=4, num_in_group=2이면 매 4개 레이어 중 마지막 2개만 오프로딩한다. 전체 레이어를 오프로딩하면 대역폭이 부족할 수 있으므로, 일부만 선택적으로 오프로딩하여 성능과 메모리의 균형을 맞춘다.

  3. Pinned 메모리 강제: Non-pinned CPU 메모리에서 non_blocking=True로 복사하면 CUDA 드라이버가 내부적으로 스트림 동기화를 수행하여, 이벤트 기반 fork 동기화가 깨진다. 이를 방지하기 위해 항상 pinned 메모리를 사용한다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글