🤗 Accelerateは、PyTorchのおかげで非常に大規模なモデルを実行する方法です

'🤗 Accelerateは、PyTorchによって大規模なモデルを実行する方法です'

大規模モデルの読み込みと実行

Meta AIとBigScienceは最近、ほとんどの一般的なハードウェアのメモリ（RAMまたはGPU）に収まらない非常に大きな言語モデルをオープンソース化しました。Hugging Faceでは、私たちの使命の一部として、それらの大きなモデルにアクセスできるようにするためのツールを開発しました。そのため、スーパーコンピュータを所有していなくても、これらのモデルを実行できるようにするためのツールを開発しました。このブログ投稿で選ばれたすべての例は、無料のColabインスタンス（制限付きのRAMとディスク容量）で実行されます。ディスク容量に余裕がある場合は、より大きなチェックポイントを選択することもできます。

ここでは、OPT-6.7Bを実行する方法を示します:

import torch
from transformers import pipeline

# これは基本的なColabインスタンスで動作します。
# もし時間がかかっても待つ時間と十分なディスク容量がある場合は、より大きなチェックポイントを選択してください！
checkpoint = "facebook/opt-6.7b"
generator = pipeline("text-generation", model=checkpoint, device_map="auto", torch_dtype=torch.float16)

# 推論を実行します
generator("More and more large language models are opensourced so Hugging Face has")

これらの引数がそれぞれ何を意味するのかについては、まもなく説明しますが、まずはPyTorchでの伝統的なモデルの読み込みパイプラインを考えてみましょう：通常は次のような手順で行われます：

1. モデルの作成
2. その重みをメモリに読み込む（通常はstate_dictと呼ばれるオブジェクトに含まれます）
3. 読み込んだ重みを作成したモデルにロードする
4. 推論用にモデルをデバイスに移動する

これまでこれでうまくいってきましたが、非常に大きなモデルではこのアプローチが難しくなります。ここで選んだモデルは67億のパラメータを持っています。デフォルトの精度では、たった1ステップ（モデルの作成）でおよそ26.8GBのRAMが必要です（float32のパラメータはメモリ上で4バイトを占有します）。これはColabのRAMにも収まりません。

• AutoTrainによる画像分類
• 日本語安定拡散
• 最適化ストーリー：ブルーム推論

次に、ステップ2ではモデルのメモリにもう1つのコピーを読み込みます（デフォルトの精度ではさらに26.8GBのRAMが必要です）。もし最大のモデル、例えばBLOOMまたはOPT-176B（いずれも1760億のパラメータを持つ）をこのように読み込もうとする場合、1.4テラバイトのCPU RAMが必要になります。これはやや過剰です！そして、全ての重みを1つのGPU（または複数のGPU）に移動するためだけにこれらすべてが行われます。

明らかに、よりスマートな方法が必要です。このブログ投稿では、AccelerateがPyTorchの機能を活用して非常に大きなモデルをロードして推論を実行する方法について説明します。これにより、メモリまたは1つのGPUに収まらない場合でも、上記のプロセスが次のように変更されます：

1. 空の（重みのない）モデルを作成する
2. 各レイヤーをどこに配置するかを決定する（複数のデバイスが利用可能な場合）
3. 重みの一部をメモリに読み込む
4. 空のモデルに重みをロードする
5. 重みを推論用のデバイスに移動する
6. すべての重みが読み込まれるまでステップ3から繰り返す

空のモデルの作成

PyTorch 1.9では、メタデバイスという新しいデバイスが導入されました。これにより、データが関連付けられていないテンソルを作成することができます。メタデバイス上では、形状があれば、CPU（またはGPU）のRAMの心配をする必要なく、任意の大きなテンソルを作成することができます。

たとえば、次のコードはColabでクラッシュします：

import torch

large_tensor = torch.randn(100000, 100000)

この大きなテンソルは4 * 10**10バイト（デフォルトの精度はFP32なので、テンソルの各要素は4バイトを占有します）つまり40GBのRAMが必要です。一方、メタデバイス上では問題ありません：

import torch

large_tensor = torch.randn(100000, 100000, device="meta")

このテンソルを表示しようとすると、PyTorchは次のように表示します：

tensor(..., device='meta', size=(100000, 100000))

前述の通り、このテンソルにはデータは関連付けられておらず、形状のみが存在します。

メタデバイス上でモデルを直接インスタンス化することもできます：

large_model = torch.nn.Linear(100000, 100000, device="meta")

ただし、既存のモデルでは、この構文により、各サブモジュールがdeviceキーワード引数を受け入れて渡すようにモデリングコード全体を書き直す必要があります。Transformersライブラリの150のモデルにはこれが実用的ではなかったため、空のモデルを自動的に生成するためのコンテキストマネージャを開発しました。

以下は、BLOOMの空のバージョンをインスタンス化する方法です:

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("bigscience/bloom")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

