본문으로 건너뛰기

[vLLM] Model Loader: 모델 가중치 로딩

들어가며

LLM의 가중치 파일은 수십 GB에 달하며, safetensors, PyTorch checkpoint, GGUF 등 다양한 포맷으로 존재한다. vLLM의 Model Loader는 이 다양한 포맷을 통합적으로 처리하고, 양자화와 텐서 병렬화를 적용하면서 가중치를 GPU에 로딩한다.

코드: vllm/model_executor/model_loader/

핵심 구조/코드 분석

BaseModelLoader: 로더의 인터페이스

class BaseModelLoader(ABC):
    def __init__(self, load_config: LoadConfig):
        self.load_config = load_config

    @abstractmethod
    def download_model(self, model_config: ModelConfig) -> None:
        raise NotImplementedError

    @abstractmethod
    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
        raise NotImplementedError

    def load_model(self, vllm_config, model_config, prefix="") -> nn.Module:
        target_device = torch.device(load_device)
        with set_default_torch_dtype(model_config.dtype):
            with target_device:
                model = initialize_model(
                    vllm_config=vllm_config,
                    model_config=model_config,
                    prefix=prefix,
                )

load_model은 3단계로 구성된다:

  1. 모델 초기화: initialize_model로 빈 모델 구조를 생성
  2. 가중치 로딩: load_weights로 가중치 파일에서 텐서를 로딩
  3. 후처리: process_weights_after_loading으로 양자화 등 후처리 적용

DefaultModelLoader: 기본 로더

class DefaultModelLoader(BaseModelLoader):
    DEFAULT_NUM_THREADS = 8

    @dataclasses.dataclass
    class Source:
        model_or_path: str
        revision: str | None
        subfolder: str | None = None
        prefix: str = ""
        fall_back_to_pt: bool = True
        allow_patterns_overrides: list[str] | None = None

    def __init__(self, load_config: LoadConfig):
        super().__init__(load_config)
        self.local_expert_ids: set[int] | None = None

Source 데이터클래스가 모델 소스의 메타데이터를 캡슐화한다. local_expert_ids는 MoE(Mixture of Experts) 모델에서 현재 워커가 담당하는 전문가 ID를 추적한다.

가중치 파일 준비

def _prepare_weights(self, model_name_or_path, subfolder, revision,
                     fall_back_to_pt, allow_patterns_overrides):
    is_local = os.path.isdir(model_name_or_path)
    load_format = self.load_config.load_format

    if load_format == "auto":
        load_format = (
            "mistral"
            if len(list_filtered_repo_files(
                model_name_or_path=model_name_or_path,
                allow_patterns=["consolidated*.safetensors"],
            )) > 0
            else "hf"
        )

    if load_format == "hf":
        allow_patterns = ["*.safetensors", "*.bin"]
    elif load_format == "safetensors":
        allow_patterns = ["*.safetensors"]
    elif load_format == "mistral":
        allow_patterns = ["consolidated*.safetensors"]
    elif load_format == "pt":
        allow_patterns = ["*.pt"]

auto 모드에서는 Mistral 공식 포맷(consolidated*.safetensors)이 있으면 Mistral 로더를, 아니면 HuggingFace 로더를 선택한다. 로컬이 아닌 경우 HuggingFace Hub에서 자동 다운로드한다.

가중치 이터레이터: 메모리 효율적 로딩

def _get_weights_iterator(self, source):
    hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(...)

    if use_safetensors:
        if self.load_config.load_format == "fastsafetensors":
            weights_iterator = fastsafetensors_weights_iterator(...)
        elif self.load_config.load_format == "instanttensor":
            weights_iterator = instanttensor_weights_iterator(...)
        else:
            if extra_config.get("enable_multithread_load"):
                weights_iterator = multi_thread_safetensors_weights_iterator(
                    hf_weights_files,
                    max_workers=extra_config.get("num_threads", self.DEFAULT_NUM_THREADS),
                )
            else:
                weights_iterator = safetensors_weights_iterator(
                    hf_weights_files,
                    local_expert_ids=self.local_expert_ids,
                )
    else:
        weights_iterator = pt_weights_iterator(hf_weights_files, ...)

가중치를 이터레이터로 반환하여 전체를 메모리에 올리지 않고 스트리밍으로 처리한다. 다양한 로딩 전략을 지원한다:

  • safetensors: 기본, 가장 일반적
  • fastsafetensors: 대규모 모델을 위한 최적화 로더
  • instanttensor: 즉시 텐서 로딩
  • multi_thread: 멀티스레드 로딩으로 I/O 병렬화
  • local_expert_ids: MoE에서 필요한 전문가의 가중치만 선택적 로딩

중복 파일 필터링

if use_safetensors:
    hf_weights_files = filter_duplicate_safetensors_files(
        hf_weights_files, hf_folder, index_file
    )
else:
    hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)

Mistral-7B-Instruct-v0.3 같은 모델에서는 샤딩된 safetensors 파일과 통합 safetensors 파일이 동시에 존재할 수 있다. 인덱스 파일을 참조하여 중복을 제거하고, 추론에 불필요한 파일(optimizer state 등)도 필터링한다.

다양한 로더 구현체

model_loader/ 디렉토리에는 다양한 전문 로더가 존재한다:

  • DefaultModelLoader: 표준 HuggingFace/safetensors 포맷
  • BitsAndBytesModelLoader: 4/8비트 양자화 로딩
  • GGUFModelLoader: llama.cpp의 GGUF 포맷 지원
  • ShardedStateLoader: 사전 샤딩된 체크포인트 로딩
  • TensorizerLoader: CoreWeave Tensorizer로 S3/GCS에서 빠른 로딩
  • DummyModelLoader: 테스트/프로파일링용 더미 가중치

왜 이 설계인가

  1. 이터레이터 기반 로딩: 가중치를 스트리밍으로 처리하여 메모리 피크를 최소화한다. 70B 모델도 가중치 전체를 CPU RAM에 올리지 않고 점진적으로 GPU에 전송할 수 있다.

  2. MoE 전문가 필터링: local_expert_ids로 해당 워커가 필요한 전문가의 가중치만 로딩한다. Expert Parallelism에서 불필요한 데이터 전송을 방지한다.

  3. 포맷 자동 감지: auto 모드에서 파일 패턴을 분석하여 최적의 로딩 전략을 자동 선택한다. 사용자가 모델 포맷을 알 필요가 없다.

  4. 멀티스레드 로딩: 대규모 모델의 경우 I/O가 병목이 될 수 있다. 멀티스레드 로딩으로 여러 가중치 파일을 동시에 읽어 로딩 시간을 단축한다.

정리

Model Loader는 vLLM의 모델 지원 범위를 결정하는 핵심 컴포넌트이다. HuggingFace Hub, GGUF, Tensorizer 등 다양한 소스와 포맷을 통합하면서, 메모리 효율성과 로딩 속도를 모두 최적화한다. 새로운 모델 포맷이 등장하면 BaseModelLoader를 상속하여 쉽게 확장할 수 있다.

댓글

관련 포스트

vLLM 의 다른글