Как работает Mixed Precision Training (FP16 + FP32 master веса)?

Краткий тезис

precision precision precision precision precision precision training|Mixed Precision Training — это техника ускорения обучения и снижения потребления памяти за счёт выполнения основных вычислений (forward/backward) в FP16 (16‑битная плавающая точка), при этом master‑веса (копия параметров модели) хранятся в FP32 для сохранения точности. Для предотвращения потери значимости градиентов (underflow) применяется loss scaling – умножение функции потерь на коэффициент перед обратным распространением. Шаг оптимизатора выполняется над FP32‑весами, после чего они конвертируются обратно в FP16 для следующей итерации. Это даёт ускорение в 2–3 раза и уменьшение использования GPU‑памяти на 30–50% по сравнению с полной FP32‑точностью, но требует аккуратной обработки операций, чувствительных к точности (например, softmax, layer normalization).

1. Что такое Mixed Precision Training (MPT)

Mixed Precision Training – это метод обучения нейронных сетей, при котором разные части вычислительного графа выполняются в разных числовых форматах. Обычно:

• FP32 (32‑бит, float32) – для master‑весов, шага оптимизатора и некоторых критических операций.
• FP16 (16‑бит, float16) – для forward‑прохода, backward‑прохода (вычисление градиентов) и хранения активаций.

Термин master‑веса (master weights) – это копия всех параметров модели, хранящаяся в FP32. Она обновляется на каждом шаге оптимизатора, а затем преобразуется в FP16 для использования в прямом и обратном проходах.

Зачем это нужно

• FP16 занимает вдвое меньше памяти, чем FP32, и позволяет использовать более быстрые тензорные ядра (Tensor Cores) на современных GPU (NVIDIA Volta, Turing, Ampere, Hopper).
• Прямое обучение в FP16 без master‑весов приводит к накоплению ошибок округления и расходимости, так как градиенты многих слоёв (особенно в глубоких сетях) становятся слишком малыми для представления в FP16.

2. Зачем нужен FP32 master‑вес

Проблема FP16 имеет ограниченный динамический диапазон (примерно 5.96×10⁻⁸ … 6.55×10⁴) и точность около 3–4 значащих цифр. При обновлении весов на шаге оптимизатора (например, SGD: w = w - lr * grad) приращение lr * grad может быть на несколько порядков меньше самого веса. В FP16 такое приращение будет округлено до нуля – веса перестанут обновляться.

Решение

• Хранить master‑веса в FP32 (высокая точность).
• На каждой итерации:
  1. Скопировать master‑веса в FP16 (конвертация).
  2. Выполнить forward/backward в FP16, получив градиенты в FP16.
  3. Применить loss scaling к градиентам (см. раздел 4).
  4. Сконвертировать градиенты обратно в FP32 (если необходимо) и применить шаг оптимизатора к FP32 master‑весам.

Таким образом, накопление ошибок округления происходит только в одном шаге, а master‑веса сохраняют полную точность.

3. Процесс forward/backward в FP16

Forward‑проход

• Входные данные (обычно FP32) преобразуются в FP16.
• Все операции (линейные слои, активации, свёртки) выполняются в FP16.
• Активации сохраняются в FP16 для backward (экономия памяти).

Backward‑проход

• Градиенты вычисляются в FP16 с использованием сохранённых активаций.
• Градиенты по весам также получаются в FP16.

Проблема underflow
Градиенты многих слоёв (особенно в начале сети) могут быть очень маленькими (например, 10⁻⁶). В FP16 минимальное представимое положительное число – 5.96×10⁻⁸, но точность вблизи нуля резко падает. Градиенты меньше ~10⁻⁷ округляются до нуля – это underflow. Для борьбы используется loss scaling.

4. Loss Scaling

Loss scaling – умножение значения функции потерь (loss) на константу scale перед вызовом backward(). После обратного распространения градиенты всех параметров оказываются умноженными на scale. Это сдвигает их в диапазон, где FP16 может их точно представить.

Алгоритм

1. Вычислить loss в FP32.
2. Умножить loss на scale (например, 2¹⁶ = 65536).
3. Выполнить backward() – градиенты теперь в FP16, но масштабированы.
4. Перед шагом оптимизатора разделить градиенты на scale (в FP32).

Динамический loss scaling (рекомендуется):

• Начать с большого scale (например, 2²⁴).
• Если на шаге возникает Inf или NaN в градиентах (переполнение FP16), scale уменьшается (например, вдвое).
• Если несколько шагов прошли без переполнения, scale увеличивается (например, вдвое каждые N шагов).

Пример кода (PyTorch AMP):

scaler = torch.cuda.amp.GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

5. Шаг оптимизатора

После backward градиенты находятся в FP16 (масштабированные). Перед применением оптимизатора их необходимо:

1. Сконвертировать в FP32 (если оптимизатор работает с FP32).
2. Разделить на scale (восстановить истинные значения).

Сам шаг оптимизатора (SGD, Adam, AdamW) выполняется над FP32 master‑весами. Это гарантирует, что обновления не теряются из‑за низкой точности FP16.

