混合精度训练的“水土不服”：从梯度爆炸到稳定收敛

在 AMD GPU 上进行大模型微调，最让人头疼的往往不是环境配置，而是训练过程中的“玄学”问题。很多从 NVIDIA 平台迁移过来的算法工程师会发现，同样的模型架构、同样的超参数，在 ROCm 环境下跑混合精度训练（AMP）时，Loss 曲线经常像过山车一样剧烈波动，甚至直接出现梯度爆炸导致训练中断。这并非硬件不行，而是不同架构对数值精度的敏感度存在差异。特别是在使用 BF16（BFloat16）时，虽然它比 FP16 拥有更宽的动态范围，但在某些算子实现和梯度累积策略上，依然需要针对性的调优。今天就来聊聊我在 LLaMA-Factory 框架下，解决 ROCm 混合精度训练稳定性问题的实战心得。

为什么 AMD 显卡上的混合精度更容易“炸”？

首先要打破一个误区：BF16 并不是万能药。虽然 BF16 的指数位与 FP32 相同，理论上能更好地保留大数值信息，减少溢出风险，但 AMD 的 CDNA 或 RDNA 架构在执行矩阵乘法（GEMM）时的累加逻辑与 NVIDIA Tensor Core 有所不同。

在实际操作中，我发现两个主要诱因：

1. 算子实现的数值误差累积：部分底层算子在 ROCm 库（如 rocBLAS）中的默认实现，为了追求极致速度，可能在中间结果的精度保持上做了妥协。当这些微小误差在深层网络中逐层传递并累积时，就会导致最终梯度的偏差被放大。
2. 损失缩放（Loss Scaling）策略不匹配：传统的动态损失缩放算法是基于 FP16 设计的，其阈值判断逻辑直接套用到 BF16 场景下可能过于激进或保守。如果缩放因子过大，梯度反传时容易溢出；过小则导致小梯度被截断为零（Underflow），模型无法收敛。

因此，在 ROCm 上开启混合精度，不能简单地照搬 CUDA 时代的配置文件，必须对精度策略进行“本地化”改造。

核心调优手段：调整缩放因子与切换精度模式

遇到 Loss 突然变成 NaN 或者 Inf，第一反应不应该是降低学习率，而是检查混合精度的配置。以下是我在多次试错后总结出的两套有效方案。

方案一：精细化调整 Loss Scaling

LLaMA-Factory 默认通常会启用动态损失缩放。在 AMD 卡上，建议先尝试手动固定缩放因子，观察训练稳定性。

如果你使用的是基于 PyTorch 原生 AMP 的后端，可以在启动脚本或配置文件中干预 GradScaler 的行为。虽然 LLaMA-Factory 封装了大部分逻辑，但我们可以通过环境变量或修改源码中的初始化参数来调整。

例如，将初始缩放因子设置得更保守一些：

# 伪代码示例：在 trainer 初始化前干预
from torch.cuda.amp import GradScaler

# 针对 ROCm 环境，适当降低初始 scale 值，避免早期溢出
scaler = GradScaler(init_scale=2**10, growth_factor=2.0, backoff_factor=0.5)

在 LLaMA-Factory 的配置文件中（如 finetune.yaml），虽然没有直接暴露 init_scale 参数，但你可以通过开启更频繁的梯度检查来间接缓解问题。如果发现训练初期就不稳定，可以尝试在命令行参数中增加 --logging_steps 的频率，以便更早捕捉异常。

方案二：果断切换至纯 BF16 或纯 FP32

如果调整缩放因子效果不佳，或者模型结构本身对精度极其敏感（如某些包含 LayerNorm 的特殊变体），最稳妥的方案是放弃动态缩放，直接使用纯 BF16 模式，甚至在关键阶段回退到 FP32。

在 LLaMA-Factory 中，这可以通过修改 compute_type 参数实现。对于 MI250/MI300 等支持原生 BF16 加速的显卡，纯 BF16 通常是性价比最高的选择，因为它避免了 FP16 那种复杂的缩放逻辑，同时保持了较快的计算速度。

修改配置文件示例：

# finetune_lora_bf16.yaml
model_name_or_path: meta-llama/Llama-3-8B
do_train: true
template: llama3
finetuning_type: lora
lora_target: all
compute_type: bf16  # 关键：强制使用 bf16，关闭自动混合精度中的 fp16 逻辑
output_dir: ./saves/llama3-lora-bf16

如果连纯 BF16 都无法收敛，且显存资源允许，最后的“大招”就是切换到 FP32。虽然这会牺牲约一半的训练速度并增加显存占用，但它能彻底消除精度带来的数值噪声。这在调试新模型架构或排查收敛问题时非常有用：

compute_type: fp32

一旦在 FP32 下确认模型能正常收敛，再逐步尝试降回 BF16，此时你就能确定问题确实出在精度而非数据或代码逻辑上。

LLaMA-Factory 中的监控与实战配置

光改配置还不够，必须建立有效的监控机制。在 ROCm 环境下训练，我强烈建议重点关注以下两个指标，它们比单纯的 Loss 值更能反映精度问题。

1. 梯度范数（Gradient Norm）

这是判断梯度爆炸最直接的指标。在 LLaMA-Factory 的日志输出中，开启详细日志后可以看到每一步的 grad_norm。

• 正常情况：梯度范数通常维持在一个相对稳定的区间（例如 0.1 到 10 之间，具体取决于模型大小）。
• 异常信号：如果某一步 grad_norm 突然飙升到几千甚至几万，紧接着 Loss 变为 NaN，那就是典型的溢出。

你可以在训练命令中加入 --plot_loss true，训练结束后查看可视化图表。如果看到梯度范数曲线有尖锐的脉冲，说明当前的精度设置无法容纳该步的梯度更新。

2. 显存占用与利用率

有时候不稳定是因为显存碎片化导致算子回退到了低效实现。使用 rocm-smi 或 rocprof 工具实时监控：

watch -n 1 rocm-smi --showmeminfo vram

如果在训练过程中显存占用剧烈跳动，可能需要检查是否开启了 gradient_checkpointing。在显存紧张时强制开启混合精度，可能会触发额外的内存交换，进而影响数值稳定性。

综合配置建议

针对大多数在 AMD 显卡上进行的 LoRA 微调任务，我推荐以下“稳健型”配置组合：

# 推荐配置：稳健优先
compute_type: bf16          # 优先使用原生 BF16
lora_alpha: 32              # 适当调整 LoRA 缩放系数
lora_dropout: 0.05          # 加入少量 Dropout 增加鲁棒性
optim: adamw_torch          # 使用 PyTorch 原生优化器，兼容性更好
gradient_accumulation_steps: 4 # 通过累积步数减小单步 Batch Size 压力

如果在上述配置下依然偶尔出现波动，可以尝试在启动命令中显式禁用某些激进的融合算子（如果有相关环境变量支持），或者将 learning_rate 稍微下调 10%-20%，给优化器更多的缓冲空间。

结语

在 AMD ROCm 生态中进行混合精度训练，本质上是一个在“性能”与“稳定性”之间寻找平衡点的过程。不要盲目迷信自动化工具的默认设置，理解 BF16 的特性，学会灵活切换精度模式，并善用梯度范数等指标进行诊断，才是解决收敛问题的关键。随着社区对 SGLang、TileLang 等底层算子的不断优化，以及 LLaMA-Factory 对 ROCm 支持的日益完善，这些“坑”正在被快速填平。作为开发者，我们既要享受异构计算带来的成本红利，也要保持对数值细节的敬畏，用扎实的实验数据去验证每一次配置的调整。

200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper