[fix] Stabilize ATOM FP8 no-eager rollout weight sync and CUDA graph lifecycle

[xysheng-AMD]

Motivation

Improve ATOM FP8 rollout stability and correctness for LumenRL DAPO long-running training with no-eager + compilation level 3.

This PR addresses several issues seen during RL training:

• FP8 rollout weights diverging from the training-side model after online weight sync.
• CUDA graph invalidation/recapture instability after rollout sleep/wake.
• HIP illegal memory access during repeated graph recapture in long runs.
• RMSNorm behavior mismatch between ATOM rollout and training-side execution.

Technical Details

This PR updates the ATOM FP8 rollout path to keep CUDA graph dependencies stable across RL weight updates and sleep/wake cycles.

Key changes:

• Align online FP8 rollout weight updates with initial ATOM FP8 weight layout.
• Correctly handle packed/fused Qwen3 weights, FP8 requantization, blockscale layout, and weight scale updates.
• Apply required post-update normalization/shuffle processing for FP8 weights.
• Add compatibility fallback for aiter versions that expose shuffle_scale instead of moe_shuffle_scale.
• Preserve CUDA graph validity in no-eager mode by keeping graph-dependent storage addresses stable:
  • Keep KV cache resident instead of releasing/reallocating it.
  • Update weights in-place via copy_.
  • Preserve parameter storage during shuffle when shape/dtype allow.
  • Avoid invalidating and recapturing CUDA graphs after every online weight update.
• Enable model-sensitive RMSNorm in ATOM rollout to match training-side behavior.

The main stability change is moving away from repeated CUDA graph recapture after weight updates. Instead, the rollout engine keeps weights/KV cache/graphs resident and refreshes values in-place, avoiding ROCm illegal memory accesses observed during long-running sleep/wake recapture.

Test Plan

Validated with LumenRL DAPO ATOM FP8 rollout runs using:

• no-eager mode
• compilation level 3
• FP8 per-block online quantization
• sleep level 2
• 8-GPU Ray colocated ATOM rollout
• periodic validation enabled

Test coverage included:

• Short smoke runs with validation every 2 steps.
• Formal long-running training with validation every 10 steps.
• Validation-after-training lifecycle checks to ensure rollout does not degrade after eval.
• Monitoring for CUDA graph recapture failures, HIP illegal memory access, eager fallback, and reward/length regressions.

Test Result

Smoke validation runs completed successfully after the lifecycle and in-place update fixes.

The updated ATOM branch was used in a formal LumenRL DAPO run and successfully passed validation checkpoints after rollout sleep/wake. In particular, the run progressed past validation after step 20 without the previous CUDA graph recapture crash pattern.

Observed:

• No fallback to eager mode.
• No post-validation generation collapse.
• No CUDA graph recapture failed after the in-place update change.
• Reward/accuracy and response lengths remained in expected ranges after validation.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1612 --add-label <label>

[yitingw1]

Please resolve conflicts, thanks.