[rollout] feat: Support Multi-stage Awake for SGLang

[hebiao064]

Co-authored with: MrAta (immrata@gmail.com) and @zhaochenyang20

Checklist Before Starting

• Search for similar PR(s).

What does this PR do?

Motivation

In RL Ecosystem which use colocate design like verl, we need to offload training model and load serving model & KV Cache frequently.

Background

• Currently SGLang is using torch_memory_saver to pause and resume.
• torch_memory_saver is a open source repo that provided easy to use api to hack cudaMalloc and cudaFree to make sure the virtual address could be consistent after pause and resume, which is critical to ensure CUDA Graph work.
• CUDA Graph is critical to make sure SGLang runs faster in decoding phases.

Here is the current behavior of VERL + SGLang

1. During Training, we have training model and optimizer state in the GPU Memory, and once training is done, we will offload optimizer state to cpu and keep the model weights in GPU, which is needed in Update Weight.
2. During Update Weight, we awake the SGLang engine, so those paused memory of Model Weights and KV Cache will come back. Then we update model from training model to serving model on the fly using the api: update_weights_in_tensor
3. After Model being updated, we delete the training model from GPU Memory.

Above design works pretty well so far, however, this would waste a big chunk of GPU Memory during rollout, which could cause a few issues we've seen so far:

• Small KV Cache: We need to use relative lower number of mem fraction ratio (e.g: 0.6), hence our KV Cache has less tokens. Given KV Cache has less tokens, we will hit RuntimeError: Prefill out of memory. Try to lower your batch size. when we try prefill large number of requests.
• Out of Memory: If we use mem fraction ratio 0.8 and run RL for 32B model on 8 H100, it will OOM during update weight

Challenge

• torch_memory_saver currently only supports Singleton, hence SGLang will pause and resume KV Cache + Weights together, they are treated as the same group of memory controlled by the singleton torch_memory_saver instance

Proposal

1. During Training, we do the same
2. During Update Weight Stage 1, we awake the model weights from SGLang and then update weights
3. During Update Weight Stage 2, we delete the training model weights from GPU Memory
4. Awake the SGLang's KV Cache

Benefit

With above feature, we can train larger model with same GPU, we can also make training/rollout more efficient given we can allocate larger KV Cache

Solution: Keep using Singleton and provide tag based pause/resume

• Support tag based resume/pause: Multi-Stage Awake: Support tag-based Resume and Pause fzyzcjy/torch_memory_saver#20
• Support Multiple Stage Awake in SGLang: Multi-Stage Awake: Support Resume and Pause KV Cache and Weights separately sgl-project/sglang#7099
• Support Multiple Stage Awake in verl: [rollout] feat: Support Multi-stage Awake for SGLang #1911

High-Level Design

Demonstrate the high-level design if this PR is complex.

Specific Changes

List the specific changes.

API

Demonstrate how the API changes if any.

Usage Example

Provide usage example(s) for easier usage.

# Add code snippet or script demonstrating how to use this 

Test

Additional Info.

• Issue Number: Fixes issue # or discussion # if any.
• Training: [Note which backend this PR will affect: FSDP, Megatron, both, or none]
• Inference: [Note which backend this PR will affect: vLLM, SGLang, both, or none]

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks.
• Add [BREAKING] to the PR title if it breaks any API.
• Update the documentation about your changes in the docs.
• New CI unit test(s) are added to cover the code path.
• Rely on existing unit tests on CI that covers the code path.

[hebiao064]

Another finding, for current GRPO training for QWEN 7B, mem-fraction 0.9 would fail with:

(WorkerDict pid=3022678) [torch_memory_saver.cpp] CUresult error  result=2 file=csrc/torch_memory_saver.cpp func=cu_mem_create line=103

Our goal is to make GRPO QWen 7B work with 0.9 or even 0.95 mem fraction by leveraging multi-stage awake

[hebiao064]

Another finding, for current GRPO training for QWEN 7B, mem-fraction 0.9 would work (peak memory goes up to 97%) on H200

When we set mem-fraction to 0.95, it will fail with:

(WorkerDict pid=3022678) [torch_memory_saver.cpp] CUresult error  result=2 file=csrc/torch_memory_saver.cpp func=cu_mem_create line=103

Our goal is make GRPO QWen 7B work with 0.95 mem fraction by leveraging multi-stage awake

[zhaochenyang20]

Great work！

[zhaochenyang20]

This feature is supported in sglang main and torch mem savor main. We will release 0.4.8 these days.