使用Pai-Megatron-Patch最佳化PyTorch版Transformer模型訓練 - Platform For AI

Pai-Megatron-Patch各加速技術的訓練效能基準評測資料，供您評估和選擇加速方案時參考。

背景資訊

以下實驗均在阿里雲 ECS 執行個體上運行。執行個體規格為 ecs.gn6e-c12g1.12xlarge，配置 48 核 CPU、368 GiB 記憶體、4 塊 NVIDIA V100 GPU，作業系統 Ubuntu 18.04 64 位元，頻寬峰值 100 Mbps。運行 nvidia-smi 確認驅動版本 440.64.00、CUDA 10.2 已就緒，4 塊 Tesla V100-SXM2 GPU 均處於空閑狀態。

| NVIDIA-SMI 440.64.00    Driver Version: 440.64.00    CUDA Version: 10.2  |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:00:07.0 Off |                    0 |
| N/A   32C    P0    41W / 300W |      0MiB / 32510MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  On   | 00000000:00:08.0 Off |                    0 |
| N/A   31C    P0    41W / 300W |      0MiB / 32510MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   2  Tesla V100-SXM2...  On   | 00000000:00:09.0 Off |                    0 |
| N/A   30C    P0    39W / 300W |      0MiB / 32510MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  Tesla V100-SXM2...  On   | 00000000:00:0A.0 Off |                    0 |
| N/A   31C    P0    40W / 300W |      0MiB / 32510MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

混合精度

實驗環境：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。原因是模型無法放入 32 GB 顯存：僅 Adam 最佳化器狀態參數就消耗 16 GB 顯存。

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py", line 351, in forward
    encoder_output = self.encoder(encoder_input,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py", line 703, in forward
    hidden_states = layer(hidden_states,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py", line 441, in forward
    self.self_attention(layernorm_output,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/transformer.py", line 264, in forward
    matmul_result = torch.baddbmm(
RuntimeError: CUDA out of memory. Tried to allocate 128.00 MiB (GPU 1; 31.75 GiB total capacity; 28.56 GiB already allocated; 84.00 MiB free; 30.19 GiB reserved in total by PyTorch)

方案	吞吐（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

ORT 計算圖最佳化相比單精確度訓練吞吐提升 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