环境搭建：从 Docker 镜像到依赖隔离

在 AMD ROCm 环境下折腾大模型微调，最劝退的往往不是算法本身，而是那令人头大的环境配置。对于急需验证模型效果的算法工程师来说，时间就是算力，我们没精力去逐行解决系统级的依赖冲突。我的建议非常直接：放弃在宿主机直接安装 PyTorch ROCm 版本的念头，直接使用官方提供的 Docker 容器。这不仅能确保驱动版本与软件栈的完美匹配，还能避免污染宿主机的 Python 环境。

我通常选用 rocm/pytorch:latest 或针对特定 MI300X 优化的镜像作为基座。启动容器时，务必加上 --device /dev/kfd --device /dev/dri 参数以确保 GPU 可见性，同时挂载代码目录和数据集目录。进入容器后，第一步是验证 ROCm 是否正常工作，运行 python -c "import torch; print(torch.cuda.is_available())"（注意：在 ROCm 中 torch 依然沿用 cuda 接口名，但底层调用的是 HIP），若返回 True 且显示显卡型号，则基础环境就绪。

接下来安装 LLaMA-Factory。虽然它支持 pip install，但在 ROCm 环境下，为了获得最佳的算子支持（特别是 Flash Attention 的 ROCm 变种），强烈建议从源码安装并指定相关环境变量：

export MAX_JOBS=4
export ROCM_PATH=/opt/rocm
pip install -e ".[torch,metrics]"

如果在编译 flash-attn 时遇到报错，通常是因为编译器找不到 HIP 头文件，此时需检查 ROCM_PATH 是否指向正确目录。一旦安装完成，运行 llamafactory-cli version 确认版本无误，我们就拥有了一个干净、可复现的微调沙箱。

避坑实录：compute_type 设置与梯度爆炸

环境跑通只是第一步，真正的挑战始于训练启动。在 NVIDIA 平台上习以为常的配置，搬到 AMD 卡上可能会引发灾难性的后果。我最深刻的一次教训是关于 compute_type 的设置。

起初，为了节省显存并加速训练，我沿用了在 H800 上的习惯，在 YAML 配置文件中将 compute_type 设置为 fp16。然而，训练刚开始几百步，Loss 瞬间变成 NaN，随后整个进程崩溃。查看日志，并没有明显的显存溢出（OOM），而是典型的梯度爆炸特征。

经过反复排查和社区 Issue 检索，问题定位到了 AMD Instinct 系列显卡（尤其是 MI250/MI300 系列）对 fp16 的数值稳定性支持与 NVIDIA A/H 系列存在差异。在某些算子实现中，ROCm 底层的 fp16 累加精度不足，导致微小梯度在反向传播时被放大。

解决方案非常明确：切换到 bf16（BFloat16）。

BF16 拥有与 FP32 相同的指数位宽，极大地提升了动态范围，能有效防止梯度溢出。修改后的配置片段如下：

# lora_finetune.yaml
compute_type: bf16  # 关键修改：弃用 fp16，改用 bf16
optim: adamw_bf16   # 配合使用支持 bf16 的优化器
lr_scheduler_type: cosine
warmup_ratio: 0.1

将 compute_type 改为 bf16 并重启训练后，Loss 曲线迅速平滑下降，收敛行为恢复正常。这个坑提醒我们：不要盲目复用 CUDA 时代的“最佳实践”，在 ROCm 环境下，BF16 往往是更稳妥的默认选择，除非你有极其特殊的理由必须使用 FP16。

多卡实战：DeepSpeed ZeRO-3 的显存魔法

单卡验证通过后，面对 70B 甚至更大参数的模型，多卡分布式训练是必经之路。AMD 的 RCCL（Rocm Communication Collectives Library）已经能够很好地替代 NCCL 进行多卡通信，而 LLaMA-Factory 对 DeepSpeed 的集成让这一过程变得相当透明。

重点在于 ZeRO-3 (Zero Redundancy Optimizer Stage 3) 的配置。在单卡显存有限的情况下（例如单卡 80GB 跑 70B 模型），ZeRO-3 通过将优化器状态、梯度和模型参数分片存储在所有卡的显存中，实现了“用空间换时间”的极致显存节省。

在我的 MI300X 八卡集群测试中，未开启 ZeRO-3 时，仅加载模型权重就已接近显存上限，根本无法进行训练。开启 ZeRO-3 后，单卡显存占用瞬间下降了约 70%，使得全量微调成为可能。

以下是我整理的一份经过实测的多卡微调配置模板，可直接复制使用：

### deepspeed_zero3.yaml
deepspeed: examples/deepspeed/ds_z3_config.json

# ds_z3_config.json 核心内容参考：
{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "none", 
      "pin_memory": true
    },
    "offload_param": {
      "device": "none",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_parition_grads": true,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

启动命令：

FORCE_TORCHRUN=1 llamafactory-cli train \
    --model_name_or_path meta-llama/Llama-3-70B-Instruct \
    --do_train \
    --finetuning_type full \
    --dataset alpaca_en_demo \
    --template llama3 \
    --deepspeed examples/deepspeed/ds_z3_config.json \
    --output_dir saves/llama3-70b/full/sft \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --learning_rate 1e-5 \
    --num_train_epochs 3 \
    --compute_type bf16 \
    --plot_loss true

在这个配置下，offload_optimizer 和 offload_param 均设为 none，意味着所有数据驻留显存以换取最快速度。如果显存依然紧张，可以将 device 改为 cpu，利用主机内存进行卸载，虽然会牺牲部分通信带宽，但能进一步突破显存限制。实测数据显示，在八卡互联环境下，ZeRO-3 不仅解决了显存瓶颈，由于减少了单卡的数据负载，通信开销也在可接受范围内，整体训练吞吐量依然保持在高位。

常见报错与快速排查手册

在 ROCm 环境下踩坑是常态，以下是几个高频报错及其“填坑”方案，建议收藏备用：

1. 报错：RuntimeError: HIP error: hipErrorNoDevice

   • 原因：容器未正确映射 GPU 设备，或当前用户无权限访问 /dev/kfd。
   • 解法：检查 Docker 启动参数是否包含 --device /dev/kfd --device /dev/dri，或在宿主机执行 chmod 666 /dev/kfd（临时方案）。

2. 报错：NCCL/RCCL initialization failed

   • 原因：多卡训练时，网卡接口识别错误或防火墙阻挡。
   • 解法：显式指定网络接口， export NCCL_SOCKET_IFNAME=eth0（替换为你的实际内网网卡名），并确保节点间端口互通。

3. 报错：Kernel launch configuration invalid

   • 原因：某些算子的 Block Size 设置超过了当前 GPU 架构的限制。
   • 解法：尝试降低 per_device_train_batch_size，或在 LLaMA-Factory 中禁用特定的融合算子（如设置 disable_flash_attn: true 进行排查）。

4. 训练 Loss 不下降或震荡

   • 原因：除了前述的 fp16 精度问题外，还可能是 Learning Rate 过大。
   • 解法：在 ROCm 上，建议初始学习率比 NVIDIA 环境略低（例如从 1e-4 降至 5e-5），并配合 Warmup 策略。

通过这套流程，从环境隔离到精度调优，再到多卡扩展，我们可以在 AMD 平台上构建出一条稳定高效的大模型微调流水线。ROCm 生态或许还在成长中，但只要掌握了正确的配置方法和思维模式，它完全能够胜任生产级的训练任务。

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