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 で浮動小數點數を表現するデータ フォーマットです。
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 に変換するという手法を取ります。
推論時における 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@stableNGC のコンテナーをお使いの場合は、各種深層學習フレームワークのコンテナーに Transformer Engine も含まれています。手動インストールの場合 Custom kernel などのビルドも必要となり依存関係が多く時間もかかるので、NGC コンテナーをお使いいただくことを推奨します。
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:24.05-py3Transformer 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.LayerNormnn.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 を行う例をご紹介します。
本チュートリアルでは以下のバージョンを使用します
- コンテナー: nvcr.io/nvidia/pytorch:24.05-py3
- ライブラリ: transformers==4.41.2 + requirements.txt
- GPU: NVIDIA L4 24GB
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 = NoneTraining と 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/ \
--fp8FP8 Training は Scaling の計算、重みの cast などの処理にオーバーヘッドがあることと、高精度の重みと FP8 の重みを両方メモリに持つ必要があるため、なるべく GPU メモリに乗る限界までバッチ サイズを上げて多くのサンプルを一度に計算したほうが速度/消費メモリ量両方の観點でメリットが得られやすくなります。モデル サイズに対して GPU メモリが不足しており、バッチ サイズ 1 でぎりぎり動くようなケースでは十分な効果が得られない場合がありますのでご注意ください。
參考: te.TransformerLayer の使用
今回の例では Linear と LayerNorm を 1:1 で置き換えていくという最もシンプルで汎用的な方法を取りましたが、より高度な使い方として、te.TransformerLayer を使用して Transformer の層全體を構成する方法があります。こちらは、モデル毎にクラスの実裝と weight の変換が必要となりますが、最適化された Attention Kernel、複數のモジュールの Fusion などが導入されており、より速度を向上させることができます。
TransformerLayer を使用した構成例については公式ドキュメントのチュートリアル (Llama2/3) で紹介されています。こちらの Blog でも、次回以降により高度な使い方として紹介させていただく予定です。
Transformer Engine のパフォーマンス向上効果
実際に Transformer Engine を適用する前と後でどの程度パフォーマンスに差が出るかを比較してみました。No TE は Transformer Engine なしの実裝、te.Linear+te.LayerNorm は今回のチュートリアルで使用した Linear と LayerNorm を置き換える実裝、te.TransformerLayer は Transformer の層全體を TransformerLayer に置き換えた実裝です。
Transformer Engine を適用することで、未適用の場合と比べて最大 1.45 倍の高速化とメモリ使用量の15% 程度の削減効果が確認できました。また、FP8 を使用しないケースにおいても、Transformer Engine の Module に置き換えただけでオリジナルの実裝より高速で省メモリになっています。
同様に、H100 (PCIe) の GPU でも比較を行いました。
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 の適用は手動で行っています。)
- Huggingface Accelerate: Low Precision Training Methods
- PyTorch Lightning: N-BIT PRECISION (INTERMEDIATE)
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=Truemodel:
## 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: FalseNeMo Framework を使用した LLM のファインチューニングのチュートリアルはこちらに公開されています。
まとめ
Transformer Engine を用いて、BERT モデルの Linear/LayerNorm 層を置き換えて FP8 Training を行う方法をご紹介しました。Hopper/Ada Lovelace の性能をフルに引き出すために、是非 FP8 を活用していただければと思います。
次回は応用編として、te.TransformerLayer を用いたより高速な実裝や、よりパフォーマンスを高めるための機能の使い方をご紹介したいと思います。
