Entraînement en précision mixte avec PyTorch (AMP)

L’entraînement en précision mixte est l’une des techniques d’optimisation des performances les plus importantes en apprentissage profond. En utilisant simultanément des nombres à virgule flottante FP32 (simple précision) et FP16 (demi-précision) pour les calculs, il peut améliorer considérablement la vitesse d’entraînement et réduire l’utilisation de la mémoire avec une perte quasi nulle de précision du modèle. Cette section détaille la technique de précision mixte automatique (AMP) dans PyTorch.

Versions applicables :Le code de cet article est écrit avec PyTorch 1.6+torch.cuda.ampAPI. PyTorch 2.4+ recommande d’utilisertorch.amp.autocastettorch.amp.GradScaler, et l’utilisation est globalement la même ; les différences seront signalées dans l’article.

1. Fondamentaux de l’entraînement en précision mixte

1.1 Pourquoi avons-nous besoin de la précision mixte

L’entraînement des modèles d’apprentissage profond implique un grand nombre d’opérations matricielles. Le FP32 traditionnel (virgule flottante 32 bits) offre une grande précision de calcul, mais occupe plus de mémoire et est plus lent. Le FP16 (virgule flottante 16 bits) est plus rapide et utilise moins de mémoire, mais sa plage numérique est plus petite et il est sujet aux problèmes de sous-dépassement de gradient.

L’idée centrale de l’entraînement en précision mixte est la suivante :Utiliser FP32 pour les opérations qui nécessitent une haute précision, et FP16 pour les opérations qui n’en nécessitent pas.. Ainsi, vous pouvez bénéficier de l’avantage de vitesse du FP16 tout en évitant les problèmes de précision.

La figure suivante montre les différences de disposition des bits des trois formats à virgule flottante—Les bits d'exposantdéterminent la plage numérique.Les bits de mantissedéterminent la précision :

1.2 Avantages de la précision mixte

Métrique	Amélioration	Description
Vitesse d'entraînement	Amélioration de 2-3x	Dépend de la prise en charge des Tensor Core du GPU
Utilisation de la mémoire	Réduit d'environ 50%	Activations et résultats intermédiaires stockés en FP16
Bande passante mémoire	Réduit d'environ 50%	Une taille de données plus petite signifie moins de transfert de données
Surcharge de communication	Réduit d'environ 50%	Volume de transfert des gradients réduit de moitié en entraînement distribué

1.3 Principe d'accélération des Tensor Core

Le Tensor Core de NVIDIA est une unité matérielle dédiée aux opérations matricielles. Il peut effectuer une opération de multiplication-accumulation matricielle 4×4 (D = A × B + C) en un seul cycle d'horloge, ce qui constitue la principale source d'accélération de l'entraînement en FP16. Comparé aux cœurs CUDA ordinaires qui nécessitent plusieurs instructions pour accomplir la même opération, le Tensor Core la compresse en une seule instruction.

Les GPU qui prennent en charge le Tensor Core incluent :

Architecture Volta(V100) — Tensor Core de première génération, prend en charge uniquement FP16
Architecture Turing(série RTX 20) — Prend en charge FP16 / INT8 / INT4
Architecture Ampere(série RTX 30, A100) — Ajout de la prise en charge de BF16 / TF32
Architecture Ada Lovelace(série RTX 40) — Ajout de la prise en charge de FP8
Architecture Hopper(H100) — Ajout du Transformer Engine FP8

Les GPU RTX grand public prennent également en charge le Tensor Core ; par exemple, les modèles RTX 3060 et supérieurs peuvent bénéficier de l'accélération AMP.

2. Utilisation de base de l'AMP PyTorch

2.1 autocast et GradScaler

L'API AMP de PyTorch comprend principalement deux composants essentiels :

autocast: Un gestionnaire de contexte qui bascule automatiquement les opérations dans sa portée en FP16 (les opérations sensibles à la précision reviennent automatiquement à FP32)
GradScaler: Ajuste dynamiquement le facteur d'échelle du gradient pour amplifier les gradients FP16 et éviter le sous-dépassement (nécessaire uniquement pour FP16 ; généralement pas nécessaire pour BF16)

