PyTorch-Training mit gemischter Präzision (AMP)

Training mit gemischter Präzision ist eine der wichtigsten Techniken zur Leistungsoptimierung im Deep Learning. Durch die gleichzeitige Verwendung von FP32- (Einfachpräzision) und FP16- (halbe Präzision) Gleitkommazahlen für Berechnungen lassen sich die Trainingsgeschwindigkeit deutlich steigern und der Speicherbedarf reduzieren, nahezu ohne Verlust der Modellgenauigkeit. Dieser Abschnitt beschreibt die Technik Automatic Mixed Precision (AMP) in PyTorch im Detail.

Anwendbare Versionen:Der Code in diesem Artikel wurde mit PyTorch 1.6+ geschriebentorch.cuda.ampAPI. PyTorch 2.4+ empfiehlt die Verwendung vontorch.amp.autocastundtorch.amp.GradScaler, und die Verwendung ist im Wesentlichen dieselbe; Unterschiede werden im Artikel erwähnt.

1. Grundlagen des Trainings mit gemischter Präzision

1.1 Warum brauchen wir gemischte Präzision

Das Training von Deep-Learning-Modellen umfasst eine große Anzahl von Matrixoperationen. Herkömmliches FP32 (32-Bit-Gleitkomma) hat eine hohe Rechenpräzision, belegt jedoch mehr Speicher und ist langsamer. FP16 (16-Bit-Gleitkomma) ist schneller und benötigt weniger Speicher, aber sein numerischer Bereich ist kleiner und es neigt zu Problemen mit Gradientenunterlauf.

Die Kernidee des Trainings mit gemischter Präzision ist:Verwenden Sie FP32 für Operationen, die eine hohe Präzision erfordern, und FP16 für Operationen, die keine hohe Präzision erfordern.. Auf diese Weise können Sie den Geschwindigkeitsvorteil von FP16 nutzen und gleichzeitig Präzisionsprobleme vermeiden.

Die folgende Abbildung zeigt die Unterschiede im Bit-Layout der drei Gleitkommaformate –Exponentenbitsbestimmen den numerischen Wertebereich.Mantissenbitsbestimmen die Genauigkeit:

1.2 Vorteile gemischter Präzision

Metrik	Verbesserung	Beschreibung
Trainingsgeschwindigkeit	2-3x Verbesserung	Hängt von der GPU-Tensor-Core-Unterstützung ab
Speicherverbrauch	Um etwa 50 % reduziert	Aktivierungen und Zwischenergebnisse in FP16 gespeichert
Speicherbandbreite	Um etwa 50 % reduziert	Kleinere Datengröße bedeutet weniger Datenübertragung
Kommunikations-Overhead	Um etwa 50 % reduziert	Gradientenübertragungsvolumen beim verteilten Training halbiert

1.3 Prinzip der Tensor-Core-Beschleunigung

NVIDIAs Tensor Core ist eine Hardwareeinheit, die auf Matrixoperationen spezialisiert ist. Er kann eine 4×4-Matrix-Multiply-Accumulate-Operation (D = A × B + C) in einem einzigen Taktzyklus abschließen, was die Hauptquelle der FP16-Trainingsbeschleunigung ist. Im Vergleich zu gewöhnlichen CUDA-Cores, die mehrere Anweisungen benötigen, um dieselbe Operation abzuschließen, komprimiert der Tensor Core sie in eine einzige Anweisung.

GPUs, die Tensor Core unterstützen, sind unter anderem:

Volta-Architektur(V100) — Tensor Core der ersten Generation, unterstützt nur FP16
Turing-Architektur(RTX 20-Serie) — Unterstützt FP16 / INT8 / INT4
Ampere-Architektur(RTX 30-Serie, A100) — BF16-/TF32-Unterstützung hinzugefügt
Ada-Lovelace-Architektur(RTX 40-Serie) — FP8-Unterstützung hinzugefügt
Hopper-Architektur(H100) — FP8-Transformer-Engine hinzugefügt

Auch Consumer-RTX-GPUs unterstützen Tensor Core; beispielsweise können Modelle ab RTX 3060 von der AMP-Beschleunigung profitieren.

2. Grundlegende Verwendung von PyTorch AMP

2.1 autocast und GradScaler

Die AMP-API von PyTorch umfasst hauptsächlich zwei Kernkomponenten:

autocast: Ein Kontextmanager, der Vorgänge innerhalb seines Gültigkeitsbereichs automatisch auf FP16 umschaltet (präzisionsempfindliche Vorgänge fallen automatisch auf FP32 zurück)
GradScaler: Passt den Gradient-Skalierungsfaktor dynamisch an, um FP16-Gradienten zu verstärken und einen Unterlauf zu verhindern (nur für FP16 erforderlich; für BF16 in der Regel nicht nötig)

