ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

10 分钟上手 AxisRL:从 Docker 环境到第一条 GRPO 训练命令的完整新手指南

10 分钟上手 AxisRL:从 Docker 环境到第一条 GRPO 训练命令的完整新手指南

【免费下载链接】axrl

AxisRL is an agentic RL post-training framework built on SGLang rollout, Megatron training, and real-world agent workflows.

项目地址:https://gitcode.com/gh_mirrors/ax/axrl
点击查看免费下载

AxisRL 是一个面向 agentic RL 的 LLM 后训练框架(post-training framework):它用 SGLang 做高吞吐 rollout、用 Megatron 做大规模分布式训练,并把多轮 agent 交互、reward 收集与权重同步串成一条完整闭环。本文带你从拉取 Docker 镜像开始,一步步跑通第一条 GRPO 训练命令,帮你快速完成 AxisRL 安装、数据准备和首次训练。

🧭 什么是 AxisRL?

传统 RLHF 大多围绕单轮问答展开,而真实的 agent 场景里,模型需要与环境多轮交互、调用工具,最后才拿到 reward。AxisRL 就是为这类agentic 强化学习后训练设计的:

  • Rollout 侧:SGLang worker 负责高吞吐推理生成;
  • 训练侧:Megatron worker 负责 PPO / GRPO 等大模型训练;
  • 中间层:AxisRL 的 driver 负责调度、环境交互、verifier 打分、样本构建和权重同步。

从这张官方架构图中可以看到一个完整循环:输入队列派发给 Rollout Actors → 与 Gym 风格环境交互 → SGLang Workers 生成 token → 打包样本送入 Output Queue → Megatron Workers 训练并更新权重 → 再同步回 rollout 侧,进入下一轮迭代。

✅ 开始前:硬件与前置条件

条件说明
操作系统Linux + NVIDIA GPU(默认 GRPO 配置约需 4 张 GPU:Megatron 侧 tp=2 × pp=2,rollout 与训练 colocate 共享显存)
运行时已安装 Docker,支持--gpus
网络可访问 Hugging Face(首次运行会自动拉取模型和数据集)

💡 显存不够?别急,后面会教你如何用命令行参数一键降低并行度、换更小的模型。

🐳 第一步:获取代码,启动 Docker 环境

AxisRL 推荐直接使用项目提供的预构建镜像,其中已包含 SGLang、Megatron Core、MagiAttention、Ray 和全部 CUDA 依赖,不用自己折腾环境。

1️⃣ 克隆代码(镜像与代码需放在同一台机器):

git clone https://gitcode.com/gh_mirrors/ax/axrl cd axrl

2️⃣ 拉取官方预构建镜像:

docker pull leejunjie/sglang-mcore:cu130-sgl0.5.14-mcore0.18-magi

3️⃣ 启动容器并把仓库目录挂载进去(镜像工作目录为/root):

docker run -it --name axrl \ --gpus all \ --net=host --ipc=host \ -v /path/to/axrl:/root/axrl \ leejunjie/sglang-mcore:cu130-sgl0.5.14-mcore0.18-magi \ bash

如果你需要自行构建镜像,可以直接参考仓库中的 cu130-sgl0.5.14-mcore0.18-magi.Dockerfile。

📦 第二步:安装 AxisRL 并准备数据

在容器内的/root/axrl目录下执行:

pip install -e .

(仓库里的 install.sh 做的就是这件事。)

然后是模型和数据。完整下载脚本 download_data.py 会拉取全部 recipe 用到的模型和数据集,体量很大(包含多 B 参数的模型)。本教程的 GRPO 示例只需要两样东西,首次运行训练时会自动按需下载,无需提前手动下载:

  • 模型:Qwen/Qwen2.5-1.5B-Instruct
  • 数据集:openai/gsm8k(train + test)

如果希望提前缓存,也可以只运行一次完整下载:

python axrl/example/download_data.py

🧩 第三步:认识 GRPO GSM8K 训练配置

每个 recipe 都有自己的配置入口。GSM8K GRPO 的配置在 axis_recipe/grpo_gsm8k/pipeline_config.yaml,几个关键字段如下:

配置项默认值含义
megatron_worker.model.nameQwen/Qwen2.5-1.5B-Instruct训练 + rollout 共用的模型
train_datasetsopenai/gsm8k/train训练数据集
grpo.loss_typegrpoGRPO 策略优化目标
online_rl_train.num_rollouts_per_conversation8每条 prompt 采样 8 条 rollout,组内对比计算 advantage
online_rl_train.model_sync_every_n_global_updates4每 4 次全局更新把新权重同步回 rollout 侧

配置背后的任务逻辑实现在 train_pipeline.py:MathEnv驱动数学题的多轮解题过程,GSM8KVerifier对最终答案打分产生 reward。

🚀 第四步:运行第一条 GRPO 训练命令

一切就绪,直接执行官方示例命令(来源:axis_recipe/grpo_gsm8k/run_train.sh):

