Transformers documentation

動的ウェイトの読み込み

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

動的ウェイトの読み込み

チェックポイントは、モデルが実行時に期待する形式と一致しない形式でシリアライズされていることがよくあります。一般的なシナリオには、以下が含まれます。

  1. 融合された重み: チェックポイントは gate_proj と up_proj の重みをそれぞれ個別に保存しますが、モデルは効率化のために融合された gate_up_proj を使用します。
  2. MoEエキスパートの統合: 個々のエキスパートの重み(experts.0.weight、experts.1.weight、…)を、単一の3次元テンソルにスタックする必要があります。
  3. 従来の命名規則: 古いチェックポイントでは異なる命名規則が使用されています(例: LayerNorm.gamma と LayerNorm.weight)。
  4. 複合モデル: ビジョン言語モデルには、それぞれ独自のチェックポイント規約を持つ2つのPreTrainedModelサブモジュールが含まれます。
  5. 量子化: 重みは、デシリアライズが必要な量子化形式で保存されている場合があります。

Dynamic weight loading は、チェックポイントのテンソルを読み込む際に、スケジュールされた可逆操作を適用することでこの問題を解決します。Transformers は、WeightConverter と WeightRenaming を通じてこれを公開しています。これらは、1 つ以上のチェックポイントキーが 1 つ以上のモデルパラメータにどのようにマッピングされるかを記述し、マッチしたテンソルに対して実行すべきコンポーザブルな ConversionOps を指定します。このアプローチは、新しい重みレイアウトに適応し、量子化された Mixture-of-Experts (MoE) をサポートし、テンソル並列性と統合できます。

このガイドでは、WeightConverterを使用してテンソルをコンバートする方法を示します。コンバージョンのマッピングはconversion_mapping.pyに定義されており、登録されたマッピングはmodel_type文字列(例: "mixtral")またはクラス名(例: "LlavaModel")をキーとして識別されます。

完全な読み込みパイプライン

すべてのモデルは動的な重み読み込みシステムを経由します。変換マッピングは、そのシステム内の任意のステップであり、モデルがクラス名またはmodel_typeに対してエントリを登録している場合にのみ有効になります。

Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model()
                                              ↓
                         ┌───────────────────────────────────────────────────────────┐
                         │  For each weight in checkpoint:                           │
                         │  1. Match renamed/processed source key to model parameter │
                         │  2. Shard the weight and send to device (async)           │
                         │  3. Collect tensors with the same source_pattern together │
                         │     (e.g. MoE experts, gate_up_proj)                      │
                         │  4. Apply dequantization/deserialization (if pre-quant)   │
                         │  5. Apply conversion (if defined)                         │
                         │  6. Apply quantization (if enabled and step 4 not used)   │
                         │  7. Set parameter on model                                │
                         └───────────────────────────────────────────────────────────┘
ステップ有効化されるタイミング
動的読み込み常に、すべてのモデルに対して
変換マッピングモデルのクラスまたはmodel_typeが_MODEL_TO_CONVERSION_PATTERNに登録されている場合のみ
TPシャーディングDistributedConfig(tp_size=N) が使用され、かつモデルに base_model_tp_plan が存在する場合のみ
デクアンタイゼーション/デシリアライズ事前量子化済みチェックポイントを読み込む場合のみ
量子化量子化設定が提供され、かつ重みが事前量子化されていない場合のみ

Denseモデル(例: Llama)

ほとんどのdenseモデルでは、チェックポイントの形式がモデルの形式と直接一致するため、変換マッピングは不要です。一部のモデルでは、リネーム(例:レガシーな命名規則)が引き続き必要になる場合があります。TPシャーディングは、有効化されている場合に引き続き適用されます。

Checkpoint:                             Model:
model.layers.0.self_attn.q_proj.weight  →  model.layers.0.self_attn.q_proj.weight
model.layers.0.self_attn.k_proj.weight  →  model.layers.0.self_attn.k_proj.weight
model.layers.0.mlp.gate_proj.weight     →  model.layers.0.mlp.gate_proj.weight
model.layers.0.mlp.up_proj.weight       →  model.layers.0.mlp.up_proj.weight
model.layers.0.mlp.down_proj.weight     →  model.layers.0.mlp.down_proj.weight

レガシーのチェックポイントでは、すべてのモデルに適用される組み込みの名称変更によって処理される、古い命名規則が使用されている場合があります。

Checkpoint:                          Model:
LayerNorm.gamma              →       LayerNorm.weight
LayerNorm.beta               →       LayerNorm.bias

MoEモデル(例: Mixtral)

MoEモデルの場合、チェックポイントの形式はモデルの形式と異なります。変換マッピングにより、個別のエキスパート重みが融合された3Dテンソルに変換され、その後にTPシャーディングが適用されます。

