Related Resources
Generative AI

Transformer Engine ではじめる FP8 Training (導入編)

2024 年 7 月 18 日
By
Reading Time: 5 minutes

Transformer Engine とは

Transformer Engine とは、Transformer モデルの学習を効率的に行うためのオープンソース ライブラリです。

  • FP8 Tensor コアを活用して行列演算 (GEMM) を行うためのモジュール
  • GPU に最適化された Attention などの Transformer を構成するモジュール
  • 各種深層学習フレームワークで、上記のモジュールを使用した Transformer のモデルの構築が容易に行える API

が含まれており、GPU における Transformer モデルの学習効率を大幅に向上させることができます。特に FP8 については、記事執筆時点では Hopper/Ada Lovelace アーキテクチャなどの最新の GPU に搭載はされているものの、深層学習フレームワークでは対応する OP がまだ実装されていない状況であるため、Transformer Engine は FP8 を活用して GPU の性能を最大限に引き出すために必須のライブラリといえます。

FP8 Training について

FP8 は、名前の通り 8bit で浮動小数点数を表現するデータ フォーマットです。

図 1. FP8 のデータ形式

FP8 を用いる利点としては次のようなものがあります。

  • 計算速度の向上
    • FP8 Tensor コアによる GEMM は FP16/BF16 の 2 倍の FLOPS
  • GPU メモリ使用量の削減/メモリ アクセス速度の向上
    • Tensor サイズが半分になる
  • 推論環境と同一の精度を用いることによる、量子化による精度低下の防止

FP8 Training を行う際は、一般的な FP16/BF16 Training と同じようにオリジナルの高精度の重みを保持し、計算時に FP8 に autocast する Mixed Precision の方式を使用します。

しかし、FP8 はダイナミックレンジがかなり狭いデータ形式となるため、FP16 の Mixed Precision のようにグローバルに単一の Scaling をするだけでは十分な精度が担保できません。

そこで、FP8 では Tensor (重みおよび activation) 毎に Scaling Factor を持ち、各 Tensor を個別にスケーリングして FP8 に変換するという手法を取ります。

図 2. FP8 Training における Tensor の Scaling

推論時における INT8 Quantization などと似たような処理ではありますが、一度キャリブレーションしておけば重みと Scaling が固定である推論時と違い、学習時は毎ステップ重みと入力のレンジが変わるため、 Scaling Factor を都度更新する必要があります。

この Tensor の Scaling と計算時の FP8 への autocast を計算のボトルネックになりにくいように効率化するためには複雑な処理が必要なのですが、これらを自動で実施してくれるのが Transformer Engine です。

Transformer Engine のインストール

Transformer Engine は Python のパッケージとして提供されているので、pip などのパッケージ マネージャーでインストールが可能です。

pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable

NGC のコンテナーをお使いの場合は、各種深層学習フレームワークのコンテナーに Transformer Engine も含まれています。手動インストールの場合 Custom kernel などのビルドも必要となり依存関係が多く時間もかかるので、NGC コンテナーをお使いいただくことを推奨します。

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:24.05-py3

Transformer Engine の使い方

Transformer Engine は、以下のようにして使うことができます。ここでは、PyTorch の API を例に説明します。

from transformer_engine.common.recipe import Format, DelayedScaling
import transformer_engine.pytorch as te
import torch

fp8_format = Format.HYBRID  # E4M3 during forward pass, E5M2 during backward pass
fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")

my_linear = te.Linear(768, 768, bias=True)

inp = torch.rand((1024, 768)).cuda()

with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
   out = my_linear(inp)

fp8_recipe (DelayedScaling) には、FP8 Training を行うための FP8 フォーマット、Scaling の更新間隔、Tensor の最大値の履歴をどれくらい持つか、などを指定します。

次に、Transformer Engine の API を使用してモデルを定義します。ここでは全結合層が 1 層のモデルを te.Linear で定義していますが、もちろん他の nn.Module を継承したクラス配下にネストしても問題ありません。実装されているモジュール一覧は公式ドキュメントに記載されていますが、よく使うものとしては以下のクラスが挙げられます。

  • te.Linear, te.LayerNorm
    • nn.Linear, nn.LayerNorm と互換のクラス
  • te.TransformerLayer
    • Transformer の 1 層全体を定義するクラス

最後に、モデルに forward する部分を te.fp8_autocast() でラップします。このコンテキスト マネージャーの中では、Transformer Engine を使用して定義したモジュールのうち、FP8 が適用可能な部分を自動で FP8 に変換して計算が行われるようになります。

