[Arist12]

Problem

There is no maintained gpt-oss-20b MoE LoRA RL launcher for AMD.

Change

Adds scripts/amd/run_gpt_oss_20b_lora.py: single-node colocated GRPO on 8× MI350X / MI355X through Megatron-Bridge, plus launcher snapshots and the LoRA recipe index entry.

The recipe uses a rank-32 per-expert adapter on gate_proj,up_proj,down_proj, math reward, trainer TP=4 with unfused BSHD attention, and two TP=4 rollout engines. Each rollout is 8 prompts × 8 samples with global batch size 64, so 20 rollouts produce 20 optimizer steps. Responses are capped at 4096 tokens.

ROCm uses SGLang Triton attention, MoE, and LoRA backends with virtual experts disabled and static memory fraction 0.2. The BF16 checkpoint is required because the Triton MoE path does not serve the MXFP4 checkpoint.

The launcher does not checkpoint; related LoRA exact-resume work is tracked in #2516 and #2661.

Corresponding dependency modifications

• SGLang / Megatron-Bridge / Megatron-LM: no changes required (Megatron-LM exercised at 8c1e05747eb6).
• aiter: ROCm/aiter#5694. Some runs of this recipe showed intermittent train_rollout_logprob_abs_diff spikes; the launcher runs without that fix.
• miles: #3263 — _setup_lora_model_via_bridge drops moe_router_use_torch_mm and moe_router_dtype, so this recipe's router GEMM runs bf16 x bf16 -> bf16 while its args table reports fp32 + torch mm. Config fidelity only.

Validation

8× MI355X, ROCm 7.2, rocm/sgl-dev:miles-rocm720-mi35x. Two full runs of this launcher, 20 rollouts / optimizer steps each, Ray job SUCCEEDED, both ~1h06m:

	aiter	train_rollout_logprob_abs_diff (min / mean / max)	rollouts > 0.15	rollout/raw_reward mean
4it2x89i	stock	0.0140 / 0.0171 / 0.0202	0	0.57
jvvnhop9	with the kernel fix	0.0150 / 0.0170 / 0.0208	0	0.56

Per-rollout series for jvvnhop9, in order:

train/train_rollout_logprob_abs_diff
0.0208 0.0185 0.0164 0.0172 0.0153 0.0166 0.0160 0.0180 0.0186 0.0153
0.0187 0.0151 0.0156 0.0161 0.0150 0.0158 0.0162 0.0204 0.0161 0.0178

rollout/raw_reward
0.5000 0.5625 0.7500 0.5625 0.6094 0.6250 0.5469 0.3906 0.4062 0.4688
0.6719 0.6094 0.6094 0.4688 0.5781 0.5156 0.5938 0.4688 0.5156 0.7656

train/train_rollout_kl 6.95e-4 – 1.53e-3 (mean 8.79e-4); response length mean 2920 with truncation
0.410; pg_clipfrac 0. Reward is in the same range as the stock run on the identical prompt
sequence, so the fix does not cost reward. No reward trend is claimed from 20 optimizer steps at
lr 1e-05.

Tests

• Focused launcher snapshot coverage: 8 passed.
• Changed-file pre-commit checks clean.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yushengsu-thu]

waiting for AMD part ci pass

[Arist12]

waiting for AMD part ci pass

Hi @yushengsu-thu , this PR alone should be recipe only, the dependent aiter PR has been merged ROCm/aiter#5694. The miles one #3263 is pending review but should be not impact the recipe. Thanks!