NVIDIA Transformer Engine: фьюзед-ядра, FP8 и BF16 для обучения трансформеров

NVIDIA Transformer Engine — une bibliothèque pour accélérer l'entraînement de modèles transformer sur les GPU Ampere et plus récents. Un tutoriel publié le 1er août 2026 sur MarkTechPost montre comment connecter TE à un modèle de type GPT, configurer la quantification FP8 et mesurer les gains en vitesse et en mémoire.

Qu'est-ce que NVIDIA Transformer Engine

NVIDIA Transformer Engine est une bibliothèque officielle pour optimiser les transformers sur les GPU de niveau Ampere et supérieur. Elle fournit des remplacements drop-in pour les couches PyTorch standard : `te.Linear`, `te.LayerNorm`, `te.LayerNormLinear`, `te.LayerNormMLP` et un bloc complet `te.TransformerLayer`. Tous les composants fonctionnent avec BF16 par défaut ; les tensor cores FP8 sont disponibles sur les cartes avec une compute capability ≥ 8.9.

Composants clés de la bibliothèque :

• `te.Linear` — couche linéaire avec support des tensor cores FP8
• `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 PyTorch pur sur les GPU sans support TE

Comment Fonctionnent FP8 et DelayedScaling

Le mode FP8 est activé via la recette `DelayedScaling` de `transformer_engine.common.recipe`. La recette gère l'historique des amax des tenseurs (`amax_history_len=16`), l'algorithme de calcul (`amax_compute_algo="max"`) et le format hybride E4M3/E5M2 : le premier format assure la précision pour les activations, le second offre une plage dynamique étendue pour les gradients.

La caractéristique clé du delayed scaling est l'accumulation de statistiques sur 16 itérations avant la mise à jour de l'échelle du tenseur. Cela stabilise l'entraînement lors des premières étapes, lorsque la distribution des activations n'est pas encore établie.

«

Transformer Engine combine des noyaux fusionnés, un 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`, une dimension cachée de 768, 12 têtes d'attention, un vocabulaire de 96 tokens et une longueur de séquence de 256. Le modèle est entraîné 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 de mémoire VRAM maximale et le temps par étape par rapport à BF16, tandis que la précision reste comparable — les auteurs le vérifient via une génération autorégressive après l'entraînement.

Exigences GPU

Transformer Engine ne fonctionne que sur les GPU avec une compute capability ≥ 8.0 (architecture Ampere). Les tensor cores FP8 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 PyTorch pur sans noyaux fusionnés.

• Minimum pour les noyaux TE : Ampere, CC ≥ 8.0 — A10, A100, RTX 30xx
• Minimum pour 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 est suffisant pour une reproduction complète ; sur T4, seul le mode de repli sans noyaux fusionnés est disponible.

Ce que Cela Signifie

NVIDIA Transformer Engine abaisse le seuil d'entrée dans l'entraînement à faible précision : au lieu d'implémenter manuellement la quantification FP8, il suffit de remplacer les couches PyTorch standard par leurs équivalents TE. Pour les équipes entraînant de grands modèles sur des clusters GPU, c'est un moyen pratique de réduire la consommation de mémoire et le temps d'itération sans retravailler l'architecture.

Questions Fréquentes

Faut-il un H100 pour Utiliser Transformer Engine ?

Non. NVIDIA Transformer Engine fonctionne sur n'importe quel GPU avec une compute capability ≥ 8.0, y compris l'A100 et la RTX 3090. Les tensor cores FP8 ne sont disponibles qu'avec CC ≥ 8.9 — sur le H100 et la RTX 4090 ; sur les autres cartes, la bibliothèque bascule automatiquement sur BF16.

Comment Installer Transformer Engine ?

Installation via pip : `pip install transformer_engine[pytorch]`. CUDA 11.8 ou supérieur et PyTorch 2.0+ sont requis.