Acumulação de gradiente

Lotes grandes produzem ativações grandes que esgotam a memória da GPU. A acumulação de gradientes permite treinar com um tamanho de lote efetivo maior, distribuindo o cálculo dos gradientes entre vários mini-lotes.

Os gradientes são acumulados ao longo de n mini-lotes antes de o otimizador atualizar os pesos. Por exemplo, com um tamanho de lote por dispositivo de 8 e 4 etapas de acumulação, o tamanho de lote efetivo é 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 a acumulação de gradiente apenas quando um lote maior não couber na memória. Ela não melhora a taxa de processamento em comparação com o treinamento com um lote grande real.

Acumule os gradientes por gradient_accumulation_steps em per_device_train_batch_size.

from transformers import TrainingArguments

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

Escala de perda

Para uma função de perda personalizada, inclua num_items_in_batch para que o Trainer divida a perda pelo número de alvos de previsão em todos os mini-lotes. Isso normaliza por tokens, em vez de por uma contagem fixa de etapas com gradient_accumulation_steps. Caso contrário, o Trainer divide a perda 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 modelos de LM causal, num_items_in_batch conta os rótulos deslocados. A perda desloca os rótulos para que a previsão na posição i tenha como alvo o token na posição i + 1, o que deixa a posição 0 de cada sequência sem um alvo. O Trainer exclui essas posições e conta sobre labels[..., 1:], de modo que o denominador corresponda ao número de alvos de previsão usados pela perda. Quando um data collator fornece shift_labels diretamente, como um collator sem padding, o Trainer conta sobre esse tensor em vez disso. Outros tipos de perda, como LM mascarado e classificação, contam o tensor de rótulos completo.

Se você criar uma subclasse de Trainer e sobrescrever compute_loss(), defina loss_is_scaled_for_ga para indicar como sua perda é normalizada: False para uma média por lote, que o Trainer então divide por gradient_accumulation_steps, ou True se você já a dividiu por 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óximas etapas

• Leia o documento uso de memória da GPU para entender o que está causando o uso de memória na GPU durante o treinamento.
• Consulte o guia Gradient checkpointing para aprender a reduzir a memória de ativações recalculando as ativações em vez de armazená-las em cache.
• Consulte o guia Mixed precision training para aprender a usar tipos de dados de menor precisão para reduzir a memória e acelerar o treinamento.
• Leia o post no blog Gradient Accumulation Fix para saber como o acúmulo de gradientes é calculado.