Sglang on CctoctoFX

[VeRL,SGLang] RL训推显存管理优化

Wed, 17 Sep 2025 11:30:12 +0800

SGLang团队的博客：https://hebiao064.github.io/rl-memory-management

Overview

上述是简化的在线RL训练流程，隐去了reference和critic model，并且用基础的reward function而非reward model来说明流程。实际上就是policy model存在的training engine和rollout engine上需要进行优化。

从简化的PPO流程开始：

for prompts, pretrain_batch in dataloader:
    # Stage 1: Rollout generation (inference)
    batch = actor.generate_sequences(prompts)
    # Stage 2: Prepare experience
    batch = reference.compute_log_prob(batch)
    batch = reward.compute_reward(batch)  # Reward function or model
    batch = compute_advantages(batch, algo_type)
    # Stage 3: Actor training
    actor_metrics = actor.update_actor(batch)

每一个iter相当于是actor model进行一次rollout再进行training，而veRL因为rollout和training共部署，所以两边可能不用version的actor model是在相同的GPU组上的，这导致了虽然资源共享但是显存管理会变得更复杂。

显存问题

训练阶段显存

FSDP（fully sharded + full activation checkpointing）下，每个GPU占据显存：

每个GPU的峰值显存：~48GB

推理阶段显存

During inference, the full model is typically loaded (not sharded):

[VeRL] AgentLoop源码走读

Thu, 14 Aug 2025 11:30:12 +0800

最近 RL sys 圈子的吴锡斌老师在 verl 上设计了将 rollout 与 tool 调用解耦的 AgentLoop，实现了自由灵活的 mutli-turn RL。在每个 AgentLoop 内部，rollout engine 只对外提供一个 token-in-token-out 的接口，而 tool 调用则通过 ToolAgentLoop 来实现。我个人比较喜欢这样解耦的设计，同时，AgentLoop 的代码结构也比较清晰。我个人学习了一次整个代码后，觉着 AgentLoop 的设计甚是不错，但是 ActorRolloutRefWorker 的历史包袱还是很重。

本文简单分析了 agent loop 的源码，并给出了一些自己的看法。

如果我们把整个 ActorRolloutRefWorker 当做一个 sgl.Engine 的话，AgentLoop 里面包装的两层 AsyncSGLangServer 和 AsyncLLMServerManager。AsyncSGLangServer 相当于在 sgl.Engine 上包装了 fastapi 成了 server，而 AsyncLLMServerManager 是在 server 上包了一层 router 做 load balance，相当于 sglang 的 router。这两层设计都是合理的，主要麻烦的是 ActorRolloutRefWorker，层层调用，最后一共经过 7 个 class 才调到 sgl.Engine，最近 verl 团队也在致力于对这块 worker class 的重构，敬请期待。最后，AgentLoopManager，AgentLoopWorker 和 AgentLoop 这三层，我觉得 AgentLoopWorker 可能未必有必要，其他两层挺合理的。

[SGLang] 后端代码速览

Wed, 13 Aug 2025 10:30:12 +0800

本文档为开发者提供 SGLang 后端代码的代码梳理，按照一个请求从输入到最后输出的顺序进行讲解。下图简要介绍了这一流程：

具体而言，请求的处理过程如下：

用户启动 Server ，初始化 FastAPI App、TokenizerManager、DetokenizerManager 和 Scheduler，每个组件运行各自的无限事件循环（infinite event loop）。
用户向 FastAPI Server 发送 /v1/chat/completions 请求，Server 通过 v1_chat_completions endpoint 将请求转发到 TokenizerManager。
v1_chat_completions 函数将请求转换为 ChatCompletionRequest，再转换为 GenerateReqInput，并调用 TokenizerManager 的 generate_request 方法。
TokenizerManager 对请求进行 tokenization，并以 Python 对象（pyobj）形式将其转发给 Scheduler，同时调用 TokenizerManager 的 _wait_one_response 方法。
Scheduler 在事件循环 event_loop_normal 中处理请求：
- Scheduler 通过 recv_requests 接收请求，调用 process_input_requests 处理输入，通过 handle_generate_request 管理生成请求的逻辑，并将其加入 waiting_queue。
- 从 waiting_queue 中，Scheduler 使用 get_next_batch_to_run 为即将处理的请求创建 ScheduleBatch。
- Scheduler 执行 run_batch 函数，将 ScheduleBatch 转换为 ModelWorkerBatch。
- Scheduler 调用 TpModelWorker 的 forward_batch_generation，等待 logits_output 和 next_token_ids。
- TpModelWorker 初始化 ForwardBatch，将其转发至 ModelRunner，并等待 logits_output。
- ModelRunner 处理 ForwardBatch，调用 forward_extend 执行模型的前向计算（forward pass）。
- 模型通过 AttentionBackend 加速生成 logits，返回给 ModelRunner，进而返回给 TpModelWorker。
- TpModelWorker 从 ModelRunner 接收 logits_output，调用 ModelRunner 的 sample 方法生成 next_token_ids，并将其发送回 Scheduler。
- Scheduler 通过 process_batch_result 处理批次结果，使用 tree_cache.cache_finished_req(req) 缓存请求，并通过 check_finished 验证完成状态。对于未完成的请求，Scheduler 继续其事件循环，直到这个请求满足结束条件；对于已完成的请求，则转发到 Scheduler 的 stream_output。
- 在 stream_output 函数中，Scheduler 处理输出，将其包装成 BatchTokenIDOut，并发送给 DetokenizerManager。
DetokenizerManager 在其事件循环中接收 BatchTokenIDOut，处理后生成 BatchStrOut 并返回给 TokenizerManager。

