feat: Sglang Rollout Refactor

[xiuhu17]

Refactor SGLang Rollout with the New Design

Overview

This PR refactors the SGLang rollout stack to support the new rollout architecture while preserving backward compatibility with the previous interface.

The new design focuses on improving scalability, reliability, and maintainability of the rollout system. In particular, it introduces support for:

• Redesign the old sglang rollout
• SGLang router
• Multi-node SGLang deployment
• Weight / KV cache / CUDA graph onload and offload
• Pause/Continue generation
• Weight check: snapshot, reset, compare for weight update check.
• async def generate_async(...) for asynchronous multi-turn rollout
• Refactored colocated weight update

The goal is to make the rollout layer production-ready for larger-scale serving and training integration, while minimizing disruption to existing callers.

---

Tests:

Thorough tests(around 75 tests):

• https://docs.google.com/document/d/1hP4zzJQ5ZaMhEnmyoWNLOx3T2E8Yj20DCmTvbn-XTWw/edit?tab=t.0

---

Summary by CodeRabbit

Release Notes

• New Features

  • Added pause and continue generation control to all generation backends
  • Introduced async generation support for streaming token output
  • Enhanced weight and memory management for optimized generation

• Improvements

  • Updated training configuration recipes for improved cluster utilization

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuki-97]

hi @xiuhu17 , thanks for the refactor and many new features! great work! I reviewed partly and left some comments.

[terrykong]

/ok to test 4d0c640

[terrykong]

SGLang Rollout Refactor — Review

Thanks for this substantial refactoring effort, @xiuhu17! The new architecture (router, multi-node support, fault tolerance, weight/KV offload) is a meaningful step toward production readiness. Below are findings grouped by priority.

Merge Conflicts

This PR currently has merge conflicts with main (mergeable: CONFLICTING). Please rebase on main and resolve conflicts before further review iterations.

Missing Copyright Headers

8 new/modified source files under nemo_rl/ are missing the NVIDIA Apache 2.0 copyright header (required per project guidelines for non-test files):

• sglang_generation.py, sglang_worker.py, fault_tolerance.py, sglang_router.py
• utils/async_utils.py, utils/http_utils.py, utils/misc.py, utils/ray_utils.py

Missing # pragma: no cover

3 @ray.remote classes are missing the coverage pragma (required for Ray actors per testing conventions):

• RouterActor in sglang_router.py:10
• Lock in utils/ray_utils.py:22
• _HttpPosterActor in utils/http_utils.py:106

Performance / Convergence Evidence

This PR rewrites the weight-streaming protocol, adds CUDA-graph/KV offload, multi-node SGLang, and a new router — but provides no quantitative evidence (throughput, refit latency, memory, convergence curves). The test link is to an external Google Doc. Could you share tokens/sec or step-time numbers (with and without the new features on a representative model) directly in the PR description?

Documentation

No docs/ updates for the user-facing additions (router config, fault tolerance, pause/continue, weight check, async generation). Per project conventions, new features should include in-repo documentation.

Test Deletions

332 lines deleted from test_grpo.py and 224 lines from test_policy_utils.py — could you clarify what coverage replaces these? They don't appear to be SGLang-specific tests.

Generated by Claude Code

[yuki-97]

mainly some style things.

[yuki-97]

hi @xiuhu17 , I finished review, sorry so many comments. could you help to solve them and the remaining previous comments? also please lmk if you need help.

besides there're some other fails will need you to fix.

1. there's a conflict with main, need to do a rebase.
2. PR title: https://github.com/NVIDIA-NeMo/RL/actions/runs/25301027605/job/74167878729?pr=2267
3. DCO: https://github.com/NVIDIA-NeMo/RL/pull/2267/checks?check_run_id=74167877728

[yuki-97]

/ok to test 2deff62

[yuki-97]

/ok to test 3979db4

[yuki-97]

/ok to test a051b72

[yuki-97]

/ok to test c620535

[yuki-97]

/ok to test 9f3948c

[yuki-97]

/ok to test 749fb1a

[yuki-97]

/ok to test 9f09507

[yuki-97]

@xiuhu17 thanks so much for the efforts!