Checkpoint:                              Model:
experts.0.w1.weight  ─┐
experts.1.w1.weight   │ MergeModulelist
...                   ├───────────────→  experts.gate_up_proj (8, hidden, 2*intermediate)
experts.0.w3.weight   │ + Concatenate
experts.1.w3.weight  ─┘

複合モデル(例:ビジョン・言語)

PreTrainedModel は、他の PreTrainedModel サブモジュールを含む場合があります。各サブモデルには独自のコンバージョンマッピングがあり、クラス名または model_type に対して登録されます。親モデルが読み込まれると、get_model_conversion_mapping がサブモデルを深さ優先順に走査し、それらのマッピングを収集して、scope_prefix を介して各トランスフォームをそのサブモジュールのパスに自動的にスコープします。

Composite model:                              Per-submodel mappings (auto-scoped):
LlavaForConditionalGeneration
├── vision_model: SiglipVisionModel    →     SiglipVisionModel mapping (scope="vision_model")
└── language_model: LlamaForCausalLM   →     LlamaForCausalLM mapping  (scope="language_model")

scope_prefix はサブモジュールのドット区切りパス("vision_model"、"language_model.model" など)です。スコープ付きの transform は、f"{scope_prefix}." で始まるキーに対してのみ発火します。パターンマッチングの前にプレフィックスが除去され、置換後に再付与されるため、各サブモデルのマッピングは、あたかもそれがルートであるかのように、サブモデルを基準として記述されます。

アーキテクチャ

システムは、core_model_loading.py で定義されるいくつかの主要なコンポーネントを中心に構築されています。

フェーズ 1 — キーごとの処理(チェックポイントのキーを反復します):

  1. transform リストを一度だけ走査します。一致したすべての WeightRenaming が発火します(例: block_sparse_moe → mlp)。また、最大 1 つの WeightConverter がそのキーを要求できます(例: experts.*.w1.weight)。
  2. ThreadPoolExecutor を使用して、シャード(TP)化してデバイスに送信を非同期で実行します。
  3. source_pattern が同じ Collect テンソルをまとめて処理します(例: すべての MoE エキスパート重み、gate + up 射影)。

フェーズ 2 — マッピングごとの処理(収集されたマッピングを反復処理します):

  1. Dequantize/deserialize(事前量子化済みチェックポイントのみ)。
  2. ConversionOps チェーンを適用: Chunk、Concatenate、MergeModulelist、Transpose など
  3. Quantize をオンザフライで実行します(事前量子化されていない場合)。
  4. モデルにパラメータを設定します。

WeightTransform

パターンマッチングとテンソルの収集を処理するベースクラス:

  • パターンのコンパイル: ソースパターンは、re.search() で照合される完全な正規表現です。* ワイルドカードは任意のインデックス可能なコンポーネントにマッチし、すべてのマッチをまとめてバッチ操作に使用します。
  • キャプチャグループとバックリファレンス: ソースパターン内のキャプチャグループは、ターゲットパターン内で\1、\2…として参照でき、リネーム後も部分文字列(レイヤーインデックスなど)を保持できます。
  • スコープ指定: scope_prefix(get_model_conversion_mapping によってサブモデルごとに自動設定されます)は、変換をそのパス以下のキーに限定します。マッチングの前にプレフィックスは除去され、置換後に再度付与されます。
  • キーの名称変更: rename_source_key() は正規表現(スコープ処理を含む)を適用し、(renamed_key, source_pattern) を返します。これにより、ローダーはどのコンバーター(該当する場合)がそのキーを要求したかを把握できます。
  • テンソル集約: add_tensor() は、解決済みのテンソル(または Future)を source_pattern ごとに蓄積し、単一の変換(例: 全 MoE エキスパートの重み)に必要なすべてのテンソルが、オペレーションチェーンの実行前にバッチとしてまとめられるようにします。
  • 可逆性: reverse_transform() はソースとターゲットのパターンを入れ替え、各操作を逆転させるため、同じリストを逆順に適用することで保存が実行されます。
  • トラッキング: was_used() は、読み込み時にトランスフォームが実際にどのキーにもマッチしたかを報告します。これにより、非双射的なリネーム(例: 既に存在する可能性のあるプレフィックスを追加する PrefixChange)を、保存時に対称的に再適用できるようになります。

WeightRenaming

WeightRenaming は、テンソル操作を伴わない純粋なキー名変更のための専用 WeightTransform です。WeightConverter とは異なり、WeightRenaming はキーを claim しないため、複数のリネームを自由にチェーンできます。ロードパスではすべてのリネームがコンバーターより前に実行され、セーブパス(逆順リスト)ではすべての逆リネームが逆コンバーターより後に実行されます — 上記の順序ルールを参照してください。

# Legacy checkpoint compatibility
WeightRenaming("LayerNorm.gamma", "LayerNorm.weight")

