DefaultModelLoader:HF Hub から GPU メモリに重みを運ぶ

源码版本v0.25.1

役割

モデルローダー (model loader) は vLLM の起動フローにおいて、モデル重みをディスクや Hugging Face Hub から GPU メモリに運ぶ部分です。DefaultModelLoader はその中で最もよく使われる実装で、safetensors / .pt / Mistral 独自形式 / npcache など多様なディスク形式を扱い、複数シャードの重みを 1 つの (name, tensor) イテレータにまとめ、モデル自身の load_weights メソッドに渡して nn.Module に順次流し込みます。BaseModelLoader の上で「ダウンロード・フィルタ・パース・イテレート」の 4 つを直列に実行します。テンソルの分割や量子化には関与しません。この 2 つは ColumnParallelLinear のようなレイヤーと QuantizationConfig が受け持ちます。

具体的な 1 回のロードは概ね次のとおりです。__init__ は LoadConfig の model_loader_extra_config をバリデーションし (default_loader.py:74-126)、enable_multithread_load、num_threads、enable_weights_track の 3 つのキーだけを許可します。download_model は _prepare_weights を呼び出し、load_format に応じて allow_patterns を選び、ローカルディレクトリなら直接 glob、リモートなら download_weights_from_hf を使います (default_loader.py:194-207)。load_weights が本当に仕事をするエントリで、まず EP フィルタ初期化を行い (default_loader.py:425-426)、get_all_weights イテレータを model.load_weights に渡し、最後に track_weights_loading で漏れがないか検証します (default_loader.py:444-445)。

設計動機

なぜローダーを Loader + Source の 2 層に分けるのか?

• 複数ソースの重み:モデル本体以外に secondary_weights を持つことがあります (例: マルチモーダルの visual encoder が別の checkpoint)。Source dataclass により各ソースは独自の prefix、fall_back_to_pt、allow_patterns_overrides を持ち (default_loader.py:49-69)、get_all_weights がこれらを 1 つのイテレータにまとめます (default_loader.py:321-340)。
• フォーマットの分岐:load_format 1 つで後の経路が決まります。auto は consolidated*.safetensors を検出すると自動で mistral に切り替え (default_loader.py:151-163)、safetensors / fastsafetensors / instanttensor は .safetensors だけを強制し、pt / npcache は従来経路に進みます。
• マルチスレッド高速化:enable_multithread_load を入れると multi_thread_safetensors_weights_iterator または multi_thread_pt_weights_iterator に切り替わり、デフォルトは 8 スレッドです (default_loader.py:278-308)。ただしマルチスレッドモードはデフォルトの lazy 戦略のみをサポートし、safetensors_load_strategy の他の戦略とは排他です (default_loader.py:116-126)。
• EP 重みフィルタ:MoE + 専門家並列時に、自 rank が担当する expert だけをロードします。_init_ep_weight_filter が local_expert_ids を計算し (default_loader.py:351-412)、後続の safetensors イテレータがこれに従って非ローカル expert テンサルをスキップしてディスク IO を節約します。
• safetensors インデックスの重複排除:Mistral-7B-Instruct-v0.3 のように sharded と consolidated の 2 つの safetensors が共存する場合、直接 glob するとすべて読んでしまい衝突します。filter_duplicate_safetensors_files が model.safetensors.index.json を使ってフィルタします (default_loader.py:217-235)。
• ロード追跡:非量子化モデルはデフォルトで enable_weights_track=True、ロード完了後に named_parameters と照合して漏れた重みを検出し、ランダム重みのまま推論が走るのを防ぎます (default_loader.py:434-445)。

主要ファイル

• DefaultModelLoader class:43-47 — クラス定義と docstring、DEFAULT_NUM_THREADS = 8。
• Source dataclass:49-69 — Source のインナーデータクラス、重みソースを 1 つ記述。
• _prepare_weights:128-242 — load_format をパース、ダウンロード、glob、重複排除、(hf_folder, hf_weights_files, use_safetensors) を返します。
• _get_weights_iterator:244-319 — load_format とマルチスレッドスイッチで weights iterator を選び、最後にテンサル名に prefix を付けます。
• get_all_weights:321-340 — primary source と secondary_weights を 1 つのジェネレータにまとめます。
• _init_ep_weight_filter:351-412 — EP モードで local_expert_ids を計算し、非ローカル expert テンサルをスキップします。
• load_weights:414-445 — ロードのエントリ、torchao 戦略切り替え → EP フィルタ → model.load_weights → 追跡検証。
• track_weights_loading:447-475 — named_parameters と loaded_weights を照合し、欠落や予想外の重みを報告します。
• weight_utils.py:1-1 — safetensors_weights_iterator、download_weights_from_hf、filter_duplicate_safetensors_files などの低レイヤーツールを提供。
• base_loader.py:1-1 — BaseModelLoader 抽象基底クラス、load_weights インターフェースを定義。

