Mixed Precision

Обучение со смешанной точностью — это метод, при котором большая часть вычислений нейронной сети выполняется в формате с плавающей точкой разрядностью 16 бит: FP16 (половинная точность) или BF16. При этом основная копия весов модели хранится в 32-битном формате (FP32), а численно чувствительные операции выполняются в FP32. 16-битные значения занимают вдвое меньше памяти и пропускной способности, чем FP32, а Tensor Cores современных GPU обрабатывают их значительно быстрее. Поэтому смешанная точность сокращает время обучения и позволяет размещать в памяти более крупные размеры пакета, сохраняя точность на уровне FP32. Описанный Micikevicius и соавторами в 2018 году метод с основными весами FP32, вычислениями FP16 и масштабированием функции потерь теперь автоматически реализуется в фреймворках глубокого обучения как автоматическая смешанная точность (AMP).

Как работает обучение со смешанной точностью#

Три компонента позволяют обратному распространению ошибки и обновлению весов при обучении с точностью 16 бит сохранять точность обучения в FP32:

1. Основные веса FP32: Оптимизатор применяет каждое обновление к копии весов в FP32. Прибавление небольшого обновления непосредственно к весу FP16 часто приводит к его округлению до нуля, поскольку FP16 хранит лишь около трёх значащих десятичных цифр.
2. Точность отдельных операций: Контекст autocast выполняет умножение матриц и свёртки в 16-битном формате, где Tensor Cores обеспечивают наибольший прирост скорости, а операции суммирования, softmax и вычисление функции потерь оставляет в FP32.
3. Масштабирование функции потерь: FP16 не может представлять числа по модулю меньше примерно 6 × 10⁻⁸, поэтому многие небольшие градиенты округлялись бы до нуля. Умножение функции потерь на коэффициент масштабирования перед обратным распространением ошибки переносит все градиенты в диапазон FP16; перед обновлением весов градиенты делят на тот же коэффициент.

В PyTorch AMP за преобразование типов для отдельных операций отвечает torch.autocast, а за масштабирование функции потерь — GradScaler:

import torch
from torch import nn

# Use FP16 on an NVIDIA GPU and BF16 on CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.bfloat16

model = nn.Linear(512, 10).to(device)  # master weights stay in FP32
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = torch.amp.GradScaler(device, enabled=dtype == torch.float16)  # loss scaling for FP16 only

x = torch.randn(64, 512, device=device)
with torch.autocast(device_type=device, dtype=dtype):
    out = model(x)  # the matrix multiply runs in 16-bit
    loss = out.float().pow(2).mean()

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Output: {out.dtype}, weights: {model.weight.dtype}, gradients: {model.weight.grad.dtype}")

На CPU фрагмент кода выводит torch.bfloat16 для выхода слоя и torch.float32 для весов и их градиентов; на GPU NVIDIA выводится torch.float16. Вычисления в 16-битном формате и основные веса FP32 сосуществуют в одном шаге обучения.

GradScaler в PyTorch динамически корректирует масштаб. Он начинает с 65,536, уменьшает масштаб вдвое и пропускает шаг, если масштабированные градиенты содержат значения inf или NaN; после 2,000 последовательных шагов без переполнения масштаб увеличивается вдвое, поэтому коэффициент стабилизируется чуть ниже уровня, при котором градиенты переполняются.

FP16, BF16 и форматы с меньшей точностью#

В FP16 для экспоненты отведено 5 бит, а для мантиссы — 10; формат обеспечивает высокую точность, но переполняется при значениях выше 65,504. В BF16 сохраняются 8 бит экспоненты FP32, поэтому диапазон достигает примерно 3.4 × 10³⁸, но для мантиссы остаётся лишь 7 бит: масштабирование функции потерь не требуется, зато округление получается более грубым. BF16 поддерживается GPU NVIDIA начиная с поколения Ampere, TPU Google и многими современными CPU; это стандартный формат обучения больших языковых моделей. Та же идея применяется к FP8 на GPU Hopper и Blackwell с помощью Transformer Engine от NVIDIA, а также к форматам FP4, например NVFP4, на Blackwell.

Смешанная точность на практике#

Основной практический риск AMP — модель или GPU, на которых в 16-битном формате возникают потери NaN. При обучении Ultralytics YOLO26 AMP включён по умолчанию с помощью amp=True: сначала сравниваются прогнозы FP32 и AMP модели YOLO26n на GPU с CUDA, а при расхождении прогнозов или использовании GPU из серии, на которой при AMP возникают потери NaN, например GTX 16, выполняется переход на FP32. Настройка amp="bf16" переключает autocast на BF16 и отключает масштабирование функции потерь.

Смешанная точность широко применяется там, где обучение ограничено объёмом памяти или временем. Команды, разрабатывающие автономные автомобили, обучают модели обнаружения объектов на миллионах кадров высокого разрешения; сокращение памяти для активаций вдвое позволяет использовать более крупные изображения и пакеты. В анализе медицинских изображений трёхмерные томограммы КТ и МРТ часто не помещаются в память GPU при FP32, а 16-битные активации делают возможным обучение моделей сегментации на целых томограммах.

Смешанная точность и связанные понятия#

• Квантование моделей: Квантование преобразует обученную модель в целочисленный формат с низкой точностью, например INT8, для ускорения инференса; смешанная точность ускоряет обучение и сохраняет веса с плавающей точкой.
• Квантование с учётом обучения: QAT имитирует округление до целых чисел во время обучения, чтобы модель сохраняла точность после развёртывания в INT8.
• Половинная точность: Половинной точностью называют сам числовой формат FP16, а смешанной точностью — метод обучения, сочетающий его с FP32. Обучение исключительно в FP16 без основных весов FP32 и масштабирования функции потерь приводит к потере небольших обновлений из-за округления, а малых градиентов — из-за исчезновения; часто такая модель не сходится.