# Module path changes
WeightRenaming(".block_sparse_moe.", ".mlp.")

PrefixChange は、パスのコンポーネント全体を削除または追加するという一般的なケース向けに、WeightRenaming をラップする上位のラッパーです。オプションの model_prefix を指定すると、その名前空間配下のキーに対してのみ操作が適用されます。

# "model.layers.bad_prefix.weight" → "model.layers.weight"
PrefixChange(prefix_to_remove="bad_prefix", model_prefix="model.layers")
# "layers.0.weight" → "model.layers.0.weight"
PrefixChange(prefix_to_add="model")

WeightConverter

WeightConverter は WeightTransform を継承し、収集されたテンソルに対して作用する ConversionOps のチェーンを備えています。サポートされている4つの基数は以下の通りです。

基数ソースパターンターゲットパターン一般的な操作
1対111Transpose、PermuteForRope
1対多1>1チャンク(例:qkv_proj をアンパックする)
many-to-one>11Concatenate、MergeModulelist(例:エキスパートの融合)
多対多>1>1明示的にサポートされている操作(例: ErnieFuseAndSplitTextVisionExperts)でのみ
WeightConverter(
    source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
    target_patterns=".experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)

WeightConverter は可逆性も備えています。reverse_transform() はソースとターゲットを入れ替え、各操作をその reverse_op に置き換えます。そのため、登録された変換マッピングは構造上双方向です(読み込みではリストをそのまま使用し、保存では逆順で使用します)。

リネームとコンバーターの順序

WeightRenaming のエントリを WeightConverter のエントリより前に配置し、それぞれのリーフを重複させないでください。リネームはコンバーターが消費するキーを正規化しますが、コンバーターが生成するリーフを対象にすることはできません。保存パスはこの前提に依存しており、マッピングを2段階で逆変換します(まずコンバーターを逆変換し、次にリネームを逆変換)。

transform リストは、チェックポイントのキーごとに1回ずつ、順番に処理されます。各キーについて:

  • 一致するすべての WeightRenaming が発火し、複数のリネームがチェーン(連鎖)して実行される場合があります。
  • WeightConverter のうち、最初に claims でそのキーに一致するものが採用されます。以降のコンバーターはスキップされます。これにより、マージ/スプリット処理においてテンソルが単一のコンバーターにルーティングされることが保証され、意図が重複する2つのコンバーターが存在することは設定ミスとみなされます。
weight_mapping = [
    WeightRenaming("^old_prefix", "encoder"),              # rename runs always
    WeightConverter(                                       # converter claims the key
        "attn.qkv_proj.weight",
        ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"],
        operations=[Chunk(dim=0)],
    ),
]
# Load:  "old_prefix.attn.qkv_proj.weight"
#   → WeightRenaming  → "encoder.attn.qkv_proj.weight"
#   → WeightConverter → "encoder.attn.{q,k,v}_proj.weight"
#
# Save: list reversed, each transform inverted:
#   → rev(WeightConverter) repacks QKV → "encoder.attn.qkv_proj.weight"
#   → rev(WeightRenaming)  fixes prefix → "old_prefix.attn.qkv_proj.weight"

変換操作

WeightConverter クラスには、from_pretrained() が呼び出された際に、チェックポイントのソーステンソルをモデルのターゲットテンソルに変換するために実行される複数の操作があります。

操作は完全に可逆です。保存すると変換が逆転し、元のチェックポイントが返されるため、異なるフレームワーク間で簡単に作業できます。

操作逆順
Chunk(dim)Concatenate(dim)
Concatenate(dim)Chunk(dim)
MergeModulelist(dim)SplitModulelist(dim)
SplitModulelist(dim)MergeModulelist(dim)
Transpose(d0, d1)Transpose(d1, d0)
PermuteForRope()PermuteForRope()
Conv3dToLinear(...)LinearToConv3d(...)

チャンク

Chunk 操作は、テンソルを指定した次元に沿って等分します。例えば、モデルが単一のテンソルではなく、Q、K、V を3つの独立したテンソルとして期待する場合などに使用します。

WeightConverter(
    "self_attn.qkv_proj",
    ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
    operations=[Chunk(dim=0)],
)

連結

Concatenate 操作は、複数のテンソルを単一のテンソルに結合します。例えば、モデルが Q、K、V を個別のテンソルではなく単一のテンソルとして期待する場合などに使用します。

WeightConverter(
    ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
    "self_attn.qkv_proj",
    operations=[Concatenate(dim=0)],
)

MergeModulelist

MergeModulelist は、2D テンソルのリストを単一の 3D テンソルに結合します。例えば、MergeModulelist と Concatenate を組み合わせることで、MoE のエキスパートをスタックし、1 つのテンソルにパックできます。

WeightConverter(
    ["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight"],
    "mlp.experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)

SplitModulelist

SplitModulelist は、3D テンソルを 2D テンソルのリストに分割します。例えば、スタックされたエキスパートを個々のエキスパートに分割できます。

WeightConverter(
    "mlp.experts.down_proj",
    "block_sparse_moe.experts.*.w2.weight",
    operations=[SplitModulelist(dim=0)],
)

PermuteForRope

PermuteForRope は、重みをインターリーブ形式から sin/cos 形式に変換します。例えば、Chunk と PermuteForRope を組み合わせることで、融合された QKV テンソルを分割し、Q と K に sin/cos RoPE 置換を適用できます。

WeightConverter(
    ["model.layers.*.self_attn.qkv_proj.weight"],
    [
        "model.layers.*.self_attn.q_proj.weight",
        "model.layers.*.self_attn.k_proj.weight",
        "model.layers.*.self_attn.v_proj.weight",
    ],
    operations=[Chunk(dim=0), PermuteForRope()],
)

転置

Transpose はテンソルの次元を入れ替えます。異なる規約間で重みのレイアウトを相互変換する際に便利です。

WeightConverter(
    source_patterns="mlp.gate.weight",
    target_patterns="mlp.text_moe.gate.weight",
    operations=[Transpose(dim0=0, dim1=1)],
)

操作のチェーン

オペレーションをチェーンすることで、複雑な変換を実行できます。オペレーションは順序通りに実行され、各オペレーションの出力が次のオペレーションの入力になります。

例: Mixtral MoEの変換

WeightConverter(
    source_patterns=[
        ".experts.*.w1.weight",  # gate_proj per expert
        ".experts.*.w3.weight",  # up_proj per expert
    ],
    target_patterns=".experts.gate_up_proj",
    operations=[
        MergeModulelist(dim=0),  # Stack all experts: (n_experts, in, out)
        Concatenate(dim=1),      # Fuse gate+up: (n_experts, in, 2*out)
    ],
)

データフロー:

Input:
  ".experts.*.w1.weight": [tensor_0, tensor_1, ..., tensor_7]  # 8 experts
  ".experts.*.w3.weight": [tensor_0, tensor_1, ..., tensor_7]  # 8 experts

After MergeModulelist(dim=0):
  ".experts.*.w1.weight": (8, 4096, 14336)  # stacked gate
  ".experts.*.w3.weight": (8, 4096, 14336)  # stacked up

After Concatenate(dim=1):
  ".experts.gate_up_proj": (8, 4096, 28672)  # fused gate_up

パターンマッチングルール

  • ソースパターンはre.search()で評価される完全な正規表現です。^はscope_prefixの除去後のスコープ付きキーの先頭にアンカーします。
  • * はインデックス単位のワイルドカードであり、同じソースパターンに一致するすべてのテンソルを収集します(正しい連結のためにチェックポイントの順序を保持します)。
  • ソースパターン内のキャプチャグループは、ターゲットパターン内で\1、\2…のように参照でき、元のキーの一部(レイヤーのインデックスなど)を保持できます。
  • スコープ付きトランスフォーム(scope_prefix が設定されている場合)は、f"{scope_prefix}." で始まるキーのみを一致させます。

変換マッピングの登録

変換リストは文字列キーに対して登録されます。キーはmodel_type(例:"mixtral")またはクラス名(例:"LlavaModel")です。

from transformers.conversion_mapping import register_checkpoint_conversion_mapping

register_checkpoint_conversion_mapping(
    "my_model_type",
    [
        WeightRenaming(".old.", ".new."),
        WeightConverter(
            ".experts.*.w1.weight",
            ".experts.gate_proj",
            operations=[MergeModulelist(dim=0)],
        ),
    ],
)

ルックアップルール

get_model_conversion_mapping が PreTrainedModel を処理する際、すべてのサブ PreTrainedModel は DFS 順で訪問されます(nn.Module.named_modules() を PreTrainedModel のインスタンスに絞り込んだもの)。各サブモデルは 2 つの段階で解決され、まずカスタムコードのフィルタが実行され、次に登録済みのマッピングが参照されます。

カスタムコードのモデルは、参照が実行される前にスキップされます。サブモデルがカスタムコードに該当するのは、is_custom_code()がTrueを返す場合です。これには、trust_remote_code=TrueでHubから読み込まれたモデルや、Transformersに存在しないモデルクラスが含まれます。カスタムコードのモデルは変換マッピングを継承しません。そのクラス名やmodel_typeがネイティブモデルと衝突する可能性があるためです。例えば、model_type="mixtral"を設定したカスタムアーキテクチャは、Mixtralの組み込みのエキスパート融合変換を引き継ぎません。

register_checkpoint_conversion_mapping を明示的に使用してカスタムコードモデルを登録し、変換マッピングを適用します。登録により、クラス名または model_type がユーザー登録マッピングのセットに追加され、カスタムコードのスキップが免除され、以下の通常のリファレンス処理にルーティングされます。

カスタムコードフィルターを通過する各サブモデルについて:

  1. クラス名による検索が最初に試行されます。次に model_type が試行されます。同じモジュールに対して両方が登録されている場合、クラス名のマッピングが優先されます。そのモジュールに対しては model_type の方が無視されます。これにより、タスクヘッド(例: LlavaForConditionalGeneration)が共有される model_type のベースライン("llava")を上書きできます。
  2. 選択されたマッピングには、scope_prefix がサブモジュールのドット区切りパス(ルートの場合は "")に設定されています。
  3. Ancestor-based deduplication マッピングを保持するかどうかを決定します:
    • 祖先パスがすでに同一の識別子(クラス名またはmodel_type)を要求している場合、サブモジュールはスキップされます。祖先の非スコープ付きまたは上位スコープのマッピングが、すでにこのサブツリーをカバーしているためです。
    • siblingのみがそれを要求している場合、サブモジュールは独自のscope_prefixを保持したまま維持されます。各siblingには、それぞれにスコープ付きのマッピングが割り当てられます。

クラス名とmodel_typeの既知リストは別々に追跡されますが、1つの注意点があります。クラス名によってモジュールが一致した場合、そのmodel_typeは既知リストに追加されません。これにより、同じmodel_typeを共有するがクラス固有のマッピングを持たない他のモジュール(例:DetrForSegmentation配下のDetrModel)が、model_typeの参照を通じて依然として到達可能になるようにするためです。

クラスベースのマッピングと model_type エイリアス

両方のスタイルが_MODEL_TO_CONVERSION_PATTERNに共存しています:

_MODEL_TO_CONVERSION_PATTERN = {
    # model_type aliases (lookup by config.model_type)
    "minimax": "mixtral",
    "qwen3_moe": "qwen2_moe",
    "mistral3": "llava",
    # class-name aliases (lookup by type(submodule).__name__)
    "PaliGemmaModel": "LlavaModel",
    "MaskFormerDetrDecoder": "DetrModel",
    ...
}

model_type が共有されているが、特定のクラスに異なる動作が必要な場合は、クラス名をキーとして使用することが推奨されます。

Tensor parallelism の統合

動的ロードシステムは、src/transformers/integrations/tensor_parallel.py で定義された TensorParallelLayer 階層を通じて、テンソル並列性 (TP) と統合されています。

TP が有効な場合、テンソルはマテリアライズ中にシャード化され、その後にシャード化されるわけではありません。つまり、各ランクは必要なテンソルの部分のみを読み込みます。

def spawn_tp_materialize(thread_pool, tensor, sharding_method, tensor_idx, device, dtype):
    def _job():
        return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
    return thread_pool.submit(_job)

利用可能な並列スタイル

スタイル重みシャード次元説明
colwise-2カラム方向: 出力特徴量がシャード化
rowwise-1行方向:入力特徴量がシャード化
packed_colwise-2融合された重み (gate_up_proj) の場合
packed_rowwise-1融合された重みについて
embedding_rowwise0語彙並列
grouped_gemm0MoE におけるエキスパート並列
sequence_parallelNone重みのシャーディングなし

パックされた重みの処理

gate_up_proj などの融合された重みについては、正しくシャードするために特別な注意が必要です:

def get_packed_weights(param, empty_param, device_mesh, rank, dim):
    """
    Interleaves gate and up shards correctly.

    Packed tensor: [G0 G1 G2 G3 | U0 U1 U2 U3]

    With TP=2:
    - Rank 0 gets: [G0 G1 | U0 U1]
    - Rank 1 gets: [G2 G3 | U2 U3]
    """

TP操作はWeightTransformに保存され、変換操作の後に適用されます:

if matched_tp_pattern := tp_plan_alt.search(renamed_key):
    tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]]
    mapping.distributed_operation = tp_layer(
        device_mesh=device_mesh,
        rank=device_mesh.get_local_rank(),
        empty_param=empty_param.clone()
    )

