Primus：在 AMD GPU 上用于大规模模型的轻量化统一训练框架

原文作者：Wen Xie, Yao Fu, Xiaoming Peng, Xiaobo Chen, Liz Li, Vidushi Goyal, Anshul Gupta

训练大规模LLM在实际落地中复杂度很高。不同框架的接口不一致，多 GPU 与分布式需要脆弱的脚本拼装，不同后端的细节差异会引入额外开销，拖慢训练迭代。Primus 通过一个简洁、与后端解耦的训练框架，帮助开发者在 AMD GPU 上更快地发起、定制与扩展训练任务。

本文将介绍如何使用Primus [1]，如何快速上手，以及它如何在兼容 ROCm 的集群与开发环境中简化规模化训练流程。我们此前已发布 Primus-Turbo 与 Primus‑SaFE 的相关内容，感兴趣可跳转阅读

Primus-Turbo 简介：在 AMD GPU 上加速 Transformer 模型的高性能库

AMD 与 Nexa AI 联手突破设备端 AI：在 NPU 上实现 SDXL-Turbo 实时图像生成

Primus [1] 是一个统一、模块化的训练框架，同时支持 Megatron-LM 与 TorchTitan 两个后端，便于在 AMD 硬件上配置、扩展与优化大模型训练。它的核心目标是统一训练接口，让Megatron-LM 与 TorchTitan 等引擎以一致的方式使用。Primus 提供基于 YAML 的配置系统、模块化后端支持、内置的预检校验，以及结构化日志工具。

重点能力：

• 基于YAML 的统一 CLI 配置

• 模块化多后端设计

• 内置预检校验，提升启动可靠性

• 结构化日志，提升问题定位效率

Primus 将不同后端的参数抽象为结构化、可组合的 YAML，使训练可以用一条命令启动。它减少重复配置，增强可复现性，并简化跨团队与跨环境协作。Primus 采用声明式设计：用 YAML 定义实验，而不是大量 CLI 参数或临时脚本，从而让训练流程更干净、可维护。

下面是一个定义Primus 训练实验的 YAML 文件，契约清晰、可复现。

work_group: AMDuser_name: rootexp_name: llama3.1_8b-pretrainworkspace: ./output
modules:  pre_trainer:    framework: megatron    config: pre_trainer.yaml
    # model to run    model: llama3.1_8B.yaml    overrides:      train_iters: 50      micro_batch_size: 2      global_batch_size: 128
      # parallel      tensor_model_parallel_size: 1      pipeline_model_parallel_size: 1
      # data      mock_date: true      train_data_path: null      valid_data_path: null      test_data_path: null

启动实验只需一条可复现的命令：

export EXP=examples/megatron/configs/llama3.1_8B-pretrain.yamlbash examples/run_pretrain.sh

Primus 与 Megatron-LM 深度集成，同时对用户屏蔽后端复杂度：

• 结构化配置映射会自动把YAML 字段转换为 Megatron 期望的 argparse.Namespace。

• 自动补丁：无需手工修改，即可调整tokenizer 类型、attention 后端、optimizer overrides 等。

  
  

这种设计让你几乎不改动上游Megatron 代码即可运行，同时获得 Primus 的简洁配置与统一启动体验。

Primus 开箱支持 Megatron-LM，包括张量并行（tensor parallelism, TP）、流水并行（pipeline parallelism, PP）、专家并行（expert parallelism, EP）等高级并行策略。它的模块化后端架构也便于扩展到更多训练引擎，例如计划支持的 TorchTitan，且都通过同一统一接口访问。每个后端实现一个通用的 launcher，封装环境准备、参数翻译与执行逻辑。Primus 不强行用同一份配置适配所有引擎，而是清晰隔离不同后端的行为差异。

这让开发者可以：

• 几乎不改动配置即可切换后端

• 避免手工重配环境

• 在不同机器与集群上一致地启动训练

  
  

后续与TorchTitan 的集成将引入 FP8 精度与 fused attention 在 AMD GPU 上的支持，同时保持统一的 CLI 与基于 YAML 的工作流。

Primus 提供强大的预检系统：验证环境变量、检查并行配置、估算模型规模与 FLOPs，在训练开始前提前发现潜在问题。适配 HPC 工作流时，还支持可选的 Slurm 作业脚本生成，便于融入分布式计算环境。启动训练前，Primus 会执行全面预检，帮助开发者确认集群环境已准备好进行大规模分布式训练。与传统静态配置校验不同。Primus 预检会主动探测运行时环境，并生成结构化诊断报告。

你将获得：

• 集群健康检查：验证集群连通性与RDMA/RCCL 通信可用性

