35. データ並列化でLLM訓練を高速化！マルチGPU活用術

スポンサーリンク

35. データ並列化でLLM訓練を高速化！マルチGPU活用術

こんにちは！今回は、複数のGPUを使ってLLM（大規模言語モデル）のトレーニングを高速化する「データ並列化（Data Parallelism）」について学んでいきます。

前回までメモリ最適化を学びましたが、今度は別の問題：トレーニング時間に目を向けます！

分散データ並列化（DDP: Distributed Data Parallel）を使用すると、複数のGPUを使ってより多くのデータを並行処理しながら、モデルの動作を同一に保つことができるんです。これは、モデルがすでに収まる場合にファインチューニングを加速する最もシンプルな方法なんですね。

はじめに：モデルが収まるがトレーニングが遅い場合

モデルが単一GPUに快適に収まると、メモリは制限要因ではなくなります。次の質問は単純です：

「トレーニングにどれくらいかかるか？」

そしてここで、コンテキストが重要になります。

トレーニング時間の重要性

10時間、さらには数日かかるトレーニング実行は、ワークフロー、予算、またはプロジェクトのタイムラインによっては完全に受け入れられるかもしれません。

しかし、より速い反復が重要になる場合：
– 1日あたり複数の実験を実行したい
– 開発サイクルを短縮したい
– より厳しい締め切りを満たす必要がある

こんなとき、トレーニング時間を削減することが重要になり始めます。そして、より多くのGPUを追加することがプロセスを加速する最も簡単な方法になるんです。

それがマルチGPUトレーニングが登場する場所なんですね！

データ並列化とは何か？

単一GPUでトレーニングすることを想像してください：

  1. GPUはモデルを通じてバッチを実行
  2. 勾配を計算
  3. 重みを更新
  4. 繰り返す

簡単ですが、大きなデータセットがある場合は遅いんですよね。

データ並列化の仕組み

データ並列化は、このプロセスを複数のGPU全体に複製します：

データの分割：
– GPU 0 → バッチAを処理
– GPU 1 → バッチBを処理
– GPU 2 → バッチCを処理
– GPU 3 → バッチDを処理

各GPUは独自のデータスライスで勾配を計算します。

同期ステップ：
1. すべてのGPUからの勾配が平均化される
2. すべてのモデルコピーが正確に同じ平均化された更新を受け取る
3. すべてのレプリカがロックステップを維持する

結果：同一のトレーニング、より短い時間！

重要な制限

これには1つの重要な制限があります：

完全なモデルは各GPUに収まる必要があります。

DDPは計算を並列化しますが、メモリ使用量を削減しません。

有効バッチサイズの自動増加

各GPUが独自のバッチを処理するため、有効バッチサイズはデバイス数で自動的に増加します。

例：
– バッチサイズ8を実行する4つのGPU
– = バッチサイズ32を実行する単一GPU

これにより、2つの即座の利点が得られます：
1. より速いトレーニング
2. より大きな有効バッチサイズ

モデルの動作を変更することなく！

勾配同期の課題

データ並列化は、すべてのGPUが同期を維持するためにのみ機能します。

勾配同期とは

各バックワードパス後、すべてのレプリカは：
1. すべての他のレプリカと勾配を交換
2. 単一の平均化された更新に同意

このステップ（勾配同期）が、すべてのモデルコピーを同一に保ちます。

通信コスト

しかし、それは無料ではありません！

10億パラメータモデルの場合：
– トレーニングステップごとに約10億の勾配値を移動
– 追加するGPUが多いほど、より多くの通信が必要
– リンクが遅いほど、計算ではなく通信により多くの時間を費やす

インターコネクトの種類

遭遇する最も一般的な2つのタイプ：

PCIe（一般的なコンピュータバス）

  • コンシューマーハードウェアで一般的
  • 比較的遅い（約16 GB/s）
  • 通信オーバーヘッドが大きい

NVLink

  • データセンターハードウェア（A100、H100など）で見られる
  • NVIDIAの高速GPU間接続
  • はるかに高速（最大600 GB/s）
  • 通信オーバーヘッドが少ない

より速いインターコネクト = 勾配がより速く同期 = 待機時間が少ない = より良いスケーリング

そのため、PCIeを使用したコンシューマーハードウェアと比較して、NVLink接続されたGPUを使用したクラウドインスタンスでより良いDDPパフォーマンスが見られるんですね。

スケーリング効率に影響する要因

