PyTorch Lightningに対応できたので、DistributedDataParallel によるマルチGPU学習を試してみた。
前回未対応だった、勾配クリッピングはconfig.yamにgradient_clip_valを記述するだけで対応できた。
また、モデルのエクスポートもon_train_end()で実装した。
DeepLearningShogi/dlshogi/ptl.py at feature/pytorch_lightning · TadaoYamaoka/DeepLearningShogi · GitHub
DDP
PyTorch Lightningで、Trainerに複数のデバイスを指定すると自動で、データパラレルでマルチGPUで並列化できる。
昨今、大規模言語モデルでは、並列化にモデル並列などの手法が使われるが、将棋AIのモデルはGPUメモリに十分のるサイズなので、データパラレルで並列化を行う。
データパラレルは、バッチを分割して、各GPUで計算を行うので、各GPUのバッチサイズ×GPU数がバッチサイズになる。
1GPUのバッチサイズが小さいとGPUの利用効率が落ちるので、それぞれのGPUで限界のバッチサイズを割り当てるのが良い。
ただし、合計のバッチサイズが大きすぎるとラージバッチ問題という、ノイズによる効果が失われて汎化性能が失われる問題が起きる。
したがって、大きなモデルで、合計のバッチサイズが抑えられる場合でしか並列化の効果は得られない。
実験
並列化による速度向上と精度を確認するため、今回は20ブロック256フィルタのモデルで実験した。
1GPUあたりのバッチサイズ固定でGPUを増やした場合と、合計のバッチサイズ固定でGPUを増やした場合で確認した。
基準のバッチサイズは4096とする。
データローダの並列化
マルチGPUで学習する場合、データローダが複数プロセスで作成される。
各プロセスでファイルからデータの読み込みが必要になるため、あらかじめ前処理済みのキャッシュファイルを作成しておくと良い。
config.yamlのデータのパラメータにcacheのパスを指定して、0エポックでfitを実行するとキャッシュを作成できる。
data: train_files: - /work/hcpe3/selfplay_pre42_resnet30x384_relu_b4096lr004-013_floodgate8_s5_po10000-01.hcpe3 val_files: - /work/hcpe3/floodgate.hcpe batch_size: 4096 val_batch_size: 4096 use_average: true use_evalfix: true temperature: 1.0 cache: cache
python ptl.py fit -c config.yaml --trainer.max_epochs 0
実験結果
| バージョン | GPU数 | バッチサイズ/GPU | 合計バッチサイズ |
|---|---|---|---|
| version_0 | 1GPU | 4096 | 4096 |
| version_1 | 2GPU | 4096 | 8192 |
| version_2 | 3GPU | 4096 | 12288 |
| version_3 | 2GPU | 2048 | 4096 |
訓練時間
GPUあたりのバッチサイズ固定でGPU数を増やすと、GPU数に反比例して訓練時間が短くなる。
全体のバッチサイズを固定して、GPU数を増やした場合、2GPUで訓練時間は0.53倍になり、少しオーバーヘッドがあるが訓練時間はGPU数に反比例して短縮される。
評価方策正解率
評価価値正解率
評価方策エントロピー
評価価値エントロピー
