Acumulación de gradientes

Los lotes grandes generan activaciones grandes que agotan la memoria de la GPU. La acumulación de gradientes permite entrenar con un tamaño de lote efectivo mayor distribuyendo el cálculo de gradientes en varios mini-lotes.

Los gradientes se acumulan a lo largo de n mini-lotes antes de que el optimizador actualice los pesos. Por ejemplo, con un tamaño de lote por dispositivo de 8 y 4 pasos de acumulación, el tamaño de lote efectivo es 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()

Use la acumulación de gradientes solo cuando un lote más grande no quepa en la memoria. No mejora el rendimiento en comparación con el entrenamiento con un lote grande real.

Acumula los gradientes durante gradient_accumulation_steps en per_device_train_batch_size.

from transformers import TrainingArguments

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

Escala de pérdida

Para una función de pérdida personalizada, incluya num_items_in_batch para que Trainer divida la pérdida por el número de objetivos de predicción en todos los mini lotes. Esto normaliza por tokens en lugar de por un número fijo de pasos con gradient_accumulation_steps. De lo contrario, Trainer divide la pérdida por 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

Para los modelos de LM causal, num_items_in_batch cuenta las etiquetas desplazadas. La pérdida desplaza las etiquetas de modo que la predicción en la posición i apunta al token en la posición i + 1, lo que deja la posición 0 de cada secuencia sin objetivo. Trainer excluye esas posiciones y cuenta sobre labels[..., 1:], de modo que el denominador coincide con el número de objetivos de predicción que utiliza la pérdida. Cuando un data collator proporciona shift_labels directamente, como un collator sin padding, Trainer cuenta sobre ese tensor en su lugar. Otros tipos de pérdida, como LM enmascarado y clasificación, cuentan el tensor de etiquetas completo.

Si creas una subclase de Trainer y sobrescribes compute_loss(), establece loss_is_scaled_for_ga para indicar cómo se normaliza tu pérdida: False para una media por lote, que Trainer luego divide entre gradient_accumulation_steps, o True si ya la dividiste entre 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

Próximos pasos

• Lea el documento Uso de la memoria de la GPU para comprender qué factores determinan el uso de la memoria en la GPU durante el entrenamiento.
• Consulte la guía de Gradient checkpointing para aprender a reducir la memoria de activaciones mediante el recálculo de activaciones en lugar de almacenarlas en caché.
• Consulte la guía de Entrenamiento con precisión mixta para aprender a usar tipos de datos de menor precisión para reducir la memoria y acelerar el entrenamiento.
• Lee el artículo de blog Gradient Accumulation Fix para aprender cómo se calcula la acumulación de gradientes.