Treinamento com precisão mista

O treinamento em precisão completa (fp32) armazena e calcula tudo em 32 bits. A precisão mista usa fp16 ou bf16 para as etapas de propagação direta e reversa, que são intensivas em computação, mantendo uma cópia dos pesos em fp32 para a atualização do otimizador. A computação fica mais rápida, a memória de pesos e ativações é reduzida e a estabilidade do treinamento é preservada.

┌─────────────────────────────────────────────────────┐
│           MIXED PRECISION TRAINING LOOP             │
│                                                     │
│  fp32 master weights ──cast──▶ fp16/bf16            │
│         ▲                          │                │
│         │                    FORWARD (autocast)     │
│         │                    matmuls in fp16/bf16   │
│         │                    reductions stay fp32   │
│         │                          │ loss           │
│         │                    LOSS SCALE ×S  ──fp16  │
│         │                          │                │
│         │                    BACKWARD               │
│         │                    grads in fp16/bf16     │
│         │                          │                │
│         │                    UNSCALE ÷S    ──fp16   │
│         │                    check inf/nan          │
│         │                    cast grads → fp32      │
│         └────────────────────────── optimizer.step  │
└─────────────────────────────────────────────────────┘

Defina bf16() ou fp16() como True para habilitar o treinamento com precisão mista. Ambos são tipos de 16 bits, mas o bf16 tem a mesma faixa de expoente que o fp32, portanto, raramente sofre estouro. Use bf16 em GPUs Ampere ou mais novas (A100, H100) e recorra ao fp16 em hardware mais antigo, como V100 ou T4.

Carregue o modelo em fp32, caso contrário, o autocast se torna inoperante. Carregar em bf16 ou fp16 não deixa uma cópia mestra em fp32 para o otimizador atualizar.

from transformers import TrainingArguments

args = TrainingArguments(..., bf16=True)
args = TrainingArguments(..., fp16=True)

Se o seu modelo for numericamente estável em bf16/fp16, você pode pular a precisão mista por completo e carregar e treinar diretamente em bf16/fp16. Isso evita a cópia dos pesos em fp32 na memória.

tf32

tf32 é um modo de computação em GPUs Ampere que usa mantissa de 10 bits para multiplicações de matrizes, em vez de 23 bits. Isso pode proporcionar um ganho de velocidade, especialmente quando combinado com bf16/fp16. O PyTorch habilita o tf32 para multiplicações de matrizes por padrão em GPUs Ampere e mais recentes, mas configurá-lo explicitamente em TrainingArguments garante que ele esteja ativo, independentemente da versão do PyTorch ou das configurações padrão do ambiente.

from transformers import TrainingArguments

args = TrainingArguments(..., bf16=True, tf32=True)

Próximas etapas

• Consulte o guia Kernels para aprender a acelerar o treinamento com kernels fundidos personalizados.
• Consulte o guia de torch.compile para aprender a compilar o passo de ida e o passo de volta para obter mais vazão.