[大規模モデル] Transformers ライブラリシングルマシンマルチカード推論 device_map

記事ディレクトリ

device_map パラメータの分析
device_map="auto" コード例
手動構成
参考文献

Hugging Face のtransformersライブラリは、GPT や ChatGLM などのモデルを自動的にロードして使用するための自動モデル (AutoModel) のモデルインスタンス化メソッドをサポートしています。AutoModel.from_pretrained()このメソッドのパラメーターはdevice_map、単一マシンのマルチカード推論を実現できます。

device_map パラメータの分析

device_mapはメソッドの重要なパラメータでありAutoModel.from_pretrained()、リソースの効果的な割り当てと利用を実現するために、モデルの各コンポーネントをどの特定のコンピューティングデバイスにロードするかを指定するために使用されます。このパラメーターは、モデルの並列トレーニングまたは分散トレーニングを実行する場合に特に役立ちます。

device_mapauto, balanced, balanced_low_0, sequential次のように、パラメータにはいくつかのオプションがあります。

“auto” 和 “balanced”: モデルをすべての GPU に均等に分割します。主に、より効率的な割り当て戦略を発見することが可能です。「バランス」パラメータの機能は安定しています。 (オンデマンドで利用可能)
“balanced_low_0”: モデルは最初の GPU を除く他の GPU に均等に分割され、最初の GPU で占有するリソースが少なくなります。このオプションは、最初の GPU で実行する必要がある関数 (反復プロセス) の生成など、最初の GPU で追加の操作を必要とするニーズに適しています。（推奨）
“sequential”: GPU 0 から最後の GPU までモデルシャードを GPU の順序で割り当てます (この場合、最後の GPU は占有されていないことがよくあり、「balanced_low_0」が最初か最後か、アンバランスな充填の違いになります)。実際の使用では、GPU 0 がビデオメモリを直接バーストします。 (推奨されません)

device_map="auto" コード例

ここでの環境は 2 つのグラフィックスカードを備えた 1 台のマシンであり、これを使用してモデルdevice_map="auto"をロードし、グラフィックスカードの使用状況を観察します。ChatGLM-6B

サンプルコードは次のとおりです。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# 加载模型
model_path = "./model/chatglm2-6b"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")

text = '什么是机器学习?'
inputs = tokenizer(text, return_tensors="pt")
print(inputs)

outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

プログラム実行前のグラフィックスカードの占有状況は次のとおりです。

0 番のグラフィックスカード自体が他のプログラムによって約 13G のビデオメモリを占有されていることがわかります。
使用device_map="auto"後のグラフィックスカードの占有状況は次のとおりです。

自動ポリシーを使用した後、グラフィックスカード 0 と 1 はそれぞれ約 6 ～ 7G 多くのビデオメモリを占有します。

手動構成

シングルカードとして構成

device = "cuda:1"
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map=device)

複数のカード用に構成する場合、
モデルの一部を最初のグラフィックスカードに配置し、他の部分を 2 番目のグラフィックスカードに配置する必要があるとします。または、モデルのコンポーネントサイズに応じてそれを適切に割り当てる必要があります。モデル。ただし、具体的なレイヤー名は実際のモデルに基づいて決定する必要があります。概念的な例を次に示します。

device = {
    
    
    "transformer.h.0": "cuda:0",  # 第一部分放在GPU 0
    "transformer.h.1": "cuda:1",  # 第二部分放在GPU 1
    # ... 根据模型结构继续分配
}
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=device)