Die folgende Abbildung zeigt den Datenfluss eines vollständigen Trainingsschritts mit AMP:

2.2 Einfaches Anwendungsbeispiel

Beispiel

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler
# Empfohlene Vorgehensweise für PyTorch 2.4+:
# from torch.amp import autocast, GradScaler

# Prüfen, ob CUDA verfügbar ist
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 unterstützt: {torch.cuda.is_bf16_supported()}")

# ── Modelldefinition ──────────────────────────────────────
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)

# Verlustfunktion und Optimierer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ── Wichtige Komponenten für Mixed-Precision-Training ──────────────────────────
# GradScaler: skaliert den Verlust, um FP16-Gradienten-Unterlauf zu vermeiden
scaler = GradScaler()

# Trainingsschleife
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()

# ── Kern: Verwende den autocast-Kontextmanager ──────
# Operationen innerhalb des autocast-Bereichs verwenden automatisch FP16
# Präzisionsempfindliche Operationen (z. B. softmax, loss) fallen automatisch auf FP32 zurück
with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)

# ── Verwende scaler für die Rückwärtspropagation ─────────────
# 1. Skaliere den Loss (multipliziere mit scale_factor)
# 2. Rückwärtspropagation (berechnet auf den skalierten Gradienten)
# 3. scaler.step hebt intern die Skalierung der Gradienten auf und prüft auf Inf/NaN
scaler.scale(loss).backward()

# Parameter aktualisieren
scaler.step(optimizer)

# Aktualisiere den Skalierungsfaktor des Scalers
scaler.update()

# Statistiken
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

# Simulierte Daten
train_loader = [
(torch.randn(32, 128), torch.randint(0, 10, (32,))) for _ in range(10)
]

# Training starten
for epoch in range(3):
loss, acc = train_epoch_amp(
model, train_loader, criterion, optimizer, scaler, device
)
print(f"Epoche {epoch+1}: Verlust={loss:.4f}, Genauigkeit={acc:.4f}")

print("Training mit gemischter Präzision abgeschlossen!")

API-Migrationshinweis:PyTorch 2.4 wirdtorch.cuda.amp.autocastals veraltet markieren; es wird empfohlen, dies zu verwendentorch.amp.autocast('cuda'). Die Verwendung ist genau dieselbe; nur der Importpfad unterscheidet sich.

2.3 Auswahl zwischen BF16 und FP16

Moderne GPUs unterstützen zwei Halbpräzisionsformate. Ihr Hauptunterschied liegt in der Zuordnungsstrategie von Exponentenbits und Mantissenbits:

Format	Exponentenbits	Mantissenbits	Numerischer Bereich	Precision	Anwendungsfall
FP16	5 Bits	10 Bits	±65504	Höher (~3.3 Stellen)	Erfordert Kompatibilität mit älteren GPUs (V100/RTX 20-Serie)
BF16	8 Bits	7 Bits	±3.4×10³⁸	Niedriger (~2.4 Stellen)	Für Stabilität (A100/RTX 30-Serie+)

Empfehlung: Wenn die Hardware BF16 unterstützt,bevorzugen Sie BF16. Der numerische Bereich ist exakt derselbe wie bei FP32, sodass Überlaufprobleme während des Trainings selten auftreten und GradScaler nicht benötigt wird.

Beispiel

# ── BF16-Ansatz (PyTorch 1.10+) ────────────────────
# BF16 erfordert keinen GradScaler, da sein numerischer Bereich derselbe ist wie bei FP32

from torch.cuda.amp import autocast

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

# Methode 2: BF16 standardmäßig global aktivieren (wenn die Hardware dies unterstützt)
# torch.backends.cuda.matmul.allow_bf16_reduced_precision = True
# torch.backends.cudnn.allow_bf16_reduced_precision = True

# ── Hardware-Unterstützung prüfen ─────────────────────────────
print(f"Hardware unterstützt BF16: {torch.cuda.is_bf16_supported()}")
print(f"Derzeit erlaubt matmul BF16: {torch.backends.cuda.matmul.allow_bf16}")
print(f"Derzeit erlaubt cuDNN BF16: {torch.backends.cudnn.allow_bf16}")

# ── Vollständiges BF16-Trainingsbeispiel (kein GradScaler erforderlich) ──────────
scaler_bf16 = None # BF16 benötigt keinen Skalierer

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)

