すべてのプロダクト
Search

このプロダクト

    Platform For AI:リファレンス：パフォーマンスベンチマーク

    ドキュメントセンター

    Platform For AI:リファレンス：パフォーマンスベンチマーク

    最終更新日:Jun 22, 2026

    このドキュメントでは、Pai-Megatron-Patch を使用して PyTorch Transformer モデルのトレーニングを最適化する方法について説明します。

    背景情報

    すべての実験は、次の仕様の Alibaba Cloud Elastic Compute Service (ECS) インスタンスで実施しました。インスタンスタイプは ecs.gn6e-c12g1.12xlarge で、48 コア CPU、368 GiB のメモリ、4 つの NVIDIA V100 GPU を搭載しています。環境は Ubuntu 18.04 64 ビットオペレーティングシステムを使用し、イメージ ID は ubuntu_18_04_x64_20G_alibase_20211227.vhd、ピーク帯域幅は 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 |
+-------------------------------+----------------------+----------------------+

    混合精度トレーニング

    実験設定：英語の Hugging Face BERT モデルの事前学習。

    • 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

    ソリューション

    スループット (サンプル/秒)

    ピークメモリ (MB)

    単精度トレーニング

    103.07 +/- 1.03

    17,025

    混合精度トレーニング

    178.15 +/- 2.10

    12,698

    GPU メモリ最適化：モデル状態分割

    実験設定：英語の Megatron GPT モデルの事前学習。

    • num-layers 24

    • hidden-size 2048

    • num-attention-heads 32

    • num-params 1,313,722,368 (13 億)

    • local-rank 4

    • seq-length 1024

    • micro-batch-size 1

    • global-batch-size 4

    最適化を行わない場合、このモデルは 32 GB の GPU に収まらず、メモリ不足エラーが発生します。たとえば、Adam オプティマイザーの状態パラメーターだけで 16 GB の GPU メモリが必要です。

    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)

    ソリューション

    スループット (サンプル/秒)

    ピークメモリ (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

    混合精度トレーニング + Fully Sharded Data Parallel (FSDP)

    NA

    NA

    混合精度トレーニング + ZeRO Stage 1

    12.88 +/- 0.10

    15,709

    混合精度トレーニング + ZeRO Stage 2

    10.27 +/- 0.08

    15,693

    混合精度トレーニング + ZeRO Stage 3

    NA

    NA

    3D ハイブリッド並列処理

    実験設定：英語の Megatron GPT モデルの事前学習。

    • num-layers 24

    • hidden-size 2048

    • num-attention-heads 32

    • num-params 1,313,722,368 (13 億)

    • local-rank 4

    • seq-length 1024

    • micro-batch-size 1

    • global-batch-size 4

    混合精度トレーニングを有効にした場合：

    テンソル並列処理

    パイプライン並列処理

    スループット (サンプル/秒)

    ピークメモリ (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 Runtime グラフ最適化

    実験設定：英語の Hugging Face BERT モデルのファインチューニング。

    • 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

    単精度トレーニングと比較して、ONNX Runtime グラフ最適化のみを使用すると、スループットが 15.6% 向上します。

    ソリューション

    スループット (サンプル/秒)

    ピークメモリ (MB)

    単精度トレーニング

    479.15 +/- 1.67

    2,112

    混合精度トレーニング

    589.66 +/- 4.79

    2,127

    ONNX Runtime グラフ最適化

    554.24 +/- 1.98

    2,430

    ONNX Runtime + 混合精度トレーニング

    614.70 +/- 8.69

    2,289