Écrêtage des gradients PyTorch : norme, ordre et AMP

Dans PyTorch, appliquez l’écrêtage des gradients après la rétropropagation et avant la mise à jour des paramètres par l’optimiseur. Si vous accumulez les gradients, terminez d’abord la fenêtre d’accumulation. Si vous utilisez un facteur d’échelle, annulez son effet sur les gradients avant de les écrêter.

Ce guide rend cet ordre concret avec quatre petits exemples sur CPU. Vous examinerez une norme, comparerez l’écrêtage par norme à l’écrêtage par valeur et observerez l’effet de deux erreurs d’ordre. Vingt contrôles locaux ont réussi avec Python 3.12.14 et PyTorch 2.15.0+cpu. Aucun jeu de données ni modèle à télécharger n’est nécessaire.

Commencer avec un gradient vérifiable à la main

L’API clip_grad_norm_ modifie les gradients en place et utilise par défaut la norme L2. La liste de paramètres définit le groupe à écrêter. Dans un modèle dont tous les paramètres sont gérés par un seul optimiseur, il s’agit généralement de model.parameters().

Ici, une perte linéaire artificielle produit le gradient [3, 4], de norme L2 égale à 5. Le seuil de 1 facilite le calcul. Ce n’est pas une recommandation valable pour tous les modèles. Chaque bloc Python ci-dessous peut être exécuté séparément dans un environnement où PyTorch est installé.

import torch

weights = torch.nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float64))
optimizer = torch.optim.SGD([weights], lr=0.1)
optimizer.zero_grad(set_to_none=True)

# Synthetic loss chosen to produce the gradient [3, 4].
loss = 3 * weights[0] + 4 * weights[1]
loss.backward()
before = torch.nn.utils.clip_grad_norm_(
    [weights], max_norm=1.0, error_if_nonfinite=True
)
print('Returned norm:', round(before.item(), 6))
print('Norm after:', round(weights.grad.norm().item(), 6))
print('Gradient:', [round(x, 6) for x in weights.grad.tolist()])
optimizer.step()
print('Weights:', [round(x, 6) for x in weights.tolist()])

Sortie observée :

Returned norm: 5.0
Norm after: 1.0
Gradient: [0.6, 0.8]
Weights: [0.94, -1.08]

La norme renvoyée vaut 5 : elle est mesurée avant l’écrêtage. Le gradient devient environ [0.6, 0.8], avec une norme légèrement inférieure à 1, arrondie à 1.0. La valeur renvoyée ne devient pas la nouvelle norme ; examinez les gradients après l’opération si vous avez besoin de cette mesure.

L’écrêtage seul ne modifie pas les paramètres. L’étape SGD suivante soustrait 0.1 fois le gradient écrêté, faisant passer les poids de [1, -1] à environ [0.94, -1.08]. Cette perte linéaire sert à rendre le calcul visible ; elle ne constitue ni un objectif d’entraînement utile ni une preuve de convergence.

Écrêter un groupe de paramètres cohérent

Nous avons aussi réparti le même gradient entre deux paramètres distincts. Leur écrêtage commun a renvoyé 5 et ramené leurs composantes à environ 0.6 et 0.8. Appeler la fonction séparément pour chaque paramètre imposerait une autre règle : chacun aurait sa propre limite.

Choisissez le groupe de paramètres explicitement. Si l’optimiseur ne met à jour qu’une partie du modèle, utilisez le groupe voulu au lieu de supposer que tous les paramètres doivent participer à cette opération. Les exemples emploient un seul optimiseur et des gradients denses.

Les contrôles ont également couvert des gradients nuls, un gradient inférieur au seuil et un paramètre dont le gradient était None. Le zéro est resté nul, [0.3, 0.4] est resté inchangé avec une limite de 1 et le gradient absent n’a pas influencé la norme calculée. Un appel réussi ne prouve pas que chaque paramètre attendu a reçu un gradient.

Distinguer l’écrêtage par norme de l’écrêtage par valeur

