Gradient Accumulation

A acumulação de gradientes é uma técnica de treino que executa vários lotes pequenos, chamados microlotes, na passagem direta e na retropropagação, soma os gradientes e só depois atualiza os pesos do modelo. A atualização aproxima um passo sobre um lote tão grande quanto o conjunto de todos os microlotes, enquanto a GPU só mantém na memória as ativações de um microlote de cada vez. Permite que um dispositivo com memória limitada treine com o tamanho de lote definido pela receita, sem reduzir a resolução das imagens nem diminuir o modelo. O tamanho efetivo do lote é o tamanho do microlote multiplicado pelo número de passos de acumulação e pelo número de dispositivos; por isso, microlotes de 4 imagens acumulados ao longo de 8 passos em 4 GPUs equivalem a um lote de 128.

Como funciona o acúmulo de gradientes#

No PyTorch e na maioria dos outros frameworks, cada chamada a backward() adiciona gradientes novos aos que já estão armazenados nos parâmetros, em vez de os substituir. A acumulação aproveita esse comportamento: executa a passagem direta e a retropropagação para cada microlote, chama o otimizador quando se atinge o número definido de microlotes, limpa os gradientes e inicia a janela seguinte:

A maioria das funções de perda calcula a média sobre as amostras de um lote; por isso, a perda de cada microlote é dividida pelo número de passos, para converter os gradientes somados numa média. Com esse escalamento, quatro microlotes de 8 produzem o mesmo gradiente que um lote de 32:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

Imprime Matches the full-batch gradient: True, com uma diferença máxima de 1.6e-07, resultante do arredondamento em vírgula flutuante ao somar os mesmos números por uma ordem diferente.

Em que difere a acumulação de um lote grande#

A equivalência mantém-se para funções de perda que calculam a média de amostras independentes, mas deixa de se verificar em três casos comuns:

• Normalização por lote: Calcula a média e a variância de cada microlote, por isso as estatísticas de 4 imagens são mais ruidosas do que as de 128.
• Sequências de comprimento variável: Quando um modelo de linguagem calcula a média da perda sobre os tokens de cada microlote, os microlotes com poucos tokens têm o mesmo peso que os mais longos.
• Recorte de gradientes: Deve ser aplicado uma única vez ao gradiente acumulado, imediatamente antes do passo do otimizador; recortar cada microlote altera o resultado.

Em outubro de 2024, a Hugging Face e a Unsloth corrigiram a incompatibilidade na contagem de tokens no ajuste fino de modelos de linguagem de grande escala (LLM), dividindo a perda somada pela contagem total de tokens de toda a janela de acumulação. A acumulação também não acelera o processamento: o mesmo número de amostras continua a passar pela rede um microlote de cada vez, e os microlotes pequenos utilizam a GPU com menos eficiência. No treino distribuído, o DistributedDataParallel do PyTorch oferece um contexto no_sync() que ignora a redução total de gradientes em todos os microlotes, exceto no último, poupando assim comunicação. Um lote efetivo maior costuma exigir uma taxa de aprendizagem mais elevada: investigadores do Facebook treinaram o ResNet-50 no ImageNet em uma hora com lotes de 8.192 imagens em 256 GPUs, aumentando linearmente a taxa de aprendizagem com o tamanho do lote e aplicando um aquecimento gradual.

Técnicas relacionadas#

• Checkpointing de gradientes: Recalcula as ativações durante a retropropagação em vez de as armazenar, poupando memória em cada microlote à custa de mais cálculo.
• Precisão mista: Armazena as ativações em formatos de 16 bits, o que permite acomodar microlotes maiores. O escalador da perda desfaz o escalamento dos gradientes uma única vez, após o último microlote.
• Paralelismo de pipeline: Divide cada lote em microlotes que percorrem etapas do modelo em diferentes GPUs e acumula os gradientes antes de um único passo do otimizador.

O treino Ultralytics YOLO26 acumula gradientes por predefinição: no modo de treino, o treinador executa um passo do otimizador a cada round(nbs / batch) lotes; por isso, o valor predefinido de batch=16 e o tamanho nominal do lote nbs=64 resultam em quatro lotes por atualização, e batch=8 resulta em oito. O número aumenta gradualmente a partir de um durante o aquecimento, e o decaimento dos pesos é escalado por batch × accumulation / nbs, pelo que os hiperparâmetros predefinidos para um lote de 64 imagens continuam a ser aplicáveis numa GPU que só consegue manter algumas imagens de cada vez.