理想的なスケーリング（2 GPU = 2倍速度、4 GPU = 4倍速度）と実際のパフォーマンスのギャップは、いくつかの要因に依存します：

  1. インターコネクト速度：NVLinkはPCIeと比較してオーバーヘッドを劇的に削減
  2. モデルサイズ：より大きなモデルは同期するデータが多い（PEFTを使用しない限り）
  3. 計算時間：より速いGPUは計算が速く終了し、通信が相対的により高価になる

簡単に言えば：
DDPはトレーニングを加速しますが、追加のGPUごとに「通信税」があります。これはトレードオフです。わずかな同期オーバーヘッドと引き換えに、より多くの計算並列化を得るんですね。

ほとんどの実用的なファインチューニングセットアップでは、速度の向上がコストをはるかに上回ります！

Hugging Face Accelerateの紹介

メカニクスとトレードオフを理解したので、次のステップは実際にそれを使用する方法です。

良いニュース：分散トレーニングコードを自分で書く必要はありません！

Hugging FaceのAccelerateライブラリが、マルチGPUトレーニングを通常のPythonスクリプトを実行するのとほぼ同じに感じさせます。

Accelerateが処理するもの

  • 複数のプロセスのスポーン
  • モデルの分散
  • 勾配の同期
  • データローダーのラップ
  • 結果の収集
  • デバイス配置の管理

最良の部分：LLMファインチューニングスクリプトをあまり変更する必要がありません！

Accelerate構成

分散トレーニングを起動する前に、Accelerateを構成する必要があります。2つのオプションがあります。

オプション1：インタラクティブ構成

インタラクティブプロンプトを実行します：

accelerate config

セットアップについて尋ねられます：

In which compute environment are you running?
> This machine

Which type of machine are you using?
> multi-GPU

How many different machines will you use?
> 1

How many processes should be used for distributed training?
> 4  # 使用したいGPU数に設定

Accelerateはこの構成を保存し、トレーニングを起動するたびに使用します。

オプション2：構成ファイルを直接作成

YAML構成ファイルを手動で作成できます。これは、異なるGPU構成間をすばやく切り替えたい場合に便利です。

2 GPUの構成例（config_2gpu.yaml）：

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_GPU
downcast_bf16: "no"
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2  # 使用するGPUの数
rdzv_backend: static
same_network: true
use_cpu: false

主要な設定：
– distributed_type: MULTI_GPU：DDPを有効にする
– num_processes: 2：使用するGPUの数
– mixed_precision: bf16：トレーニング構成に一致

Hugging Face TrainerでAccelerateを使用

QLoRAのような量子化モデルをファインチューニングする場合、Accelerator APIを直接広範囲に使用する必要はありません。代わりに、Hugging FaceのTrainerがAccelerateとシームレスに統合し、分散トレーニングの複雑さのほとんどを処理します。

1. Accelerateで起動

# デフォルト構成を使用（accelerate configから）
accelerate launch train_qlora.py

# または構成ファイルを明示的に指定
accelerate launch --config_file config_2gpu.yaml train_ddp_accelerate.py

このコマンドは自動的に：
– 利用可能なGPUを検出
– GPUごとに1つのプロセスをスポーン
– 分散通信をセットアップ
– 環境変数を構成（LOCAL_RANK、WORLD_SIZEなど）

手動構成は不要です。Accelerateはハードウェアを読み取り、正しいことを行います！

2. 重要：量子化モデルのデバイスマッピング

マルチGPUトレーニングで4-bit量子化を使用する場合、各プロセスは割り当てられたGPUでモデルをロードする必要があります。

処理方法：

from accelerate import Accelerator

# Acceleratorを初期化してプロセス情報を取得
accelerator = Accelerator()

# ローカルプロセスインデックスに基づいてデバイスマッピングでモデルをロード
model, tokenizer = setup_model_and_tokenizer(
    cfg,
    use_4bit=True,
    use_lora=True,
    device_map=accelerator.local_process_index
)

このアプローチが機能する理由：

accelerator.local_process_indexプロパティは、このプロセスが使用すべきGPUを返します：
– プロセス0 → local_process_index=0 → GPU 0でモデルをロード
– プロセス1 → local_process_index=1 → GPU 1でモデルをロード
– そして続く…

これをdevice_mapとして直接渡すと、モデルローダーに伝えます：「このプロセスのGPU Nにすべてのモジュールを配置してください。」

3. DDPの勾配チェックポイント

勾配チェックポイントを使用した分散トレーニングには、非再入可能実装が必要です：

args = TrainingArguments(
    output_dir=output_dir,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},  # DDPに不可欠
    # ... その他の引数
)

use_reentrant=Falseがないと、「variable marked ready twice」エラーが発生します。これは、DistributedDataParallelと互換性のある新しいチェックポイント実装を使用するようPyTorchに伝えるんですね。