量子化の統合

Quantization は、チェックポイントがすでに量子化されているかどうかによって、ロードパイプラインに2つの方法で組み込まれます。

  • 事前量子化済みチェックポイント: 量子化器は、量子化済みテンソルをデシリアライズする WeightConverter インスタンス(get_weight_conversions() を通じて)を提供します。不要なキャストを避けるため、チェックポイントの dtype は保持されます。
  • オンザフライ量子化: 量子化器は、変換オペレーションの後に適用される量子化操作を提供し、重みがロードされる際に量子化を行います。

get_model_conversion_mapping の末尾で、update_weight_conversions(...) を使用して、クオンタイザーは変換リスト全体を書き換えることもできます。例えば、FP8 デクオンタイザーは、すべての既存のコンバーターに Fp8Dequantize オペレーションを先頭に追加することで、エキスパートマージ / 連結オペレーションがエキスパートごとの構造をフラットにする前に、ブロックごとのスケールが適用されるようにします。

高速かつ効率的なモデルの読み込み

モデルの読み込みは、ローダーが操作に必要なテンソルを把握し、そのマテリアライゼーションを遅延評価でスケジュールするため、より高速かつメモリ使用量が少ないです。

ローダーはチェックポイントを1回だけスキャンして、パターンマッチを発見し、テンソルを収集します。それらをFutureオブジェクトとして保存し、GILをブロックせずに非同期で読み込むためにスレッドプールに送信します。スレッドが利用可能になった時点で、パラメータの読み込みが開始されます。