これはどのモデルでも機能しますが、直接使用できないシェルが返されます。一部の操作はメタデバイスで実装されていますが、すべての操作はまだ実装されていません。たとえば、上記で定義したlarge_modelは入力と共に使用できますが、BLOOMモデルは使用できません。使用しても、出力はメタデバイスのテンソルとなり、結果の形状は取得できますが、それ以上の情報は得られません。

さらなる作業として、PyTorchチームは新しいFakeTensorクラス上で作業しています。これは、メタデバイス上のテンソルのようなものですが、デバイス情報（形状とdtypeに加えて）も持っています。

各重みの形状を知っているため、事前学習済みのテンソルを完全にロードした場合にそれらがどれだけのメモリを消費するかを知ることができます。そのため、モデルをCPUとGPUに分割する方法についての決定を下すことができます。

デバイスマップの計算

事前学習済みの重みをロードする前に、それらを配置する場所を知る必要があります。これにより、重みを正しい場所に配置するたびにCPUのRAMを解放することができます。これは、メモリ内でどれだけのスペースを占有するかを計算するため、空のモデル上のメタデバイスで実行できます。

Accelerateは、空のモデルから自動的にデバイスマップを決定するための関数を提供しています。これにより、利用可能なすべてのGPUの使用を最大化し、次にCPUのRAMを使用し、最後にディスクオフロードに適合しない重みを示します。OPT-13bを使用して詳細を見てみましょう。

from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

device_map = infer_auto_device_map(model)

これにより、モジュールまたは重みをデバイスにマッピングする辞書が返されます。たとえば、Titan RTXが1つ搭載されたマシンでは、次のようになります:

{'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10.self_attn': 0,
 'model.decoder.layers.10.activation_fn': 0,
 'model.decoder.layers.10.self_attn_layer_norm': 0,
 'model.decoder.layers.10.fc1': 'cpu',
 'model.decoder.layers.10.fc2': 'cpu',
 'model.decoder.layers.10.final_layer_norm': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18.self_attn': 'cpu',
 'model.decoder.layers.18.activation_fn': 'cpu',
 'model.decoder.layers.18.self_attn_layer_norm': 'cpu',
 'model.decoder.layers.18.fc1': 'disk',
 'model.decoder.layers.18.fc2': 'disk',
 'model.decoder.layers.18.final_layer_norm': 'disk',
 'model.decoder.layers.19': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

Accelerateは、埋め込みおよびデコーダーの9番目のブロックまでがすべてGPU（デバイス0）に収まると評価し、10番目のブロックの一部はCPUに配置する必要があります。また、17番目のレイヤーまでの重みもCPUに配置する必要があります。次に、18番目のレイヤーはCPUとディスクの両方に分割され、その後のレイヤーはすべてディスクにオフロードする必要があります。

ただし、このデバイスマップを後で使用するとエラーが発生します。なぜなら、このモデルを構成するレイヤーには残差接続（ブロックの入力がブロックの出力に追加される）があるため、特定のレイヤーのすべての要素は同じデバイス上にある必要があるからです。これをAccelerateに伝えるために、no_split_module_classesキーワード引数で分割しないモジュールのリストを渡すことができます:

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])

これにより、次の結果が返されます。

'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

各レイヤーは常に同じデバイス上にあります。

Transformersでは、from_pretrained()メソッドまたはpipelineでdevice_mapを使用する場合、同じデバイスに残すブロックのクラスは自動的に提供されるため、心配する必要はありません。 device_mapには次のオプションがあります（複数のGPUがある場合にのみ関連します）：

• "auto"または"balanced"：Accelerateは重みを均等に分割して各GPUを均等に使用します。
• "balanced_low_0"：Accelerateは重みを均等に分割し、最初のGPUには可能な限り少ない重みが含まれるようにします（generate関数を使用してモデルの出力で作業する場合などに便利です）。
• "sequential"：AccelerateはGPUを順番に埋めます（最後のGPUは使用されない場合があります）。

また、辞書形式のdevice_mapを自分で渡すこともできます（レイヤー/モジュール名からデバイスへのマッピング）。

最後に、受け取るdevice_mapの結果は、選択したdtypeに依存することに注意してください（異なる浮動小数点数の型は異なるスペースを使用します）。 dtype="float16"を指定すると、異なる結果が得られます：

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")

この精度では、モデルをレイヤー21までGPUに収めることができます：