4. 有効バッチサイズの理解

per_device_train_batch_size = 4
gradient_accumulation_steps = 4
num_gpus = 4

有効バッチサイズ：
4（デバイスあたり）× 4（GPU）× 4（累積）= 64

動作の流れ：
1. 各GPUはステップあたり4サンプルを処理
2. 勾配は同期前に4ステップにわたって累積
3. すべてのGPUがall-reduceを実行して勾配を平均化
4. モデルは重み更新あたり64サンプルの有効バッチを見る

5. 結果：ほぼ線形スケーリング

SAMSumデータセットでQLoRAを使用してLlama 3.1 8Bをファインチューニングする実験結果：

トレーニング時間：
– 1 GPU: 87.6分
– 2 GPU: 44.6分（1.96倍速、98.2%効率）
– 4 GPU: 22.6分（3.87倍速、96.8%効率）

モデル品質（ROUGE-1スコア）：
– 1 GPU: 53.1%
– 2 GPU: 52.1%
– 4 GPU: 53.0%

ほぼ線形スケーリングは優れた並列効率を示します！モデル品質は構成全体で一貫しており、分散トレーニングでの異なるランダムバッチ順序による小さな変動があるだけですね。

6. 重複出力の処理

各GPUは独自のPythonプロセスを実行するため、print文は複数回実行されます。

解決策：

def print_rank0(msg):
    """メインプロセスからのみ出力"""
    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
        print(msg)

これはTrainerの組み込みロギング（損失、メトリクス）には必要ありません。これはすでに調整されています。しかし、カスタムprint文の場合、LOCAL_RANKをチェックすることでコンソールをきれいに保てます！

重要なポイント

Trainer + Accelerateでは、分散トレーニングはほぼ透明です。

主な落とし穴：
1. 量子化モデルのデバイス配置：accelerator.local_process_indexを使用
2. 勾配チェックポイント構成：use_reentrant=False

これらを正しく設定すると：
– 2 GPUで2倍の速度
– 4 GPUで4倍の速度
– モデル品質の損失なし

全体像：これまでの結果のまとめ

すべてのモデルがどのように積み重なるか：

ファインチューニングされたLlama 3.2 3B：
– すべてのROUGEメトリクスでファインチューニングされたGPT-4o-miniに近い
– ROUGE-1で4.12ポイント以内
– ROUGE-2で3.3ポイント以内
– ROUGE-Lで3.17ポイント以内

3Bモデルは1Bバリアントを上回る：
– ROUGE-1で4.27ポイント改善
– ROUGE-2で7.3ポイント改善

Accelerateを使用したデータ並列化により、この大きなモデルを複数のGPU全体でトレーニングでき、一貫したモデル品質を維持しながらほぼ線形の速度向上（4 GPUで4倍）を達成できました！

この改善は、効率的なマルチGPUトレーニングを使用したより大きなモデルへのスケーリングが、インフラストラクチャとデータを完全に制御しながら、管理されたAPIソリューションに匹敵するパフォーマンスをもたらす方法を強調しています。

次のステップ

DDPは非常に具体的な問題を解決します：

「モデルが単一GPUに収まる – どうすればより速くトレーニングできるか？」

でも、モデルが収まらなくなると、ボトルネックが変わります。速度はもはや問題ではありません。メモリです。

その時点で、データ並列化は役に立ちません。なぜなら、すべてのGPUがまだモデルの完全なコピーを保持する必要があるからです。

解決策はシャードトレーニングです：
– パラメータ、勾配、オプティマイザー状態が複製ではなくデバイス全体に分割される
– それがFSDPやDeepSpeed ZeROのようなテクニックが設計されている理由

次回は、これらのアプローチがメモリプレッシャーをどのように削減し、単一GPUが保持できるよりもはるかに大きなモデルをトレーニングすることを可能にするかを見ていきます！

まとめ

今回の記事では、データ並列化（DDP）を使ったマルチGPUトレーニングについて学びました。

重要なポイント：

  1. DDPは計算を並列化：メモリではなく速度のために
  2. ほぼ線形スケーリング：4 GPUで約4倍の速度
  3. Accelerateが複雑さを隠す：数行の変更でマルチGPU対応
  4. 重要な設定：デバイスマッピングと勾配チェックポイント
  5. 通信コスト：インターコネクトの種類が効率に影響

モデルが単一GPUに収まる場合、DDPは速度を上げる最もシンプルで効果的な方法です！

次回は、モデルが収まらない場合のシャードトレーニング（FSDP/ZeRO）について学んでいきましょう！

コメント

タイトルとURLをコピーしました