Aller au contenu principal

Module 7 — GPU et précision mixte

Jusqu'ici, l'entraînement tournait sur processeur. Une époque sur Fashion-MNIST prend une minute environ, ce qui reste acceptable pour prototyper. À la première image de 224 par 224 pixels ou au premier ResNet, le CPU ne suffit plus. Ce module explique comment déplacer proprement un entraînement sur GPU et comment lui appliquer la précision mixte, qui divise souvent la durée par deux sans dégrader la qualité.

Un device unique, référencé partout​

La règle est simple : on décide du device au début du script, et tout — modèle, entrées, cibles — vit dessus.

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

Cette ligne permet au même code de tourner sur portable sans GPU et sur serveur GPU sans rien changer. On la voit systématiquement en tête des scripts d'entraînement bien écrits.

Déplacer le modèle et les tenseurs​

to(device) copie un tenseur ou déplace un module. Sur nn.Module, c'est récursif : tous les sous-modules et paramètres suivent.

modele = ReseauFashionMNIST().to(device)

for images, cibles in train_loader:
    images = images.to(device, non_blocking=True)
    cibles = cibles.to(device, non_blocking=True)
    # ... suite habituelle

non_blocking=True a un effet uniquement quand la mémoire hôte est épinglée (pin_memory=True dans le DataLoader) : le transfert CPU→GPU se fait alors en parallèle du calcul en cours. Sans épingle, l'argument est ignoré silencieusement.

L'erreur qui bloque l'entraînement au premier lot

RuntimeError: Expected all tensors to be on the same device signale qu'un tenseur est resté sur CPU alors que le reste est sur GPU. Cause la plus fréquente : oublier .to(device) sur les cibles, ou créer un tenseur dans le forward sans préciser le device (torch.zeros(3) au lieu de torch.zeros(3, device=x.device)).

Le transfert de données est un vrai goulot​

Un GPU moderne calcule beaucoup plus vite qu'il ne reçoit ses entrées. Trois leviers rendent le transfert transparent.

  1. pin_memory=True dans le DataLoader réserve la mémoire hôte pour un transfert direct, plus rapide.
  2. num_workers > 0 délègue la préparation à des processus fils pendant que le GPU calcule le lot précédent.
  3. non_blocking=True dans to(device) autorise le recouvrement effectif calcul/transfert.

Sans ces trois éléments, on voit un GPU coûteux à 30 % d'utilisation qui attend ses données. La mesure de débit ci-dessous rend le phénomène visible.

import time

def mesurer_debit(loader, modele, device, nb_lots=100):
    modele.eval()
    torch.cuda.synchronize() if device.type == "cuda" else None
    debut = time.perf_counter()
    exemples = 0
    with torch.no_grad():
        for i, (x, y) in enumerate(loader):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            _ = modele(x)
            exemples += x.size(0)
            if i + 1 >= nb_lots:
                break
    torch.cuda.synchronize() if device.type == "cuda" else None
    duree = time.perf_counter() - debut
    return exemples / duree

torch.cuda.synchronize() est indispensable avant et après la mesure : les appels CUDA sont asynchrones, sans quoi on mesurerait le temps de programmation et non le temps de calcul.

La précision mixte, en deux mots​

Un GPU récent calcule beaucoup plus vite en float16 (ou bfloat16) qu'en float32. La précision mixte alterne les deux : les opérations massives (matmul, convolutions) tournent en 16 bits, les cumulateurs et les paramètres restent en 32 bits pour préserver la stabilité numérique. Gains typiques : 1,5 à 3 fois plus rapide, avec la même exactitude finale.

torch.cuda.amp fournit les deux outils nécessaires : autocast qui choisit automatiquement la précision à utiliser opération par opération, et GradScaler qui empêche les gradients trop petits d'être arrondis à zéro.

autocast : le contexte qui change la précision​

from torch.cuda.amp import autocast, GradScaler

modele = ReseauFashionMNIST().to(device)
optimiseur = torch.optim.AdamW(modele.parameters(), lr=1e-3)
scaler = GradScaler(enabled=(device.type == "cuda"))

for images, cibles in train_loader:
    images = images.to(device, non_blocking=True)
    cibles = cibles.to(device, non_blocking=True)

    optimiseur.zero_grad(set_to_none=True)

    with autocast(device_type="cuda", dtype=torch.float16, enabled=(device.type == "cuda")):
        logits = modele(images)
        perte = criterion(logits, cibles)

    scaler.scale(perte).backward()
    scaler.step(optimiseur)
    scaler.update()