HuggingFace BERT を例にした Transformer Engine の適用

次に、実際の既存のモデル コードに対して Transformer Engine を適用する例として、HuggingFace Transformers ライブラリを用いて BERT の Masked Language Modeling を行う例をご紹介します。

本チュートリアルでは以下のバージョンを使用します

GPU については Hopper および Ada Lovelace のアーキテクチャであれば実行可能です。コンシューマー向けの GeForce RTX 40 シリーズでもお試しいただくことが可能ですが、メモリ容量などが異なるためバッチ サイズなどを適宜調整して実施してみてください。

今回の例では、transformers==4.41.2 に含まれる run_mlm_no_trainer.py をベースに適用していきます。

Docker コンテナーの起動

チュートリアル用のディレクトリを作成し、PyTorch の docker コンテナーを起動します。以降のステップはコンテナー内で作業していきます。

mkdir te-examples
cd te-examples

docker run --rm -it --gpus all --shm-size 4g -v ${PWD}:/workspace -w /workspace nvcr.io/nvidia/pytorch:24.05-py3 bash

ソース コードのダウンロードとライブラリ インストール

HuggingFace Transformers のレポジトリから examples のコードと requirements.txt をダウンロードし、依存ライブラリをコンテナー内にインストールします。

wget https://raw.githubusercontent.com/huggingface/transformers/v4.41.2/examples/pytorch/language-modeling/requirements.txt
wget https://raw.githubusercontent.com/huggingface/transformers/v4.41.2/examples/pytorch/language-modeling/run_mlm_no_trainer.py

pip install transformers==4.41.2 -r requirements.txt

一部モジュールの Transformer Engine への置き換え

まず、既存の BERT モデルで使用されているモジュールを Transformer Engine のモジュールに置き換えていくためのパッチ用のコード (te_patch.py) を作成します。ここでは既存のライブラリに含まれるモデルを書き換えるため、モデル作成後にパッチを当てる方法を取っていますが、ご自身でモデルの実装コードを記述する場合は最初から te.Linear/te.LayerNorm を使用して実装することができます。

from io import BytesIO

from torch import nn
import transformer_engine.pytorch as te

def remove_extra_state_from_state_dict(self, destination, prefix, local_metadata):
   """
   HFのsave_pretrained()メソッドがBytesIO型を保存できないため、保存前に削除するhook
   """
   for key in list(destination.keys()):
       if key.endswith('._extra_state') and isinstance(destination[key], BytesIO):
           del destination[key]


def patch_linear_norm(model):
   for name, module in model.named_children():
       if isinstance(module, nn.Linear):
           # Tensor Coreの制約のため次元が16の倍数である必要がある
           if any(p % 16 != 0 for p in module.weight.shape):
               return
           has_bias = module.bias is not None
           te_module = te.Linear(
               module.in_features, module.out_features, bias=has_bias,
               params_dtype=module.weight.dtype
           )
           te_module.weight.copy_(module.weight)
           if has_bias:
               te_module.bias.copy_(module.bias)
           te_module._register_state_dict_hook(remove_extra_state_from_state_dict)

           setattr(model, name, te_module)

       elif isinstance(module, nn.LayerNorm):
           te_module = te.LayerNorm(
               module.normalized_shape[0], eps=module.eps,
               params_dtype=module.weight.dtype
           )
           te_module.weight.copy_(module.weight)
           te_module.bias.copy_(module.bias)
           te_module._register_state_dict_hook(remove_extra_state_from_state_dict)

           setattr(model, name, te_module)

       else:
           patch_linear_norm(module)

このコードでは、再帰的にモデルの nn.Linear と nn.LayerNorm を探索し、te.Linear と te.LayerNorm に置き換えています。この 2 つのクラスは PyTorch 標準のモジュールと互換パラメーターを持っているため、 weight/bias をコピーすることで同じように動作させることができます。

なお、state_dict の保存時に .extra_state を削除する hook を追加していますが、HuggingFace Transformers の v4.41.2 現在では、save_pretrained() メソッドで Transformer Engine の FP8 関連の metadata (BytesIO 型) を保存しようとするとエラーになってしまうため、state_dict に保存しないように対応しています。この点はいずれバージョンアップで解消されるかもしれません。

こちらの関数を作成したら、run_mlm_no_trainer.py で te_patch モジュールを読み込み、モデル初期化後に作成した関数を適用することで Transformer Engine 対応のモデルに変換することができます。

