參考：性能基準評測
本文介紹如何使用Pai-Megatron-Patch優(yōu)化PyTorch版Transformer模型訓練。
背景信息
以下所有實驗結(jié)果均在阿里云ECS服務器上進行，實驗使用的ECS配置信息如下。
混合精度
實驗環(huán)境：英文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 |
分布式顯存優(yōu)化：模型狀態(tài)切分
實驗環(huán)境：英文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原生的分布式數(shù)據(jù)并行會導致出現(xiàn)OOM，導致OOM的關(guān)鍵原因是模型無法放在32G的顯卡上，因為Adam優(yōu)化器的狀態(tài)參數(shù)就消耗16G顯存。
方案 | 吞吐 （samples/s） | Peak Memory （MB） |
無加速技術(shù) | OOM | OOM |
混合精度訓練 | 9.57 +/- 0.26 | 25061 |
混合精度訓練 + oss模型狀態(tài)切分 | 6.02 +/- 0.06 | 22077 |
混合精度訓練 + oss/sdp模型狀態(tài)切分 | 7.01 +/- 0.07 | 17113 |
混合精度訓練 + fsdp模型狀態(tài)切分 | NA | NA |
混合精度訓練 + Zero-1 | 12.88 +/- 0.10 | 15709 |
混合精度訓練 + Zero-2 | 10.27 +/- 0.08 | 15693 |
混合精度訓練 + Zero-3 | NA | NA |
3D混合并行
實驗環(huán)境：英文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計算圖優(yōu)化
實驗環(huán)境：英文huggingface bert 微調(diào)
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計算圖優(yōu)化 | 554.24 +/- 1.98 | 2430 |
ORT+混合精度 | 614.70 +/- 8.69 | 2289 |