Trois nouveautés :

  • autocast enveloppe la passe avant et le calcul de la perte, pas la rétropropagation.
  • scaler.scale(perte).backward() multiplie la perte par un facteur d'échelle avant backward, pour que les gradients en float16 ne s'annulent pas.
  • scaler.step(optimiseur) dévoile le gradient, vérifie l'absence d'infini ou de NaN, puis appelle step uniquement s'ils sont sains.
  • scaler.update() ajuste le facteur d'échelle : il l'augmente quand tout va bien, le divise par deux quand un gradient déborde.

set_to_none=True sur zero_grad remplace les tenseurs par None plutôt que par des zéros, ce qui libère de la mémoire et est légèrement plus rapide.

bfloat16 sur les GPU récents

Sur Ampere (A100, RTX 30xx) et plus récent, bfloat16 évite le besoin de GradScaler grâce à sa plage exponentielle identique à float32. Écrire autocast(dtype=torch.bfloat16) et supprimer le GradScaler simplifie la boucle. float16 reste le défaut sur les GPU plus anciens.

Vérifier que la précision mixte n'altère pas la qualité​

C'est le contrôle qu'on oublie et qu'il faut toujours faire : la précision mixte est censée être neutre, elle ne l'est pas toujours si le modèle contient des opérations numériquement instables.

Procédure : entraîner deux versions identiques, l'une en float32, l'autre en précision mixte, sur trois à cinq époques, comparer les exactitudes de validation. Un écart de plus de 0,5 point signale un problème, souvent lié à une couche personnalisée ou à une perte qui opère sur de très petits nombres.

Aperçu de torch.compile​

Depuis PyTorch 2, torch.compile prend un nn.Module et le compile à la volée en un graphe optimisé, sans changer votre code.

modele = ReseauFashionMNIST().to(device)
modele = torch.compile(modele, mode="default")

C'est l'équivalent moderne du traçage de TensorFlow présenté au cours 08, sans le fardeau du traçage explicite. Les gains typiques vont de 10 % à 3 fois plus rapide selon le modèle. En 2026, il fonctionne bien sur les architectures standards ; les architectures très dynamiques (contrôle de flux dépendant des données) peuvent encore poser problème. À essayer systématiquement, à retirer si l'entraînement ne démarre plus.

Une boucle GPU + AMP + compile, prête à l'emploi​

def entrainer_gpu_amp(modele, train_loader, val_loader, nb_epoques=10):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    modele = modele.to(device)
    if device.type == "cuda":
        modele = torch.compile(modele)

    criterion = nn.CrossEntropyLoss()
    optimiseur = torch.optim.AdamW(modele.parameters(), lr=1e-3, weight_decay=1e-4)
    planificateur = torch.optim.lr_scheduler.CosineAnnealingLR(optimiseur, T_max=nb_epoques)
    scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))

    for epoque in range(nb_epoques):
        modele.train()
        for x, y in train_loader:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            optimiseur.zero_grad(set_to_none=True)
            with torch.cuda.amp.autocast(device_type="cuda", dtype=torch.float16,
                                          enabled=(device.type == "cuda")):
                perte = criterion(modele(x), y)
            scaler.scale(perte).backward()
            scaler.step(optimiseur)
            scaler.update()
        planificateur.step()

Sur un GPU d'entrée de gamme récent, une époque Fashion-MNIST passe de la minute à quelques secondes.

En résumé​

  • Un device unique décidé en tête de script ; to(device) sur le modèle est récursif, sur les entrées et cibles il ne l'est pas.
  • Le trio pin_memory + num_workers + non_blocking rend le transfert transparent ; sans lui, un GPU coûteux attend ses données.
  • autocast + GradScaler apportent la précision mixte en une poignée de lignes ; sur GPU récent, bfloat16 dispense du GradScaler.
  • torch.compile est un ajout d'une ligne qui peut accélérer significativement ; toujours vérifier qu'il n'altère ni la qualité, ni la stabilité de l'entraînement.

Le module suivant s'attaque à un problème très pragmatique : sauvegarder l'entraînement à intervalles réguliers et le reprendre à l'octet près après une interruption.