执行摘要
- 一句话:补齐 AMD 平台 Qwen3-4B LoRA RL 启动器
- 推荐动作:值得精读。若团队需要维护或扩展 AMD / ROCm 上的 LoRA RL 训练,这是当前唯一入口;其「平台差异收敛在启动器层」的做法——环境变量桥接(HIP → CUDA 可见设备)、SGLang Triton 后端固定、以及「配方 + 快照 + 文档」三件套的组织方式——是值得复用的模式。验证数据(logprob 差、KL、截断率、AIME 趋势)可作为后续 AMD 训练健康度对比的参考基线,但不应据此断言训练收益。
功能与动机
PR body 开头即点明动机:There is no maintained Qwen3-4B LoRA RL launcher for AMD。此前 AMD 用户只能从 CUDA 配方自行移植,缺少可维护、可复现的入口;同时 ROCm 上 SGLang 的 CUDA-only rollout 后端不可用,必须在启动器层把 attention 与 LoRA 后端固定为 Triton 并做 HIP/CUDA 设备环境桥接,才能把已验证的 Megatron-Bridge LoRA 配方带到 AMD 平台。
实现拆解
- 新增启动器入口:
scripts/amd/run_qwen3_4b_lora.py 定义 ScriptArgs(继承 U.ExecuteTrainConfig),以 dataclass 声明全部可调项:hardware 限定为 auto / MI350X / MI355X,num_gpus_per_node 允许覆盖可见卡数;LoRA 超参(rank 64、alpha 32、dropout 0.0、target-modules all-linear、lora-base-cpu-backup)沿用已验证的 Qwen3-4B Megatron LoRA 配方,rollout 形状(8 prompt × 8 样本 = 64 全局 batch)与 8192 token 响应窗口对齐既有 Qwen3 launcher。
- ROCm 环境适配:
_set_rocm_environment() 设置 RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES 与 RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES,阻止 Ray 启动 worker 时重设设备可见性;把 HIP_VISIBLE_DEVICES 同步到 CUDA_VISIBLE_DEVICES,保证上层按 CUDA 习惯读取设备;再设 NCCL_NVLS_ENABLE=0 避开 execute_train 的 NVIDIA 拓扑探测与 ROCm 不支持的 NVLink SHARP。
- 命令组装与提交:
_execute() 按 checkpoint / LoRA / rollout / eval / optimizer / GRPO / SGLang / misc 八组参数块拼接训练命令,其中 SGLang 块固定 --sglang-lora-backend triton 与 --sglang-attention-backend triton,并设置 --sglang-chunked-prefill-size 4096、--sglang-mem-fraction-static 0.4;最终调用 U.execute_train 提交 Ray job,--colocate 让 actor 与 rollout 共享 8 卡。
- CLI 入口划分:
prepare 只下载模型(Qwen/Qwen3-4B)与数据集(dapo-math-17k、可选 aime-2024);train 假设输入已在节点上直接训练;full-train 合并两者,便于一键跑通。
- 测试与文档配套:
tests/manual/launch_scripts/test_py_launch_scripts.py 把新脚本加入 _HARDWARE_A_RECORDING_REPRESENTS(固定 MI355X);新增 tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/ 下 prepare / train / full_train 三个命令录制快照,作为启动器行为回归基准;docs/advanced/lora.md 的 LoRA 支持矩阵补充 AMD 启动器条目并注明 Triton 后端。作者在 8× MI355X、ROCm 7.2、torch 2.9.1+rocm7.2.0 上完成 100 个 optimizer 步骤,train/train_rollout_logprob_abs_diff 均值 0.01232(低于 0.03 门限),AIME-2024 avg@16 从 0.5625 提升至峰值 0.6833。
关键文件:
scripts/amd/run_qwen3_4b_lora.py(模块 启动器;类别 source;类型 core-logic;符号 ScriptArgs, _set_rocm_environment, _resolve_num_gpus, _download_inputs): 本 PR 的核心:新增 Qwen3-4B GRPO LoRA 启动器,集中处理 AMD 硬件声明、ROCm 环境变量桥接、SGLang Triton 后端固定与训练命令组装,是 AMD LoRA RL 的维护入口。
tests/manual/launch_scripts/test_py_launch_scripts.py(模块 启动测试;类别 test;类型 test-coverage;符号 _HARDWARE_A_RECORDING_REPRESENTS): 把新启动器注册到 _HARDWARE_A_RECORDING_REPRESENTS(固定为 MI355X),让快照测试以正确的硬件假设录制命令序列。
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/full_train.txt(模块 命令快照;类别 test;类型 test-fixture): full-train 入口的完整命令录制快照(109 行),覆盖 HF 下载、Ray 启动与 train.py 全量参数,是启动器行为回归的主要依据。
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/train.txt(模块 命令快照;类别 test;类型 test-fixture): train 入口的命令录制快照(92 行),锁定训练参数序列,防止默认值漂移导致启动器行为悄悄变化。
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/prepare.txt(模块 命令快照;类别 test;类型 test-fixture): prepare 入口的下载命令录制快照(16 行),校验模型与数据集下载路径与仓库选择。
docs/advanced/lora.md(模块 文档;类别 docs;类型 documentation): LoRA 支持矩阵补充 AMD 启动器条目并注明 Triton attention/LoRA 后端,让用户能直接定位 AMD 支持状态。
关键符号:ScriptArgs, _set_rocm_environment, _resolve_num_gpus, _download_inputs, _get_wandb_args, _execute, prepare, train, full_train
关键源码片段
scripts/amd/run_qwen3_4b_lora.py
本 PR 的核心:新增 Qwen3-4B GRPO LoRA 启动器,集中处理 AMD 硬件声明、ROCm 环境变量桥接、SGLang Triton 后端固定与训练命令组装,是 AMD LoRA RL 的维护入口。
# scripts/amd/run_qwen3_4b_lora.py(新增启动器核心,节选)
# 平台差异被收敛在启动器层:参数声明、ROCm 环境适配与命令组装都在这里完成,
# 避免把 AMD 特例写进核心训练代码。
@dataclass
class ScriptArgs(U.ExecuteTrainConfig):
hardware: Literal["auto", "MI350X", "MI355X"] = "auto"
num_gpus_per_node: int | None = None
# LoRA 超参沿用已在 CUDA 上验证的 Qwen3-4B Megatron LoRA 配方
lora_rank: int = 64
lora_alpha: int = 32
lora_dropout: float = 0.0
target_modules: str = "all-linear"
lora_base_cpu_backup: bool = True
# Rollout 形状:8 个 prompt × 8 个样本 = 一次 64 样本更新
num_rollout: int = 100
rollout_batch_size: int = 8
n_samples_per_prompt: int = 8
# 8192 token 响应窗口与已维护的 Qwen3 launcher 对齐
rollout_max_response_len: int = 8192
global_batch_size: int = 64
lr: float = 2e-5
# 每个 GPU 一个 4B rollout engine
rollout_num_gpus_per_engine: int = 1
sglang_mem_fraction_static: float = 0.4
def _set_rocm_environment() -> None:
"""在 Ray 启动前后保持调用方的 ROCm 设备选择。"""
# 阻止 Ray 在启动 worker 时重设 HIP/CUDA 可见设备
os.environ.setdefault("RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", "1")
os.environ.setdefault("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "1")
# 把 HIP_VISIBLE_DEVICES 同步到 CUDA_VISIBLE_DEVICES,
# 让上层逻辑仍按 CUDA 习惯读取设备编号
if hip_visible_devices := os.environ.get("HIP_VISIBLE_DEVICES"):
os.environ["CUDA_VISIBLE_DEVICES"] = hip_visible_devices
# execute_train 会触发 NVIDIA 拓扑探测;ROCm 上禁用不支持的 NVLink SHARP
os.environ.setdefault("NCCL_NVLS_ENABLE", "0")
def _execute(args: ScriptArgs) -> None:
_set_rocm_environment()
hardware, num_gpus = _resolve_num_gpus(args)
# LoRA 参数块:lora-base-cpu-backup 让基座权重常驻 CPU,缓解显存压力
lora_args = (
f"--lora-rank {args.lora_rank} "
f"--lora-alpha {args.lora_alpha} "
f"--lora-dropout {args.lora_dropout} "
f'--target-modules "{args.target_modules}" '
)
if args.lora_base_cpu_backup:
lora_args += "--lora-base-cpu-backup "
# 本启动器与 CUDA 配方最大的差异点:ROCm 没有 CUDA-only 的 SGLang rollout 后端,
# 必须把 attention 与 LoRA 后端固定为 Triton
sglang_args = (
f"--rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine} "
f"--sglang-mem-fraction-static {args.sglang_mem_fraction_static} "
"--sglang-chunked-prefill-size 4096 "
"--sglang-lora-backend triton "
"--sglang-attention-backend triton "
)
# 单节点 colocated:actor 与 rollout 共享 8 卡,训练与采样在同一 Ray 集群内完成
misc_args = (
"--train-backend megatron "
f"--actor-num-nodes {args.num_nodes} "
f"--actor-num-gpus-per-node {num_gpus} "
f"--num-gpus-per-node {num_gpus} "
"--colocate "
)
# 其余参数块(checkpoint / rollout / eval / optimizer / GRPO / wandb)在此省略,
# 最终统一交给通用命令工具层提交 Ray job;config 对象与参数字符串共同描述一次训练
train_args = " ".join(
[ckpt_args, lora_args, rollout_args, eval_args, optimizer_args,
grpo_args, _get_wandb_args(args), sglang_args, misc_args, args.extra_args]
)
U.execute_train(
train_args=train_args,
config=args,
num_gpus_per_node=num_gpus,
megatron_model_type=args.megatron_model_type,
megatron_path=args.megatron_path,
)
评论区精华
维护者 yushengsu-thu 在 review 中提问 Should this PR add snapshots?(指向 @yueming-yuan),关注新启动器是否应随命令录制快照一起合入;最终版本已包含 prepare / train / full_train 三个快照文件,且 yushengsu-thu 随后给出 APPROVED,说明该疑问在合并前得到解决。另 claude[bot] 提示本 PR 来自 fork、自动 review 被禁用,维护者可评论 @claude review 触发一次性 review;实际未触发,由维护者人肉审核并批准。
- 是否应把启动器快照纳入本 PR (question): 最终版本已包含 prepare.txt、train.txt、full_train.txt 三个快照文件,且 yushengsu-thu 随后给出 APPROVED,说明该疑问在合并前得到解决。
- fork PR 自动 review 被禁用 (other): 未触发额外 bot review,由维护者 yushengsu-thu 人肉审核并批准合并。
风险与影响
- 风险:
- 快照对工具链版本敏感:
tests/snapshots 冻结了极长且细节级的 train.py 参数序列(含全部模型超参与 --wandb-key 'frozen-wandb-api-key'),一旦 train.py 默认参数演进或 SGLang 参数改名,快照即失效,需要同步更新。
- ROCm 环境变量依赖:
_set_rocm_environment() 依赖 Ray 的 RAY_EXPERIMENTAL_NOSET_* 行为与 SGLang 在 ROCm 上的设备读取约定,相关版本升级后可能失效且难以在 CPU CI 捕获。
- 验证范围有限:仅 8× MI355X 单节点 colocated 实测,MI350X 与多节点 / 跨节点场景、ROCm 7.2 之外版本未覆盖;colocated 模式下 actor 与 rollout 同机显存竞争的行为未展开讨论。
- 评估信号强度不足:AIME-2024 只有 30 道独立题目,avg@16 的涨跌统计功效有限,作者自己也声明这是次要信号而非效果声明。
- 无 GPU CI 覆盖:该启动器只进入
tests/manual 层,PR CI 不实际跑 ROCm 训练,回归只能依赖快照与人工验证。
- 影响:对 AMD / ROCm 用户,这是开箱即用的 Qwen3-4B LoRA RL 起点,prepare / train / full-train 三步 CLI 大幅降低上手成本;对系统,核心训练库零改动(+469/-1 中仅 1 行删除为文档表格),训练路径无回归面,新增快照扩大了 launcher 回归保护;对团队,填补 AMD launcher 维护空白,后续 AMD LoRA 迭代可在该脚本上延续,LoRA 支持矩阵也让用户能直接定位 AMD 状态。整体影响范围为中低,聚焦新增入口与配套资产。
- 风险标记:ROCm 环境变量依赖, 快照对工具链版本敏感, 新入口无 GPU CI 覆盖, 仅单机 colocated 验证, AIME 评估样本量小
关联脉络
- PR #2701 feat(ci): declare supported GPU generations on every CUDA test: 本 PR 在 tests/manual/launch_scripts/test_py_launch_scripts.py 的 _HARDWARE_A_RECORDING_REPRESENTS 中注册新启动器并固定 MI355X,延续了 #2701 建立起的「测试显式声明硬件」体系;同方向的 #3184 也在为遗漏的 CUDA 测试补硬件声明。
- PR #3184 ci: declare hardware on the two CUDA tests that landed after #2701: 与 #2701 同属 CI 硬件声明体系的补充,体现仓库对测试硬件显式化的持续收紧;与本 PR 的硬件映射注册属同一维护方向。
参与讨论