La figure suivante montre le flux de données d'une étape d'entraînement complète avec AMP :

2.2 Exemple d'utilisation de base

Exemple

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler
# Approche recommandée pour PyTorch 2.4+ :
# from torch.amp import autocast, GradScaler

# Vérifier si CUDA est disponible
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"BF16 pris en charge : {torch.cuda.is_bf16_supported()}")

# ── Définition du modèle ──────────────────────────────────────
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, 256),
nn.ReLU(),
nn.Linear(256, 10)
)

def forward(self, x):
return self.net(x)

model = SimpleModel().to(device)

# Fonction de perte et optimiseur
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ── Composants clés pour l’entraînement en précision mixte ──────────────────────────
# GradScaler : met à l’échelle la perte pour éviter le sous-dépassement de gradient FP16
scaler = GradScaler()

# Boucle d'entraînement
def train_epoch_amp(model, loader, criterion, optimizer, scaler, device):
model.train()
total_loss = 0
correct = 0
total = 0

for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)

optimizer.zero_grad()

# ── Cœur : utiliser le gestionnaire de contexte autocast ──────
# Les opérations dans la portée autocast utilisent automatiquement FP16
# Les opérations sensibles à la précision (par ex., softmax, loss) reviennent automatiquement à FP32
with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)

# ── Utiliser le scaler pour la rétropropagation ─────────────
# 1. Mettre à l'échelle la perte (multiplier par scale_factor)
# 2. Rétropropagation (calculée sur les gradients mis à l'échelle)
# 3. scaler.step annule en interne la mise à l'échelle des gradients et vérifie la présence d'Inf/NaN
scaler.scale(loss).backward()

# Mettre à jour les paramètres
scaler.step(optimizer)

# Mettre à jour le facteur d'échelle du scaler
scaler.update()

# Statistiques
total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)

return total_loss / total, correct / total

# Données simulées
train_loader = [
(torch.randn(32, 128), torch.randint(0, 10, (32,))) for _ in range(10)
]

# Démarrer l'entraînement
for epoch in range(3):
loss, acc = train_epoch_amp(
model, train_loader, criterion, optimizer, scaler, device
)
print(f"Époque {epoch+1}: Perte={loss:.4f}, Exactitude={acc:.4f}")

print("Entraînement en précision mixte terminé !")

Note de migration de l'API :PyTorch 2.4 marqueratorch.cuda.amp.autocastcomme obsolète ; il est recommandé d'utilisertorch.amp.autocast('cuda'). L'utilisation est exactement la même ; seul le chemin d'importation diffère.

2.3 Choisir entre BF16 et FP16

Les GPU modernes prennent en charge deux formats de demi-précision. Leur différence fondamentale réside dans la stratégie d'allocation des bits d'exposant et des bits de mantisse :

Format	Bits d'exposant	Bits de mantisse	Plage numérique	Precision	Cas d'utilisation
FP16	5 bits	10 bits	±65504	Plus élevé (~3.3 chiffres)	Nécessite une compatibilité avec les GPU plus anciens (V100/RTX 20 series)
BF16	8 bits	7 bits	±3.4×10³⁸	Plus faible (~2.4 chiffres)	Pour la stabilité (A100/RTX 30 series+)

Recommandation : si le matériel prend en charge BF16,préférez BF16. Sa plage numérique est exactement la même que celle de FP32, de sorte que les problèmes de débordement se produisent rarement pendant l’entraînement, et GradScaler n’est pas nécessaire.

Exemple

# ── approche BF16 (PyTorch 1.10+) ────────────────────
# BF16 ne nécessite pas GradScaler car sa plage numérique est la même que FP32

from torch.cuda.amp import autocast

# Méthode 1 : Spécifier dtype dans autocast
with autocast(device_type='cuda', dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)

# Méthode 2 : Activer globalement BF16 par défaut (si le matériel le prend en charge)
# torch.backends.cuda.matmul.allow_bf16_reduced_precision = True
# torch.backends.cudnn.allow_bf16_reduced_precision = True

