本文介紹如何使用Pai-Megatron-Patch最佳化PyTorch版Transformer模型訓練。
背景資訊
以下所有實驗結果均在阿里雲ECS伺服器上進行，實驗使用的ECS配置資訊如下。
混合精度
實驗環境：英文huggingface 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
方案 | 吞吐 （samples/s） | Peak Memory （MB） |
單精確度訓練 | 103.07 +/- 1.03 | 17025 |
混合精度訓練 | 178.15 +/- 2.10 | 12698 |
分布式顯存最佳化：模型狀態切分
實驗環境：英文megatron gpt預訓練
num-layers 24
hidden-size 2048
num-attention-heads 32
num-params 1313722368（13億）
local-rank 4
seq-length 1024
micro-batch-size 1
global-batch-size 4
使用Pytorch原生的分布式資料並行會導致出現OOM，導致OOM的關鍵原因是模型無法放在32G的顯卡上，因為Adam最佳化器的狀態參數就消耗16G顯存。
方案 | 吞吐 （samples/s） | Peak Memory （MB） |
無加速技術 | OOM | OOM |
混合精度訓練 | 9.57 +/- 0.26 | 25061 |
混合精度訓練 + oss模型狀態切分 | 6.02 +/- 0.06 | 22077 |
混合精度訓練 + oss/sdp模型狀態切分 | 7.01 +/- 0.07 | 17113 |
混合精度訓練 + fsdp模型狀態切分 | NA | NA |
混合精度訓練 + Zero-1 | 12.88 +/- 0.10 | 15709 |
混合精度訓練 + Zero-2 | 10.27 +/- 0.08 | 15693 |
混合精度訓練 + Zero-3 | NA | NA |
3D混合并行
實驗環境：英文megatron gpt預訓練
num-layers 24
hidden-size 2048
num-attention-heads 32
num-params 1313722368（13億）
local-rank 4
seq-length 1024
micro-batch-size 1
global-batch-size 4
開啟混合精度訓練下：
運算元拆分 | 流水並行 | 吞吐 （samples/s） | Peak Memory （MB） |
1 | 1 | 9.63 +/- 0.29 | 25061 |
2 | 1 | 7.59 +/- 0.14 | 11300 |
4 | 1 | 6.16 +/- 0.06 | 5673 |
1 | 2 | 8.46 +/- 0.17 | 12375 |
1 | 4 | 8.03 +/- 0.12 | 8141 |
2 | 2 | 7.37 +/- 0.11 | 6211 |
4 | 4 | 6.24 +/- 0.08 | 5673 |
ORT計算圖最佳化
實驗環境：英文huggingface 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%：
方案 | 吞吐 （samples/s） | Peak Memory （MB） |
單精確度訓練 | 479.15 +/- 1.67 | 2112 |
混合精度訓練 | 589.66 +/- 4.79 | 2127 |
ORT計算圖最佳化 | 554.24 +/- 1.98 | 2430 |
ORT+混合精度 | 614.70 +/- 8.69 | 2289 |