clip_grad_value_ borne chaque composante entre l’opposé du seuil et le seuil positif. Appliquons cette opération au gradient initial [3, 4] :

import torch

weights = torch.nn.Parameter(torch.zeros(2, dtype=torch.float64))
(3 * weights[0] + 4 * weights[1]).backward()
torch.nn.utils.clip_grad_value_([weights], clip_value=1.0)
print('Value-clipped gradient:', weights.grad.tolist())
print('Its L2 norm:', round(weights.grad.norm().item(), 6))

Sortie observée :

Value-clipped gradient: [1.0, 1.0]
Its L2 norm: 1.414214

Les deux composantes respectent maintenant la limite de 1, mais leur norme L2 commune vaut environ 1.414214. Dans le premier exemple, l’écrêtage par norme applique un facteur commun et conserve le rapport entre les composantes. L’écrêtage par valeur fait passer ce rapport de 3:4 à 1:1. Choisissez l’opération qui correspond à la règle recherchée.

Terminer l’accumulation avant l’écrêtage

Considérons deux micro-lots de même poids, dont les gradients non normalisés sont [6, 0] et [0, 8]. Diviser chaque perte artificielle par deux donne un gradient accumulé de [3, 4]. Écrêtez ce gradient final une seule fois, pour la prochaine étape de l’optimiseur.

La comparaison suivante inclut volontairement l’option incorrecte consistant à écrêter après chaque rétropropagation :

import torch

def accumulated_gradient(clip_after_each):
    weights = torch.nn.Parameter(torch.zeros(2, dtype=torch.float64))
    # Two equally weighted synthetic microbatches.
    for coefficients in ([6.0, 0.0], [0.0, 8.0]):
        loss = (weights * weights.new_tensor(coefficients)).sum() / 2
        loss.backward()
        if clip_after_each:
            torch.nn.utils.clip_grad_norm_(
                [weights], 1.0, error_if_nonfinite=True
            )
    if not clip_after_each:
        torch.nn.utils.clip_grad_norm_(
            [weights], 1.0, error_if_nonfinite=True
        )
    return [round(x, 6) for x in weights.grad.tolist()]

print('Clip once at the end:', accumulated_gradient(False))
print('Clip after each backward:', accumulated_gradient(True))

Sortie observée :

Clip once at the end: [0.6, 0.8]
Clip after each backward: [0.242535, 0.970142]

Dans le second cas, l’écrêtage réduit la première contribution avant l’ajout de la seconde. La direction finale diffère de celle obtenue par un seul écrêtage à la fin, même si les deux normes finales sont proches de 1. Placer l’opération dans la boucle de micro-lots peut donc modifier la mise à jour que vous vouliez calculer.

La division par deux suppose deux contributions de même poids. Des tailles de lots inégales, des cibles masquées ou une dernière fenêtre incomplète exigent une normalisation adaptée. Notre guide d’accumulation des gradients, en anglais présente la boucle complète et la gestion des fenêtres. Traitez le choix de normalisation séparément de celui du seuil.

Annuler la mise à l’échelle avant l’écrêtage avec GradScaler

Les exemples PyTorch de précision mixte placent scaler.unscale_(optimizer) après toutes les rétropropagations de l’étape et avant l’écrêtage. Appelez cette méthode une fois par optimiseur et par étape, puis utilisez scaler.step(optimizer) et scaler.update(). Conservez le même facteur pendant une fenêtre d’accumulation.

Cette démonstration autonome sur CPU vérifie cet ordre avec un facteur initial de 8. Elle utilise des tenseurs float64 sans autocast ; elle ne teste ni l’arithmétique FP16 ni l’exécution CUDA.

import torch

# CPU demonstration of scaling order, without autocast or FP16.
weights = torch.nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float64))
optimizer = torch.optim.SGD([weights], lr=0.1)
scaler = torch.amp.GradScaler('cpu', init_scale=8.0)
optimizer.zero_grad(set_to_none=True)
loss = 3 * weights[0] + 4 * weights[1]
scaler.scale(loss).backward()
print('Scaled norm:', round(weights.grad.norm().item(), 6))