# ── Vérifier la prise en charge matérielle ─────────────────────────────
print(f"Le matériel prend en charge BF16 : {torch.cuda.is_bf16_supported()}")
print(f"Le matmul actuel autorise BF16 : {torch.backends.cuda.matmul.allow_bf16}")
print(f"Le cuDNN actuel autorise BF16 : {torch.backends.cudnn.allow_bf16}")

# ── Exemple complet d'entraînement BF16 (sans GradScaler nécessaire) ──────────
scaler_bf16 = None # BF16 n'a pas besoin de scaler

for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()

with autocast(device_type='cuda', dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)

# Rétropropager directement, aucun scaler nécessaire
loss.backward()
optimizer.step()

3. Conseils avancés et optimisation

3.1 Mise à l'échelle dynamique de la perte

Le mécanisme central de GradScaler estla mise à l'échelle dynamique de la perte—il ajuste automatiquement le facteur d'échelle en fonction de l'état de l'entraînement, empêchant le sous-dépassement des gradients tout en réduisant progressivement la mise à l'échelle lorsque l'entraînement se stabilise afin de minimiser la perte de précision :

GradScaler gère automatiquement la boucle de rétroaction ci-dessus ; vous n'avez qu'à ajuster finement son comportement via des paramètres :

Exemple

from torch.cuda.amp import GradScaler

# Paramètres personnalisés de GradScaler
scaler = GradScaler(
init_scale=2**16, # Facteur d'échelle initial, par défaut 65536
growth_factor=2.0, # Multiplicateur de croissance du facteur d'échelle, par défaut 2.0
backoff_factor=0.5, # Multiplicateur de repli du facteur d'échelle, par défaut 0.5
growth_interval=2000, # Nombre d'étapes consécutives réussies avant l'augmentation
enabled=True # Indique si activé (peut être basculé dynamiquement)
)

# Flux de travail :
# 1. Échelle initiale = 65536
# 2. Si une étape contient Inf/NaN → échelle × 0.5 (réduire), ignorer cette étape
# 3. Si 2000 étapes consécutives ne contiennent pas de Inf/NaN → échelle × 2.0 (augmenter)
# 4. Toujours trouver l'échelle maximale utilisable dans une plage sûre

# Voir le facteur d'échelle actuel
print(f"Facteur d'échelle actuel : {scaler.get_scale()}")

# Déterminer si le scaler considère la dernière étape comme réussie
print(f"Réussi récemment : {scaler._found_inf.item() == 0}")

3.2 AMP avec accumulation de gradient

L'accumulation de gradient simule une taille de lot plus grande avec une mémoire limitée. Lors de l'utilisation d'AMP, notez que la perte de chaque sous-lot doit être divisée par le nombre d'étapes d'accumulation ; sinon, le gradient final sera amplifié :

Exemple

# ── Accumulation de gradient + précision mixte ────────────────
accumulation_steps = 8

scaler = GradScaler()
model.train()

for batch_idx, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device)

with autocast(device_type='cuda'):
outputs = model(inputs)
# Point clé : diviser la perte de chaque sous-lot par le nombre d'étapes d'accumulation
# Ainsi, le gradient accumulé est équivalent au gradient moyen d'un grand lot
loss = criterion(outputs, labels) / accumulation_steps

# Accumuler les gradients mis à l'échelle (remarque : ne pas effacer ici)
scaler.scale(loss).backward()

# Mettre à jour les paramètres tous les accumulation_steps lots
if (batch_idx + 1) % accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

# ── Gestion des gradients à la fin lorsqu'ils représentent moins d'un cycle d'accumulation complet ──────────
remainder = len(train_loader) % accumulation_steps
if remainder != 0:
# Les gradients accumulés doivent être remis à l'échelle en fonction du nombre réel d'étapes
# Approche simple : exécuter quand même l'étape ; les gradients seront légèrement plus grands, mais l'impact est généralement minimal
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

3.3 AMP en validation et inférence

