본문으로 건너뛰기

[vLLM] CPU/XPU Worker: 비NVIDIA 하드웨어 워커

들어가며

vLLM은 NVIDIA GPU 없이도 CPU에서 LLM 추론을 실행할 수 있다. vllm/v1/worker/cpu_worker.py의 CPUWorker는 GPU Worker의 인터페이스를 유지하면서, CPU 환경에 맞는 최적화와 제약을 구현한다.

핵심 구조/코드 분석

CPUWorker 클래스

class CPUWorker(Worker):
    def __init__(self, vllm_config, local_rank, rank, distributed_init_method,
                 is_driver_worker=False):
        super().__init__(vllm_config, local_rank, rank, distributed_init_method,
                        is_driver_worker=is_driver_worker)
        self.parallel_config.disable_custom_all_reduce = True

        profiler_config = vllm_config.profiler_config
        if profiler_config.profiler == "torch":
            worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
            self.profiler = TorchProfilerWrapper(
                profiler_config, worker_name=worker_name,
                local_rank=self.local_rank, activities=["CPU"],
            )

GPU Worker를 상속하되, 두 가지를 변경한다:

  1. disable_custom_all_reduce = True: CPU 환경에서는 NCCL 기반 커스텀 all-reduce를 사용할 수 없다.
  2. 프로파일러 activities를 ["CPU"]로 설정한다.

디바이스 초기화

def init_device(self):
    def check_preloaded_libs(name: str):
        ld_preload_list = os.environ.get("LD_PRELOAD", "")
        if name not in ld_preload_list:
            logger.warning(
                "%s is not found in LD_PRELOAD. "
                "For best performance, please follow the section "
                "`set LD_PRELOAD` in "
                "https://docs.vllm.ai/en/latest/getting_started/installation/cpu/",
                name,
            )

    if sys.platform.startswith("linux"):
        check_preloaded_libs("libtcmalloc")
        if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
            check_preloaded_libs("libiomp")

CPU 추론 성능에 필수적인 라이브러리들을 검증한다:

  • libtcmalloc: Google의 고성능 메모리 할당자. 기본 glibc malloc보다 멀티스레드 환경에서 훨씬 빠르다.
  • libiomp: Intel OpenMP 라이브러리. x86 CPU에서 병렬 연산 성능을 극대화한다. ARM에서는 불필요하므로 x86에서만 검사한다.

스레드 수 고정

def skip_set_num_threads(x: int):
    logger.warning(
        "CPU backend doesn't allow to use "
        "`torch.set_num_threads` after the thread binding, skip it."
    )

torch.set_num_threads = skip_set_num_threads

CPU 백엔드는 스레드 바인딩 후 torch.set_num_threads를 호출하면 성능이 저하된다. 이를 방지하기 위해 함수 자체를 no-op으로 교체한다. 대담한 몽키패칭이지만, CPU 추론 성능을 위해 필수적이다.

분산 환경 초기화

os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(":")[-1]
init_worker_distributed_environment(
    self.vllm_config, self.rank, self.distributed_init_method,
    self.local_rank, current_platform.dist_backend,
)
set_random_seed(self.model_config.seed)
self.model_runner: CPUModelRunner = CPUModelRunner(self.vllm_config, torch.device("cpu"))

VLLM_DIST_IDENT는 allreduce 공유 메모리의 고유 식별자로, 분산 초기화 메서드의 포트 번호를 사용한다. dist_backend는 CPU 환경에서 gloo를 사용한다.

Sleep Mode 비활성화

def sleep(self, level: int = 1) -> None:
    logger.warning("sleep mode is not supported on CPU, ignore it.")
    pass

def wake_up(self, tags: list[str] | None = None) -> None:
    logger.warning("wake_up is not supported on CPU, ignore it.")
    pass

CPU에는 GPU 메모리 오프로딩 개념이 없으므로, sleep/wake_up을 no-op으로 구현한다. 경고 로그만 남기고 무시한다.

왜 이 설계인가

  1. GPU Worker 상속: CPUWorker는 GPUWorker의 서브클래스다. 스케줄러, 엔진 코어 등 상위 레이어가 Worker 인터페이스에 의존하므로, 동일한 인터페이스를 유지하면서 CPU 특화 로직만 오버라이드한다. 코드 중복을 최소화하면서 호환성을 유지하는 전략이다.

  2. LD_PRELOAD 검증: CPU 추론에서 메모리 할당자와 스레딩 라이브러리의 선택이 성능에 2-3배 차이를 만들 수 있다. 경고로 안내하되 강제하지는 않아, 비표준 환경에서도 실행은 가능하다.

  3. torch.set_num_threads 몽키패치: CPU 백엔드는 numactl이나 taskset으로 코어를 명시적으로 바인딩하는 것이 일반적이다. 이 바인딩 후에 PyTorch나 다른 라이브러리가 set_num_threads를 호출하면 바인딩이 깨져서 NUMA 노드 간 메모리 접근으로 성능이 급락한다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글