Back

/ 5 min read

Transformersでbig modelsの読み込み

transformersを使ってLLMの推論をさせる時にGPUを活用したいが、十分なメモリ量がない場合、学習済みモデルを読み込む際にとりあえず device_map='auto'を指定して推論させることがあるが、いまいちパフォーマンスが出ないことがあるので内部の仕組みがどのように行われるのかを調べる。

accelerateによるdevice_mapの推定

transformersでモデルを読み込む際は一般に以下のように記述することが多い。

import transformers


tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct", token="foo")
model = transformers.AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3.1-8B-Instruct,
    device_map="auto", # or 'cpu' or 'cuda' etc...
    token="foo",
)

ここで device_map="auto"を指定することでモデルのパラメータをどのデバイスで読み込むかを実際のメモリ量をもとに自動で推定し、ロードしてくれる機構となっている。

この機構はaccelerateでサポートされていて、どのように割り当てを決めているかなどは以下のドキュメントが詳しい。

https://huggingface.co/docs/accelerate/usage_guides/big_modeling

基本的にまず利用できるGPUメモリに可能な限り読み込んで、その後CPU -> ディスクという順序で読み込み先が決まる。

実際に各パラメータがどこで読み込まれることになるのかは accelerate.infer_auto_device_map を使って知ることができる。

with accelerate.init_empty_weights():  # meta dataのみ読み込む
    model = transformers.AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")
accelerate.infer_auto_device_map(model)


OrderedDict([('model.embed_tokens', 0),
             ('model.layers.0', 0),
             ('model.layers.1', 0),
             ('model.layers.2', 0),
             ('model.layers.3', 0),
             ('model.layers.4', 0),
             ('model.layers.5', 0),
             ('model.layers.6', 0),
             ('model.layers.7', 0),
             ('model.layers.8.self_attn', 0),
             ('model.layers.8.input_layernorm', 'cpu'),
             ('model.layers.8.post_attention_layernorm', 'cpu'),
             ('model.layers.9', 'cpu'),
             ('model.layers.10', 'cpu'),
             ('model.layers.11', 'cpu'),
             ('model.layers.12', 'cpu'),
             ('model.layers.13', 'cpu'),
             ('model.layers.14', 'cpu'),
             ('model.layers.15', 'cpu'),
             ('model.layers.16', 'cpu'),
             ('model.layers.17', 'cpu'),
             ('model.layers.18', 'cpu'),
             ('model.layers.19', 'cpu'),
             ('model.layers.20', 'cpu'),
             ('model.layers.21', 'cpu'),
             ('model.layers.22', 'cpu'),
             ('model.layers.23', 'cpu'),
             ('model.layers.24', 'cpu'),
             ('model.layers.25', 'cpu'),
             ('model.layers.26', 'cpu'),
             ('model.layers.27', 'cpu'),
             ('model.layers.28', 'cpu'),
             ('model.layers.29', 'cpu'),
             ('model.layers.30', 'cpu'),
             ('model.layers.31', 'cpu'),
             ('model.norm', 'cpu'),
             ('model.rotary_emb', 'cpu'),
             ('lm_head', 'cpu'),
             ('model.layers.8.mlp', 'cpu')])

現状の空きがあるメモリ量からどこにロードするかの推論結果を得られることがわかる。 accelerate.init_empty_weightsではpytorchのmeta deviceを使って実際のパラメータの値は読み込むことなしにメタデータのみ取得して割り当ての推論に使っている。

例えば、半精度でモデルを読み込むとGPUに読み込めるパラメータ量がさらに増えていることが確認できる。

with accelerate.init_empty_weights():
    model = transformers.AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct", torch_dtype=torch.float16)
accelerate.infer_auto_device_map(model)


OrderedDict([('model.embed_tokens', 0),
             ('model.layers.0', 0),
             ('model.layers.1', 0),
             ('model.layers.2', 0),
             ('model.layers.3', 0),
             ('model.layers.4', 0),
             ('model.layers.5', 0),
             ('model.layers.6', 0),
             ('model.layers.7', 0),
             ('model.layers.8', 0),
             ('model.layers.9', 0),
             ('model.layers.10', 0),
             ('model.layers.11', 0),
             ('model.layers.12', 0),
             ('model.layers.13', 0),
             ('model.layers.14', 0),
             ('model.layers.15', 0),
             ('model.layers.16', 0),
             ('model.layers.17', 0),
             ('model.layers.18', 0),
             ('model.layers.19', 0),
             ('model.layers.20', 0),
             ('model.layers.21.self_attn', 0),
             ('model.layers.21.input_layernorm', 'cpu'),
             ('model.layers.21.post_attention_layernorm', 'cpu'),
             ('model.layers.22', 'cpu'),
             ('model.layers.23', 'cpu'),
             ('model.layers.24', 'cpu'),
             ('model.layers.25', 'cpu'),
             ('model.layers.26', 'cpu'),
             ('model.layers.27', 'cpu'),
             ('model.layers.28', 'cpu'),
             ('model.layers.29', 'cpu'),
             ('model.layers.30', 'cpu'),
             ('model.layers.31', 'cpu'),
             ('model.norm', 'cpu'),
             ('model.rotary_emb', 'cpu'),
             ('lm_head', 'cpu'),
             ('model.layers.21.mlp', 'cpu')])

モデルが一つのGPUに載るくらい十分に小さい場合はこのように表示される。

with accelerate.init_empty_weights():
    model = transformers.AutoModelForCausalLM.from_pretrained("neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8")
accelerate.infer_auto_device_map(model)


OrderedDict([('', 0)])

accelerateによれば、レイヤーを通過するごとにGPU, CPU間でパラメータを移動させているということであるが、そこのオーバーヘッドが気になるところではある。筆者の環境だと一部のパラメータをGPUに載せるよりCPU onlyで利用する方が推論速度が出るケースもあり、必要に応じてdeviceを手動設定するのが良いのかもしれない。

Each time an input is passed through a layer, it is sent from the CPU to the GPU (or disk to CPU to GPU), the output is calculated, and the layer is removed from the GPU going back down the line. While this adds some overhead to inference, it enables you to run any size model on your system, as long as the largest layer fits on your GPU.