AMP est également recommandé pour la validation et l'inférence. Comme la rétropropagation et la mise à l'échelle des gradients ne sont pas nécessaires, le code est plus simple :

Exemple

# ── Utilisation d'autocast pour l'inférence ──────────────────────────
@torch.no_grad()
def inference_amp(model, inputs, device):
model.eval()
# L'inférence utilise FP16, aucun GradScaler n'est nécessaire
with autocast(device_type='cuda'):
outputs = model(inputs.to(device))
return outputs

# ── inference_mode est plus rapide que no_grad ──────────────
# inference_mode désactive davantage la surcharge de traçage, adapté à l'inférence pure
@torch.inference_mode()
def inference_fast(model, inputs, device):
model.eval()
with autocast(device_type='cuda'):
outputs = model(inputs.to(device))
return outputs

# ── Exemple d'inférence par lots ───────────────────────────────
def batch_inference(model, dataloader, device):
model.eval()
all_outputs = []

with torch.inference_mode():
for inputs in dataloader:
with autocast(device_type='cuda'):
outputs = model(inputs.to(device))
all_outputs.append(outputs.cpu())

return torch.cat(all_outputs, dim=0)

4. Problèmes courants et solutions

4.1 Entraînement instable

Problème	Cause	Solution
La perte devient NaN / Inf	Débordement de gradient (échelle trop grande ou taux d'apprentissage trop élevé)	Réduire`init_scale`, ajouter un écrêtage du gradient, abaisser le taux d'apprentissage
La perte ne diminue pas / oscille	Sous-dépassement de gradient (échelle trop petite, gradients tronqués à 0)	Augmentez`init_scale`, vérifier les statistiques des gradients, essayer BF16
Baisse significative de la précision	Certaines couches (par ex., LayerNorm, Softmax) sont sensibles à FP16	Conserver manuellement ces couches en FP32 (voir 4.2)
Instabilité en fin d'entraînement	Incompatibilité d'échelle en raison de changements dans les plages de valeurs des paramètres du modèle	Augmenter de manière appropriée`growth_interval`, rendant les ajustements d'échelle plus conservateurs

4.2 Contrôle manuel de la précision

Certaines opérations subissent une perte de précision importante en FP16 et doivent être spécifiées manuellement pour utiliser FP32. L'autocast de PyTorch dispose d'une protection intégrée pour ces opérations, mais les opérations personnalisées peuvent nécessiter une gestion manuelle :

Exemple

# ── Méthode 1 : Encapsuler une fonction de perte qui force le FP32 ────────
class FP32Loss(nn.Module):
"""Encapsulation de fonction de perte qui force le calcul en FP32"""
def __init__(self, base_criterion):
super().__init__()
self.base_criterion = base_criterion

def forward(self, input, target):
# Convertir explicitement en FP32, en dehors de l'influence d'autocast
return self.base_criterion(input.float(), target.float())

# ── Méthode 2 : désactiver localement dans autocast ──────────────
with autocast(device_type='cuda'):
outputs = model(inputs)

# Pour les opérations sensibles à la précision, désactiver temporairement autocast
with autocast(enabled=False):
loss = criterion(outputs.float(), labels)

# ── Méthode 3 : conserver des couches spécifiques en FP32 dans le modèle ────────
class CustomModel(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64), # BN est sensible à la précision
nn.ReLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
)
self.classifier = nn.Linear(128, 10)

def forward(self, x):
x = self.features[0](x) # Conv : FP16 (contrôlé par autocast)
x = self.features[1](x.float()) # BN : forcer FP32
x = self.features[2](x) # ReLU : FP16
x = self.features[3](x)
x = self.features[4](x.float()) # BN : forcer FP32
x = self.features[5](x)
x = x.mean(dim=[2, 3]) # Pooling moyen global
x = self.classifier(x)
return x

Astuce :L'autocast de PyTorch conserve automatiquement les opérations suivantes en FP32 : softmax, log_softmax, cross_entropy, layer_norm, batch_norm, etc. Généralement, seuls les opérateurs personnalisés nécessitent une gestion manuelle.

4.3 Écrêtage des gradients et AMP

