NVIDIA/TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
What it solves
Transformer Engine (TE) は、Transformer モデル（LLM、MoE、マルチモーダルモデルなど）を数百億のパラメータにスケールアップさせる際の、高いメモリと計算需要を満たす課題を解決します。低精度数値フォーマットを活用することで、訓練および推論の両方においてメモリ利用率を低減し、スループットを向上させます。
How it works
TE は、Transformer アーキテクチャ向けに高度に最適化されたビルディングブロックと融合カーネル (fused kernels) を提供します。自動混合精度 API を実装しており、ユーザーが既存のフレームワークのコードに低精度フォーマットをシームレスに統合できる仕組みを提供します。
主な技術的能力は以下の通りです：
- Low-Precision Support: Hopper, Ada, および Blackwell GPU での 8 ビット浮動小数点 (FP8) のネイティブサポート、および Blackwell GPU での MXFP8 と NVFP4 によるさらなる効率向上。
- Framework Integration: PyTorch と JAX (Flax) 用の Python API、および他のディープラーニング・ライブラリと統合するためのフレームワークに依存しない C++ API を提供します。
- Automated Scaling: 内部的に FP8 訓練に必要なスケーリング・ファクターを管理し、ユーザーの混合精度プロセスを簡素化します。
- Optimized Kernels: Mixture-of-Experts (MoE) や、さまざまな並列化戦略 (tensor, sequence, context) に向けた融合操作と最適化を含みます。
Who it’s for
NVIDIA GPU (Ampere アーキテクチャ以降) を使用し、ハードウェアの効率と訓練速度を最大化したい、大規模な Transformer モデルを構築・訓練する研究者やエンジニア向けです。
Highlights
- Broad Precision Support: 最新の NVIDIA GPU で FP8, MXFP8, NVFP4 をサポートし、Ampere 以降で FP16/BF16 をサポートします。
- Easy Integration: FP8 サポートを備えた Transformer レイヤーを構築するための、すぐに使えるモジュールを提供します。
- C++ API: ディープラーニング・ライブラリの開発者向けに、フレームワークに依存しない C++ ライブラリを提供します。
- Performance Optimizations: 融合カーネル、および PyTorch における FlashAttention-2 と FlashAttention-3 のサポートを特徴としています。