transformerengine

Transformer Engine（TE）は、NVIDIA GPU上のTransformerモデルを高速化するためのNVIDIAライブラリです。中心的な機能は低精度計算であり、Hopper、Ada、Blackwell GPUでの8ビット浮動小数点（FP8）に加え、BlackwellではMXFP8およびNVFP4フォーマットをサポートし、学習と推論の両方で性能向上とメモリ使用量の削減を目的としています。また、Ampere以降のアーキテクチャでFP16およびBF16にわたる最適化も行います。 このライブラリは、一般的なTransformerアーキテクチャ向けに高度に最適化されたビルディングブロックと、フレームワーク固有のコードで使用できる自動混合精度のようなAPIを提供します。フレームワークに依存しないC++ APIも含まれており、他の深層学習ライブラリがTransformer向けのFP8サポートを追加できるようになっています。TEモジュールは内部的にFP8学習に必要なスケーリングファクターと関連値を維持し、ユーザーの混合精度ワークフローを簡素化します。 プロジェクトが挙げるハイライトには、FP8サポートを備えたTransformerレイヤーを構築するための使いやすいモジュール、融合カーネル最適化、Hopper、Ada、BlackwellでのFP8サポート、BlackwellでのMXFP8およびNVFP4サポート、Ampere以降でのFP16/BF16にわたる最適化が含まれます。 使用例はPyTorchとJAX/Flax向けに示されています。PyTorchでは、ユーザーはtransformer_engine.pytorchをインポートし、選択したFP8フォーマットでDelayedScalingなどのレシピを作成し、フォワードパスをte.autocastでラップします。JAX/Flaxでは、te_flaxモジュールとレシピを用いた同様のautocastの使用法が示されています。より完全なチュートリアルとして、はじめにガイドへのリンクがあります。 インストールオプションには、NGC Dockerコンテナ（推奨）、PyTorchおよび/またはJAX向けのextras付きpipパッケージ、PyTorch統合向けのconda-forgeパッケージ、ソースビルドが含まれます。システム要件では、Blackwell、Hopper、Grace Hopper/Blackwell、Ada、Ampereハードウェア、公式OSとしてのLinux（限定的なWSL2サポート）、CUDA 12.1以上（Blackwellでは12.8以上）、cuDNN 9.12以上、C++17対応のGCC 9以上またはClang 10以上、Python 3.12推奨が挙げられています。FP8機能にはコンピュートケイパビリティ8.9以上が必要です。CUDA_PATH、CUD_PATH、CXX、NVTE_FRAMEWORK、MAX_JOBS、NVTE_CUDA_ARCHSなどのビルド関連環境変数が文書化されています。 READMEでは、PyTorchにおけるFlashAttention-2およびFlashAttention-3のサポートが説明されており、両方が存在する場合はFlashAttention-3が優先され、FlashAttention-2のコンパイルはメモリを大量に消費する可能性があることが注記されています。トラブルシューティングセクションでは、ABI互換性のインポートエラー、ヘッダーやライブラリの欠落、ビルドリソースの問題、詳細なビルドログ、cuDNNサブライブラリの読み込み失敗を含む仮想環境の問題、JAX FFI登録エラーが扱われています。 v1.7の破壊的変更が文書化されています。PyTorchのパディングマスクの定義が変更され、Trueは位置を含めるのではなくマスクアウトすることを意味するようになり、フレームワーク間でマスクのセマンティクスが統一されました。収束に関する注記では、テストされた構成においてFP8とMXFP8はBF16の学習損失曲線と有意な差を示さず、下流のLLMタスクで検証されたことが述べられており、Mosaic Composer、Alibaba Pai、Megatron CoreなどのフレームワークにわたるMPT-1.3B、Llama2-7B、LLM-8B、MPT-13B、MoE-16B、Llama2-70Bなどのモデルが挙げられています。 記載されている統合には、DeepSpeed、Hugging Face Accelerate、Lightning、MosaicML Composer、NVIDIA JAX Toolbox、NVIDIA Megatron-LM、NVIDIA NeMo Megatron Bridge、Amazon SageMaker Model Parallel Library、Levanter、GPT-NeoX、Hugging Face Nanotronが含まれます。このプロジェクトはApache-2.0ライセンスであり、CONTRIBUTINGガイドを通じた貢献を歓迎しています。