Что такое mixed precision training?

Mixed Precision Training

Mixed Precision Training — это техника оптимизации обучения глубоких нейронных сетей, которая использует разные типы данных (разные precision) для разных операций. Обычно используют float32 и float16 одновременно.

Основная идея

Традиционно нейронные сети обучаются в float32 (32-битные числа с плавающей точкой). Это даёт высокую точность, но требует много памяти и вычислений.

Mixed Precision использует стратегию:

• Forward pass и некоторые вычисления: float16 (16-битные) — быстро и экономно
• Градиенты и обновление весов: float32 — для численной стабильности

Результат: на 2-3x раз меньше памяти, на 1.5-3x раз быстрее, БЕЗ потери качества модели.

Числовые основы

float32 (single precision):

• 32 бита: 1 знак + 8 экспонент + 23 мантисса
• Диапазон: примерно 10 в степени -38 до 10 в степени 38
• Точность: примерно 7 десятичных цифр

float16 (half precision):

• 16 бит: 1 знак + 5 экспонент + 10 мантисса
• Диапазон: примерно 10 в степени -4 до 10 в степени 4 (узкий)
• Точность: примерно 3-4 десятичных цифры
• Проблема: gradient underflow (градиенты становятся нулём)

bfloat16 (brain float):

• 16 бит: 1 знак + 8 экспонент + 7 мантисса
• Диапазон: такой же как float32
• Точность: примерно 3 цифры
• Разработан Google для Deep Learning
• Более стабилен чем float16

Проблема: Gradient Underflow

При обратном распространении gradients часто становятся очень малыми числами. В float16 они могут стать нулём — это underflow.

# Пример underflow
gradient_float32 = 1e-5  # нормально
gradient_float16 = 1e-5  # может стать 0 (слишком маленькое)

Решение: Loss Scaling

Loss Scaling — это техника избежать underflow:

1. Умножить loss на большое число перед backward pass
2. Вычислить градиенты (они будут больше, не underflow)
3. Разделить градиенты на то же число перед обновлением весов

loss_scale = 1024
loss_scaled = loss * loss_scale  # Увеличить градиенты
loss_scaled.backward()           # Backward в float16

# Разделить градиенты перед обновлением
for param in model.parameters():
    if param.grad is not None:
        param.grad.div_(loss_scale)  # Вернуть нормальный масштаб

Как работает в PyTorch

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())
scaler = GradScaler()  # Автоматический Loss Scaling

for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        
        # Forward pass в float16 внутри autocast
        with autocast(dtype=torch.float16):
            outputs = model(batch['input'])
            loss = criterion(outputs, batch['target'])
        
        # Backward в float16
        scaler.scale(loss).backward()
        
        # Optimizer step с автоматическим scaling
        scaler.step(optimizer)
        scaler.update()  # Обновить scale для следующей итерации

Компоненты Mixed Precision

Forward pass - float16/bfloat16 для скорости вычисления Loss computation - float16 для уменьшения памяти Gradients - float16/bfloat16 с loss scaling Gradient updates - float32 для численной стабильности Model weights - float32 для precision при сходимости

Преимущества

1. Меньше памяти — float16 использует 2x меньше памяти чем float32

   • Можно больший batch size
   • Можно большие модели на одном GPU

2. Быстрее обучение — tensor core операции (NVIDIA) оптимизированы для float16

   • На A100: 2.5x speedup
   • На V100: 1.5-2x speedup

3. Нет потери качества — float32 для весов сохраняет convergence

4. Распределённое обучение — меньше bandwidth при синхронизации gradients

Ограничения

1. Не все операции — некоторые слои (normalization) нужны в float32
2. Требует новые GPU — старые GPU (pre-Volta) поддерживают плохо
3. Нужна настройка — loss scale нужно выбрать правильно
4. Может быть нестабильно — если loss scale слишком большой, может быть overflow

Когда использовать Mixed Precision

• Нужна максимальная скорость — training больших моделей
• Ограничена памяти — GPU с малой памятью
• Много экспериментов — быстрее пробовать разные конфигурации
• Распределённое обучение — много серверов, нужна оптимизация

Когда избегать

• Требуется экстремальная точность — например, обучение Transformer для машинного перевода
• Старые архитектуры — некоторые модели чувствительны
• Отладка — в float32 ошибки более очевидны

Современные подходы

• Automatic Mixed Precision (AMP) — PyTorch autocast выбирает тип автоматически
• DeepSpeed ZeRO — распределённое mixed precision
• TensorFlow mixed_precision API — встроенная поддержка
• bfloat16 everywhere — Google TPU использует bfloat16 по умолчанию

Mixed Precision Training — это must-know техника для любого Data Scientist, работающего с большими моделями. Это дало возможность обучать GPT-3 и другие гигантские модели вообще.