Accelerate documentation
分散推論
分散推論
分散推論は、2つのカテゴリに分類できます。
- モデル全体を各GPUにロードし、バッチのチャンクを一度に各GPUのモデルコピーに送信する
- モデルの一部を各GPUに読み込み、一度に1つの入力のみを処理する
このガイドでは、最初の括弧について説明します。2つ目については、Big Model Inferenceを参照してください。
バッチのチャンクを自動的に各ロードされたモデルに送信する
これは最もメモリ集約的なソリューションであり、ある時点で各GPUがモデルの完全なコピーをメモリに保持する必要があります。
通常、この処理を行う際、ユーザーはモデルをCPUから読み込むために特定のデバイスに送信し、その後、各プロンプトを別のデバイスに移動します。
diffusersライブラリを使った基本的なパイプラインは、次のようになります。
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)その後、特定のプロンプトに基づいて推論を実行します。
def run_inference(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
pipe.to(rank)
if torch.distributed.get_rank() == 0:
prompt = "a dog"
elif torch.distributed.get_rank() == 1:
prompt = "a cat"
result = pipe(prompt).images[0]
result.save(f"result_{rank}.png")ランクを確認して送信するプロンプトを判断する必要があるため、やや煩雑になる点に気づくでしょう。
ユーザーは、Accelerate を使用して Accelerator でデータローダーを準備することで、このようなタスクの管理を簡単にできるのではないかと考えるかもしれません。（詳しくは、クイックツアーの関連セクションをご覧ください）
管理できるか？できます。ただし、不要な余分なコードを追加する点については、その通りです。
Accelerateを使用することで、Accelerator.split_between_processes()コンテキストマネージャー（PartialStateおよびAcceleratorStateにも存在します）を用いて、このプロセスを簡素化できます。
この関数は、渡されたデータ（プロンプト、テンソルのセット、過去のデータの辞書など）を自動的にすべてのプロセスに分割し、必要に応じてパディングを適用して、すぐに使用できる状態にします。
このコンテキストマネージャーを使用して、上記の例を書き直してみましょう。
import torch
from accelerate import PartialState # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)
# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
result = pipe(prompt).images[0]
result.save(f"result_{distributed_state.process_index}.png")その後、コードを起動するには、Accelerateを使用できます。
accelerate config を使用して設定ファイルを生成した場合は:
accelerate launch distributed_inference.py
特定の設定ファイルを使用する場合:
accelerate launch --config_file my_config.json distributed_inference.py
または、設定ファイルを作成せずに2つのGPUで起動したい場合は:
Note: システムに基づいて値が推測されるという警告が表示されます。これらを削除するには、
accelerate config defaultを実行するか、accelerate configを実行して設定ファイルを作成してください。
accelerate launch --num_processes 2 distributed_inference.py
これにより、このデータを分割するために必要な定型コードを、ごく少数のコード行に簡単に削減できました。
では、プロンプトとGPUの数が一致しない場合はどうでしょうか。例えば、プロンプトが3つあるのにGPUが2つしかない場合です。
コンテキストマネージャーの下では、最初のGPUが最初の2つのプロンプトを受け取り、2番目のGPUが3番目のプロンプトを受け取るため、すべてのプロンプトが分割され、オーバーヘッドは不要になります。
しかし、すべての GPUの結果に対して何らかの処理（すべてを集約して後処理を実行するなど）を行いたい場合はどうすればよいでしょうか。
apply_padding=True を渡すことで、プロンプトのリストが同じ長さにパディングされ、余分なデータは最後のサンプルから取得されるようになります。これにより、すべての GPU が同じ数のプロンプトを持つため、結果を集約できます。
これは、結果の集約など、各デバイス上のデータの長さが同じである必要があるアクションを実行する場合にのみ必要です。基本的な推論では、これはいりません。
例えば:
import torch
from accelerate import PartialState # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)
# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
result = pipe(prompt).images最初のGPUではプロンプトが["a dog", "a cat"]となり、2番目のGPUでは["a chicken", "a chicken"]となります。
最後のサンプルは直前のものと同じになるため、除外してください。
LLMとの使用方法など、より複雑な例はこちらでご覧いただけます。
Update on GitHub