[VeRL] Multi-Turn RL训练源码走读（2）

Sun, 03 Aug 2025 17:30:12 +0800

在 Part 1 中，我们介绍了 verl 的初始化过程，我们进一步介绍 verl 的训练过程，包括rollout部分、make experience部分以及training部分。

在 GRPO 中，单个 step 包含四个阶段：load data -> rollout -> make experience -> update model。区别于前一节的详述，本节会使用伪代码结合源码的方式进行阐述。

  flowchart LR
subgraph W2["Initialize"]
WP[Process Data] --> A
direction TB D1[Data Prepare] --> A
A[TaskRunner] --> B1[RayPPOTrainer]
B1 --> Workers

    subgraph Workers["Workers"]
        direction TB
                WA[ActorRolloutWorker] --> WD[FSDP Engine]
        WB[CriticWorker] --> WD
        WC[RewardModelWorker] --> WD
        WD --> WE[SGLang Engine]
    end
    
    Workers --> C1[Hybrid Engine]
end 

subgraph W3["Train Loop"]
    direction TB
    E[DataLoader] --> RolloutBox
    
    subgraph RolloutBox["Rollout"]
        F1[Prepare Data] --> F2[SGLang Async Rollout]
        F2 --> F3[Multi-turn Chat Process]
    end
    
    RolloutBox --> ExpBox
    
    subgraph ExpBox["Make Experience"]
        G1[Recompute Log Probs] --> G2[Compute Reward]
        G2 --> G3[Compute Advantage]
    end
    
    ExpBox --> UpdateBox
    
    subgraph UpdateBox["Train The Model"]
        H1[Load FSDP Model Weight] --> H2[Compute Gradient]
        H2 --> H3[Weights Update]
        H3 --> H4[Sync Weights]
    end
    
    UpdateBox --> E
end

W2 --> W3

数据加载与预处理

verl 通过 DataProto 和 RLHFDataset 来实现数据处理。具体来说，在 main_ppo.py 中，我们观察这个函数：

[VeRL] Multi-Turn RL训练源码走读（1）

Sun, 03 Aug 2025 15:30:12 +0800

该part主要聚焦相关模块初始化部分

还是以 verl 出发，分析其 end to end mutli-turn RL 训练的全过程。整体上，我希望覆盖所有重要的 class 以及函数，更细粒度的代码不再展开。

为了前后内容的一致性，基于 76f63cffa5 的 commit 进行分析。

虽然本文以分析 verl 的代码为主，写完之后我才意识到，系统设计问题是非常通用的。诸如“log probs 重计算”，“Rollout Engine 显存管理”等等系统设计，是各大 RL 框架都需要考虑的核心问题。

此外因为最近在学习SGLang的实现，本文的推理后端选择的是SGLang展开分析。

整个训练的示意图如下，我们会具体展开每个部分。

  flowchart LR
subgraph W2["Initialize"]
WP[Process Data] --> A
direction TB D1[Data Prepare] --> A
A[TaskRunner] --> B1[RayPPOTrainer]
B1 --> Workers

    subgraph Workers["Workers"]
        direction TB
                WA[ActorRolloutWorker] --> WD[FSDP Engine]
        WB[CriticWorker] --> WD
        WC[RewardModelWorker] --> WD
        WD --> WE[SGLang Engine]
    end
    
    Workers --> C1[Hybrid Engine]
end

subgraph W3["Train Loop"]
    direction TB
    E[DataLoader] --> RolloutBox
    
    subgraph RolloutBox["Rollout"]
        F1[Prepare Data] --> F2[SGLang Async Rollout]
        F2 --> F3[Multi-turn Chat Process]
    end
    
    RolloutBox --> ExpBox
    
    subgraph ExpBox["Make Experience"]
        G1[Recompute Log Probs] --> G2[Compute Reward]
        G2 --> G3[Compute Advantage]
    end
    
    ExpBox --> UpdateBox
    
    subgraph UpdateBox["Train The Model"]
        H1[Load FSDP Model Weight] --> H2[Compute Gradient]
        H2 --> H3[Weights Update]
        H3 --> H4[Sync Weights]
    end
    
    UpdateBox --> E
end

W2 --> W3

数据预处理

以 GSM8K 为例，预处理脚本是 examples/data_preprocess/gsm8k_multiturn_w_tool.py。整个脚本只做了经典的 huggingface datasets mapping，核心逻辑如下：