Instantiating a big model

非常に大規模な事前学習済みモデルを使用する場合、RAMの使用量を最小限に抑えることは課題の1つです。通常のPyTorchのワークフローは次のとおりです：

1. ランダムな重みを持つモデルを作成します。
2. 事前学習済みの重みをロードします。
3. これらの事前学習済みの重みをランダムなモデルに配置します。

ステップ1と2の両方がメモリにモデルの完全なバージョンを必要とし、ほとんどの場合は問題ありませんが、モデルのサイズが数ギガバイトになると、これらの2つのコピーをRAMから排除することができなくなる可能性があります。さらに悪いことに、分散トレーニングを実行するためにtorch.distributedを使用している場合、各プロセスは事前学習済みモデルをロードし、これらの2つのコピーをRAMに保存します。

このガイドでは、Transformersがこの問題に対処するために提供するソリューションを探ります。なお、これは現在も開発が進行中の分野であり、将来、ここで説明されているAPIがわずかに変更される可能性があることに注意してください。

Sharded checkpoints

バージョン4.18.0から、10GBを超えるサイズのモデルチェックポイントは自動的に複数の小さな部分に分割されます。model.save_pretrained(save_dir)を実行する際に1つの単一のチェックポイントを持つ代わりに、いくつかの部分的なチェックポイント（それぞれのサイズが<10GB）と、パラメータ名をそれらが格納されているファイルにマップするインデックスが生成されます。

max_shard_sizeパラメータでシャーディング前の最大サイズを制御できるため、例として通常サイズのモデルと小さなシャードサイズを使用します。従来のBERTモデルを使用してみましょう。

from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-cased")

もしsave_pretrained()を使用して保存する場合、新しいフォルダが2つのファイルを含む形で作成されます: モデルの設定情報とその重み情報です。

>>> import os
>>> import tempfile

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir)
...     print(sorted(os.listdir(tmp_dir)))
['config.json', 'pytorch_model.bin']

最大シャードサイズを200MBに設定します：

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     print(sorted(os.listdir(tmp_dir)))
['config.json', 'pytorch_model-00001-of-00003.bin', 'pytorch_model-00002-of-00003.bin', 'pytorch_model-00003-of-00003.bin', 'pytorch_model.bin.index.json']

モデルの設定の上に、3つの異なる重みファイルと、index.jsonファイルが見られます。これは私たちのインデックスです。 このようなチェックポイントは、from_pretrained()メソッドを使用して完全に再ロードできます：

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     new_model = AutoModel.from_pretrained(tmp_dir)

主要な利点は、大規模なモデルの場合、上記のワークフローのステップ2において、各チェックポイントのシャードが前のシャードの後にロードされ、RAMのメモリ使用量をモデルのサイズと最大のシャードのサイズを合わせたものに制限できることです。

内部では、インデックスファイルが使用され、どのキーがチェックポイントに存在し、対応する重みがどこに格納されているかを判断します。このインデックスは通常のJSONファイルのように読み込むことができ、辞書として取得できます。

>>> import json

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     with open(os.path.join(tmp_dir, "pytorch_model.bin.index.json"), "r") as f:
...         index = json.load(f)

>>> print(index.keys())
dict_keys(['metadata', 'weight_map'])

メタデータには現時点ではモデルの総サイズのみが含まれています。 将来的には他の情報を追加する予定です：

>>> index["metadata"]
{'total_size': 433245184}

重みマップはこのインデックスの主要な部分であり、各パラメータ名（通常はPyTorchモデルのstate_dictで見つかるもの）をその格納されているファイルにマップします：

>>> index["weight_map"]
{'embeddings.LayerNorm.bias': 'pytorch_model-00001-of-00003.bin',
 'embeddings.LayerNorm.weight': 'pytorch_model-00001-of-00003.bin',
 ...

直接モデル内でfrom_pretrained()を使用せずに、 シャーディングされたチェックポイントをロードしたい場合（フルチェックポイントの場合にmodel.load_state_dict()を使用するように行う方法）、load_sharded_checkpoint()を使用する必要があります：

>>> from transformers.modeling_utils import load_sharded_checkpoint

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     load_sharded_checkpoint(model, tmp_dir)

Low memory loading

シャードされたチェックポイントは、上記のワークフローのステップ2におけるメモリ使用量を削減しますが、 低メモリの環境でそのモデルを使用するために、Accelerateライブラリに基づいた当社のツールを活用することをお勧めします。

詳細については、以下のガイドをご覧ください：Accelerateを使用した大規模モデルの読み込み