[SGLang] Model Loader: 가중치 로딩 인프라와 최적화

2026년 4월 11일수정: 2026년 4월 11일

들어가며

LLM 서빙에서 모델 로딩은 서버 시작 시간의 대부분을 차지한다. 수십 GB의 가중치를 디스크에서 GPU로 효율적으로 전송하는 것은 시스템 성능의 핵심이다. SGLang은 Safetensors 멀티스레드 로딩, 메모리 매핑, 텐서 병렬 분산 등 다양한 최적화를 제공한다.

이 글에서는 python/sglang/srt/model_loader/loader.py와 weight_utils.py를 중심으로 Model Loader의 설계를 분석한다.

로더 계층 구조

SGLang은 다양한 로딩 형식을 지원하기 위해 추상 기반 클래스를 사용한다.

┌────────────────────┐
│  BaseModelLoader   │  (추상 기반)
├────────────────────┤
│  download_model()  │
│  load_model()      │
└────────┬───────────┘
         │
    ┌────┴──────────────────────┐
    │                           │
┌──────────────────┐  ┌─────────────────────┐
│ DefaultModelLoader│  │ DummyModelLoader    │
│ (Safetensors/PT) │  │ (테스트용 더미 가중치)│
└──────────────────┘  └─────────────────────┘

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_model(self, *, model_config, device_config) -> nn.Module:
        raise NotImplementedError

DefaultModelLoader: 가중치 준비

DefaultModelLoader는 모델 가중치 파일을 찾고, 적절한 이터레이터를 선택한다.

가중치 파일 탐색

_prepare_weights()는 로컬 또는 HuggingFace Hub에서 가중치 파일을 찾는다.

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

    if load_format == LoadFormat.AUTO:
        allow_patterns = ["*.safetensors", "*.bin"]
    elif load_format == LoadFormat.SAFETENSORS:
        use_safetensors = True
        allow_patterns = ["*.safetensors"]
    elif load_format == LoadFormat.PT:
        allow_patterns = ["*.pt"]

    if not is_local:
        hf_folder = download_weights_from_hf(
            model_name_or_path, self.load_config.download_dir,
            allow_patterns, revision,
        )
    else:
        hf_folder = model_name_or_path

    return hf_folder, hf_weights_files, use_safetensors

Safetensors를 우선 사용하되, .pt 파일로 폴백할 수 있다.

가중치 이터레이터 선택

_get_weights_iterator()에서 포맷에 맞는 이터레이터를 선택한다.

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

    if use_safetensors:
        if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
            weights_iterator = fastsafetensors_weights_iterator(hf_weights_files)
        elif use_multithread:
            weights_iterator = buffered_multi_thread_safetensors_weights_iterator(
                hf_weights_files,
                max_workers=extra_config.get("num_threads", self.DEFAULT_NUM_THREADS),
                disable_mmap=weight_loader_disable_mmap,
            )
        else:
            weights_iterator = safetensors_weights_iterator(
                hf_weights_files, disable_mmap=weight_loader_disable_mmap
            )
    else:
        weights_iterator = pt_weights_iterator(hf_weights_files)

기본 스레드 수는 8개이다.

class DefaultModelLoader(BaseModelLoader):
    DEFAULT_NUM_THREADS = 8

양자화 설정 로딩

모델 로딩 전에 양자화 설정을 확인한다.

def _get_quantization_config(model_config, load_config):
    model_class, _ = get_model_architecture(model_config)
    packed_modules_mapping = getattr(model_class, "packed_modules_mapping", {})

    if model_config.quantization is not None:
        quant_config = get_quant_config(
            model_config, load_config, packed_modules_mapping
        )
        # GPU Capability 검증
        major, minor = get_device_capability()
        capability = major * 10 + minor
        if capability < quant_config.get_min_capability():
            raise ValueError(
                f"The quantization method {model_config.quantization} "
                "is not supported for the current GPU."
            )
        return quant_config
    return None

device_loading_context: CPU ↔ GPU 전환

가중치를 CPU에서 로딩한 후 GPU로 옮기는 컨텍스트 매니저이다.

@contextmanager
def device_loading_context(module, target_device):
    if target_device.type == "cpu":
        yield module
        return

    original_infos = {}
    # CPU 파라미터를 GPU로 이동
    for name, p in module.named_parameters():
        if p.device.type == "cpu":
            device_data = p.data.to(target_device)
            original_infos[name] = dict(
                device=p.device,
                original_data=p.data,
                device_data=device_data,
            )
            p.data = device_data

    try:
        yield module
    finally:
        # 원래 디바이스로 복원
        for name, p in module.named_parameters():
            if name in original_infos:
                original_data = original_infos[name]["original_data"]
                original_data.copy_(p.data.to(original_data.device))
                p.data = original_data

weight_utils.py: 저수준 유틸리티

파일 잠금

여러 프로세스가 동시에 같은 모델을 다운로드하는 것을 방지한다.

def get_lock(model_name_or_path, cache_dir=None, suffix=""):
    lock_dir = cache_dir or temp_dir
    model_name = model_name_or_path.replace("/", "-")
    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
    lock_file_name = hash_name + model_name + suffix + ".lock"
    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
    return lock

mode=0o666으로 여러 사용자가 잠금을 공유할 수 있다.

Safetensors 중복 파일 필터링

분할 Safetensors 파일과 통합 파일이 공존하는 경우 충돌을 방지한다.

hf_weights_files = filter_duplicate_safetensors_files(
    hf_weights_files, hf_folder, index_file
)

MTP 가중치 필터링

Multi-Token Prediction 모델에서 특정 draft layer만 로딩한다.

@classmethod
def _filter_mtp_weights(cls, weights_iterator, prefix, draft_model_idx):
    filtered_weights = []
    for name, tensor in weights_iterator:
        match = cls._MTP_PATTERN.match(name)
        if match is not None:
            idx = int(match.group(1))
            if idx != draft_model_idx:
                continue
            new_name = name.replace(match.group(), "model.mtp.layers.0.")
        else:
            new_name = name
        filtered_weights.append((prefix + new_name, tensor))
    return tuple(filtered_weights)

로딩 형식 요약

SGLang이 지원하는 로딩 형식을 정리한다.

형식	파일 패턴	특징
AUTO	`.safetensors`, `.bin`	자동 감지
SAFETENSORS	`*.safetensors`	메모리 매핑, 빠른 로딩
FASTSAFETENSORS	`*.safetensors`	fastsafetensors 라이브러리
PT	`*.pt`	PyTorch 기본 형식
NPCACHE	`*.bin`	NumPy 캐시
MISTRAL	`consolidated*.safetensors`	Mistral 모델 전용
DUMMY	-	테스트용 랜덤 가중치

설계 근거: 멀티스레드 로딩

가중치 로딩은 I/O 바운드이므로 멀티스레드가 효과적이다. buffered_multi_thread_safetensors_weights_iterator는 여러 스레드가 파일을 병렬로 읽고, 버퍼에 텐서를 쌓아 순차 소비한다.

Thread 1 ─→ file_1.safetensors ─→ ┐
Thread 2 ─→ file_2.safetensors ─→ ├→ Buffer ─→ weight_loader
Thread 3 ─→ file_3.safetensors ─→ ┘

참고

SGLang 의 다른글

이전글 [SGLang] Piecewise CUDA Graph: 분할 그래프 컴파일 전략
현재글 : [SGLang] Model Loader: 가중치 로딩 인프라와 최적화
다음글 [SGLang] torch.compile & Inductor: PyTorch 컴파일러 통합