システムで他の重いプロセスが実行されている場合、マルチスレッドはロードを高速化させるどころか、逆に遅くしてしまう可能性があります。この場合は、環境変数 HF_DEACTIVATE_ASYNC_LOAD=1 を設定して、重みをシーケンシャルにロードしてください。

デフォルトでは、非同期パラメータの読み込みに4スレッドを使用します。これにより、読み込みシナリオやハードウェアを問わず、最適なバランスが得られます。処理は主にI/Oに依存しますが、アクセラレータのハードウェアや読み込み時に必要なdtypeによっては、dtypeがシリアライズされたものとは異なる場合、CPU/GPUに依存する可能性があります(この場合、追加のコピー操作が必要です)。

非同期読み込みと同期読み込み

def spawn_materialize(thread_pool, tensor, device, dtype) -> Future | Callable:
    def _job():
        return _materialize_copy(tensor, device, dtype)

    if thread_pool is not None:
        return thread_pool.submit(_job)  # Async: returns Future
    else:
        return _job  # Sync: returns Callable (deferred execution)

同期読み込みは、次の場合に使用されます:

  • 環境変数 HF_DEACTIVATE_ASYNC_LOAD=1 が設定されています。
  • ディスクオフロードが有効になっています(メモリ制約により、シーケンシャルな読み込みが必要です)。
  • オンザフライ量子化が有効になっています(量子化ステップより前にワーカースレッドが先行することを回避します)。