scaler.unscale_(optimizer)  # Once, after all backward calls for this step.
before = torch.nn.utils.clip_grad_norm_(
    [weights], 1.0, error_if_nonfinite=True
)
print('Unscaled norm before clipping:', round(before.item(), 6))
print('Clipped norm:', round(weights.grad.norm().item(), 6))
scaler.step(optimizer)
scaler.update()
print('Weights:', [round(x, 6) for x in weights.tolist()])

Sortie observée :

Scaled norm: 40.0
Unscaled norm before clipping: 5.0
Clipped norm: 1.0
Weights: [0.94, -1.08]

La norme mise à l’échelle vaut 40. L’annulation de l’échelle rétablit 5, puis l’écrêtage produit les mêmes poids que l’exemple initial. Dans un test distinct et volontairement incorrect, nous avons écrêté le gradient encore multiplié par le facteur à une norme de 1, puis annulé l’échelle. Sa norme finale valait environ 0.125, soit un huitième de la limite voulue.

Les tests ont aussi confirmé qu’un second appel à unscale_ avant l’étape provoque une erreur. Notre guide FP16, en anglais explique les questions de précision et de plage numérique qui motivent la précision mixte. Le test CPU présenté ici établit l’ordre des opérations, pas les performances ou le comportement numérique d’un GPU.

Traiter séparément les gradients non finis

Les appels d’écrêtage par norme utilisent error_if_nonfinite=True. Dans nos tests, des gradients contenant NaN, l’infini positif ou l’infini négatif ont déclenché une exception avant l’étape suivante de l’optimiseur, et les poids sont restés inchangés. Il s’agit d’un arrêt immédiat sur erreur. L’écrêtage n’a pas réparé les valeurs invalides.

Dans ce cas, examinez le lot, le calcul de la perte, les réglages de précision et les opérations précédentes. Conservez les informations nécessaires pour reproduire l’échec. Réduire le seuil ne permet pas, à lui seul, de comprendre l’origine d’une valeur non finie.

Avec un facteur d’échelle, ce choix d’arrêt immédiat signifie aussi que l’exception interrompt l’exemple avant que le scaler puisse omettre l’étape selon son comportement habituel. Si la boucle doit reprendre en sautant une mise à jour non finie, prévoyez explicitement ce fonctionnement ; le code ci-dessus ne l’implémente pas.

Choisir le seuil à partir du comportement observé

Commencez par enregistrer la norme renvoyée avant écrêtage, le seuil, la perte, le taux d’apprentissage et la fréquence d’application de l’écrêtage. Conservez aussi les réglages de réduction de la perte et d’accumulation : changer leur échelle modifie le sens d’un seuil fixe.

Comparez plusieurs seuils selon le résultat recherché, par exemple la perte de validation, en contrôlant les autres conditions de l’expérience. Si l’écrêtage intervient presque à chaque étape, vérifiez si le seuil est volontairement restrictif ou si un autre réglage demande une correction. Ces petits exemples n’établissent aucun seuil optimal universel.

Avec SGD sans momentum ni décroissance des poids, la mise à jour de notre exemple est le gradient écrêté multiplié par le taux d’apprentissage. Ne généralisez pas cette relation à tous les optimiseurs : Adam utilise des statistiques glissantes des gradients ; le momentum et la décroissance des poids peuvent aussi modifier la mise à jour. Une limite de norme du gradient ne borne pas universellement le déplacement des paramètres à la même valeur.

Les 20 contrôles couvrent les gradients denses sur CPU, les blocs affichés, l’ordre de l’accumulation, les erreurs non finies et la mise à l’échelle sur CPU. Ils ne valident ni les gradients creux, ni l’entraînement distribué ou réparti en fragments, ni autocast sur CUDA, ni la précision d’un modèle. L’environnement a émis un avertissement d’initialisation lié à l’absence de NumPy ; ces exemples utilisant uniquement des tenseurs ont abouti sans cette bibliothèque.