Treinamento de precisão mista: autocasting, escalonamento de perda e limites numéricos

O que é

O treinamento de precisão mista combina formatos de ponto flutuante em uma mesma execução de treinamento. As operações matriciais mais caras podem usar FP16 ou BF16, enquanto as operações sensíveis continuam em FP32. A precisão mista automática escolhe os formatos segundo regras por operação e por dispositivo; isso não equivale a converter todos os tensores para meia precisão. Os benefícios dependem do modelo, do hardware e da carga de trabalho. [1][2]

Como funciona

O autocasting escolhe a precisão de cada operação. O escalonamento de perda (loss scaling) resolve um problema distinto: gradientes FP16 muito pequenos podem se tornar zero. Multiplicar a perda por um fator de escala amplia os gradientes no backward; remover esse fator antes da atualização do otimizador preserva a magnitude pretendida. Um scaler dinâmico detecta gradientes não finitos, ignora as atualizações do otimizador afetadas e ajusta a escala. Ele não consegue tornar válido todo modelo numericamente instável. [1][2]

Fluxo de trabalho prático

Estabeleça uma referência em FP32, com dados e regras de avaliação fixos. Em um treinamento comum em PyTorch no modo eager, envolva o forward pass e o cálculo da perda em torch.autocast e execute o backward fora desse contexto. Com FP16, use torch.amp.GradScaler quando apropriado. Não converta manualmente o modelo inteiro só para habilitar o autocasting. [1]

Se você fizer clipping de gradientes, desfaça a escala (unscale) antes do clipping. Chame os métodos step e update do scaler na ordem documentada e depois zere os gradientes para a próxima atualização. Salve o estado do scaler junto com os checkpoints do modelo e do otimizador. [2] Com acúmulo de gradientes, mantenha uma única escala em todo o batch efetivo. A execução compilada e as operações customizadas exigem verificações de precisão próprias, específicas de cada versão.

Exemplo ilustrativo

Considere um classificador cujas convoluções dominam o tempo de execução, mas cuja perda final precisa de maior faixa numérica. O autocasting pode executar as convoluções elegíveis em precisão menor e avaliar em FP32 uma perda sensível que seja suportada. Se um passo produzir gradientes infinitos, o scaler pode ignorar essa atualização; isso é um evento de diagnóstico, não uma evidência de que o treinamento convergiu. O exemplo descreve um comportamento, não um ganho de velocidade medido.

Limites e verificação

FP16 e BF16 têm faixas numéricas diferentes. Um modelo treinado para BF16 pode sofrer overflow em FP16, cujo maior valor finito é 65.504. Aumentar a escala da perda não corrige um cálculo de forward que já estoura. O BF16 costuma dispensar o escalonamento de gradientes típico do FP16, mas ainda introduz arredondamento e exige validação. [1]

Compare as métricas de validação e as trajetórias da perda com a referência em FP32, inspecione valores não finitos e atualizações ignoradas, e meça o pico de memória e a vazão (throughput) após o warm-up. Modelos pequenos ou limitados pela entrada de dados podem ganhar pouco. [2] Registre a precisão, o hardware, as versões de software e as diferenças observadas. Este artigo não relata nenhum experimento local.

Fontes

[1] PyTorch: Automatic Mixed Precision package [2] PyTorch: Automatic Mixed Precision recipe