PyTorch 混合精度训练底层机理：FP16 vs BF16 vs FP8 数值动态范围与溢出防御

在大模型与深度学习训练工程中，混合精度训练（Mixed Precision Training） 是将 GPU 显存占用减半、矩阵乘法算力翻倍（打满 Tensor Core）的核心底层基石。

然而，许多开发者在从传统的单精度（FP32）切换到低精度浮点数时，经常被各种诡异的数值溢出（Numeric Overflow / Underflow）与 NaN / Inf 梯度爆炸折磨得苦不堪言：

• 为什么在 FP16 模式下训练，Loss 会在几百步后突然变成 NaN，而换成 BF16 后却能稳如磐石地收敛？
• 为什么 FP16 必须死死绑定 GradScaler（损失缩放），而 BF16 却完全不需要？
• 最新英伟达 Hopper / Blackwell 架构力推的 FP8（E4M3 vs E5M2） 又是如何将吞吐再次推向极限的？

本文从 IEEE 754 底层二进制比特位结构 出发，深度解构 FP16、BF16 与 FP8 的数值物理极限与溢出防御实战。

1. 浮点数据格式的底层二进制位结构剖析

[IEEE 754 浮点数物理比特布局 (Sign 符号位 | Exponent 指数位 | Mantissa 尾数位)]:

1. FP32 (单精度基准, 32-bit):
   [S: 1b] | [Exponent: 8b] | [Mantissa: 23b] ──> 动态范围: ~10^38, 极高精度

2. FP16 (传统半精度, 16-bit):
   [S: 1b] | [Exponent: 5b] | [Mantissa: 10b] ──> 动态范围: 仅 ~6.5x10^4 (极窄！极易溢出！)

3. BF16 (Bfloat16 谷歌脑神经格式, 16-bit):
   [S: 1b] | [Exponent: 8b] | [Mantissa: 7b]  ──> 动态范围: 与 FP32 100% 相同 (~10^38)！

4. FP8-E4M3 (前向计算格式, 8-bit):
   [S: 1b] | [Exponent: 4b] | [Mantissa: 3b]  ──> 专注于更高精度 (Max: 448)

5. FP8-E5M2 (反向梯度格式, 8-bit):
   [S: 1b] | [Exponent: 5b] | [Mantissa: 2b]  ──> 专注于更大动态范围 (Max: 57344)

2. 核心浮点格式数值特性全景对比矩阵

3. PyTorch 生产级自动混合精度（AMP）与防御实战

在 PyTorch 2.x 中，推荐根据硬件环境自适应切换最优精度：

import torch
import torch.nn as nn
from typing import Tuple

class PrecisionSafeTrainer:
    def __init__(self, model: nn.Module, optimizer: torch.optim.Optimizer):
        self.model = model
        self.optimizer = optimizer
        
        # 1. 硬件自适应精度判定
        if torch.cuda.is_bf16_supported():
            # A100/H100/4090 首选 BF16 (最稳健，完全免除 GradScaler 开销！)
            self.dtype = torch.bfloat16
            self.scaler = None
            print("[AMP Mode] 硬件支持原生 BF16！开启零缩放开销混合精度训练！")
        else:
            # 旧版显卡 (V100/T4) 回退至 FP16，并必须挂载 GradScaler 动态防下溢
            self.dtype = torch.float16
            self.scaler = torch.cuda.amp.GradScaler(
                init_scale=65536.0,
                growth_factor=2.0,
                backoff_factor=0.5,
                growth_interval=2000
            )
            print("[AMP Mode] 旧版显卡回退至 FP16，已挂载动态 GradScaler 防线！")

    def train_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> float:
        self.optimizer.zero_grad()
        
        if self.dtype == torch.bfloat16:
            # BF16 极简前向与反向
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                outputs = self.model(inputs)
                loss = nn.functional.cross_entropy(outputs, targets)
                
            loss.backward()
            # 梯度范数裁剪 (防梯度爆炸)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.optimizer.step()
            return loss.item()
        else:
            # FP16 严格缩放前向与反向
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                outputs = self.model(inputs)
                loss = nn.functional.cross_entropy(outputs, targets)
                
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
            self.scaler.step(self.optimizer)
            self.scaler.update()
            return loss.item()

4. 不同精度格式在 7B 模型训练中的稳定性实测对比

实测数据表明：BF16 凭借与 FP32 相同的 8 位指数宽度，在保持显存减半、吞吐提升 2.2 倍的同时，彻底告别了 FP16 容易发生的数值溢出闪崩！

5. 架构师混合精度避坑铁律

1. Ampere (A100/3090) 及以上架构强制使用 BF16：在大模型预训练与微调中，一律废弃 FP16，无脑选用 torch.bfloat16；
2. Softmax 与 LayerNorm 保持 FP32 计算：在 Transformer Block 中，Attention Softmax 与归一化算子内部必须强制转换为 FP32 进行累加，防止由于尾数位较少导致的精度损失。