Mixed Precision Training

Mixed Precision Training (Micikevicius et al., 2018) kết hợp nhiều định dạng số thực trong cùng một quá trình huấn luyện, nhằm tận dụng tốc độ của Tensor Cores mà vẫn giữ ổn định số học. Khác với quantization (thường cho inference), mixed precision được dùng khi đang huấn luyện và phụ thuộc chặt vào kiến trúc GPU (xem thêm tối ưu GPU).

1. Master Weights FP32. Ta giữ một bản sao FP32 của trọng số để tích luỹ các bước cập nhật nhỏ:

Đầu mỗi step ta cast sang FP16 để forward/backward; cuối step ta unscale gradient và cập nhật bản FP32.

2. Loss Scaling. Gradient thực tế thường nằm trong khoảng 1e-9 tới 1e-3. FP16 chỉ phủ ~6e-5 tới 65504. Ta nhân loss lên S trước backward để dịch phân phối gradient lên dải FP16:

3. FP16 / BF16 Compute. Forward và backward chạy ở 16 bit trên Tensor Cores. Accumulator (phép cộng) vẫn ở FP32 — tức là D = A·B + C với A, B ở FP16 và C, D ở FP32. Cách này giữ được độ chính xác khi tổng hàng nghìn phần tử.

import torch
from torch import nn
from torch.amp import autocast, GradScaler
from torch.utils.data import DataLoader

device = "cuda"
model = MyTransformer().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss()

# GradScaler duy nhất cho toàn bộ quá trình train
scaler = GradScaler()

for epoch in range(num_epochs):
    for batch in DataLoader(train_ds, batch_size=32):
        optimizer.zero_grad(set_to_none=True)

        # 1) Forward + loss ở FP16 (hoặc BF16)
        with autocast(device_type=device, dtype=torch.float16):
            logits = model(batch.input_ids)
            loss = criterion(logits, batch.labels)

        # 2) Nhân loss lên S rồi backward (gradient đã được scale)
        scaler.scale(loss).backward()

        # 3) Unscale + clip gradient nếu cần
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # 4) Step: nếu overflow -> skip, sau đó update scale
        scaler.step(optimizer)
        scaler.update()

Cái bẫy cần tránh. Một số op rất nhạy với FP16 và nên giữ ở FP32: softmax (có exp), layernorm (trừ trung bình), loss cross-entropy (log), một số op reduce-sum dài. PyTorch autocast xử lý phần lớn các trường hợp này, nhưng khi viết kernel custom cần chú ý.