• GPU 诊断：检查GPU 可用性、拓扑与内存带宽

• 网络性能基准：测量跨节点的有效带宽，定位潜在瓶颈

  
  

结果会自动整理为可读的PDF 报告，便于在训练前审阅系统就绪情况、与团队共享诊断，快速定位性能瓶颈或配置异常节点。Primus 适合在集群与 CI 流水线中进行稳健、面向生产的部署。

分布式训练的调试通常很难。Primus 引入结构化日志，把输出按实验与 rank 组织，便于在大规模任务中快速定位问题、还原行为。分布式训练常会在多 rank 与多节点上产生分散且庞大的日志，难以高效分析。Primus 的分层日志体系按实验、模块、rank 与严重级别组织日志，提升分析与问题跟踪效率。

日志目录采用清晰的规范，如下所示：

{workspace}/└── {work_group}/{user_name}/{exp_name}/logs/pre_trainer/    ├── rank-0/    │   ├── debug.log    │   ├── error.log    │   ├── info.log    │   └── warning.log    ├── rank-1/    │   └── ...    ...    ├── rank-7/

这让开发人员可以拥有：

• 日志根目录遵循workspace 层级

  
  

全部日志输出路径为：

{workspace}/{work_group}/{user_name}/{exp_name}/logs/

注意：这可确保日志按团队、用户、实验隔离。

• 模块级隔离

    o 每个模块（如pre_trainer）独立写日志，避免冲突并增强模块化。

• Rank 分离

    o 每个rank 写自己的日志（rank-0、rank-1 等），便于分布式中的细粒度定位。

• 日志级别分流

    o 每个rank 会按严重级别输出（debug.log、info.log、warning.log、error.log），便于检索或聚合关键问题。

这种结构让你可以：

• 将错误精确追溯到具体rank 或模块

• 对比不同worker 的行为

• 更好地归档、审计或跨团队共享日志

  
  

通过将日志与用户定义的元数据（work_group、user_name、exp_name）对齐，Primus 让分布式训练不仅可复现，也具备规模化的可观测性与可诊断性。

你可以通过两种方式立即试用Primus：Docker 交互模式或快速开始模式。

设置Docker

使用官方ROCm Megatron Docker 镜像，保持一致的运行环境：

docker pull docker.io/rocm/megatron-lm:v25.7_py310

交互模式

进入容器，在容器内执行训练：

# Launch the containerdocker run -it --device /dev/dri --device /dev/kfd --device /dev/infiniband --network host --ipc host --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged -v $HOME:$HOME -v  $HOME/.ssh:/root/.ssh --shm-size 128G --name primus_training_env rocm/megatron-lm:v25.7_py310
# Install dependenciespip install -r requirements.txt
# Export your HF_TOKEN in the workspaceexport HF_TOKEN=<your_hftoken>
# Launch training(e.g llama3.1_8B)EXP=examples/megatron/configs/llama3.1_8B-pretrain.yaml bash ./examples/run_pretrain.sh

快速开始模式

无需进入Docker 容器。本地克隆仓库，设置配置并直接运行。

1. 克隆仓库

克隆仓库并安装依赖：

# Clone with submodulesgit clone -b v0.1.0-rc1 --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.gitcd Primus
# Or initialize submodules if already clonedgit submodule update --init --recursive
# Install Python dependenciespip install -r requirements.txt

2. 启动训练

使用run_local_pretrain.sh 脚本本地启动训练：

# Export DOCKER IMAGEexport DOCKER_IMAGE=rocm/megatron-lm:v25.7_py310
# Export your HF_TOKEN in the workspaceexport HF_TOKEN=<your_hftoken>
# Example for megatron llama3.1_8Bexport EXP=examples/megatron/configs/llama3.1_8B-pretrain.yaml bash ./examples/run_local_pretrain.sh

本文展示了如何使用Primus 简化 AMD ROCm 上的大模型训练。Primus [1] 让训练实验更易配置、更安全启动、更易调试。从在 AMD GPU 上的大规模预训练到常规微调，Primus 提供一致、可靠的工作流。无论是微调 7B 模型，还是在 AMD GPU 上扩展预训练，Primus 都能让你的迭代更快、更安心。

[1] Primus：https://github.com/AMD-AIG-AIMA/Primus

第三方内容由其权利方直接授权，并非由AMD 授权。所有链接的第三方内容按“现状”提供，不附带任何形式的担保。使用该等第三方内容完全由你自行决定，AMD 在任何情况下均不对你使用第三方内容承担责任。使用第三方内容的风险由你自行承担，并对由此产生的任何损害自行负责。