Обучение с смешанной точностью

Обучение с полной точностью (fp32) хранит и вычисляет все данные в 32 битах. Смешанная точность использует fp16 или bf16 для вычислительно интенсивных прямого и обратного проходов, при этом сохраняя копию весов в fp32 для обновления оптимизатора. Вычисления становятся быстрее, уменьшается потребление памяти для весов и активаций, а стабильность обучения сохраняется.

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

Установите bf16() или fp16() в True, чтобы включить обучение с смешанной точностью. Оба типа имеют 16 бит, но у bf16 диапазон экспоненты такой же, как у fp32, поэтому переполнение происходит практически никогда. Используйте bf16 на GPU Ampere и новее (A100, H100), а на более старом оборудовании, таком как V100 или T4, откатывайтесь на fp16.

Загрузите модель в fp32, иначе autocast не будет выполнять никаких действий. Загрузка в bf16 или fp16 не оставляет мастер-копии в fp32, от которой оптимизатор мог бы обновлять веса.

from transformers import TrainingArguments

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

Если ваша модель численно стабильна в bf16/fp16, вы можете полностью отказаться от смешанной точности и загружать и обучать модель непосредственно в bf16/fp16. Это позволяет избежать создания копии весов в fp32 в памяти.

tf32

tf32 — это режим вычислений на GPU Ampere, который использует 10-битную мантиссу для матричных умножений вместо 23 бит. Это может обеспечить ускорение, особенно в сочетании с bf16/fp16. PyTorch по умолчанию включает tf32 для матричных умножений на GPU Ampere и новее, но явное указание этого параметра в TrainingArguments гарантирует его активацию независимо от версии PyTorch или настроек окружения по умолчанию.

from transformers import TrainingArguments

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

Дальнейшие шаги

• См. руководство Kernels, чтобы узнать, как ускорить обучение с помощью пользовательских слиянных ядер.
• См. руководство по torch.compile, чтобы узнать, как компилировать прямой и обратный проходы для повышения пропускной способности.