from te_patch import patch_linear_norm
with torch.no_grad():
    patch_linear_norm(model)

fp8_autocast の使用

FP8 を使用するためには、モデルの forward pass を te.fp8_autocast() のコンテキスト マネージャーでラップする必要があります。これは、FP16/BF16 における torch.amp.autocast() のようなもので、この領域で Transformer Engine のモジュールが呼び出された際は、FP8 で計算可能な部分が自動的に FP8 に変換されます。

ここでは run_mlm_no_trainer.py を書き換えていきます。まず、コマンドライン引数に FP8 の使用フラグを追加します。

parser.add_argument(
    '--fp8',
    action="store_true",
    help="Whether to use fp16 (mixed) precision instead of 32-bit",
)

次に、main() 関数内の最初 (args の parse が終わった後) に FP8 の Scaling の設定を行う recipe を作成します。

import transformer_engine.pytorch as te
from transformer_engine.common import recipe

if args.fp8:
    fp8_recipe = recipe.DelayedScaling()
else:
    fp8_recipe = None

Training と Evaluation の loop 部分を fp8_autocast の context manager で囲みます。te.fp8_autocast() は forward pass のみを囲う必要があります。loss.backward() や optimizer.step() は囲わないように注意してください。

for step, batch in enumerate(active_dataloader):
    with accelerator.accumulate(model):
        with ( # ここのwithステートメントを追加
            torch.cuda.amp.autocast(dtype=torch.bfloat16),
            te.fp8_autocast(enabled=args.fp8, fp8_recipe=fp8_recipe)
        ):
            outputs = model(**batch)
for step, batch in enumerate(eval_dataloader):
    with torch.no_grad():
        with ( # ここのwithステートメントを追加
            torch.cuda.amp.autocast(dtype=torch.bfloat16),
            te.fp8_autocast(enabled=args.fp8, fp8_recipe=fp8_recipe)
        ):
            outputs = model(**batch)

モデルや Optimizer を FP32 で初期化している場合、FP8 で計算可能な GEMM 以外の部分の計算は FP32 のままになってしまいますので、フレームワークの amp も同時に併用することをおすすめします。上記の例では torch.cuda.amp.autocast を使用して BF16 の Mixed Precision も同時に適用しています。

実行

コードの修正が完了したら早速学習を走らせてみましょう。環境変数 NVTE_DEBUG=1 をつけて実行すると、内部で FP8 が使用されているかどうかを標準出力で確認することができます。実際の学習速度やメモリ使用量は使用するモデルのサイズや GPU に依存しますので、バッチ サイズ等を調整して 16bit の時と 16bit+8bit のときで消費メモリ量や速度を比較してみてください。

python run_mlm_no_trainer.py \
   --model_name_or_path google-bert/bert-large-cased \
   --dataset_name wikitext \
   --dataset_config_name wikitext-2-raw-v1 \
   --per_device_train_batch_size 24 \
   --per_device_eval_batch_size 24 \
   --output_dir outputs/ \
   --fp8

FP8 Training は Scaling の計算、重みの cast などの処理にオーバーヘッドがあることと、高精度の重みと FP8 の重みを両方メモリに持つ必要があるため、なるべく GPU メモリに乗る限界までバッチ サイズを上げて多くのサンプルを一度に計算したほうが速度/消費メモリ量両方の観点でメリットが得られやすくなります。モデル サイズに対して GPU メモリが不足しており、バッチ サイズ 1 でぎりぎり動くようなケースでは十分な効果が得られない場合がありますのでご注意ください。

参考: te.TransformerLayer の使用

今回の例では Linear と LayerNorm を 1:1 で置き換えていくという最もシンプルで汎用的な方法を取りましたが、より高度な使い方として、te.TransformerLayer を使用して Transformer の層全体を構成する方法があります。こちらは、モデル毎にクラスの実装と weight の変換が必要となりますが、最適化された Attention Kernel、複数のモジュールの Fusion などが導入されており、より速度を向上させることができます。

TransformerLayer を使用した構成例については公式ドキュメントのチュートリアル (Llama2/3) で紹介されています。こちらの Blog でも、次回以降により高度な使い方として紹介させていただく予定です。

図 3. TransformerLayer の構造

Transformer Engine のパフォーマンス向上効果