# Direkte Backpropagation, kein Skalierer erforderlich
loss.backward()
optimizer.step()

3. Fortgeschrittene Tipps und Optimierung

3.1 Dynamische Verlustskalierung

Der Kernmechanismus von GradScaler istdynamische Verlustskalierung—er passt den Skalierungsfaktor automatisch basierend auf dem Trainingszustand an, verhindert Gradienten-Underflow, während er die Skalierung allmählich reduziert, wenn sich das Training stabilisiert, um Präzisionsverlust zu minimieren:

GradScaler verwaltet die obige Rückkopplungsschleife automatisch; Sie müssen sein Verhalten nur über Parameter feinjustieren:

Beispiel

from torch.cuda.amp import GradScaler

# Benutzerdefinierte GradScaler-Parameter
scaler = GradScaler(
init_scale=2**16, # Initialer Skalierungsfaktor, Standardwert 65536
growth_factor=2.0, # Wachstumsmultiplikator des Skalierungsfaktors, Standardwert 2.0
backoff_factor=0.5, # Fallback-Multiplikator für den Skalierungsfaktor, Standard 0.5
growth_interval=2000, # Anzahl aufeinanderfolgender erfolgreicher Schritte vor dem Wachstum
enabled=True # Ob aktiviert (kann dynamisch umgeschaltet werden)
)

# Arbeitsablauf:
# 1. Anfängliche Skalierung = 65536
# 2. Wenn ein Schritt Inf/NaN enthält → Skalierung × 0.5 (verringern), diesen Schritt überspringen
# 3. Wenn 2000 aufeinanderfolgende Schritte kein Inf/NaN haben → Skalierung × 2.0 (erhöhen)
# 4. Immer den maximal nutzbaren Skalierungsfaktor innerhalb eines sicheren Bereichs finden

# Aktuellen Skalierungsfaktor anzeigen
print(f"Aktueller Skalierungsfaktor: {scaler.get_scale()}")

# Ermitteln, ob scaler den letzten Schritt als erfolgreich betrachtet
print(f"Kürzlich erfolgreich: {scaler._found_inf.item() == 0}")

3.2 AMP mit Gradientenakkumulation

Gradientenakkumulation simuliert eine größere Batch-Größe bei begrenztem Speicher. Bei der Verwendung von AMP ist zu beachten, dass der Verlust jedes Teil-Batches durch die Anzahl der Akkumulationsschritte geteilt werden muss; andernfalls wird der endgültige Gradient verstärkt:

Beispiel

# ── Gradientenakkumulation + gemischte Präzision ────────────────
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)
# Wichtig: Teile den Verlust jedes Teil-Batches durch die Akkumulationsschritte
# Auf diese Weise entspricht der akkumulierte Gradient dem durchschnittlichen Gradienten eines großen Batches
loss = criterion(outputs, labels) / accumulation_steps

# Akkumuliere die skalierten Gradienten (Hinweis: hier nicht zurücksetzen)
scaler.scale(loss).backward()

# Parameter alle accumulation_steps Batches aktualisieren
if (batch_idx + 1) % accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

# ── Behandlung von Gradienten am Ende, die weniger als einen vollständigen Akkumulationszyklus ausmachen ──────────
remainder = len(train_loader) % accumulation_steps
if remainder != 0:
# Die akkumulierten Gradienten müssen gemäß der tatsächlichen Anzahl an Schritten neu skaliert werden
# Einfacher Ansatz: Schritt trotzdem ausführen; die Gradienten sind etwas größer, aber die Auswirkung ist normalerweise minimal
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

3.3 AMP in Validierung und Inferenz

AMP wird auch für Validierung und Inferenz empfohlen. Da Backpropagation und Gradientenskalierung nicht benötigt werden, ist der Code einfacher:

Beispiel

# ── autocast für Inferenz verwenden ──────────────────────────
@torch.no_grad()
def inference_amp(model, inputs, device):
model.eval()
# Inferenz verwendet FP16, kein GradScaler erforderlich
with autocast(device_type='cuda'):
outputs = model(inputs.to(device))
return outputs

# ── inference_mode ist schneller als no_grad ──────────────
# inference_mode deaktiviert weiteren Tracing-Overhead, geeignet für reine Inferenz
@torch.inference_mode()
def inference_fast(model, inputs, device):
model.eval()
with autocast(device_type='cuda'):
outputs = model(inputs.to(device))
return outputs

# ── Batch-Inferenz-Beispiel ───────────────────────────────
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. Häufige Probleme und Lösungen

4.1 Instabiles Training