{'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.21': 0,
 'model.decoder.layers.22': 'cpu',
 ...
 'model.decoder.layers.37': 'cpu',
 'model.decoder.layers.38': 'disk',
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

各重みがどこに配置されるべきかを知ったので、モデル内に事前学習済みの重みを逐次的に読み込むことができます。

シャーディングされた状態辞書

従来、PyTorchモデルは、パラメータ名から重みへのマップを含む1つのファイルに保存されます。このマップは通常state_dictと呼ばれます。以下は、PyTorchの保存と読み込みに関するドキュメントの抜粋です：

# モデルの重みを保存する
torch.save(my_model.state_dict(), 'model_weights.pth')

# それらを再度ロードする
new_model = ModelClass()
new_model.load_state_dict(torch.load('model_weights.pth'))

これは、10億パラメータ以下のモデルには非常に適していますが、より大きなモデルでは、これはRAMに非常に負荷がかかります。BLOOMモデルには1760億のパラメータがあります。スペースを節約するためにbfloat16で重みが保存されていても、全体として352GBを表します。このモデルを訓練したスーパーコンピュータは、この量のメモリを利用できるかもしれませんが、推論にこれを必要とすることは現実的ではありません。

これが、Hugging Face Hubの大規模モデルが1つの大きなファイルで保存されず、複数のファイルで共有される理由です。たとえば、BLOOMモデルのページに移動すると、pytorch_model_xxxxx-of-00072.binという72個のファイルがあることがわかります。各ファイルにはモデルの一部の重みが含まれています。この形式を使用すると、メモリに1つのシャード（部分）の状態辞書を読み込み、重みをモデルに入れて、適切なデバイスに移動してからこの状態辞書の一部を破棄し、次のシャードに進むことができます。モデル全体を収容するために十分なRAMが必要とされるのではなく、最大のチェックポイントパートを取得するために十分なRAMが必要です。これをシャードと呼びます。たとえば、BLOOMの場合は7.19GBです。

私たちは、BLOOMの分割されたチェックポイントと呼ばれる複数のファイルに保存されたチェックポイントを呼びます。そして、それらの形式を次のように標準化しています：

• 1つのファイル（pytorch_model.bin.index.jsonと呼ばれる）には、メタデータとパラメータ名からファイル名へのマップが含まれており、各重みがどこにあるかを示しています。
• 他のすべてのファイルは、標準的なPyTorchのステート辞書であり、全体ではなくモデルの一部を含んでいます。インデックスファイルの内容は、こちらで確認できます。

モデルにこのような分割されたチェックポイントをロードするには、さまざまなシャードをループで処理するだけです。Accelerateは、Hubのリポジトリをクローンしている場合にこれを行うためのload_checkpoint_in_modelという関数を提供しています。または、Transformersのfrom_pretrainedメソッドを直接使用することもできます。このメソッドは、ダウンロードとキャッシュを処理します：

import torch
from transformers import AutoModelForCausalLM

# エラーが発生します
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16)

自動的に計算されたデバイスマップにより、GPUとCPUのRAMが十分でないために一部の重みをディスクにオフロードする必要がある場合、以下のエラーが表示されます：

ValueError: The current `device_map` had weights offloaded to the disk. Please provide an 
`offload_folder` for them.

このエラーを解決するために、次の引数を追加してください：

import torch
from transformers import AutoModelForCausalLM

# Colab上でRAMが不足します
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)

注意してください。CPUのオフロードに加えてディスクのオフロードも必要な非常に大きなモデルをロードしようとしている場合、チェックポイントの最後のシャードがロードされる際にRAM不足になる可能性があります。なぜなら、CPU上に残っているモデルの一部がスペースを取るからです。その場合は、オプションのoffload_state_dict=Trueを使用して、重みがすべて処理された後にCPU上にあるモデルの一部を一時的にオフロードし、RAMに再度ロードしてください。

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)

これでColabに収まるようになりますが、予測を生成しようとすると、使用可能なすべてのRAMをほぼ使い切ってしまい、RAM不足になります。使用可能なモデルを得るためには、さらに1つのレイヤーをディスク上にオフロードする必要があります。前のセクションで計算されたdevice_mapを少し改変して、from_pretrained呼び出しに渡すことで、これを行うことができます：

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
device_map["model.decoder.layers.37"] = "disk"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map=device_map, offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)

複数のデバイスで分割されたモデルの実行

最後に触れていない最後の部分は、Accelerateがモデルを複数のGPU、CPU RAM、およびディスクフォルダに分散して実行する方法です。これは非常にシンプルに、フックを使用して実現されています。

フックは、各forward呼び出しの直前に実行される関数を追加するPyTorchのAPIです

直接使用することはできませんが、これと同じ考え方を取り入れています。モデルがロードされると、dispatch_model関数は、各モジュールとサブモジュールにフックを追加し、各forwardパスの前後に実行されます。これらのフックは以下の処理を行います：

• モジュールのすべての入力が重みと同じデバイスにあることを確認します。
• 重みがCPUにオフロードされている場合、forwardパスの前にそれらをGPU 0に移動し、その後すぐにCPUに戻します。
• 重みがディスクにオフロードされている場合、RAMにロードしてからforwardパスの前にGPU 0に移動し、そのメモリを解放します。

以下のビデオで、全体のプロセスを要約しています：

この方法では、GPU RAMとCPU RAMが十分でなくてもモデルを読み込んで実行することができます。必要なのはディスク容量（そしてたくさんの忍耐力！）だけです。この解決策は、複数のGPUを持っている場合にはかなり単純です（クレバーなパイプライン並列処理はなく、単純にGPUを順次使用するだけです）。それでも、BLOOMにはかなり良い結果をもたらし、より小規模なセットアップでもモデルを実行することができます（ただし、より遅くなります）。

大規模モデルの推論の高速化について詳しくは、ドキュメントを参照してください。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!