transformersのdevice_map="auto"について

概要

https://huggingface.co/docs/transformers/pipeline_tutorial#using-pipeline-on-large-models-with-accelerate

によると、transformersのpipeline実行時に device_map="auto" を渡すと、大規模なモデルでも効率よく実行してくれるとのことです。

内部的にどういう動作をしているのか気になったので調べてみました。

accelerateについて

!pip install accelerate

前提として上記の動作をさせるにはaccelerateをインストールする必要があります。
効率化はaccelerateを適応することによって行われるようです。

https://huggingface.co/docs/accelerate/index

accelerateを適応することでpytorchのコードを

• CPU only
• multi-CPU on one node (machine)
• multi-CPU on several nodes (machines)
• single GPU
• multi-GPU on one node (machine)
• multi-GPU on several nodes (machines)
• TPU
• FP16/BFloat16 mixed precision
• FP8 mixed precision with Transformer Engine
• DeepSpeed support (Experimental)
• PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
• Megatron-LM support (Experimental)

に対応させることができます。

具体的には以下のように学習に必要なオブジェクトをすべてacceleratorに渡します。

model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)

accelerator.prepare

accelerator.prepareが何をしているのか少しコードを読んでみます。

https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1147

self.distributed_type
によって大きく分岐しています。

https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/utils/dataclasses.py#L144-L165

DistributedTypeはこのような種類があり、どの環境かによって挙動が大きく変わりそうです。
MEGATRON_LMなどのマルチノード環境で動かすタイプにも対応していますが、そのあたりもよしなにやってくれるっぽいです。

まとめ

それぞれのデータやパラメータをどのノードのどのGPUに乗せるのか、などの分散処理は掘り下げるともっと面白そうなのですが一旦ここまでにしておきます。