実際に Transformer Engine を適用する前と後でどの程度パフォーマンスに差が出るかを比較してみました。No TE は Transformer Engine なしの実装、te.Linear+te.LayerNorm は今回のチュートリアルで使用した Linear と LayerNorm を置き換える実装、te.TransformerLayer は Transformer の層全体を TransformerLayer に置き換えた実装です。

図 4. Transformer Engine を適用する前後のパフォーマンス比較

Transformer Engine を適用することで、未適用の場合と比べて最大 1.45 倍の高速化とメモリ使用量の15% 程度の削減効果が確認できました。また、FP8 を使用しないケースにおいても、Transformer Engine の Module に置き換えただけでオリジナルの実装より高速で省メモリになっています。

同様に、H100 (PCIe) の GPU でも比較を行いました。

図 5. H100 (PCIe) GPU における Transformer Engine を適用する前後のパフォーマンス比較

H100 では Transformer Engine 未適用と比べて最大 1.57 倍の高速化、21% のメモリ節約となっており、L4 より効果が大きくなっています。これは、バッチ サイズを大きく取れたことにより FP8 の計算/メモリオーバーヘッド以上の削減効果が得られたこと、Hopper と Ada Lovelace で kernel の実装が異なることなどが要因として考えられます。

Transformer Engine 統合済みライブラリの利用

今回は手動でモデルを書き換えて Transformer Engine の適用を行いましたが、Transformer Engine が統合されたライブラリを使うことで、モデルの書き換えや fp8_autocast の適用を簡単に行うことが出来ます。

HuggingFace Accelerate や PyTorch Ligtning では、precision を指定するオプションによって Transformer Engine の適用を自動で行ってくれます。詳細については、各ライブラリのドキュメントをご覧ください。

(HuggingFace Accelerate については、v0.29.3 現在、FP16/BF16 と FP8 を同時に適用することができません。そのため、本記事のチュートリアルで使用している Transformers の examples でも accelerate を使うことはできますが、BF16/FP8 の Mixed Precision の適用は手動で行っています。)

NeMo Framework (Megatron-Core backend) をお使いの場合は、Transformer Engine を完全に組み込んだモデル実装が含まれているため、CLI または YAML のオプションで Transformer Engine と FP8 を有効にするだけで最適化された高速な実装を使用することができます。Tensor Parallel/Pipeline Parallel などの分散学習の仕組みにも対応していますので、LLM のトレーニングを行う際は NeMo Framework の使用をおすすめします。

python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py \
  --config-path /opt/NeMo-Megatron-Launcher/launcher_scripts/conf/training/gpt3 \
  --config-name 1b_improved \
  ...
  model.transformer_engine=True \
  model.fp8=True
model:
  ## Transformer Engine
  transformer_engine: True
  fp8: True               # enables fp8 in TransformerLayer forward
  fp8_e4m3: False         # sets fp8_format = recipe.Format.E4M3
  fp8_hybrid: True        # sets fp8_format = recipe.Format.HYBRID
  fp8_margin: 0           # scaling margin
  fp8_interval: 1         # scaling update interval
  fp8_amax_history_len: 1024 # Number of steps for which amax history is recorded per tensor
  fp8_amax_compute_algo: max # 'most_recent' or 'max'. Algorithm for computing amax from history
  fp8_wgrad: True
  ub_tp_comm_overlap: False
  tp_comm_atomic_ag: False
  tp_comm_atomic_rs: False

NeMo Framework を使用した LLM のファインチューニングのチュートリアルはこちらに公開されています。

まとめ

Transformer Engine を用いて、BERT モデルの Linear/LayerNorm 層を置き換えて FP8 Training を行う方法をご紹介しました。Hopper/Ada Lovelace の性能をフルに引き出すために、是非 FP8 を活用していただければと思います。

次回は応用編として、te.TransformerLayer を用いたより高速な実装や、よりパフォーマンスを高めるための機能の使い方をご紹介したいと思います。

関連情報

Tags

Conversational AI / NLP | Generative AI | General | Hopper | Tutorial | LLMs | technical walkthrough | Training AI Models | Transformers

About the Authors

About Kotaro Yamamoto
山本 小太郎は、NVIDIA のシニア Deep Learning ソリューション アーキテクトとして、特に AI/ML 分野において、顧客への技術支援を行っています。NVIDIA に入社する前は、システム開発企業でクラウド基盤の開発や自然言語処理技術の業務適用、AI スタートアップで文字認識モデルの開発と推論システムの最適化に従事していました。現在は NVIDIA NIM を活用した LLM の推論などを中心に、AI ワークロードにおける GPU の活用を推進しています。