モデルの読み込み

Transformers は、1 行のコードで利用可能な多くの事前学習済みモデルを提供しています。これには、モデルクラスと from_pretrained() メソッドが必要です。

Hugging Face Hub に保存されているモデルの重みと設定をダウンロードして読み込むには、from_pretrained() を呼び出します。

The from_pretrained() method loads weights stored in the safetensors file format if they’re available. Traditionally, PyTorch model weights are serialized with the pickle utility which is known to be insecure. Safetensor files are more secure and faster to load.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

このガイドでは、モデルの読み込み方法、モデルを読み込むためのさまざまな方法、非常に大きなモデルにおけるメモリ問題の解決策、およびカスタムモデルの読み込み方法について説明します。

モデルと設定

すべてのモデルには、隠れ層数、語彙サイズ、活性化関数などの特定の属性を持つconfiguration.pyファイルがあります。また、各層内で実行される層と数学的演算を定義するmodeling.pyファイルも存在します。modeling.pyファイルは、configuration.py内のモデル属性を受け取り、それに基づいてモデルを構築します。この時点で、あなたはランダムな重みを持つモデルを手にしていますが、意味のある結果を出力するには訓練が必要です。

アーキテクチャとはモデルの骨格を指し、チェックポイントとは特定のアーキテクチャに対するモデルの重みを指します。例えば、BERTはアーキテクチャであり、google-bert/bert-base-uncasedはチェックポイントです。モデルという用語は、アーキテクチャとチェックポイントを指すために互換的に使用されます。

読み込めるモデルには、主に2つのタイプがあります。

AutoModelやLlamaModelのように、隠れ状態を出力する最小構成のモデル。
head が付与されたモデル。例えば AutoModelForCausalLM や LlamaForCausalLM など、特定のタスクを実行するために使用されます。

モデルクラス

事前学習済みモデルを取得するには、重みをモデルに読み込む必要があります。これは、Hugging Face Hubまたはローカルディレクトリからの重みを引数として受け取るfrom_pretrained()を呼び出すことで行います。

モデルクラスには2種類あり、AutoModelクラスと、モデル固有のクラスがあります。

AutoModel

モデル固有のクラス

大規模モデル

大規模な事前学習済みモデルを読み込むには、大量のメモリが必要です。読み込みプロセスには以下が含まれます。

ランダムな重みでモデルを作成する
事前学習済み重みの読み込み
モデルに事前学習済み重みを配置すること

モデルの重み（ランダム初期化と事前学習済み）の2コピーを保持するのに十分なメモリが必要ですが、ハードウェアによってはこれが不可能な場合があります。分散学習環境では、各プロセスが事前学習済みモデルを読み込むため、さらに困難になります。

Transformers は、高速な初期化、シャード化されたチェックポイント、Accelerate の Big Model Inference 機能、および低ビット幅データ型へのサポートにより、これらのメモリ関連の課題を軽減します。

シャード化されたチェックポイント

save_pretrained() は、50GB を超えるチェックポイントを自動的にシャード化します。これにより、大規模モデルにおけるシャード数を抑え、ファイル管理を簡素化できます。

パラメータは並列で読み込まれ、ピークメモリはモデルサイズにのみ依存します。シャーディングを行う前のチェックポイントの最大サイズを設定するには、save_pretrained() で max_shard_size を使用します。

動的ウェイト変換を必要とするモデルのメモリ使用量は、モデルサイズと、1回の変換における最大のパラメータサイズに依存します。これは通常、mixture-of-experts (MoE) モデルに当てはまり、メモリ使用量はモデルサイズに1レイヤーあたりのエキスパート数を加えたものになります。モデルの読み込み方法の詳細については、動的ウェイトローダーのガイドを参照してください。

save_pretrained() は、パラメータ名をそれぞれのシャードファイルにマッピングするインデックスファイルも作成します。このインデックスには、metadata と weight_map という2つのキーが含まれます。

import json
import os
import tempfile

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="50GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

metadataには、モデルの総サイズが保存されます。

index["metadata"]
{'total_size': 28966928384}

weight_map は、各パラメータを対応するシャードファイルにマッピングします。

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

大規模モデルの推論

この機能を使用するには、Accelerate v0.9.0 と PyTorch v1.9.0 以降がインストールされていることを確認してください！

