AMP

Что такое AMP

AMP (Automatic Mixed Precision) — подход к обучению, при котором часть вычислений выполняется в пониженной точности (FP16 или BF16), а критичные операции — в полной (FP32). PyTorch автоматически решает, какую точность использовать для каждой операции.

Зачем нужна смешанная точность

BF16 — тот же диапазон, что FP32, но меньше точность дробной части. Не требует loss scaling. Доступен на Ampere+ (A100, H100).

Loss Scaling (для FP16)

Проблема FP16: малые градиенты (~10⁻⁸) обнуляются (underflow). Решение:
1. Умножить loss на большое число (scale factor) перед backward
2. Градиенты получаются пропорционально большими → не теряются
3. Перед optimizer.step() поделить обратно

BF16 не требует loss scaling — для A100/H100 рекомендуется BF16.

Пример

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    with autocast(dtype=torch.float16):
        output = model(batch)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Эффект

AMP — одна из первых оптимизаций, которую стоит включить. Практически бесплатное ускорение.