Prhub

#3183 [AMD][LoRA] Add a Qwen3-4B LoRA RL launcher

原始 PR 作者 Arist12 合并时间 2026-09-15 05:33 文件变更 6 提交数 2 评论 1 代码增减 +469 / -1

执行摘要

补齐 AMD 平台 Qwen3-4B LoRA RL 启动器

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 平台。

值得精读。若团队需要维护或扩展 AMD / ROCm 上的 LoRA RL 训练,这是当前唯一入口;其「平台差异收敛在启动器层」的做法——环境变量桥接(HIP → CUDA 可见设备)、SGLang Triton 后端固定、以及「配方 + 快照 + 文档」三件套的组织方式——是值得复用的模式。验证数据(logprob 差、KL、截断率、AIME 趋势)可作为后续 AMD 训练健康度对比的参考基线,但不应据此断言训练收益。

讨论亮点

维护者 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;实际未触发,由维护者人肉审核并批准。

实现拆解

  1. 新增启动器入口: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。
  2. 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。
  3. 命令组装与提交:_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 卡。
  4. CLI 入口划分:prepare 只下载模型(Qwen/Qwen3-4B)与数据集(dapo-math-17k、可选 aime-2024);train 假设输入已在节点上直接训练;full-train 合并两者,便于一键跑通。
  5. 测试与文档配套: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 启动器 added 8.89
tests/manual/launch_scripts/test_py_launch_scripts.py 启动测试 modified 3.51
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/full_train.txt 命令快照 added 4.02
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/train.txt 命令快照 added 3.56
tests/snapshots/launch_scripts/py/scripts/amd/run_qwen3_4b_lora.py/prepare.txt 命令快照 added 2.56
docs/advanced/lora.md 文档 modified 1.72

关键符号

ScriptArgs _set_rocm_environment _resolve_num_gpus _download_inputs _get_wandb_args _execute prepare train full_train

关键源码片段

scripts/amd/run_qwen3_4b_lora.py core-logic

本 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,
    )

评论区精华

是否应把启动器快照纳入本 PR question

维护者 yushengsu-thu 在 review 中提问 `Should this PR add snapshots?`,并指向 @yueming-yuan,意在确认 launch_scripts 的命令录制快照(tests/snapshots/launch_scripts/...)是否应随新启动器一起合入。

结论:最终版本已包含 prepare.txt、train.txt、full_train.txt 三个快照文件,且 yushengsu-thu 随后给出 APPROVED,说明该疑问在合并前得到解决。 · 已解决

fork PR 自动 review 被禁用 other

claude[bot] 说明本 PR 来自 fork,自动化 review 被禁用;维护者可评论 `@claude review` 触发一次性 review。

结论:未触发额外 bot review,由维护者 yushengsu-thu 人肉审核并批准合并。 · 已解决

风险与影响

  1. 快照对工具链版本敏感:tests/snapshots 冻结了极长且细节级的 train.py 参数序列(含全部模型超参与 --wandb-key 'frozen-wandb-api-key'),一旦 train.py 默认参数演进或 SGLang 参数改名,快照即失效,需要同步更新。
  2. ROCm 环境变量依赖:_set_rocm_environment() 依赖 Ray 的 RAY_EXPERIMENTAL_NOSET_* 行为与 SGLang 在 ROCm 上的设备读取约定,相关版本升级后可能失效且难以在 CPU CI 捕获。
  3. 验证范围有限:仅 8× MI355X 单节点 colocated 实测,MI350X 与多节点 / 跨节点场景、ROCm 7.2 之外版本未覆盖;colocated 模式下 actor 与 rollout 同机显存竞争的行为未展开讨论。
  4. 评估信号强度不足:AIME-2024 只有 30 道独立题目,avg@16 的涨跌统计功效有限,作者自己也声明这是次要信号而非效果声明。
  5. 无 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 评估样本量小

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论