NVIDIA Transformer Engine : noyaux fusionnés, FP8 et BF16 pour l'entraînement des transformers

NVIDIA Transformer Engine est une bibliothèque qui accélère l'entraînement des modèles transformers sur les GPU Ampere et plus récents. Publié le 1er août 2026 sur MarkTechPost, un tutoriel montre comment intégrer TE à un modèle de type GPT, configurer la quantification FP8 et mesurer les gains de vitesse et de mémoire.

Qu'est-ce que NVIDIA Transformer Engine ?

NVIDIA Transformer Engine est la bibliothèque officielle d'optimisation des transformers pour les GPU de génération Ampere et supérieure. Elle fournit des remplaçants directs (drop-in) des couches PyTorch standard : `te.Linear`, `te.LayerNorm`, `te.LayerNormLinear`, `te.LayerNormMLP`, ainsi que le bloc complet `te.TransformerLayer`. Tous les composants utilisent BF16 par défaut ; les FP8 tensor cores sont disponibles sur les cartes dont la compute capability est ≥ 8.9.

Principaux composants de la bibliothèque :

• `te.Linear` — couche linéaire compatible avec les FP8 tensor cores
• `te.LayerNormLinear` — combine LayerNorm et projection dans un seul noyau
• `te.LayerNormMLP` — bloc MLP fusionné avec normalisation
• `te.TransformerLayer` — bloc transformer complet : attention, FFN, normalisation
• Repli automatique sur du PyTorch pur lorsque le GPU ne prend pas en charge TE

Comment fonctionnent le FP8 et DelayedScaling ?

Le mode FP8 s'active via la recette `DelayedScaling` du module `transformer_engine.common.recipe`. Cette recette contrôle l'historique des maxima des tenseurs (`amax_history_len=16`), l'algorithme de calcul (`amax_compute_algo="max"`) et le format hybride E4M3/E5M2 : le premier assure la précision pour les activations, le second offre une plage dynamique étendue pour les gradients.

La particularité du « delayed scaling » : la bibliothèque accumule des statistiques sur 16 itérations avant de mettre à jour l'échelle du tenseur. Cela stabilise l'entraînement durant les premières étapes, lorsque la distribution des activations n'est pas encore établie.

«

Transformer Engine combine des noyaux fusionnés, la mise à l'échelle FP8 et des chemins d'exécution optimisés pour Ampere, ce qui en fait un choix pratique pour l'entraînement de grands modèles », indique la documentation officielle de NVIDIA.

Dans le tutoriel, les auteurs construisent un modèle compact de type GPT, `MiniGPT_TE` : quatre blocs `te.TransformerLayer`, un espace de dimension 768, 12 têtes d'attention, un vocabulaire de 96 tokens et une longueur de séquence de 256. Le modèle s'entraîne sur des données synthétiques déterministes. Les mesures montrent que le mode FP8 sur un GPU de classe H100 réduit la consommation maximale de VRAM et la durée d'un pas d'entraînement par rapport à BF16, tandis que la précision reste comparable — les auteurs la vérifient par génération autorégressive après l'entraînement.

Quels GPU sont requis ?

Transformer Engine ne fonctionne que sur les GPU dont la compute capability est ≥ 8.0 (architecture Ampere). Les FP8 tensor cores sont disponibles à partir de CC ≥ 8.9 : Ada Lovelace (RTX 40xx) et Hopper (H100/H200). Sur les cartes T4 et plus anciennes, la bibliothèque se replie sur du PyTorch pur, sans noyaux fusionnés.

• Minimum pour les noyaux TE : Ampere, CC ≥ 8.0 — A10, A100, RTX 30xx
• Minimum pour le FP8 : Ada Lovelace / Hopper, CC ≥ 8.9
• Installation : `pip install transformer_engine[pytorch]`
• Dépendances : CUDA 11.8+, PyTorch 2.0+

Selon les auteurs du tutoriel, un GPU A100 ou L4 dans Google Colab suffit pour reproduire l'ensemble ; sur T4, seul le mode de repli, sans noyaux fusionnés, est disponible.

Ce que cela change

NVIDIA Transformer Engine abaisse le seuil d'entrée dans l'entraînement en basse précision : plutôt que d'implémenter soi-même la quantification FP8, il suffit de remplacer les couches PyTorch standard par leurs équivalents TE. Pour les équipes qui entraînent de grands modèles sur des clusters de GPU, c'est un moyen concret de réduire la consommation de mémoire et la durée des itérations sans retoucher l'architecture.

Questions fréquentes

Faut-il un H100 pour utiliser Transformer Engine ?

Non. NVIDIA Transformer Engine fonctionne sur tout GPU dont la compute capability est ≥ 8.0, y compris les A100 et RTX 3090. Les FP8 tensor cores ne sont disponibles qu'à partir de CC ≥ 8.9 — sur H100 et RTX 4090 ; sur les autres cartes, la bibliothèque bascule automatiquement sur BF16.

Comment installer Transformer Engine ?

L'installation se fait via pip : `pip install transformer_engine[pytorch]`. CUDA 11.8 ou supérieur et PyTorch 2.0+ sont requis.