Dokumen ini menjelaskan cara menggunakan Pai-Megatron-Patch untuk mengoptimalkan pelatihan model Transformer PyTorch.
Informasi latar belakang
Semua eksperimen dilakukan pada instance ECS Alibaba Cloud dengan spesifikasi berikut: tipe instance ecs.gn6e-c12g1.12xlarge, CPU 48-core, memori 368 GiB, dan empat GPU NVIDIA V100. Lingkungan menggunakan sistem operasi Ubuntu 18.04 64-bit dengan ID image ubuntu_18_04_x64_20G_alibase_20211227.vhd serta bandwidth puncak 100 Mbps. Perintah nvidia-smi dijalankan untuk memastikan bahwa driver versi 440.64.00 dan CUDA 10.2 telah siap, serta keempat GPU Tesla V100-SXM2 dalam keadaan idle.
| 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 |
+-------------------------------+----------------------+----------------------+Pelatihan presisi campuran
Pengaturan eksperimen: Pre-training model BERT Hugging Face bahasa Inggris.
-
num-layers 12
-
hidden-size 768
-
num-attention-heads 12
-
num-params 110.106.428
-
local-rank 4
-
seq-length 512
-
micro-batch-size 16
-
global-batch-size 64
|
Solusi |
Throughput (samples/s) |
Peak memory (MB) |
|
Pelatihan presisi tunggal |
103,07 +/- 1,03 |
17.025 |
|
Pelatihan presisi campuran |
178,15 +/- 2,10 |
12.698 |
Optimasi memori GPU: Partisi state model
Pengaturan eksperimen: Pre-training model Megatron GPT bahasa Inggris.
-
num-layers 24
-
hidden-size 2048
-
num-attention-heads 32
-
num-params 1.313.722.368 (1,3 miliar)
-
local-rank 4
-
seq-length 1024
-
micro-batch-size 1
-
global-batch-size 4
Tanpa optimasi, model ini tidak dapat muat pada GPU 32 GB dan menyebabkan error kehabisan memori. Sebagai contoh, parameter state pengoptimal Adam saja membutuhkan memori GPU sebesar 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)|
Solusi |
Throughput (samples/s) |
Peak memory (MB) |
|
Tanpa optimasi |
OOM |
OOM |
|
Pelatihan presisi campuran |
9,57 +/- 0,26 |
25.061 |
|
Pelatihan presisi campuran + Optimizer State Sharding (OSS) |
6,02 +/- 0,06 |
22.077 |
|
Pelatihan presisi campuran + OSS / Sharded Data Parallel (SDP) |
7,01 +/- 0,07 |
17.113 |
|
Pelatihan presisi campuran + Fully Sharded Data Parallel (FSDP) |
NA |
NA |
|
Pelatihan presisi campuran + ZeRO Stage 1 |
12,88 +/- 0,10 |
15.709 |
|
Pelatihan presisi campuran + ZeRO Stage 2 |
10,27 +/- 0,08 |
15.693 |
|
Pelatihan presisi campuran + ZeRO Stage 3 |
NA |
NA |
Paralelisme hibrid 3D
Pengaturan eksperimen: Pre-training model Megatron GPT bahasa Inggris.
-
num-layers 24
-
hidden-size 2048
-
num-attention-heads 32
-
num-params 1.313.722.368 (1,3 miliar)
-
local-rank 4
-
seq-length 1024
-
micro-batch-size 1
-
global-batch-size 4
Dengan pelatihan presisi campuran diaktifkan:
|
Tensor parallelism |
Pipeline parallelism |
Throughput (samples/s) |
Peak memory (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 |
Optimasi graf ONNX Runtime
Pengaturan eksperimen: Fine-tuning model BERT Hugging Face bahasa Inggris.
-
num-layers 12
-
hidden-size 768
-
num-attention-heads 12
-
num-params 110.106.428
-
local-rank 4
-
seq-length 512
-
micro-batch-size 16
-
global-batch-size 64
Dibandingkan dengan pelatihan presisi tunggal, penggunaan optimasi graf ONNX Runtime saja meningkatkan throughput sebesar 15,6%.
|
Solusi |
Throughput (samples/s) |
Peak memory (MB) |
|
Pelatihan presisi tunggal |
479,15 +/- 1,67 |
2.112 |
|
Pelatihan presisi campuran |
589,66 +/- 4,79 |
2.127 |
|
Optimasi graf ONNX Runtime |
554,24 +/- 1,98 |
2.430 |
|
ONNX Runtime + Pelatihan presisi campuran |
614,70 +/- 8,69 |
2.289 |