データフロー

呼び出し側 (通常は worker の load_model) はまず DefaultModelLoader を __init__ し、次に load_weights(model, model_config) を呼びます。関数本体の最初のステップは quantization == "torchao" に応じて safetensors ロード戦略を調整すること、その後 _init_ep_weight_filter で EP フィルタ集合を計算し、get_all_weights が各 Source を _get_weights_iterator に渡して 1 つのジェネレータにまとめ、最後に model.load_weights が自身で消費します:

@instrument(span_name="Load weights")
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
    if model_config.quantization == "torchao":
        quant_config = get_quant_config(model_config, self.load_config)
        if (
            hasattr(quant_config, "is_checkpoint_torchao_serialized")
            and quant_config.is_checkpoint_torchao_serialized
            and torchao_version_at_least("0.15.0")
        ):
            self.load_config.safetensors_load_strategy = "torchao"

    self._init_ep_weight_filter(model_config)

    loaded_weights = model.load_weights(self.get_all_weights(model_config, model))

    self.counter_after_loading_weights = time.perf_counter()
    logger.info_once(
        "Loading weights took %.2f seconds",
        self.counter_after_loading_weights - self.counter_before_loading_weights,
    )
    # We only enable strict check for non-quantized models
    # that have loaded weights tracking by default.
    default_enable_weights_track = (
        model_config.quantization is None and loaded_weights is not None
    )

(default_loader.py:414-438) model.load_weights(...) こそが本当にテンサルを nn.Module にコピーする場所です。vLLM の各モデル実装はこのメソッドをカスタマイズし、テンサル名に基づいて対応する weight_loader (例: ColumnParallelLinear.weight_loader_v2) にマッチさせます。シャーディングも量子化もそのレイヤーで行われます。詳細は テンサル並列線形レイヤー と 量子化レイヤーのロード を参照してください。

境界と失敗

• load_format が未知:hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache のリストにない値は直接 raise ValueError します (default_loader.py:183-184)。
• 重みファイルが見つからない:glob で 1 つもファイルが見つからないときは Cannot find any model weights を投げます (default_loader.py:237-240)。
• マルチスレッド + safetensors_load_strategy の非互換:enable_multithread_load=True はデフォルトの lazy 戦略のみをサポートし、他の戦略と組み合わせると raise ValueError します (default_loader.py:118-126)。
• extra_config の型エラー:model_loader_extra_config が dict でない、キーがホワイトリストにない、num_threads が正整数でない、これらは即座にエラーになります (default_loader.py:79-110)。
• EPLB 時は EP フィルタをスキップ:enable_eplb が有効なとき、冗長 expert スロットは別 rank の論理 expert を指すため、フィルタするとこれらのスロットが漏れてしまいます。そのためこの区間はそのまま return します (default_loader.py:370-375)。
• 量子化モデルは追跡検証をスキップ:量子化メソッドが process_weights_after_loading や uses_meta_device を持つ場合、checkpoint に対応する scale がないことがあります。これらのパラメータは追跡集合から外されます (default_loader.py:453-462)。

まとめ

DefaultModelLoader は「ファイルを探す・ダウンロードする・テンサルをイテレートする」の 3 つを確実に行い、具体的なシャーディング / 量子化 / 逆量子化はすべて具体的なレイヤーと量子化メソッドに任せます。テンサル並列線形レイヤー との接続点は model.load_weights 内部でテンサル名ごとに weight_loader を引くことで、量子化レイヤーのロード との接続点は quant_config.get_quant_method が各レイヤーに LinearMethodBase を取り付けることで行われます。ローダー自身はそれらを意識しません。

公式資料: vLLM 文档 · README