Accumulation du gradient

Les lots volumineux génèrent de grandes activations qui épuisent la mémoire du GPU. L’accumulation de gradients permet d’entraîner avec une taille de lot effective plus grande en répartissant le calcul des gradients sur plusieurs mini-lots.

Les gradients s'accumulent sur n mini-lots avant que l'optimiseur ne mette à jour les poids. Par exemple, avec une taille de lot par appareil de 8 et 4 étapes d'accumulation, la taille de lot effective est de 32.

Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()

Utilisez l’accumulation de gradients uniquement lorsqu’un lot plus volumineux ne tient pas en mémoire. Elle n’améliore pas le débit par rapport à l’entraînement avec un véritable lot de grande taille.

Accumuler les gradients sur gradient_accumulation_steps à travers per_device_train_batch_size.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)

Mise à l'échelle de la perte

Pour une fonction de perte personnalisée, incluez num_items_in_batch afin que Trainer divise la perte par le nombre de cibles de prédiction sur l’ensemble des mini-lots. Cela normalise par jetons plutôt que par un nombre fixe d’étapes avec gradient_accumulation_steps. Sinon, Trainer divise la perte par gradient_accumulation_steps.

import torch.nn.functional as F

def compute_loss(outputs, labels, num_items_in_batch=None):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

Pour les modèles de LM causaux, num_items_in_batch compte les étiquettes décalées. La perte décale les étiquettes de sorte que la prédiction à la position i cible le jeton à la position i + 1, ce qui laisse la position 0 de chaque séquence sans cible. Trainer exclut ces positions et compte sur labels[..., 1:], de sorte que le dénominateur correspond au nombre de cibles de prédiction utilisées par la perte. Lorsqu'un collateur de données fournit directement shift_labels, comme un collateur sans bourrage, Trainer compte sur ce tenseur à la place. D'autres types de perte, comme le LM masqué et la classification, comptent le tenseur d'étiquettes complet.

Si vous sous-classez Trainer et surchargez compute_loss(), définissez loss_is_scaled_for_ga pour indiquer comment votre perte est normalisée : False pour une moyenne par lot, que Trainer divise ensuite par gradient_accumulation_steps, ou True si vous l’avez déjà divisée par num_items_in_batch.

class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch

Étapes suivantes

• Consultez le document utilisation de la mémoire GPU pour comprendre ce qui entraîne l'utilisation de la mémoire sur le GPU pendant l'entraînement.
• Consultez le guide Gradient checkpointing pour apprendre à réduire la mémoire d’activation en recalculant les activations au lieu de les mettre en cache.
• Consultez le guide Entraînement en précision mixte pour apprendre à utiliser des types de données de précision inférieure afin de réduire la mémoire et d’accélérer l’entraînement.
• Lisez le billet de blog Gradient Accumulation Fix pour comprendre comment l'accumulation des gradients est calculée.