PyTorch et TensorFlow on GPU

Pourquoi ce chapitre est le plus rentable du parcours

Un GPUGPUIAProcesseur graphique parallélisant massivement les calculs matriciels ; indispensable à l'entraînement et à l'inférence des modèles de deep learning.Voir dans le glossaire récent délivre sa puissance annoncée seulement si le code qui l'exploite respecte un ensemble de contraintes précises : minimiser les synchronisations, éviter les copies mémoire inutiles, utiliser les bons formats numériques. La plupart des pipelines PyTorch ou TensorFlow que l'on rencontre en production tournent à 20 ou 30 % de l'efficacité théorique du matériel, non pas parce que le modèle est mal conçu, mais parce que le code autour du modèle multiplie les allers-retours CPU-GPU et les kernels non fusionnés.

Ce chapitre descend sous l'API haut niveau pour montrer ce qui se passe réellement sur le device : comment l'autograd construit et rejoue un graphe, comment écrire une opération CUDA personnalisée quand l'API standard ne suffit plus, comment les compilateurs (TorchInductor côté PyTorch, XLA côté TensorFlow) transforment un graphe d'opérations en kernels fusionnés, et comment la précisionprécisionIAProportion des alertes émises par un modèle qui sont justifiées. Elle s'oppose au rappel : améliorer l'une dégrade l'autre.Voir dans le glossaire mixte réduit la bande passante mémoire sans casser la convergence.

Aucune de ces techniques ne remplace un profilage réel. nsys et nvprof (ou leur équivalent torch.profiler) doivent confirmer chaque gain supposé. Une optimisation qui n'apparaît pas dans une trace GPU n'existe pas.

Autograd sous CUDA : ce qui se passe réellement sur le device

PyTorch construit son graphe de calcul dynamiquement, opération par opération, pendant le forward(). Chaque tensor CUDA porte une référence vers la fonction qui l'a produit (grad_fn) ; ce graphe n'est ni sérialisé ni optimisé au moment de sa construction, il sert uniquement à orchestrer le backward().

Ce point a une conséquence directe sur la performance : construire le graphe a un coût CPU non négligeable (allocation d'objets Python, empilement des grad_fn), pendant que l'exécution des kernels a lieu de façon asynchrone sur le GPU. Sur un modèle avec beaucoup de petites opérations (LSTM déroulé manuellement, boucles Python sur des time-steps), le CPU devient le goulot d'étranglement : le GPU termine ses kernels avant que le CPU n'ait fini de soumettre les suivants, et l'utilisation SM chute entre deux appels.

Le backward() rejoue ce graphe en sens inverse en soumettant les kernels gradient correspondants. Chaque opération CUDA custom doit donc fournir une implémentation forward et une fonction de gradient explicite via torch.autograd.Function, faute de quoi PyTorch ne sait pas comment la différencier.

class ScaledReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.save_for_backward(x)
        ctx.alpha = alpha
        return torch.clamp(x, min=0) * alpha

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[x < 0] = 0
        return grad_input * ctx.alpha, None

ctx.save_for_backward retient le tensor jusqu'à l'appel du backward(). Sur une boucle d'entraînement qui accumule des Function sans jamais appeler .backward() (debug, inférenceinférenceIAUtilisation d'un modèle déjà entraîné sur une donnée nouvelle. Peu coûteuse à l'unité mais répétée à chaque requête, elle constitue le coût récurrent d'exploitation.Voir dans le glossaire mal isolée), la mémoire GPU grossit silencieusement jusqu'au CUDA out of memory. Enveloppez systématiquement l'inférence dans torch.no_grad() ou torch.inference_mode().

Écrire un custom op : cpp_extension et ATen

Quand une opération n'existe pas dans l'API PyTorch, ou que sa version native multiplie les kernels alors qu'une fusion manuelle ferait mieux, on écrit une extension C++/CUDA. torch.utils.cpp_extension compile ce code à la volée (JIT) ou via un setup.py classique, et l'expose comme un module Python normal.

ATen (A Tensor library) est la bibliothèque C++ sous-jacente qui implémente les opérations sur Tensor — c'est elle que le code Python appelle en coulisses via les bindings. Écrire un custom op consiste à manipuler directement des at::Tensor en C++, avec accès aux pointeurs device et aux flux CUDA.

// fused_bias_relu.cu
#include <torch/extension.h>

__global__ void bias_relu_kernel(float* x, const float* bias, int n, int c) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        int channel = idx % c;
        float v = x[idx] + bias[channel];
        x[idx] = v > 0.f ? v : 0.f;
    }
}

torch::Tensor fused_bias_relu(torch::Tensor x, torch::Tensor bias) {
    auto n = x.numel();
    auto c = bias.numel();
    const int threads = 256;
    const int blocks = (n + threads - 1) / threads;
    bias_relu_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(), bias.data_ptr<float>(), n, c);
    return x;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_bias_relu", &fused_bias_relu, "Fused bias + ReLU (CUDA)");
}

