このトピックでは、Pai-Megatron-Patchを使用してPyTorch変圧器のトレーニングを最適化する場合のトレーニングパフォーマンスのベンチマークについて説明します。
混合精度トレーニング
トレーニング環境: 英語でHugging Faceによって開発されたBERTの事前トレーニング環境。
num-レイヤー12
隠しサイズの768
num-注意-ヘッド12
num-params 110106428
ローカルランク4
seq-length 512
マイクロバッチサイズ16
global-batch-size 64
解決策 | スループット (サンプル /秒) | ピークメモリ (MB) |
単精度トレーニング | 103.07 +/- 1.03 | 17,025 |
混合精度トレーニング | 178.15 +/- 2.10 | 12,698 |
分散メモリ最適化: モデル状態分割
トレーニング環境: 英語でメガトロンによって開発されたGPTの事前トレーニング環境。
num-layers 24
hidden-size 2048
num-attention-heads 32
num-params 1313722368 (1.3 billion)
local-rank 4
seq-length 1024
micro-batch-size 1
global-batch-size 4
PyTorchネイティブ分散データ並列処理により、メモリ不足 (OOM) 例外が発生する場合があります。 これは、モデルを32 GBメモリに基づいて実行できないためです。 Adam Optimizerの設定では、最大16 GBのメモリを消費できます。 
解決策 | スループット (サンプル /秒) | ピークメモリ (MB) |
加速技術なし | OOM | OOM |
混合精度トレーニング | 9.57 +/- 0.26 | 25,061 |
混合精度トレーニング、およびOptimizer state Sharding (OSS) を使用したモデル状態分割 | 6.02 +/- 0.06 | 22,077 |
混合精度トレーニング、およびOSSまたはSharded Data Parallel (SDP) を使用したモデル状態分割 | 7.01 +/- 0.07 | 17,113 |
混合精度トレーニング、およびFSDP (Fully Sharded Data Parallel) を使用したモデル状態分割 | N/A | N/A |
混合精度トレーニング、およびZero Redundancy optimizer (ZeRO) を使用したオプティマイザの状態分割 | 12.88 +/- 0.10 | 15,709 |
混合精度トレーニング、およびZeROを使用したオプティマイザ状態と勾配の分割 | 10.27 +/- 0.08 | 15,693 |
混合精度トレーニング、およびZeROを使用したオプティマイザの状態、勾配、およびパラメータの分割 | N/A | N/A |
3D並列処理
トレーニング環境: 英語でメガトロンによって開発されたGPTの事前トレーニング環境。
num-layers 24
hidden-size 2048
num-attention-heads 32
num-params 1313722368 (1.3 billion)
local-rank 4
seq-length 1024
micro-batch-size 1
global-batch-size 4
次の表に、3D並列処理と混合精度トレーニングを同時に有効にする場合のベンチマークを示します。
演算子の分割 | パイプライン並列性 | スループット (サンプル /秒) | ピークメモリ (MB) |
1 | 1 | 9.63 +/- 0.29 | 25,061 |
2 | 1 | 7.59 +/- 0.14 | 11,300 |
4 | 1 | 6.16 +/- 0.06 | 5,673 |
1 | 2 | 8.46 +/- 0.17 | 12,375 |
1 | 4 | 8.03 +/- 0.12 | 8,141 |
2 | 2 | 7.37 +/- 0.11 | 6,211 |
4 | 4 | 6.24 +/- 0.08 | 5,673 |
ONNXランタイムのグラフ最適化
トレーニング環境: 英語でHugging Faceによって開発されたBERTの微調整。
num-layers 12
hidden-size 768
num-attention-heads 12
num-params 110106428
local-rank 4
seq-length 512
micro-batch-size 16
global-batch-size 64
次の表に、グラフのパフォーマンスを15.6% で最適化する場合のベンチマークを示します。
解決策 | スループット (サンプル /秒) | ピークメモリ (MB) |
単精度トレーニング | 479.15 +/- 1.67 | 2,112 |
混合精度トレーニング | 589.66 +/- 4.79 | 2,127 |
ONNXランタイムのグラフ最適化 | 554.24 +/- 1.98 | 2,430 |
ONNXランタイムのグラフ最適化と混合精度トレーニング | 614.70 +/- 8.69 | 2,289 |