from_pretrained() is supercharged with Accelerate’s Big Model Inference feature.

Big Model Inference は、PyTorch の meta デバイス上に モデルのスケルトン を作成します。meta デバイスは実際のデータを保存せず、メタデータのみを保持します。

事前学習済み重みの読み込み時にのみ、ランダム初期化された重みが作成されます。これにより、メモリ上にモデルのコピーを2つ同時に保持することを回避できます。最大メモリ使用量は、モデルのサイズのみです。

デバイスマップの設計で、デバイスの配置について詳しく学んでください。

Big Model Inferenceの2つ目の機能は、モデルのスケルトンにおける重みの読み込みとディスパッチの仕組みに関係しています。モデルの重みは、利用可能なすべてのデバイスにディスパッチされます。まず最速のデバイス（通常はGPU）から開始し、残りの重みはより低速なデバイス（CPUやハードドライブ）へオフロードされます。

これらの2つの機能を組み合わせることで、大規模な事前学習済みモデルのメモリ使用量と読み込み時間を削減できます。

Big Model Inference を有効にするには、device_map を "auto" に設定してください。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

device_map でレイヤーを手動でデバイスに割り当てることもできます。モデルのパラメータ全体をデバイスにマッピングする必要がありますが、レイヤー全体が同じデバイスにある場合、そのレイヤーのすべてのサブモジュールの配置を詳細に指定する必要はありません。

hf_device_map 属性にアクセスして、モデルがデバイス間でどのように分散されているかを確認できます。

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

ディスクオフロード

モデルが GPU と CPU RAM の合計容量を超えて大きい場合、Big Model Inference は残りの重みをディスクにオフロードします。"disk" にマッピングされる重みが存在する場合は、offload_folder をディレクトリに設定してください。Transformers はオフロードされた重みをこのディレクトリに保存するためです。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    offload_folder="./offload",
)

max_memory を使用して、各デバイスが保持する容量の上限を設定します。Transformers はデバイスを速度の順に埋め、収まらない分はディスクに送信します。以下の例では、GPU を 10GB、CPU を 30GB に制限しているため、余った重みは offload_folder に保持されます。

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
)

max_memory={} を設定すると、GPU メモリおよび CPU メモリが利用できないことを指定し、モデル全体をディスクにオフロードするよう強制できます。

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={},
    offload_folder="./offload",
)

ディスクオフロードは、速度を犠牲にしてメモリを節約する手法です。フォワードパスのたびにディスクから重みを読み込むのは、CPUメモリやGPUメモリから読み込むよりも遅くなります。フルディスクオフロードにはaccelerate>1.14.0が必要です。

モデルのデータ型

dtype 引数は、モデルの重みを初期化するために使用される PyTorch の dtype を制御します。デフォルトでは、Transformers は config.json から dtype またはレガシーの torch_dtype の値を読み取って重みをロードします。config.json にこれらの値のいずれも含まれていない場合、Transformers はチェックポイント内の最初の浮動小数点重みの dtype を使用します。

特定のdtypeを渡すことで、デフォルトをオーバーライドできます。

import torch
from transformers import AutoModelForCausalLM

# specific dtype
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype=torch.float16)

AutoConfig は、ゼロからインスタンス化されるモデルに対して dtype を受け取ります。

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", dtype=torch.float16)
model = AutoModel.from_config(my_config)

カスタムモデル

カスタムモデルは、Transformers の設定クラスとモデリングクラスを基盤として構築され、AutoClass API をサポートし、from_pretrained() で読み込まれます。違いは、モデリングコードが Transformers 由来でない点です。

カスタムモデルを読み込む際は、特に注意を払ってください。Hubではすべてのリポジトリに対してマルウェアスキャンが実施されていますが、悪意のあるコードを誤って実行しないよう、引き続き注意が必要です。

カスタムモデルを読み込むには、from_pretrained() で trust_remote_code=True を設定してください。

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

セキュリティをさらに強化するため、変更された可能性のあるモデルコードを読み込まないよう、特定のrevisionからカスタムモデルを読み込みます。commitハッシュは、モデルのcommit履歴からコピーできます。

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

詳細については、モデルのカスタマイズガイドを参照してください。

Update on GitHub