マテリアライゼーションの流れ

1. Checkpoint iteration (Phase 1):
   - For each key, walk the transform list once
   - Submit materialization job to ThreadPoolExecutor
   - Job returns Future (async) or Callable (sync)
   - Collect into the matching WeightConverter / WeightRenaming

2. Per-mapping processing (Phase 2, one mapping at a time):
   - materialize_tensors() waits for this mapping's Futures only
   - Apply conversion operations chain (self.operations)
   - Apply quantization operation (if on-the-fly)
   - Set parameters on model
   - Delete realized tensors immediately

3. Cleanup:
   - Thread pool shutdown (with cancel_futures=True for interrupts)

メモリ効率

重みをコンバートする際、コンバーターは、まだ読み込まれていない必要なテンソルがすべてマテリアライズされるまで待機します。例えば、MergeModulelist操作では、マージする前にModuleList内のすべての重みが読み込まれている必要があります。

テンソルの連結には一時コピーが必要であるため、MergeModulelist や Concatenate のような操作は、変換中に元のテンソルの2倍のメモリを必要とします。マージが完了すると、メモリに残るのは結果のテンソルのみです。理論上の最悪ケースのメモリピークは、モデルサイズに、最大の MergeModulelist または Concatenate 操作に必要なテンソルを加えたものです。

この最悪ケースは、要求の高い変換が実行される前に他のすべてのパラメータが読み込まれた場合のみ発生します。このトリガーとなるシナリオは2つあります。

  1. 要求の厳しい変換処理に入る前に、すべてのパラメータが非同期で読み込まれます(スレッドプールが変換キューよりも高速でした)。
  2. 最も要求の厳しい変換が最後です。

例えば、各層のエキスパートにMergeModulelistを使用するMoEモデルの場合、理論上の最悪ケースのメモリピークは、モデルサイズに1層分のエキスパートを加えた値になります。

これらの最悪ケースは稀です。実際のメモリピークは、モデルサイズに近い値に留まる傾向があります。

可逆性

このシステムは、逆変換を伴うモデルの保存をサポートしており、ラウンドトリップでの保存/読み込みを可能にします。保存は2つのフェーズで実行されます。まず逆変換コンバーターが実行され(各テンソルは最大1つに一致します)、次に上記の順序ルールに従って逆リネームが実行されます。

def revert_weight_conversion(model, state_dict):
    """Applies reverse conversions for saving."""
    weight_conversions = getattr(model, "_weight_conversions", None)

    # Reverse all transforms
    reverse_weight_conversion = [
        conversion.reverse_transform() for conversion in weight_conversions
    ]

    # Apply in reverse
    for first_param_name, reversed_converter in conversion_mapping.items():
        realized_value = reversed_converter.convert(first_param_name, model=model)

