大規模モデル推論

Accelerate が提供する最大の機能の一つがBig Model Inferenceです。これにより、グラフィックカードに完全に収まらないモデルでも推論を実行できます。

このチュートリアルでは、AccelerateとHugging FaceエコシステムでBig Model Inferenceを使用する方法を示します。

Accelerate

PyTorchモデルを読み込むための典型的なワークフローを以下に示します。ModelClassは、お使いのデバイス（mps、cuda、またはxpu）のGPUメモリを超えるモデルです。

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

Big Model Inferenceでは、最初のステップとしてinit_empty_weightsコンテキストマネージャーを使用してモデルの空のスケルトン（骨格）を初期化します。これはメモリを必要としません。なぜならmy_modelは「パラメータなし」だからです。

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

次に、推論用のモデルに重みが読み込まれます。

load_checkpoint_and_dispatch() メソッドは、空のモデル内にチェックポイントをロードし、各レイヤーの重みを利用可能なすべてのデバイスに分散します。この際、より高速なデバイス（GPU、MPS、XPU、NPU、MLU、SDAA、MUSA）から優先的に配置し、その後、より低速なデバイス（CPUおよびハードドライブ）に配置します。

device_map="auto" を設定すると、まず GPU の利用可能な空き領域すべてに、次に CPU に、それでもメモリが不足している場合は最後にハードドライブ（最も遅い選択肢）に自動的に配置されます。

独自のデバイスマップの設計方法の詳細については、デバイスマップの設計ガイドを参照してください。

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

特定のレイヤーの「チャンク」を分割すべきでない場合は、no_split_module_classesに渡してください（詳細はこちらを参照）。

モデルの重みは、state_dictがメモリに収まらない場合など、メモリを節約するために複数のチェックポイントに分割（シャーディング）することもできます（詳細はこちらを参照）。

モデルが完全にディスパッチされたので、推論を実行できます。

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

入力データが各層を通過するたびに、CPUからGPU（またはディスクからCPU、さらにGPU）に送信され、出力が計算された後、その層はGPUから解放されて処理ラインの後ろへ戻ります。これにより推論にオーバーヘッドが加わりますが、最大の層がGPUに収まる限り、システム上で任意のサイズのモデルを実行できるようになります。

複数のGPU、いわゆる「モデル並列性」を利用することは可能ですが、同時にアクティブなGPUは常に1つだけです。これにより、GPUは前のGPUから出力が送信されるまで待機することになります。torchrunやaccelerate launchなどのツールではなく、Pythonでスクリプトを通常通り起動してください。

以下に、Big Model Inferenceの完全な例を示します。

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Hugging Faceエコシステム

Transformers や Diffusers など、Hugging Face エコシステム内の他のライブラリも、from_pretrained コンストラクタで Big Model Inference をサポートしています。

Big Model Inferenceを有効にするには、from_pretrainedにdevice_map="auto"を追加するだけです。

例えば、Big Model Inferenceを使用して、Big Sciencesの110億パラメータモデルであるT0ppを読み込みます。

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

モデルを読み込んだ後、前述の空のinitとスマートディスパッチのステップが実行され、モデルはマシンのすべてのリソースを活用する準備が完全に整います。これらのコンストラクタを通じて、torch_dtypeパラメータを指定してモデルを低精度で読み込むことで、メモリをさらに節約することもできます。

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)

次のステップ

Big Model Inference のより詳細な説明については、コンセプトガイドをぜひご覧ください。