L'écrêtage des gradients doit êtrescaler.step()avant, etscaler.unscale_()après exécution. En effet, l'écrêtage doit porter surles valeurs réelles du gradientplutôt que sur le gradient mis à l'échelle :

Exemple

# ── Ordre correct d'écrêtage des gradients ──────────────────────────
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()

with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)

# Étape 1 : mettre à l'échelle la perte et rétropropager
scaler.scale(loss).backward()

# Étape 2 : Annuler la mise à l'échelle des gradients (restaurer les gradients de l'échelle aux valeurs réelles)
# Doit être appelé avant l'écrêtage et l'étape
scaler.unscale_(optimizer)

# Étape 3 : Écrêter les gradients réels
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Étape 4 : Mettre à jour les paramètres (vérifie en interne si les gradients contiennent Inf/NaN)
scaler.step(optimizer)

# Étape 5 : Mettre à jour le facteur d'échelle
scaler.update()

# ── Erreur courante : écrêter sans unscale_ ──────────────
# Faux ! L'écrêtage est appliqué aux gradients amplifiés, et le seuil est affecté par l'échelle
# scaler.scale(loss).backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # ← Faux
# scaler.step(optimizer)

5. Comparaison des performances et bonnes pratiques

5.1 Benchmark de performance

Exemple

import time

def benchmark_training(model, train_loader, device, use_amp=True,
dtype=torch.float16, num_iterations=100):
"""Comparer la vitesse d'entraînement de FP32 et AMP"""
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler(enabled=use_amp and dtype == torch.float16)

# ── Échauffement (les 10 premières étapes ne sont pas chronométrées) ────────────────────
for i, (inputs, labels) in enumerate(train_loader):
if i >= 10:
break
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()

if use_amp:
with autocast(device_type='cuda', dtype=dtype):
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

# ── Test formel ─────────────────────────────────
torch.cuda.synchronize()
start = time.time()

for i, (inputs, labels) in enumerate(train_loader):
if i >= num_iterations:
break
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()

if use_amp:
with autocast(device_type='cuda', dtype=dtype):
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

torch.cuda.synchronize()
elapsed = time.time() - start

return elapsed / num_iterations

# ── Lancer la comparaison ─────────────────────────────────────
model = SimpleModel()

fp32_time = benchmark_training(model, train_loader, device, use_amp=False)
fp16_time = benchmark_training(model, train_loader, device, use_amp=True, dtype=torch.float16)

print(f"Temps moyen FP32 : {fp32_time*1000:.2f} ms/batch")
print(f"Temps moyen FP16 : {fp16_time*1000:.2f} ms/batch")
print(f"Accélération : {fp32_time / fp16_time:.2f}x")

# Si BF16 est pris en charge, testez-le également
if torch.cuda.is_bf16_supported():
bf16_time = benchmark_training(
model, train_loader, device, use_amp=True, dtype=torch.bfloat16
)
print(f"Temps moyen BF16 : {bf16_time*1000:.2f} ms/batch")
print(f"Accélération BF16 : {fp32_time / bf16_time:.2f}x")

5.2 Comparaison de l'utilisation de la mémoire

Exemple

def compare_memory(model_class, train_loader, device):
"""Comparer l'utilisation de la mémoire GPU de FP32 et AMP"""
def reset():
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

# ── Test de mémoire FP32 ────────────────────────────
réinitialiser()
model_fp32 = model_class().to(device)
optimizer_fp32 = optim.Adam(model_fp32.parameters())

for inputs, labels in list(train_loader)[:5]:
inputs, labels = inputs.to(device), labels.to(device)
optimizer_fp32.zero_grad()
outputs = model_fp32(inputs)
loss = nn.CrossEntropyLoss()(outputs, labels)
loss.backward()
optimizer_fp32.step()

fp32_peak = torch.cuda.max_memory_allocated() / 1024**2

# ── Test de mémoire AMP ─────────────────────────────
réinitialiser()
model_amp = model_class().to(device)
optimizer_amp = optim.Adam(model_amp.parameters())
scaler = GradScaler()