ロード時に使用された変換のリストは、モデル上で_weight_conversionsとしてキャッシュされます(実際に発火したエントリのみが保持されるため、PrefixChangeのような非双射的な名前付け変更は、対称的に正しく再適用されます)。モデルがfrom_pretrainedなしでインスタンス化された場合(したがって_weight_conversionsを持っていない場合)、revert_weight_conversionはget_model_conversion_mappingを介してマッピングを再計算するフォールバックに切り替わり、そこからPrefixChangeを除外します(元のチェックポイントにプレフィックスがあったかどうかを判定できないためです)。

ターゲットパターンには、逆方向の処理のために処理が必要な正規表現の要素が含まれる場合があります:

def process_target_pattern(pattern: str) -> tuple[str, str | None]:
    """
    - Removes `^` and `$` anchors
    - Removes negative lookahead/lookbehind
    - Detects capturing groups, replaces with \1
    """

実際の例

MixtralスタイルのMoE

チェックポイント形式:

model.layers.0.block_sparse_moe.experts.0.w1.weight  # gate per expert
model.layers.0.block_sparse_moe.experts.0.w2.weight  # down per expert
model.layers.0.block_sparse_moe.experts.0.w3.weight  # up per expert
...
model.layers.0.block_sparse_moe.experts.7.w1.weight

モデル形式:

model.layers.0.mlp.experts.gate_up_proj  # (8, 4096, 28672)
model.layers.0.mlp.experts.down_proj     # (8, 14336, 4096)

変換マッピング(conversion_mapping.pyから):

"mixtral": [
    WeightRenaming(".block_sparse_moe.", ".mlp."),
    WeightConverter(
        source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
        target_patterns=".experts.gate_up_proj",
        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    ),
    WeightConverter(
        source_patterns=[".experts.*.w2.weight"],
        target_patterns=".experts.down_proj",
        operations=[MergeModulelist(dim=0)],
    ),
],

複合ビジョン言語モデル(Gemma3)

Gemma3 は、ベースモデルとヘッドモデルで異なるプレフィックスロジックを用いる典型的な構成になっています。2 つのエントリが登録されています(ここでは、共有の "llava" / "LlavaModel" リストを指すエイリアスを通じて):

# model_type: applies to Gemma3ForConditionalGeneration / Gemma3ForSequenceClassification
"gemma3": "llava"           # → adds "model." in front of language_model / vision_tower / ...

# class:    applies to the inner Gemma3Model only
"Gemma3Model": "LlavaModel"  # → minimal rename inside the already-prefixed namespace

Gemma3ForConditionalGenerationのDFS探索:

  1. Root — クラスの検索がミス(Gemma3ForConditionalGeneration は登録されていません)、model_type の検索がヒット("gemma3")→ "llava" のプレフィックス書き換え変換が スコープなしで追加され、language_model、vision_tower、multi_modal_projector が model.* の下に配置されます。
  2. 内部の model: Gemma3Model — クラス検索がヒット(Gemma3Model → LlavaModel)→ LlavaModel のマッピングのみが "model" にスコープを限定して適用されます。はるかに広い "gemma3"(="llava")プレフィックスのリネームはここでは再適用されません。これはまさに望ましい動作です。内部モデルは、プレフィックスがすでに正しい名前空間に存在しているためです。

LlavaファミリーのVLM(PaliGemma、InternVL、Mistral3、…)はすべて同じペア形状を使用しています。model_typeエントリはヘッドモデルのプレフィックス処理を担い、クラスエントリは内部のベースモデルのマッピングを最小限に保ちます。

ヘッド固有のオーバーライドによるより深いネスト(DETR)

DetrForSegmentation は、2つのネストされたレベルの上に、ヘッド固有のオーバーライドが重ねられた構造を示しています。

DetrForSegmentation                  (class registered: segmentation-only renames)
├── detr: DetrForObjectDetection     (no mapping; just walked through)
│   └── model: DetrModel             (class registered: shared base transforms)
│       └── backbone, encoder, ...
└── mask_head, bbox_attention        (head-specific weights)

2つのマッピングが関与します(登録されたキーごとに1つ):

"DetrModel":           [WeightRenaming("backbone.conv_encoder", "backbone"), ...]      # shared base
"DetrForSegmentation": [WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), ...] # head-specific

DFS 探索:

  1. ルートであるDetrForSegmentationはクラスで一致し、セグメンテーション用に追加されたunscopedが名前変更されます。
  2. detr: DetrForObjectDetection は登録されていません → トランスフォームは追加されません。DFS はその内部に続行します。
  3. detr.model: DetrModel はクラスで一致 → scope_prefix="detr.model" を付けてベースのトランスフォームが追加されます。

そのため、DetrForObjectDetectionには独自のマッピングは不要です。そのサブツリー(DetrModel)に登録されている唯一のマッピングは、自動的に正しいパスにスコープが設定されるためです。