AXRL_OUTPUT_DIR_NAME=grpo_gsm8k \ bash axis_recipe/grpo_gsm8k/run_train.sh \ --online_rl_train.max_global_updates=4

逐行拆解这条命令:

部分作用
AXRL_OUTPUT_DIR_NAME=grpo_gsm8k指定输出目录名,结果默认写入~/axrl-data/outputs/grpo_gsm8k/
run_train.sh先由 generate_configs.py 生成 YAML,再启动训练 pipeline
--online_rl_train.max_global_updates=4只跑 4 个全局更新就结束——一次 10 分钟级的小规模 smoke run,先验证环境没问题

脚本还会把完整日志 tee 到~/axrl-data/outputs/grpo_gsm8k/run-train-pipeline-grpo-GSM8K.log,并用 TensorBoard 记录 loss、reward 等指标。启动后框架会先跑一次 GSM8K test 集的 baseline 评估(eval_on_start: true),然后进入 rollout → 训练 → 权重同步的循环。看到 reward 指标随更新推进,就说明你的第一条 GRPO 训练命令跑通了 🎉

🔁 训练过程中发生了什么?

每一轮迭代大致如下:

  1. Rollout:8 个 rollout actor 从 GSM8K 采样 prompt,SGLang worker 生成模型回答;
  2. 打分:verifier 按最终答案的正确性给每条轨迹计算 reward;
  3. 训练:GRPO 以组内 reward 均值作为 baseline 计算 advantage,Megatron 执行策略梯度更新;
  4. 同步:每 4 次全局更新,新权重同步回 SGLang,下一轮 rollout 用新模型——如此循环。

⚙️ 进阶:不改 YAML,用命令行覆盖参数

AxisRL 的参数覆盖语法是--path.to.field=value,绝大多数配置项都能在命令行直接改,优先级为YAML < 命令行 < 环境变量(详见 config-parsing.md)。

显存紧张时的最小化示例:

bash axis_recipe/grpo_gsm8k/run_train.sh \ --megatron_worker.model.name=Qwen/Qwen3-0.6B-Base \ --megatron_worker.tp_size=1 \ --megatron_worker.pp_size=1 \ --online_rl_train.max_global_updates=4

跑通 GSM8K 后,还可以换用其他 recipe:

Recipe入口脚本场景
GSM8K PPOaxis_recipe/ppo_gsm8k/run_train.sh带 value worker 的 PPO 数学训练
Search-R1axis_recipe/search_r1/run_train.sh检索增强的多轮搜索 RL
Black-Box RLaxis_recipe/blackbox_rl/README.md通过 OpenAI 兼容代理接入 OpenHands/E2B 黑盒 harness

📈 进阶能力:rollout / trainer 一致性分析

大规模训练中最隐蔽的问题往往来自 rollout 与训练两侧执行路径的细微差异(tokenization、routing、packing 等)。AxisRL 内置了 mismatch analysis 与 spike replay:对比两侧 token 级 KL 差异、在 loss 尖峰前自动保存快照用于重放复现。下图就是开启 R3 routing replay 前后,KL mismatch 分布的对比——橙色曲线整体显著更平稳:

❓ 常见问题(FAQ)

Q1:显存不足、GPU 数量不够怎么办?用命令行参数降低并行度并换小模型,如--megatron_worker.tp_size=1 --megatron_worker.pp_size=1 --megatron_worker.model.name=Qwen/Qwen3-0.6B-Base。

Q2:结果和日志在哪里看?默认输出目录为~/axrl-data/outputs/<AXRL_OUTPUT_DIR_NAME>/,其中logs/存放 worker 日志,根目录有 tee 出的完整运行日志与 TensorBoard 事件文件。

Q3:训练报错找不到模型或数据集?确认机器能访问 Hugging Face;或提前运行python axrl/example/download_data.py缓存,再把 YAML 中的name字段指向本地路径。

Q4:只想快速验证环境,不想跑完整训练?加--online_rl_train.max_global_updates=4即可把训练压缩到几个更新,几分钟内结束。

📚 下一步

  • 中文文档总览:docs/README-cn.md
  • 技术博客(深入架构与工程实践):docs/axrl-blog-cn.md
  • 配置解析完整规则:docs/config-parsing.md
  • 一致性测试与回归用例:tests/

至此,你已经完成了 AxisRL 从 Docker 环境搭建到第一条 GRPO 训练命令的全部流程。接下来可以试着把max_global_updates调大、换用 PPO 或 Search-R1 recipe,开始你自己的 agentic RL 实验吧!

【免费下载链接】axrl

AxisRL is an agentic RL post-training framework built on SGLang rollout, Megatron training, and real-world agent workflows.

项目地址:https://gitcode.com/gh_mirrors/ax/axrl
点击查看免费下载
上一篇:终极 Visual C++ 运行库 AIO 解决方案:一键解决所有 DLL 错误问题
下一篇:Windows电脑内存清理终极指南:如何快速释放内存提升性能

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进