Problem	Ursache	Lösung
Loss wird NaN / Inf	Gradientenüberlauf (Scale zu groß oder Lernrate zu hoch)	Reduzieren`init_scale`, Gradient Clipping hinzufügen, Lernrate senken
Loss nimmt nicht ab / oszilliert	Gradientenunterlauf (Scale zu klein, Gradienten auf 0 abgeschnitten)	Erhöhen`init_scale`, Gradientenstatistiken prüfen, BF16 ausprobieren
Deutlicher Präzisionsabfall	Bestimmte Schichten (z. B. LayerNorm, Softmax) sind empfindlich gegenüber FP16	Diese Schichten manuell in FP32 belassen (siehe 4.2)
Instabilität im späten Training	Scale-Fehlanpassung aufgrund von Änderungen der Wertebereiche der Modellparameter	Entsprechend erhöhen`growth_interval`, wodurch Skalenanpassungen konservativer ausfallen

4.2 Manuelle Präzisionssteuerung

Einige Operationen erleiden unter FP16 erheblichen Präzisionsverlust und müssen manuell für die Verwendung von FP32 angegeben werden. PyTorchs autocast bietet integrierten Schutz für diese Operationen, aber benutzerdefinierte Operationen erfordern möglicherweise manuelle Behandlung:

Beispiel

# ── Methode 1: Eine Loss-Funktion wrappen, die FP32 erzwingt ────────
class FP32Loss(nn.Module):
"""Wrapper für eine Verlustfunktion, der FP32-Berechnung erzwingt"""
def __init__(self, base_criterion):
super().__init__()
self.base_criterion = base_criterion

def forward(self, input, target):
# Explizit nach FP32 konvertieren, außerhalb des Einflusses von autocast
return self.base_criterion(input.float(), target.float())

# ── Methode 2: Lokales Deaktivieren innerhalb von autocast ──────────────
with autocast(device_type='cuda'):
outputs = model(inputs)

# Für präzisionsempfindliche Operationen autocast vorübergehend deaktivieren
with autocast(enabled=False):
loss = criterion(outputs.float(), labels)

# ── Methode 3: Bestimmte Layer im Modell in FP32 behalten ────────
class CustomModel(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64), # BN ist präzisionsempfindlich
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 (von autocast gesteuert)
x = self.features[1](x.float()) # BN: FP32 erzwingen
x = self.features[2](x) # ReLU: FP16
x = self.features[3](x)
x = self.features[4](x.float()) # BN: FP32 erzwingen
x = self.features[5](x)
x = x.mean(dim=[2, 3]) # Globales Average-Pooling
x = self.classifier(x)
return x

Tipp:PyTorchs autocast hält die folgenden Operationen automatisch in FP32: softmax, log_softmax, cross_entropy, layer_norm, batch_norm, usw. Normalerweise müssen nur benutzerdefinierte Operatoren manuell behandelt werden.

4.3 Gradientenclipping und AMP

Gradientenclipping mussscaler.step()vorher undscaler.unscale_()nach der Ausführung. Das liegt daran, dass Clipping aufechten Gradientenwertenstatt auf dem skalierten Gradienten operieren muss:

Beispiel

# ── Korrekte Reihenfolge beim Gradientenclipping ──────────────────────────
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)

# Schritt 1: Loss skalieren und Backpropagation durchführen
scaler.scale(loss).backward()

# Schritt 2: Gradienten entskalieren (Gradienten von der Skalierung auf die wahren Werte zurücksetzen)
# Muss vor Clipping und Schritt aufgerufen werden
scaler.unscale_(optimizer)

# Schritt 3: Die wahren Gradienten clippen
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Schritt 4: Parameter aktualisieren (prüft intern, ob Gradienten Inf/NaN enthalten)
scaler.step(optimizer)

# Schritt 5: Den Skalierungsfaktor aktualisieren
scaler.update()

# ── Häufiger Fehler: Clipping ohne unscale_ ──────────────
# Falsch! Das Clipping wird auf die verstärkten Gradienten angewendet, und der Schwellenwert wird durch die Skalierung beeinflusst
# scaler.scale(loss).backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # ← Falsch
# scaler.step(optimizer)

5. Leistungsvergleich und bewährte Verfahren

5.1 Leistungsbenchmark

Beispiel

import time

