Entrenamiento con precisión mixta

El entrenamiento con precisión completa (fp32) almacena y calcula todo en 32 bits. La precisión mixta utiliza fp16 o bf16 para las pasadas de propagación hacia adelante y hacia atrás, que son intensivas en cómputo, mientras que mantiene una copia de los pesos en fp32 para la actualización del optimizador. El cómputo es más rápido, la memoria de pesos y activaciones se reduce y se preserva la estabilidad del entrenamiento.

┌─────────────────────────────────────────────────────┐
│           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  │
└─────────────────────────────────────────────────────┘

Establezca bf16() o fp16() en True para habilitar el entrenamiento con precisión mixta. Ambos son tipos de 16 bits, pero bf16 tiene el mismo rango de exponente que fp32, por lo que casi nunca se desborda. Use bf16 en GPUs Ampere o más nuevas (A100, H100) y recurra a fp16 en hardware más antiguo como V100 o T4.

Cargue el modelo en fp32, de lo contrario autocast no tendrá efecto. Cargarlo en bf16 o fp16 deja sin una copia maestra en fp32 para que el optimizador realice las actualizaciones.

from transformers import TrainingArguments

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

Si tu modelo es numéricamente estable en bf16/fp16, puedes omitir por completo la precisión mixta y cargar y entrenar directamente en bf16/fp16. Esto evita la copia de los pesos en memoria en fp32.

tf32

tf32 es un modo de cómputo en GPUs Ampere que utiliza una mantisa de 10 bits para las multiplicaciones de matrices en lugar de 23 bits. Esto puede proporcionarte una aceleración, especialmente cuando se combina con bf16/fp16. PyTorch habilita tf32 para las multiplicaciones de matrices de forma predeterminada en GPUs Ampere y más recientes, pero configurarlo explícitamente en TrainingArguments garantiza que esté activo independientemente de la versión de PyTorch o de los valores predeterminados del entorno.

from transformers import TrainingArguments

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

Próximos pasos

• Consulte la guía de Kernels para aprender a acelerar el entrenamiento con kernels fusionados personalizados.
• Consulte la guía de torch.compile para aprender a compilar el paso hacia adelante y el paso hacia atrás para obtener un rendimiento adicional.