Fusionner l'addition du biaisbiaisIARégularité correctement apprise dans des données qui ne représentent pas la réalité visée, ou qui enregistrent des décisions passées avec leurs préjugés. Changer d'algorithme ne le corrige pas.Voir dans le glossaire et la ReLUfonction d'activationIAOpération non linéaire appliquée en sortie d'un neurone. Sans elle, empiler des couches serait inutile : une succession d'opérations linéaires reste équivalente à une seule.Voir dans le glossaire en un seul kernel évite d'écrire puis de relire le tensor intermédiaire en mémoire HBM. Sur un tensor de plusieurs millions d'éléments, cela économise un aller-retour mémoire complet — souvent plus déterminant que le gain de calcul lui-même, la bande passante mémoire étant le facteur limitant sur la plupart des opérations element-wise.

Côté build, deux options : load() pour une compilation JIT pratique en développement, ou une extension packagée avec setup.py et CUDAExtension pour la production, qui évite de recompiler à chaque import.

torch.compile et TorchInductor : compiler plutôt qu'interpréter

Depuis PyTorch 2.x, torch.compile() capture le graphe d'exécution via TorchDynamo (qui intercepte le bytecode Python sans réécrire le modèle), le transforme en graphe FX d'opérations ATen, puis le passe à TorchInductor, qui génère des kernels Triton fusionnés — ou du code C++ vectorisé pour la partie CPU.

model = MyModel().cuda()
compiled_model = torch.compile(model, mode="max-autotune")

for batch in loader:
    out = compiled_model(batch)
    loss = criterion(out, target)
    loss.backward()
    optimizer.step()

L'intérêt principal d'Inductor est la fusion de kernels : au lieu d'exécuter dix opérations element-wise comme dix kernels séparés (dix lectures et dix écritures en mémoire globale), il les fusionne en un seul kernel qui ne lit et n'écrit qu'une fois. C'est le même principe que le custom op manuel ci-dessus, mais automatisé et appliqué à l'ensemble du graphe.

Le premier passage dans une fonction compilée déclenche la compilation elle-même ("warm-up"), visible comme un pic de latence anormal sur la première itération. Ne mesurez jamais le temps d'exécution sur le premier batch après un torch.compile : mesurez à partir de la troisième ou quatrième itération, une fois le cache de kernels stabilisé.

Un piège classique est le recompilation storm : si les formes (shapes) des tensors d'entrée varient d'un batch à l'autre (séquences de longueur variable, batch size dynamique), Dynamo recompile un nouveau graphe à chaque forme rencontrée, ce qui peut rendre torch.compile plus lent qu'un mode eager. Le paramètre dynamic=True indique à Dynamo de généraliser les dimensions variables plutôt que de spécialiser à chaque forme, au prix d'une optimisation un peu moins agressive.

TensorFlow et XLA : la compilation par graphe statique

TensorFlow adopte une philosophie différente. Même en mode eager par défaut, il propose XLA (Accelerated Linear Algebra), un compilateur qui prend un sous-graphe TensorFlow et le transforme en code machine optimisé pour le device cible, GPU ou TPU.

@tf.function(jit_compile=True)
def train_step(x, y):
    with tf.GradientTape() as tape:
        preds = model(x, training=True)
        loss = loss_fn(y, preds)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

jit_compile=True force XLA sur la fonction annotée. Comme TorchInductor, XLA fusionne les opérations et élimine les tensors intermédiaires inutiles ; contrairement à Inductor qui génère du Triton, XLA génère du code spécifique au backend via son propre pipeline HLO (High Level Optimizer).

XLA impose des formes statiques à la compilation. Un tf.function avec jit_compile=True appelé sur des formes changeantes déclenche une recompilation à chaque nouvelle forme, exactement comme Dynamo. La pratique standard consiste à paddiquer les séquences à une longueur fixe (ou à un petit nombre de longueurs discrètes, une technique dite de bucketing) plutôt que de laisser les formes varier librement.

Précision mixte automatique (AMP)

La précision mixte consiste à exécuter la majorité des calculs en FP16 ou BF16 — deux fois moins d'octets à déplacer, kernels Tensor Core dédiés bien plus rapides — tout en conservant certaines opérations sensibles (réductions, normalisations, la mise à jour des poids) en FP32 pour la stabilité numérique.

scaler = torch.cuda.amp.GradScaler()

for x, y in loader:
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