После обновления master‑веса снова конвертируются в FP16 для следующей итерации (обычно это делается автоматически при следующем forward в autocast).

6. Ускорение и экономия памяти

Почему master‑веса не съедают всю экономию?

• Master‑веса хранятся только для параметров, которые обновляются (обычно это все веса модели). Активации же занимают основную память во время обучения – их сокращение вдвое даёт основной выигрыш.
• Для inference master‑веса не нужны – модель можно сохранить в FP16, получив двукратное сжатие.

7. Проблемы и ограничения

7.1 Операции, требующие FP32

Некоторые операции в FP16 дают нестабильные результаты:

• Softmax – экспоненты могут переполниться (Inf) или обнулиться.
• Layer Normalization – вычисление среднего и дисперсии в FP16 теряет точность.
• Cross‑entropy loss – логарифм от очень малых вероятностей даёт Inf.

Решение PyTorch AMP автоматически выполняет эти операции в FP32 внутри контекста autocast, если они помечены как fp32 в списке ops.

7.2 Underflow градиентов для очень глубоких сетей

Даже с loss scaling градиенты в ранних слоях могут быть слишком малы. Помогает:

• Использование BF16 (bfloat16) – имеет такой же динамический диапазон, как FP32, но меньшую точность.
• Gradient checkpointing (уменьшение памяти, но не underflow).

7.3 Не все GPU поддерживают Tensor Cores в FP16

Старые GPU (Pascal и ранее) не имеют тензорных ядер – ускорение от FP16 будет меньше, но экономия памяти остаётся.

8. Сравнение с другими форматами

BF16 (bfloat16) – предпочтителен для обучения больших моделей (LLM), так как не требует loss scaling и реже вызывает underflow. Однако не все GPU поддерживают BF16 (требуется Ampere или новее).

9. Реализация на практике (PyTorch AMP)

Automatic Mixed Precision (AMP) – встроенный механизм PyTorch, автоматизирующий выбор формата для каждой операции.

Основные компоненты

• torch.cuda.amp.autocast – контекстный менеджер, внутри которого forward выполняется в FP16/BF16 (в зависимости от устройства).
• torch.cuda.amp.GradScaler – реализует динамический loss scaling.

Пример полного цикла обучения

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = torch.cuda.amp.GradScaler()

for epoch in range(epochs):
    for batch in dataloader:
        inputs, targets = batch
        optimizer.zero_grad()
        
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)
        
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

Важные нюансы

• autocast должен охватывать только forward и вычисление loss.
• scaler.step(optimizer) автоматически распаковывает градиенты (делит на scale) и обновляет веса.
• scaler.update() корректирует scale на основе наличия Inf/NaN.

10. Когда применять Mixed Precision Training

• Обучение больших моделей (LLM, Vision Transformers, GANs) – ускорение в 2–3 раза критично.
• Ограниченная память GPU – позволяет увеличить batch size или размер модели.
• Fine‑tuning – часто используется по умолчанию (например, в Hugging Face Trainer).

Когда не стоит

• Модель очень мала (ускорение незаметно, а накладные расходы на конвертацию могут даже замедлить).
• Используются операции, не поддерживающие FP16 (например, пользовательские CUDA‑ядра).
• GPU без тензорных ядер (ускорение минимально, но экономия памяти остаётся).

11. Связь с обучением больших моделей

Mixed Precision Training – обязательный компонент для обучения современных LLM (GPT, LLaMA, BLOOM). Без него невозможно уместить модель с миллиардами параметров в память GPU. В сочетании с:

• Gradient checkpointing (уменьшение памяти активаций).
• Distributed Data Parallel (DDP) или Fully Sharded Data Parallel (FSDP).
• LoRA / QLoRA (адаптеры в FP16, base model в FP16 или 4‑бит).

MPT позволяет эффективно использовать Tensor Cores и снижать время обучения с недель до дней.

Пет-проект для закрепления

Задача Обучить небольшую модель (например, ResNet‑18 на CIFAR‑10) с использованием Mixed Precision и без, сравнить скорость, потребление памяти и точность.

Инструменты PyTorch, torch.cuda.amp, nvidia‑smi (мониторинг памяти), time.

Шаги:

1. Написать скрипт обучения в FP32 (без AMP).
2. Написать скрипт обучения с AMP (autocast + GradScaler).
3. Замерить время эпохи, пиковое использование GPU‑памяти (через torch.cuda.max_memory_allocated()).
4. Сравнить финальную точность (accuracy) на тестовом наборе.
5. Попробовать отключить loss scaling и наблюдать расходимость.

Ожидаемый результат

• Ускорение в 1.5–2 раза.
• Снижение памяти на 30–40%.
• Точность в пределах 0.1–0.5% от FP32 (при правильном scaling).

Дополнительно

• Визуализировать гистограмму градиентов до и после scaling.
• Попробовать BF16 (если GPU поддерживает) и сравнить.

Связь с другими вопросами

Навигация

• Предыдущий: 468
• Следующий: 470
• Индекс: 00. Индекс разборов