混合精度训练的“水土不服”从梯度爆炸到稳定收敛在 AMD GPU 上进行大模型微调最让人头疼的往往不是环境配置而是训练过程中的“玄学”问题。很多从 NVIDIA 平台迁移过来的算法工程师会发现同样的模型架构、同样的超参数在 ROCm 环境下跑混合精度训练AMP时Loss 曲线经常像过山车一样剧烈波动甚至直接出现梯度爆炸导致训练中断。这并非硬件不行而是不同架构对数值精度的敏感度存在差异。特别是在使用 BF16BFloat16时虽然它比 FP16 拥有更宽的动态范围但在某些算子实现和梯度累积策略上依然需要针对性的调优。今天就来聊聊我在 LLaMA-Factory 框架下解决 ROCm 混合精度训练稳定性问题的实战心得。为什么 AMD 显卡上的混合精度更容易“炸”首先要打破一个误区BF16 并不是万能药。虽然 BF16 的指数位与 FP32 相同理论上能更好地保留大数值信息减少溢出风险但 AMD 的 CDNA 或 RDNA 架构在执行矩阵乘法GEMM时的累加逻辑与 NVIDIA Tensor Core 有所不同。在实际操作中我发现两个主要诱因算子实现的数值误差累积部分底层算子在 ROCm 库如 rocBLAS中的默认实现为了追求极致速度可能在中间结果的精度保持上做了妥协。当这些微小误差在深层网络中逐层传递并累积时就会导致最终梯度的偏差被放大。损失缩放Loss Scaling策略不匹配传统的动态损失缩放算法是基于 FP16 设计的其阈值判断逻辑直接套用到 BF16 场景下可能过于激进或保守。如果缩放因子过大梯度反传时容易溢出过小则导致小梯度被截断为零Underflow模型无法收敛。因此在 ROCm 上开启混合精度不能简单地照搬 CUDA 时代的配置文件必须对精度策略进行“本地化”改造。核心调优手段调整缩放因子与切换精度模式遇到 Loss 突然变成 NaN 或者 Inf第一反应不应该是降低学习率而是检查混合精度的配置。以下是我在多次试错后总结出的两套有效方案。方案一精细化调整 Loss ScalingLLaMA-Factory 默认通常会启用动态损失缩放。在 AMD 卡上建议先尝试手动固定缩放因子观察训练稳定性。如果你使用的是基于 PyTorch 原生 AMP 的后端可以在启动脚本或配置文件中干预GradScaler的行为。虽然 LLaMA-Factory 封装了大部分逻辑但我们可以通过环境变量或修改源码中的初始化参数来调整。例如将初始缩放因子设置得更保守一些# 伪代码示例在 trainer 初始化前干预fromtorch.cuda.ampimportGradScaler# 针对 ROCm 环境适当降低初始 scale 值避免早期溢出scalerGradScaler(init_scale2**10,growth_factor2.0,backoff_factor0.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.yamlmodel_name_or_path:meta-llama/Llama-3-8Bdo_train:truetemplate:llama3finetuning_type:loralora_target:allcompute_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-n1rocm-smi--showmeminfovram如果在训练过程中显存占用剧烈跳动可能需要检查是否开启了gradient_checkpointing。在显存紧张时强制开启混合精度可能会触发额外的内存交换进而影响数值稳定性。综合配置建议针对大多数在 AMD 显卡上进行的 LoRA 微调任务我推荐以下“稳健型”配置组合# 推荐配置稳健优先compute_type:bf16# 优先使用原生 BF16lora_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_sourceAIpaper