Accumulation des gradients dans PyTorch : un exemple testé

L’accumulation des gradients combine les gradients de plusieurs petits lots avant que l’optimiseur mette à jour le modèle. Avec PyTorch, chaque micro-lot passe par une propagation avant et une rétropropagation, puis un seul appel à optimizer.step() applique la mise à jour du groupe. Vous augmentez ainsi le nombre d’exemples qui contribuent à une mise à jour tout en traitant moins d’exemples à la fois.

Le calcul doit rester cohérent : la perte doit utiliser le bon dénominateur, le dernier groupe incomplet doit être traité et les gradients doivent être conservés entre les micro-lots. Ce guide propose une petite boucle d’entraînement qui couvre ces cas, puis précise les adaptations nécessaires pour la précision mixte et plusieurs GPU.

Comment fonctionne l’accumulation des gradients dans PyTorch ?

PyTorch ajoute les gradients de chaque rétropropagation aux attributs .grad existants des paramètres. Il ne les efface pas automatiquement. Une boucle d’accumulation exploite ce fonctionnement : remettre les gradients à zéro au début d’une fenêtre, appeler backward() pour chaque micro-lot, puis mettre à jour les paramètres une fois. Le guide PyTorch sur la remise à zéro des gradients explique ce mécanisme.

Avec des micro-lots de même taille, la taille nominale du lot effectif est :

Taille du micro-lot × étapes d’accumulation × processus de parallélisme de données.

Par exemple, quatre exemples par micro-lot et trois étapes d’accumulation donnent 12 exemples par mise à jour sur un processus. Deux processus qui traitent des exemples différents avec ces mêmes réglages donnent un lot global nominal de 24. Ces nombres concernent les fenêtres complètes ; une dernière fenêtre plus courte contient moins d’exemples.

Consignez séparément les micro-lots traités, les exemples ou tokens valides traités et les mises à jour de l’optimiseur. Modifier le nombre d’étapes d’accumulation change le nombre de mises à jour par époque, sauf si un autre réglage compense ce changement.

Normaliser la perte sur la fenêtre réelle

Si chaque micro-lot contient le même nombre d’exemples de poids égal, diviser sa perte moyenne par le nombre de micro-lots de la fenêtre donne la même contribution à chaque exemple. En revanche, utiliser systématiquement le nombre d’étapes configuré ne convient pas lorsque la dernière fenêtre est plus courte.

Des micro-lots de tailles différentes demandent un autre ajustement. Une moyenne calculée sur quatre exemples ne doit pas recevoir le même poids qu’une moyenne sur un seul exemple. Le code ci-dessous additionne les pertes de chaque micro-lot et les divise par le nombre total d’exemples de la fenêtre. Il traite ainsi les tailles inégales et la dernière fenêtre incomplète.

Ce dénominateur convient à une classification avec un indice de classe par exemple, sans pondération des classes ni cibles ignorées. Pour un entraînement par tokens, comptez les cibles de prédiction valides. Avec des classes pondérées ou des étiquettes ignorées, suivez les règles de réduction de la fonction de perte ; CrossEntropyLoss documente ces distinctions. Ne remplacez pas un nombre de tokens par un nombre de séquences sans vérifier l’objectif.

Un exemple d’accumulation des gradients avec PyTorch

Enregistrez ce code dans accumulation_demo.py, dans un environnement où PyTorch est installé. La démonstration autonome utilise des données synthétiques sur CPU. Le chargeur fournit des tenseurs CPU ; la boucle conserve une fenêtre de lots en mémoire hôte et transfère un micro-lot à la fois vers l’appareil du modèle.

from itertools import islice
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


def train_epoch(model, loader, optimizer, accumulation_steps=3):
    if accumulation_steps < 1:
        raise ValueError('accumulation_steps must be positive')
    model.train()
    device = next(model.parameters()).device
    criterion = nn.CrossEntropyLoss(reduction='sum')
    batches = iter(loader)
    update_sizes = []

    while True:
        window = list(islice(batches, accumulation_steps))
        if not window:
            break
        examples = sum(y.size(0) for _, y in window)
        optimizer.zero_grad(set_to_none=True)

        for x, y in window:
            x, y = x.to(device), y.to(device)
            loss = criterion(model(x), y) / examples
            loss.backward()

        optimizer.step()
        update_sizes.append(examples)
    return update_sizes


if __name__ == '__main__':
    torch.manual_seed(7)
    x = torch.randn(19, 5)
    y = torch.randint(0, 3, (19,))
    loader = DataLoader(TensorDataset(x, y), batch_size=4,
                        shuffle=False, drop_last=False)
    model = nn.Linear(5, 3)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    print(train_epoch(model, loader, optimizer))

Exécutez le fichier avec :

python accumulation_demo.py

Le résultat indique le nombre d’exemples ayant contribué à chaque mise à jour :

[12, 7]

La première mise à jour utilise trois micro-lots de quatre exemples. La seconde utilise deux micro-lots de quatre et trois exemples. Chaque perte est divisée par le nombre réel d’exemples de sa fenêtre : les sept derniers exemples sont donc inclus avec la pondération prévue pour une perte moyenne.

Ce regroupement utilise de la mémoire hôte supplémentaire pour rendre le dénominateur explicite. Conservez les données du chargeur sur CPU et tenez compte de la taille de la fenêtre. Pour des données en flux ou des entrées particulièrement volumineuses, prévoyez un regroupement borné qui conserve la même normalisation. Charger toute la fenêtre sur le GPU compromettrait le gain mémoire recherché.

Ce que les tests de l’exemple établissent

Nous avons exécuté la démonstration avec PyTorch 2.15.0+cpu. Dix comparaisons supplémentaires ont porté sur SGD avec momentum et AdamW : fenêtres complètes ou incomplètes, micro-lots inégaux, jeu de données plus petit qu’un micro-lot et accumulation sur une seule étape. Avec un réseau déterministe sans couches stochastiques ni dépendance entre exemples d’un lot, les gradients accumulés et les paramètres mis à jour correspondaient au calcul sur des lots concaténés, à chaque mise à jour et dans la tolérance numérique fixée.

Ces comparaisons utilisaient des tenseurs CPU en float64. Elles vérifient le calcul et les limites des mises à jour, sans mesurer les économies de mémoire GPU, le débit, la précision mixte ou l’exécution distribuée.

Respecter les limites de chaque mise à jour

Appeler zero_grad() à l’intérieur de la boucle des micro-lots effacerait les contributions précédentes. Y appeler step() modifierait les paramètres avant la fin de la fenêtre. Dans les deux cas, la comparaison recherchée avec une seule mise à jour sur un lot plus grand ne tient plus.

Évitez aussi d’additionner les tenseurs de perte des micro-lots pour reporter toutes les rétropropagations si vous souhaitez libérer rapidement les activations sauvegardées. L’exemple appelle backward séparément. Une accumulation ordinaire entre de nouvelles propagations avant ne nécessite pas retain_graph=True ; consultez les règles de durée de vie du graphe de backward.

Définissez l’unité du programme de taux d’apprentissage avant de l’ajouter à la boucle. S’il est défini par mise à jour de l’optimiseur, faites-le avancer à cette limite, plutôt qu’après chaque micro-lot. S’il dépend des époques ou des validations, conservez ce fonctionnement. Consignez le nombre de mises à jour pour qu’un changement d’accumulation ne modifie pas silencieusement la montée initiale ou la décroissance du taux.

Précision mixte et écrêtage des gradients

En FP16 avec un gradient scaler, gardez les gradients à l’échelle et conservez ce facteur pendant toute la fenêtre. Après la dernière rétropropagation, retirez l’échelle une fois si un écrêtage est nécessaire, écrêtez les gradients accumulés, puis appelez les méthodes step et update du scaler. Des gradients non finis peuvent entraîner l’annulation de la mise à jour. Suivez l’exemple officiel d’accumulation avec AMP.

Pour l’installation et les choix de précision, consultez notre guide PyTorch sur A100. L’exemple CPU présenté ici ne valide aucune configuration CUDA ou AMP.

Plusieurs GPU et synchronisation DDP

Avec DistributedDataParallel, no_sync() peut reporter la synchronisation des premiers micro-lots. Incluez la propagation avant et la rétropropagation dans ce contexte, puis exécutez la dernière paire hors du contexte. Synchronisez avant la mise à jour de l’optimiseur. La documentation DDP précise ce fonctionnement.

La formule simple du lot global suppose des contributions égales entre processus. Des nombres d’exemples ou de tokens valides différents exigent une normalisation entre processus, en tenant compte de la moyenne des gradients effectuée par DDP. Coordonnez les limites des fenêtres entre les rangs. L’exemple à un seul processus ne constitue pas une boucle distribuée prête à l’emploi.

Quand l’accumulation diffère d’un lot physique plus grand

Notre test d’équivalence exclut volontairement les couches qui relient les exemples d’un lot. La normalisation par lots calcule ses statistiques sur le mini-lot courant : le diviser peut donc modifier les sorties et les statistiques suivies. Dropout et d’autres opérations aléatoires peuvent aussi changer la trajectoire numérique. Une même taille de lot effectif ne garantit pas les mêmes résultats d’entraînement.

Contrôlez l’ordre des données, les poids initiaux, la précision, les réglages de l’optimiseur et le nombre de mises à jour pour comparer les essais. Le guide de reproductibilité PyTorch rappelle qu’une même graine ne garantit pas des résultats identiques entre plateformes ou versions. Choisissez des tolérances et des mesures de validation adaptées.

L’accumulation réduit les besoins liés aux activations d’un grand lot physique, mais ne supprime ni les paramètres, ni les gradients, ni l’état de l’optimiseur. Elle répète aussi les propagations avant et arrière pour chaque mise à jour. Mesurez la mémoire et le nombre d’exemples ou de tokens valides traités par seconde avant de choisir un réglage. Davantage d’étapes d’accumulation ne garantit pas un entraînement plus rapide.

Vérifications avant un entraînement plus long

• Vérifiez que chaque exemple prévu contribue à une mise à jour, y compris dans la dernière fenêtre partielle.
• Comparez la pondération de la perte à une petite référence sur un lot concaténé.
• Consignez séparément les mises à jour, les micro-lots et le volume de données.
• Placez l’écrêtage et le programme de taux aux limites prévues par votre entraînement.
• Comparez la qualité de validation, la mémoire et le débit.
• Enregistrez les checkpoints après une mise à jour terminée, sauf si votre implémentation sauvegarde aussi les gradients d’une accumulation inachevée.

Si votre outil d’entraînement gère déjà l’accumulation, vérifiez ses réglages avant d’ajouter une deuxième boucle. Notre guide du fine-tuning LoRA montre où ce paramètre intervient. Commencez par un petit essai contrôlé qui vérifie le comportement attendu des mises à jour.