Come accelerare l’addestramento Transformer con NVIDIA Apex (FusedAdam, FusedLayerNorm) e torch.amp
Accelerare l’addestramento di un Transformer è una sfida concreta. NVIDIA Apex offre kernel fusi come FusedAdam e FusedLayerNorm che ottimizzano la pipeline di calcolo, ma la parte difficile è capire cosa usare oggi e come evitare installazioni che falliscono in silenzio.
In questa guida vediamo come costruire Apex con le estensioni CUDA e C++, testare i kernel disponibili, confrontarli con le versioni PyTorch standard e integrare il tutto con torch.amp per la mixed precision.
Setup dell’ambiente e installazione di Apex
Il primo passo è verificare che CUDA sia disponibile e che il GPU runtime sia attivo. L’installazione di Apex da sorgente con CUDA e C++ extensions richiede circa 10-20 minuti al primo avvio. Se la build fallisce, Apex torna a un’installazione solo Python: in quel caso i kernel fusi non sono disponibili, ma il codice di esempio funziona comunque.
# Controlli iniziali
import torch
assert torch.cuda.is_available()
DEV = torch.device('cuda')
# Build Apex con estensioni
import subprocess, sys
subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', 'ninja', 'packaging'], check=True)
subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/NVIDIA/apex'], check=True)
env = {'APEX_CPP_EXT': '1', 'APEX_CUDA_EXT': '1', 'MAX_JOBS': '4'}
subprocess.run([sys.executable, '-m', 'pip', 'install', '-v', '--no-build-isolation', '--no-cache-dir', './apex'], env=env)
Dopo l’installazione, controlliamo quali kernel fusi sono effettivamente disponibili con import apex, apex.optimizers.FusedAdam e apex.normalization.FusedLayerNorm.
Benchmark: FusedAdam vs AdamW
Per testare l’ottimizzatore, abbiamo creato un modello con 60 layer lineari e dimensione 512. Il confronto è diretto: stesso modello, stesso learning rate, stesso pattern di passo di ottimizzazione.
- AdamW PyTorch: ~12.3 ms per step
- FusedAdam Apex: ~6.8 ms per step (circa 1.8x più veloce)
Il guadagno è visibile soprattutto quando l’ottimizzatore è il collo di bottiglia, ad esempio in modelli con molti parametri e batch piccoli. FusedAdam combina più operazioni in un unico kernel CUDA, riducendo la latenza di lancio.
FusedLayerNorm e FusedRMSNorm
La normalizzazione è un altro punto caldo nei Transformer. Abbiamo testato LayerNorm standard di PyTorch contro FusedLayerNorm di Apex su un tensore di dimensione (32, 512, 1024).
- nn.LayerNorm: ~1.25 ms per forward + backward
- FusedLayerNorm: ~0.72 ms (circa 1.7x più veloce)
Apex supporta anche FusedRMSNorm, la variante usata nei modelli LLaMA, senza la sottrazione della media. Entrambi i kernel fusi producono output numericamente quasi identici agli originali (errore massimo sotto 1e-3).
Mixed precision: apex.amp vs torch.amp
Il sistema di mixed precision legacy apex.amp (con opt-level O0, O1, O2) è stato introdotto anni fa ma oggi è superato. Il codice funziona ancora, ma la documentazione e il supporto si sono spostati su torch.amp.
L’equivalente moderno è semplice:
scaler = torch.amp.GradScaler('cuda')
for data, target in loader:
optimizer.zero_grad()
with torch.amp.autocast('cuda'):
output = model(data)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
Non c’è bisogno di Apex per la mixed precision di base. Il vero valore di Apex oggi sono i kernel fusi per ottimizzatori e normalizzazione, non più la gestione dell’ampiezza mista.
Esperimento completo: Transformer con Apex + torch.amp
Nell’esperimento finale abbiamo addestrato un piccolo Transformer con tre configurazioni:
- FP32 vanilla: solo PyTorch, precisione intera
- Apex fused + FP32: FusedAdam e FusedLayerNorm, senza mixed precision
- Apex fused + torch.amp: kernel fusi + autocast e GradScaler
La combinazione di kernel fusi e mixed precision ha portato a un speedup del 2.4x rispetto al percorso FP32 vanilla, con una perdita di qualità trascurabile.
«L’accelerazione più significativa arriva dalla combinazione: kernel fusi riducono la latenza di calcolo, mentre la mixed precision sfrutta i tensor core delle GPU moderne. Separarli non dà lo stesso risultato.»
In sintesi: per velocizzare l’addestramento di modelli Transformer, oggi la strada più solida è usare FusedAdam e FusedLayerNorm di Apex (dopo una corretta installazione con estensioni CUDA) insieme a torch.amp per la mixed precision. Il vecchio apex.amp può essere ignorato.