def benchmark_training(model, train_loader, device, use_amp=True,
dtype=torch.float16, num_iterations=100):
"""Vergleicht die Trainingsgeschwindigkeit von FP32 und 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)

# ── Warm-up (erste 10 Schritte werden nicht gemessen) ────────────────────
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()

# ── Formaler Test ─────────────────────────────────
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

# ── Vergleich durchführen ─────────────────────────────────────
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"FP32-Durchschnittszeit: {fp32_time*1000:.2f} ms/batch")
print(f"FP16-Durchschnittszeit: {fp16_time*1000:.2f} ms/batch")
print(f"Beschleunigung: {fp32_time / fp16_time:.2f}x")

# Wenn BF16 unterstützt wird, auch testen
if torch.cuda.is_bf16_supported():
bf16_time = benchmark_training(
model, train_loader, device, use_amp=True, dtype=torch.bfloat16
)
print(f"BF16-Durchschnittszeit: {bf16_time*1000:.2f} ms/batch")
print(f"BF16-Beschleunigung: {fp32_time / bf16_time:.2f}x")

5.2 Vergleich der Speichernutzung

Beispiel

def compare_memory(model_class, train_loader, device):
"""Vergleiche den GPU-Speicherverbrauch von FP32 und AMP"""
def reset():
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

# ── FP32-Speichertest ────────────────────────────
Zurücksetzen()
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

# ── AMP-Speichertest ─────────────────────────────
Zurücksetzen()
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"FP32 maximaler Speicherverbrauch: {fp32_peak:.1f} MB")
print(f"AMP maximaler Speicherverbrauch: {amp_peak:.1f} MB")
print(f"Speichereinsparung: {(fp32_peak - amp_peak) / fp32_peak * 100:.1f}%")

compare_memory(SimpleModel, train_loader, device)

5.3 Zusammenfassung der Best Practices

Szenario	Empfohlene Konfiguration	Grund
A100 / H100 / RTX 40-Serie	BF16, kein GradScaler erforderlich	Gleicher numerischer Bereich wie FP32, am stabilsten
V100 / RTX 20-Serie	FP16 + GradScaler	Hardware unterstützt BF16 nicht; ein Scaler ist erforderlich, um Überlauf zu vermeiden
RTX 30-Serie	BF16 bevorzugt, FP16 als Alternative	Ampere-Architektur unterstützt BF16
Training großer Modelle (begrenzter GPU-Speicher)	AMP + Gradientenakkumulation + Gradienten-Checkpointing	Die Kombination aller drei maximiert die GPU-Speicherauslastung
Inferenz-Deployment	autocast + inference_mode	Kein Scaler erforderlich, schnellste Inferenz

PyTorch 2.0+ hat AMP integriert intorch.compilees, das automatisch Mixed-Precision-Optimierung anwendet. Bei Verwendung vontorch.compile(model)ihm bestimmt der Compiler automatisch, welche Operationen für FP16 geeignet sind.

6. Kombination mit anderen Optimierungstechniken

6.1 AMP + torch.compile

Beispiel

# ── Kombinierte Verwendung mit PyTorch 2.0+ ───────────────────────
model = model.to(device)

# Methode 1: Zuerst kompilieren, dann mit AMP trainieren
# torch.compile führt Operatoren automatisch zusammen und optimiert den Speicherzugriff
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()

# Methode 2: TF32 aktivieren (zusätzliche Beschleunigung auf Ampere+-Architekturen)
# TF32 verwendet 19-Bit-Präzision, mit Geschwindigkeit nahe FP16 und Genauigkeit nahe FP32
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True # Beschleunigt Faltung, wenn die Eingabegröße fest ist

6.2 AMP + Verteiltes Training

Beispiel

# ── Verteiltes Training + gemischte Präzision ────────────────────────
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 synchronisiert Gradienten während des Backward-Passes automatisch
scaler.scale(loss).backward()

# Gradient-Clipping (zuerst Unscale durchführen)
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

scaler.step(optimizer)
scaler.update()

dist.destroy_process_group()

Zusammenfassung

Training mit gemischter Präzision ist eineStandardpraxis。

Wichtige Punkte:

Präzisionsformat wählen: Bevorzuge BF16, wenn die Hardware es unterstützt, andernfalls verwende FP16 + GradScaler
Autocast verstehen: Es verwaltet den Präzisionswechsel automatisch, und in den meisten Fällen ist kein manueller Eingriff erforderlich
GradScaler verstehen: nur für FP16 erforderlich, verhindert Gradienten-Unterlauf durch dynamische Skalierung
Auf die Reihenfolge des Clippings achten: unscale → clip → step → update, die Reihenfolge darf nicht vertauscht werden
Kombinierte Optimierung nutzen: AMP lässt sich nahtlos mit torch.compile, Gradientenakkumulation und verteiltem Training kombinieren

Andere Erweiterungen