トレーニングの分散

TRLによるマルチGPUトレーニング

TRLのトレーナーは、🤗 Accelerateを使用して、複数のGPUやノードにまたがる分散学習を可能にします。そのために、まず以下を実行して🤗 Accelerateの設定ファイルを作成します。

accelerate config

マルチGPU / マルチノード構成に基づいて質問に答えてください。その後、以下を実行して分散学習を開始できます。

accelerate launch train.py

examplesフォルダには、テンプレートとして使用できる設定ファイルも提供しています。これらのテンプレートを使用するには、ジョブの起動時に設定ファイルのパスを渡すだけです。例:

accelerate launch --config_file examples/accelerate_configs/multi_gpu.yaml train.py <SCRIPT_ARGS>

これにより、ワークロードが利用可能なすべてのGPUに自動的に分散されます。

内部では、🤗 AccelerateはGPUごとに1つのモデルを作成します。各プロセスは:

独自のデータバッチを処理します
そのバッチの損失と勾配を計算します
すべてのGPU間で勾配の更新を共有します

multi gpu

実効バッチサイズは次のように計算されます: $\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$

複数のGPUにスケーリングする際にバッチサイズを一定に保つには、per_device_train_batch_sizeとgradient_accumulation_stepsを適切に更新してください。

例として、これらの設定は同等であり、同じ結果を生成する必要があります。

GPUの数	デバイスごとのバッチサイズ	勾配蓄積ステップ	コメント
1	32	1	メモリ使用量は高くなる可能性がありますが、トレーニングは高速です
1	4	8	メモリ使用量は低くなりますが、トレーニングは遅くなります
8	4	1	マルチGPUで両方の利点を最大限に活用

GPUごとに1つのモデルを配置するとメモリ使用量が高くなるため、大規模モデルやメモリが少ないGPUでは実行できない場合があります。そのような場合は、モデルシャーディング、Zero Redundancy Optimizer、混合精度学習、CPUまたはNVMeへのオフロードなどの最適化を提供するDeepSpeedを活用できます。詳細については、DeepSpeed Integrationガイドを参照してください。

非常に長いシーケンスでの学習については、専用のガイドがあります: 100万トークン超の学習。

マルチノードトレーニング

単一のマシンに十分なGPUがない場合、TRLは🤗 Accelerateを使用して、複数のマシン（ノード）にわたってトレーニングをスケーリングできます。

Accelerateの設定

マルチノード学習用のaccelerate設定ファイル（例：multi_node.yaml）を作成します。主なフィールドは以下の通りです：

compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
num_machines: 2
machine_rank: 0  # 0 for main node, 1 for second node
main_process_ip: 10.0.0.1  # IP of rank 0 node
main_process_port: 29500
num_processes: 16  # total processes across nodes
mixed_precision: bf16
use_cpu: false
same_network: true

num_processesを、すべてのノードにまたがるGPUの総数に合わせて調整してください。

10.0.0.1 を、ランク 0（メイン）ノードの実際の IP アドレスに置き換えてください。

ローンチ

オプション1: 手動での起動（非HPC）

各ノードで手動で以下を実行します:

# Node 0 (main node)
accelerate launch --config_file multi_node.yaml --machine_rank 0 train.py

# Node 1
accelerate launch --config_file multi_node.yaml --machine_rank 1 train.py

オプション2: SLURM による起動 (HPC クラスタ)

SLURMジョブスケジューラーを使用するクラスターでは、ジョブスクリプト（例: slurm_job.sh）を作成します。

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

srun accelerate launch --config_file multi_node.yaml train.py

次に、ジョブを送信します:

sbatch slurm_job.sh

SLURM は、要求されたすべてのノードと GPU にわたってトレーニングを自動的に分散し、srun がマルチノード通信に必要な環境変数を設定します。

SLURMの主要なディレクティブ:

--nodes=2: 2つのコンピュートノードを要求します
--gpus-per-node=8: ノードごとに8つのGPUを割り当てます(合計16個)
--job-name: ジョブキューでの追跡用のラベル

distributed_type: DEEPSPEED を設定し、deepspeed_config ブロックを追加することで、マルチノードと DeepSpeed を組み合わせることができます。DeepSpeed 統合ガイドを参照してください。

さらに詳しく

Accelerate: 起動スクリプト
Accelerate: 例のコレクション
SLURM Workload Manager ドキュメント - クラスタのジョブスケジューリング用

Update on GitHub