クラスキーのエイリアスは、追加の登録なしでベースのマッピングを再利用します。"MaskFormerDetrDecoder": "DetrModel" を設定すると、MaskFormer のデコーダーが、自身のクラス名で同じトランスフォームを適用するようになります。

カスタムオペレーション(ERNIE 4.5 VL MoE)

組み込みの操作が不十分な場合は、ConversionOps のサブクラスをカスタムで作成できます。例えば、ERNIE 4.5 VL MoE では、共有エキスパートのリストをテキストとビジョンのモダリティ間で分割する必要がありますが、これは単一の組み込み操作では処理できません。カスタム操作 ErnieFuseAndSplitTextVisionExperts は、2つのターゲットキー間でエキスパートを分割し、再スタックします。

"ernie4_5_vl_moe": [
    WeightRenaming("vision_model", "vision_tower"),
    WeightConverter(
        source_patterns=["experts.*.down_proj.weight"],
        target_patterns=[
            "text_moe.experts.down_proj",
            "vision_moe.experts.down_proj",
        ],
        operations=[ErnieFuseAndSplitTextVisionExperts(stack_dim=0, concat_dim=1)],
    ),
],

カスタムオペレーションは、ラウンドトリップの保存/読み込みをサポートするために、convert() と reverse_op プロパティを実装する必要があります。

モデルタイプの別名

多くのモデルが変換パターンを共有しています:

_MODEL_TO_CONVERSION_PATTERN = {
    "mixtral": "mixtral",
    "minimax": "mixtral",
    "qwen2_moe": "qwen2_moe",
    "deepseek_v2": "qwen2_moe",
    "deepseek_v3": "qwen2_moe",
    "qwen3_moe": "qwen2_moe",
    "olmoe": "qwen2_moe",
    ...
}

動的読み込みのビルディングブロックの再利用

動的ウェイトの読み込みは、フルモデルのチェックポイントに限定されません。チェックポイントのキーがパラメータにどのようにマッピングされるかを記述し、ターゲットのモジュールが存在することを保証できる限り、同じ構成要素を使って任意のウェイトのセットを読み込むことができます。

大まかに言うと、この契約は次のような形になります。

  1. モデルのネームスペースを準備します。 読み込みたいモジュールやパラメータが存在し、マッピングで対象とする名前になっていることを確認してください。アダプターの場合、読み込み前にアダプターモジュールが存在するようにinject_adapter_in_model(...)を呼び出します。カスタムヘッドや追加モジュールの場合は、まずモデル上でインスタンス化してください。
  2. Describe how to map weights. 変換/名前変更リストを構築します(例えば、 _build_peft_weight_mapping(...)) を使用する WeightConverter または WeightRenaming. チェックポイントのキーを、モデルのネームスペースに合わせて変換、分割、マージ、または名前変更する方法をここで定義します。主に以下の3つのことができます:
    • コンバーターのリストに操作を追加します。これらは、WeightConverter に収集された重み以外のすべての重みに適用されます。一般的には、WeightRenaming 操作であるべきです。
    • 各コンバーターの操作リストに操作を追加します。Quantizationでは、任意のWeightConverterの操作リストの後に量子化操作を追加するだけです。
    • カスタムな操作への置き換え / マッピング: peftではこのように動作します。例えば mixtral の Concatenate 操作を PeftConcatenate (PEFT で定義されています) に置き換えます。これにより、アダプターチェックポイントを読み込む際に、連結対象の重みが収集され、peft 用に適切に整形されます。
  3. Load + finalize + report. コアローダーを使用して変換を実行し、テンソルを設定してから、最終処理と結果のログ出力を行います。具体的には、このフローは以下のようになります:
    • LoadStateDictConfig(...) と _load_pretrained_model(...) を使用して、読み込みと変換を行います。
    • _finalize_load_state_dict(...) を呼び出し、meta に残っている欠落または不一致のテンソルを移動して初期化し、重みを束ねます。
    • log_state_dict_report(...) を使用して、欠落キー/予期しないキー/不一致キー(および変換エラー)を報告します。

これらの API は、カスタムコードやカスタム重み形式の処理を可能にするだけでなく、transformers API の最高かつ最も効率的な重みの読み込み、シャーディング、そして優れた使いやすさの恩恵を受けることを保証するために公開されています。

主要ファイルの参照

ファイル目的
src/transformers/core_model_loading.pyコアの読み込みロジック、WeightConverter、WeightRenaming、ConversionOps
src/transformers/conversion_mapping.py組み込みのマッピングとサブモデルごとの合成(get_model_conversion_mapping)
src/transformers/integrations/tensor_parallel.pyTPシャードリングのクラスとユーティリティ
src/transformers/quantizers/base.py量子化フックとベースクラス
Update on GitHub