CoreWeave Tensorizer로 모델 로드하기

vLLM은 CoreWeave의 Tensorizer를 사용한 모델 로딩을 지원해요. 디스크, HTTP/HTTPS 엔드포인트, S3 엔드포인트로 직렬화(serialize)된 vLLM 모델 텐서를 런타임에 GPU로 직접 매우 빠르게 역직렬화(deserialize)할 수 있어요. 이를 통해 Pod 스타트업 시간과 CPU 메모리 사용량을 크게 줄일 수 있고, 텐서 암호화(encryption)도 지원해요.

출처: 문서

본문

vLLM은 Tensorizer를 모델 로딩 기계(mechanism)에 완전히 통합해요. 다음은 vLLM에서 Tensorizer를 쓰기 시작하는 간단한 개요예요.

Tensorizer 설치

pip install vllm[tensorizer]

기본 개념

Tensorizer로 모델을 로드하려면 먼저 모델이 Tensorizer로 직렬화되어 있어야 해요. 이 과정을 예시 스크립트(examples/features/tensorize_vllm_model.py)가 처리해줘요. facebook/opt-125m을 직렬화한 후 추론에 로드하는 기본 예시를 살펴볼게요.

vLLM 모델 직렬화

Tensorizer로 모델을 직렬화하려면 필요 CLI 인자와 함께 예시 스크립트를 호출해요. 스크립트의 docstring이 CLI 인자를 매우 상세히 설명해줘요. 여기서는 S3 버킷 s3://my-bucket에 모델을 직렬화해 저장한다고 가정하고 docstring의 예시 중 하나를 직접 사용할게요:

python examples/features/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   serialize \
   --serialized-directory s3://my-bucket \
   --suffix v1

이렇게 하면 모델 텐서가 s3://my-bucket/vllm/facebook/opt-125m/v1에 저장돼요. 텐서화된 모델에 LoRA 어댑터를 적용하려면 위 명령에 LoRA 어댑터의 HF id를 전달할 수 있고, 그 아티팩트도 함께 저장돼요:

python examples/features/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
   --serialized-directory s3://my-bucket \
   --suffix v1

Tensorizer로 모델 서빙

모델을 원하는 곳에 직렬화했다면 vllm serve나 LLM 진입점으로 모델을 로드할 수 있어요. 모델을 저장한 디렉터리를 LLM()과 vllm serve의 model 인자로 전달하면 돼요. 앞서 LoRA 어댑터와 함께 저장한 텐서화된 모델을 서빙하는 예시:

vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
    --load-format tensorizer \
    --enable-lora

LLM() 사용 시:

from vllm import LLM
llm = LLM(
    "s3://my-bucket/vllm/facebook/opt-125m/v1",
    load_format="tensorizer",
    enable_lora=True,
)

Tensorizer 구성 옵션

tensorizer의 직렬화/역직렬화 핵심 객체는 각각 TensorSerializer와 TensorDeserializer예요. 이들에게 임의 kwargs를 전달해 직렬화·역직렬화 프로세스를 구성하려면, model_loader_extra_config의 키로 serialization_kwargs와 deserialization_kwargs를 제공하면 돼요. 두 객체의 모든 파라미터에 대한 전체 docstring은 tensorizer의 serialization.py 파일에서 볼 수 있어요.

직렬화 시 CPU 동시성을 TensorSerializer 초기화기의 limit_cpu_concurrency 파라미터로 제한하는 예시:

python examples/features/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
   --serialized-directory s3://my-bucket \
   --serialization-kwargs '{"limit_cpu_concurrency": 2}' \
   --suffix v1

로딩 프로세스를 TensorDeserializer로 커스터마이즈하는 예시 — 역직렬화 시 동시 리더 수를 초기화기의 num_readers 파라미터로 model_loader_extra_config를 통해 제한:

vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
    --load-format tensorizer \
    --enable-lora \
    --model-loader-extra-config '{"deserialization_kwargs": {"num_readers": 2}}'

LLM() 사용 시:

from vllm import LLM
llm = LLM(
    "s3://my-bucket/vllm/facebook/opt-125m/v1",
    load_format="tensorizer",
    enable_lora=True,
    model_loader_extra_config={"deserialization_kwargs": {"num_readers": 2}},
)

CoreWeave Tensorizer로 모델 로드하기