Acúmulo de gradientes: batches efetivos sem perder a normalização correta

O que é

O acúmulo de gradientes processa vários microbatches antes de atualizar os parâmetros do modelo. Os gradientes de todos eles contribuem para um único passo do otimizador, o que permite um batch efetivo maior quando um batch físico não cabe na memória do dispositivo. No PyTorch, as chamadas de backward somam-se aos gradientes já existentes nos parâmetros até que esses gradientes sejam zerados. [1][2]

Como funciona

Para microbatches do mesmo tamanho e uma perda média sobre exemplos com o mesmo peso, divida cada perda pelo número de microbatches da janela de acúmulo. A soma desses gradientes resulta então no gradiente da perda média combinada. Se os tamanhos forem desiguais, é preciso ponderar pelo número real de exemplos que contribuem. Em perdas por token, conte os tokens válidos, isto é, não mascarados; pesos de classe ou outras formas de redução podem exigir um denominador diferente. Calcular a média das médias dos microbatches sem esse cuidado altera o objetivo. Essas regras de normalização decorrem da derivação da média ponderada pretendida.

Fluxo de trabalho prático

Zere os gradientes no início de cada janela, execute forward e backward para cada microbatch já normalizado e faça então um único passo do otimizador. Depois, zere os gradientes para a janela seguinte. Passagens forward separadas normalmente não exigem manter os grafos de computação anteriores. [2] Trate explicitamente a última janela, quando incompleta. Com precisão mista, mantenha a escala do gradiente constante durante todo o acúmulo, faça o unscale uma única vez após o último backward e aplique o clipping ao gradiente combinado antes do passo. Atualize o scaler nesse limite. [1] Um scheduler configurado por atualização do otimizador só deve avançar quando houver de fato uma atualização.

Exemplo ilustrativo

Suponha que um dispositivo processe quatro microbatches de oito exemplos. O batch efetivo tem 32 exemplos, de modo que a média de cada microbatch recebe peso 1/4. Uma última janela com 8, 8, 8 e 2 exemplos contém 26 exemplos, e os pesos correspondentes são 8/26, 8/26, 8/26 e 2/26. Isso é aritmética, não um resultado de treinamento medido.

Limites e verificação

A batch normalization usa estatísticas de cada microbatch físico, portanto o acúmulo não reproduz automaticamente um único batch grande. [3] Operações aleatórias e a ordem de soma em ponto flutuante também podem alterar os resultados. Aplicar clipping a cada microbatch separadamente geralmente muda a atualização combinada. O método reduz o armazenamento simultâneo de ativações; não elimina os pesos do modelo nem os estados do otimizador, e as passagens adicionais podem aumentar o tempo de execução.

Antes de uma execução longa, compare uma atualização acumulada com uma atualização de batch combinado em um modelo pequeno e determinístico, sem camadas que dependam do batch. Verifique a normalização da perda, as diferenças de gradiente dentro de uma tolerância adequada, a contagem de atualizações e o comportamento da janela parcial. Essas são verificações recomendadas, não testes realizados para este artigo.

Fontes

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d