4、混合精度训练与推理：FP16、BF16、TF32的原理、优缺点，以及如何通过混合精度减少显存占用

大模型训练与推理的显存瓶颈，本质上是对“数值精度”与“存储/计算开销”之间平衡的极致追求。FP32（32位浮点）长期作为默认精度，但其对显存和带宽的消耗在大模型时代已成为不可承受之重。混合精度技术并非简单地“用低精度替代高精度”，而是基于不同数值格式的动态范围与精度特性，在保持模型收敛质量的前提下，将大部分张量运算切换到低精度格式，从而显著降低显存占用并提升计算吞吐。本章将深入剖析FP16、BF16、TF32三种核心低精度格式的原理与适用场景，并给出工程落地的完整方法论。

上图展示了三种低精度格式在混合精度体系中的定位与决策路径。FP16、BF16、TF32并非互相替代的关系，而是分别针对训练与推理中不同环节的数值特征而设计。FP16在A100等GPU上拥有Tensor Core加速，但动态范围窄；BF16牺牲尾数位换取与FP32相同的指数范围，解决了梯度下溢问题；TF32则专为NVIDIA Ampere架构的矩阵运算设计，在计算加速与精度之间取得折中。理解三者的差异，是配置混合精度策略的前提。

4.1 核心原理：三种低精度格式的数值表示与硬件映射

浮点数的精度与范围由符号位、指数位和尾数位的分配决定。FP16采用1位符号、5位指数、10位尾数，其最大表示范围为65504，最小正规格化数为6.1e-5。在训练大模型时，梯度值往往远小于1，FP16极易发生下溢（underflow）——即数值小于最小可表示范围而被舍入为0，导致梯度消失。BF16则采用1位符号、8位指数、7位尾数，其指数范围与FP32完全一致（最大约3.4e38），从根本上规避了下溢风险，但尾数位仅7位，精度相对较低。TF32是NVIDIA专为Tensor Core设计的格式，在A100及后续架构中，TF32模式将FP32输入截断为1位符号、8位指数、10位尾数进行矩阵乘累加，但输出仍以FP32累加，从而在硬件层面实现接近FP32的精度与FP16的吞吐。

从硬件映射角度看，FP16和BF16在Tensor Core上均支持双倍于FP32的吞吐量（例如A100的FP16/BF16算力为312 TFLOPS，而FP32为19.5 TFLOPS）。TF32在A100上的算力为156 TFLOPS，介于两者之间。这意味着，即使TF32不减少显存占用，它也能通过计算加速间接缩短训练时间，从而减少因长时间训练导致的显存碎片化风险。在推理场景中，FP16和BF16的显存节省效果直接且显著——模型权重和激活值均减半，这对于部署在显存受限的GPU（如消费级RTX 3090 24GB）上的大模型至关重要。

4.2 判断方法与操作流程：如何选择精度格式并配置混合精度

选择精度格式需考虑三个维度：模型规模、硬件架构、训练稳定性。对于参数量超过10B的模型，FP16的梯度下溢问题会显著加剧，此时BF16是更安全的选择——尤其在使用AdamW优化器时，其梯度二阶矩估计对数值范围极为敏感。若硬件为NVIDIA Ampere或Hopper架构，TF32可作为卷积和矩阵乘的默认计算格式，但需注意TF32不减少显存占用，仅加速计算。对于推理场景，若模型已收敛且无需继续训练，FP16通常足够，因为推理时前向传播的数值范围相对稳定。

实际操作中，PyTorch的torch.cuda.amp模块提供了自动混合精度（AMP）功能。其核心机制是：维护一个FP32的主权重副本，在前向传播时使用FP16/BF16计算，反向传播时梯度以低精度存储，但在优化器更新前转换为FP32。配置流程如下：