for inputs, labels in list(train_loader)[:5]:
inputs, labels = inputs.to(device), labels.to(device)
optimizer_amp.zero_grad()
with autocast(device_type='cuda'):
outputs = model_amp(inputs)
loss = nn.CrossEntropyLoss()(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer_amp)
scaler.update()

amp_peak = torch.cuda.max_memory_allocated() / 1024**2

print(f"Mémoire de pointe FP32 : {fp32_peak:.1f} MB")
print(f"Mémoire de pointe AMP : {amp_peak:.1f} MB")
print(f"Économies de mémoire : {(fp32_peak - amp_peak) / fp32_peak * 100:.1f}%")

compare_memory(SimpleModel, train_loader, device)

5.3 Résumé des bonnes pratiques

Scénario	Configuration recommandée	Raison
A100 / H100 / série RTX 40	BF16, pas besoin de GradScaler	Même plage numérique que FP32, le plus stable
V100 / série RTX 20	FP16 + GradScaler	Le matériel ne prend pas en charge BF16 ; un scaler est nécessaire pour éviter le débordement
Série RTX 30	BF16 de préférence, FP16 en alternative	L'architecture Ampere prend en charge BF16
Entraînement de grands modèles (mémoire GPU limitée)	AMP + accumulation de gradient + checkpointing de gradient	La combinaison des trois maximise l'utilisation de la mémoire GPU
Déploiement d'inférence	autocast + inference_mode	Aucun scaler nécessaire, inférence la plus rapide

PyTorch 2.0+ a intégré AMP danstorch.compilecelui-ci, qui applique automatiquement l'optimisation de précision mixte. Lors de l'utilisationtorch.compile(model)de celui-ci, le compilateur détermine automatiquement quelles opérations conviennent au FP16.

6. Combinaison avec d'autres techniques d'optimisation

6.1 AMP + torch.compile

Exemple

# ── Utilisation combinée avec PyTorch 2.0+ ───────────────────────
model = model.to(device)

# Méthode 1 : Compiler d'abord, puis entraîner avec AMP
# torch.compile fusionne automatiquement les opérateurs et optimise les accès mémoire
model_compiled = torch.compile(model, mode="reduce-overhead")

scaler = GradScaler()
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()

with autocast(device_type='cuda'):
outputs = model_compiled(inputs)
loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# Méthode 2 : Activer TF32 (accélération supplémentaire sur les architectures Ampere+)
# TF32 utilise une précision de 19 bits, avec une vitesse proche de FP16 et une exactitude proche de FP32
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True # Accélère la convolution lorsque la taille d'entrée est fixe

6.2 AMP + Entraînement distribué

Exemple

# ── Entraînement distribué + précision mixte ────────────────────────
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)

def train_ddp_amp(rank, world_size):
setup(rank, world_size)
device = torch.device(f"cuda:{rank}")

model = SimpleModel().to(device)
model = DDP(model, device_ids=[rank])

optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

for inputs, labels in train_loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)

optimizer.zero_grad()

with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)

# DDP synchronise automatiquement les gradients lors de la rétropropagation
scaler.scale(loss).backward()

# Écrêtage du gradient (il faut d'abord annuler la mise à l'échelle)
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

scaler.step(optimizer)
scaler.update()

dist.destroy_process_group()

Résumé

L'entraînement en précision mixte est unepratique standard。

Points clés :

Choisissez le format de précision: préférez BF16 lorsque le matériel le prend en charge, sinon utilisez FP16 + GradScaler
Comprenez autocast: il gère automatiquement le changement de précision, et dans la plupart des cas, aucune intervention manuelle n'est nécessaire
Comprendre GradScaler: requis uniquement pour FP16, empêche l'underflow des gradients via une mise à l'échelle dynamique
Faites attention à l'ordre de clipping: unscale → clip → step → update, l'ordre ne peut pas être inversé
Tirez parti de l'optimisation combinée: AMP peut être combiné de manière transparente avec torch.compile, l'accumulation de gradients et l'entraînement distribué

Autres extensions