What is mixed precision training (FP16, BF16)?

Mixed precision training is a technique for training neural networks that uses different floating-point formats, such as FP16 (16-bit float) and BF16 (bfloat16), instead of the standard FP32 (32-bit float). The goal is to accelerate training and reduce memory usage without significant loss of accuracy.

FP16 and BF16 allow for faster and more efficient computations on modern GPUs, as they require less memory and bandwidth. However, due to lower precision, there can be issues with numerical stability, so formats are usually mixed: core computations are in FP16/BF16, while some critical operations are in FP32.

An example of using mixed precision in PyTorch with automatic mixed precision:

model = ...
optimizer = ...
scaler = torch.cuda.amp.GradScaler()

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

This approach achieves a balance between performance and accuracy, which is especially important in distributed training and working with large models.