# PyTorch 2.x 自动混合精度示例
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 用于FP16梯度缩放，BF16通常不需要
for data, target in dataloader:
    optimizer.zero_grad()
    with autocast(dtype=torch.float16):  # 或 torch.bfloat16
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

对于BF16，由于动态范围与FP32一致，通常无需GradScaler进行梯度缩放。但需注意，BF16的精度损失可能导致训练曲线出现轻微震荡，此时可适当降低学习率或增加warmup步数。在推理阶段，可使用torch.inference_mode()配合model.half()将模型整体转换为FP16，但需验证BatchNorm层的统计量是否已冻结。

4.3 案例拆解：LLaMA-7B 微调中的显存优化实战

假设我们在一张NVIDIA A100 40GB GPU上微调LLaMA-7B模型。初始配置使用FP32，显存占用分布如下：模型权重14GB（7B×2字节），梯度14GB，优化器状态（AdamW的一阶和二阶矩）28GB，激活值约8GB，总计约64GB——远超40GB显存。直接训练必然OOM。

我们采用混合精度策略进行优化。第一步，将模型权重和梯度切换为FP16，显存占用降至14GB（权重）+14GB（梯度）+28GB（优化器状态，仍为FP32）+4GB（激活值FP16）=60GB，仍超限。第二步，引入梯度检查点（gradient checkpointing），将激活值显存从4GB降至1GB，总计57GB。第三步，使用8-bit优化器（如bitsandbytes的AdamW8bit），将优化器状态从28GB降至7GB，总计约36GB——成功放入40GB显存。

实际训练中，我们观察到FP16的梯度缩放因子（scale factor）在训练初期频繁触发溢出（scale值降至1），导致loss震荡。分析发现，LLaMA-7B的某些层（如最后的lm_head）梯度幅值可达1e3，而FP16最大值为65504，看似安全，但AdamW的二阶矩估计（梯度平方）可达1e6，在FP16下直接溢出。解决方案是：将lm_head层单独保持FP32计算，其余层使用FP16。这一调整后，训练稳定，最终loss收敛曲线与纯FP32基线几乎重合。

该案例表明，混合精度并非简单的“减半”操作，而是需要结合模型结构、优化器类型和硬件特性进行精细调优。FP16与BF16在显存节省上效果相同，但BF16在训练稳定性上更优，尤其适合LLaMA这类深层Transformer模型。

4.4 常见误区与总结复盘

误区一：认为FP16一定比BF16好。FP16的尾数位更多（10位 vs 7位），在数值精度上略高，但其动态范围窄，容易导致梯度下溢。对于大模型，BF16的稳定性优势往往超过精度优势。实践中，若使用FP16出现loss不降或NaN，应首先尝试BF16。

误区二：忽略主权重的FP32副本。有些初学者直接将模型整体转换为FP16进行训练，这会导致权重更新时累积误差，最终模型发散。必须保留FP32主权重，低精度仅用于前向和反向计算。

误区三：TF32与FP16混为一谈。TF32不减少显存占用，它只是计算加速格式。在显存不足时，TF32无法提供帮助；但在计算密集的卷积层，TF32可提升约2倍吞吐，且无需修改代码（设置torch.backends.cuda.matmul.allow_tf32 = True）。

误区四：推理时盲目使用FP16。对于包含LayerNorm或Softmax的模型，FP16的精度损失可能导致输出分布偏移。建议在推理前使用校准数据集验证FP16输出的最大误差，若超过1e-3，则考虑使用BF16或混合精度推理（部分层保持FP32）。

总结而言，混合精度是当前大模型训练与推理的必备技术。FP16、BF16、TF32各有其适用边界：FP16适合小规模模型或激活值敏感的层；BF16适合大规模模型和梯度易下溢的场景；TF32适合计算密集但显存充足的卷积网络。工程实践中，应结合显存监控（如nvidia-smi）和训练曲线，动态调整精度策略。最终目标是：在显存约束下，以最低的精度损失换取最大的吞吐提升。