如何使用官方 ROCm 镜像在 AMD GPU 上启动 SWIFT 训练？

如何使用官方 ROCm 镜像在 AMD GPU 上启动 SWIFT 训练【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift如果你有一台或多台配备 AMD GPU 的机器想用 ms-swiftSWIFT在上面跑模型训练官方文档提供了一条完整路径拉取适配了 ROCm 生态的 ms-swift Docker 镜像启动容器并确认 PyTorch 能识别 AMD GPU然后在容器内执行megatron sft等训练命令。本文依据 AMD GPU 支持文档 整理适用于在 AMD GPU 上首次搭建 SWIFT 训练环境并跑通单机训练的场景多机训练的关键环境变量也一并给出。拉取官方 ROCm 镜像并启动容器文档给出的基础环境命令如下IMAGE_NAME是适配了 AMD ROCm 生态的 ms-swift 镜像tag 中已固定 ROCm 7.2.0、PyTorch 2.10.0、vLLM 0.18.1、SWIFT 4.1.0 版本组合IMAGE_NAMEamdagi/modelscope:ubuntu22.04-rocm7.2.0-py312-torch2.10.0-vllm0.18.1-modelscope1.35.1-swift4.1.0 docker pull ${IMAGE_NAME} CONTAINER_NAMEswift_test docker run -it --networkhost --ipchost --privileged --group-add video \ --device/dev/dri --device/dev/kfd \ --shm-size 512G --ulimit memlock-1 \ --security-opt seccompunconfined --cap-add SYS_PTRACE \ --name ${CONTAINER_NAME} \ ${IMAGE_NAME} \ /bin/bash几个与 AMD 平台直接相关的参数需要保留--device/dev/dri和--device/dev/kfd是把 AMD GPU 设备节点映射进容器--group-add video保证容器内用户有权限访问 GPU 设备--shm-size 512G和--ulimit memlock-1是为大规模显存通信预留的共享内存与锁定内存上限。如果后续需要运行更新版本的 ms-swift文档建议使用 pip 升级或基于源码安装并添加--no-deps选项以避免自动升级其他依赖引起的问题。启动训练前先做两项环境检查进入容器后先确认 PyTorch 正确识别 AMD GPUpython -c import torch;print(torch.cuda.is_available()) # output: True文档给出的预期输出是True。若输出不是True说明容器或宿主机侧的 GPU 设备没有正确映射应回到上面的docker run参数排查。然后检查 GPU 拓扑连接及 NUMA 亲和性这在决定并行度切分时有用rocm-smi --showtopo该命令输出 GPU 之间的 Weight/Hops/Link Type 矩阵以及各 GPU 所属的 NUMA Node文档中示例为一台 8 卡机器GPU 0–3 在 NUMA 0、GPU 4–7 在 NUMA 1。另外可以用rocm-smi或rocm-smi -u --showmeminfo vram查看 GPU 利用率与显存占用训练过程中用它确认卡确实被用起来了。单机训练Megatron-SWIFT 全量微调 Qwen3.5-35B-A3B环境检查通过后文档给出的训练示例是用 Megatron-Swift 对Qwen/Qwen3.5-35B-A3B做全参数 SFT单机 8 卡。AMD GPU 单卡显存较大文档建议借此联合调优以下几项来提升吞吐并行度调优TP/PP/EP 等尽量减小并行切分带来的通信开销优先级 PP/EP TP显存允许时关闭 optimizer CPU offload--optimizer_cpu_offload false显存允许时调整 activation/gradient checkpointing--recompute_granularity none或--recompute_granularity selective配合--recompute_modules做细粒度控制MoE 模型建议设置export NVTE_USE_GROUPED_GEMM_TRITON1使用 triton 实现的 grouped gemm kernel带有 GatedDeltaNet 结构的模型建议设置USE_MCORE_GDN1使用 mcore 的实现版本为避免某些 AMD GPU 上可能出现的问题、保证性能更稳定建议export HSA_NO_SCRATCH_RECLAIM1。完整的单机训练命令文档原样给出output_dir默认写入当前目录下的megatron_output/Qwen3.5-35B-A3Bexport HSA_NO_SCRATCH_RECLAIM1 export NVTE_USE_GROUPED_GEMM_TRITON1 output_dir${PWD}/megatron_output/Qwen3.5-35B-A3B mkdir -p ${output_dir} current_time$(date %Y.%m.%d-%H.%M.%S) log_file${output_dir}/1node_full_megatron_Qwen3.5-35B-A3B_${current_time}.log PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True \ NPROC_PER_NODE8 \ MAX_PIXELS1003520 \ VIDEO_MAX_PIXELS50176 \ FPS_MAX_FRAMES12 \ SKIP_MULTIMODAL_MTP_VALIDATION1 \ USE_MCORE_GDN1 \ megatron sft \ --model Qwen/Qwen3.5-35B-A3B \ --dataset AI-ModelScope/LongAlpaca-12k \ --save_safetensors true \ --load_from_cache_file true \ --tuner_type full \ --add_non_thinking_prefix true \ --split_dataset_ratio 0.01 \ --tensor_model_parallel_size 1 \ --pipeline_model_parallel_size 1 \ --expert_model_parallel_size 8 \ --sequence_parallel true \ --moe_permute_fusion true \ --moe_grouped_gemm true \ --moe_shared_expert_overlap true \ --moe_aux_loss_coeff 1e-6 \ --moe_expert_capacity_factor 2 \ --micro_batch_size 1 \ --global_batch_size 8 \ --recompute_granularity selective \ --recompute_modules core_attn mlp moe \ --gradient_accumulation_fusion false \ --num_train_epochs 500 \ --group_by_length true \ --finetune true \ --freeze_llm false \ --freeze_vit false \ --freeze_aligner false \ --cross_entropy_loss_fusion true \ --lr 1e-5 \ --lr_warmup_fraction 0.05 \ --min_lr 1e-6 \ --weight_decay 0.1 \ --adam_beta2 0.95 \ --eval_steps 500 \ --save_steps 500 \ --save_total_limit 10 \ --logging_steps 1 \ --max_length 16384 \ --dataloader_num_workers 8 \ --dataset_num_proc 8 \ --no_save_optim true \ --no_save_rng true \ --optimizer_cpu_offload false \ --attention_backend flash \ --padding_free false \ --output_dir ${output_dir} \ 21 | tee ${log_file}命令中NPROC_PER_NODE8对应 8 卡单机--expert_model_parallel_size 8与之匹配换用其他卡数或换用非 MoE 模型时需要按文档前面的并行度调优原则自行调整这些并行参数。训练日志会被tee写入output_dir下的${log_file}权重按--save_steps 500的步长保存到output_dir这是判断训练是否持续推进的直接依据。文档还提供了 GRPO 强化学习的单机样例megatron rlhf --rlhf_type grpo配合 vLLM colocate 模式如果你要做的是 RL 训练而非 SFT可参考同一文档中的 2.2 节命令其前置环境变量与单机 SFT 相同。扩展到多机训练多机训练时文档要求在每台机器上先导出以下环境变量以 2 节点为例再执行与单机相同的训练脚本主体export NNODES2 # 此处以 2 节点为例 export NODE_RANK0 # 主节点设置为 0从节点设置为 1 export MASTER_ADDRMASTER_NODE_IP # 根据主节点 ip 设置 export MASTER_PORT29500 # 设置通信端口 export NCCL_SOCKET_IFNAMEens50f1np1 # 根据机器实际通信网口名设置可通过 ifconfig 查看 export GLOO_SOCKET_IFNAMEens50f1np1 # 根据机器实际通信网口名设置可通过 ifconfig 查看 export NCCL_IB_HCAmlx5_0,mlx5_1,mlx5_2,mlx5_3 # 根据实际IB网卡名设置可通过 ibv_devices 查看 export NCCL_IB_GID_INDEX3其中MASTER_NODE_IP需要替换为实际主节点的 IPNCCL_SOCKET_IFNAME/GLOO_SOCKET_IFNAME中的ens50f1np1和NCCL_IB_HCA中的网卡名都是文档示例值需用ifconfig和ibv_devices查出本机实际值后替换。已知问题与限制文档明确列出的 AMD 平台注意事项强化学习训练如果 RL 训练使用 vLLM 作为推理引擎需要 vLLM 0.11.0并建议使用 ROCm 7.0 或官方镜像以避免 sleep mode memory leak 问题。Ray Megatron 多 GPU/Node 训练使用 Ray Megatron 而非torchrun方式启动时不要设置CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES等变量以避免冲突。MoE 模型训练建议增加环境变量NVTE_USE_GROUPED_GEMM_TRITON1和参数--gradient_accumulation_fusion false以避免偶发的 GPU 卡死问题。另外常见问题文档 中关于如何在 ROCm/MI300X 上使用 megatron-swift 训练 Qwen3-Omni的回答也指向同一份 AMD GPU 支持文档可以按本文的路径操作。【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考