autocast choisit automatiquement, opération par opération, le dtype le plus adapté : convolutions et matmuls passent en FP16/BF16, softmax et normalisations restent en FP32. Le GradScaler compense le problème spécifique au FP16 : sa plage dynamique réduite fait que des gradients très petits sous-flottent (underflow) et deviennent zéro. Le scaler multiplie la loss par un facteur avant le backward(), puis divise les gradients par ce même facteur avant l'étape d'optimisation — ce qui préserve les gradients faibles dans la plage représentable du FP16.

Sur les GPU récents (Ampere et plus), BF16 a la même plage dynamique que FP32 (mêmes bits d'exposant) mais une précision réduite, ce qui évite le besoin d'un GradScaler — au prix d'une précision de mantisse plus faible que le FP16.

Éviter les copies host-device : le vrai goulot d'étranglement

Le transfert PCIe entre mémoire CPU (host) et mémoire GPU (device) est plusieurs ordres de grandeur plus lent que la bande passante interne du GPU (HBM). Un .to("cuda") ou .cpu() placé au mauvais endroit dans une boucle d'entraînement peut annuler tous les gains de compilation et de précision mixte.

Les sources les plus fréquentes de copies inutiles :

• Logging à chaque itération — appeler .item() ou .cpu().numpy() sur une loss ou une métrique force une synchronisation CUDA complète, car le CPU doit attendre que tous les kernels en attente se terminent avant de lire la valeur.
• DataLoader sans mémoire épinglée — sans pin_memory=True, chaque batch transite par de la mémoire paginable, ce qui empêche un transfert DMA asynchrone et bloque le CPU pendant la copie.
• Indexation avancée avec des tensors CPU — indexer un tensor GPU avec un tenseur d'indices resté sur CPU force une copie implicite de l'index vers le device, souvent invisible dans le code.
• Conversion .numpy() dans une fonction de perte custom — toute bibliothèque hors PyTorch/TensorFlow (SciPy, scikit-learn) appelée dans la boucle chaude ramène les donnéesdonnéesIAEnsemble d'informations structurées ou non utilisées pour entraîner, évaluer ou alimenter un modèle. La qualité, la quantité et la représentativité des données sont les facteurs décisifs pour les performances en apprentissage automatique.Voir dans le glossaire sur CPU.

pin_memory=True sur le DataLoader couplé à non_blocking=True sur le .to("cuda") permet un transfert asynchrone qui se chevauche avec le calcul du batch précédent. Le gain n'est réel que si le DataLoader utilise plusieurs workers (num_workers > 0) : sinon, le CPU reste de toute façon le facteur limitant pendant la préparation du batch suivant.

Synchronisations implicites à traquer

Certains appels, en apparence anodins, forcent une synchronisation CPU-GPU complète et cassent le pipeline asynchrone :

• print(tensor) ou toute conversion vers une chaîne
• if tensor.item() > seuil: dans une boucle de contrôle
• torch.cuda.synchronize() explicite laissé par du code de debug
• les assertions qui inspectent une valeur de tensor (assert not torch.isnan(loss).any()) exécutées à chaque itération plutôt qu'en mode debug seulement

Ces synchronisations ne provoquent pas d'erreur — le code reste correct — mais elles annulent le chevauchement calcul/transfert et rendent chaque itération strictement séquentielle. Un profil torch.profiler qui montre des blancs réguliers sur la timeline GPU, alignés avec des appels .item(), en est la signature caractéristique.

Checklist avant mise en production

1. Le DataLoader utilise pin_memory=True, num_workers > 0, et les transferts .to(device, non_blocking=True).
2. Aucune synchronisation implicite (.item(), print) dans la boucle chaude ; le logging des métriques est agrégé et déplacé hors device toutes les N itérations seulement.
3. torch.compile (ou jit_compile=True sous TensorFlow) est activé, avec des formes d'entrée stabilisées ou du bucketing pour éviter les recompilations.
4. AMP est actif avec le format adapté au matériel (BF16 sur Ampere/Hopper si la stabilité prime, FP16 + GradScaler sinon).
5. Un custom op n'est écrit qu'après avoir vérifié, profil à l'appui, qu'aucune combinaison d'opérations natives fusionnées par le compilateur n'atteint la même performance.
6. Le profil GPU (torch.profiler, Nsight Systems) confirme une occupation SM soutenue, sans trous réguliers correspondant à des copies ou synchronisations.

Ce qu'il faut retenir

La performance sur GPU n'est jamais uniquement une question de puissance de calcul brute : c'est une question d'orchestration entre CPU et GPU, de format numérique, et de fusion de kernels. torch.compile et XLA automatisent une grande partie de cette orchestration, mais aucun compilateur ne corrige un pipeline de données mal conçu ou des synchronisations implicites disséminées dans le code d'entraînement. Le réflexe à acquérir est de toujours vérifier, par le profilage, que l'optimisation appliquée se traduit effectivement par une occupation GPU plus dense — et non de se fier à la théorie.