TRL SGLang Support

[zhaochenyang20]

什么都不用考虑了，就一天，这事情做不完就不做了

[zhaochenyang20]

1. https://docs.sglang.ai/backend/send_request.html
2. https://docs.sglang.ai/backend/openai_api_completions.html
3. https://docs.sglang.ai/backend/native_api.html
4. https://docs.sglang.ai/backend/offline_engine_api.html

@renxinx

[zhaochenyang20]

任何问题，首先问 Grok，然后查 google，比如 huggingface blog，一小时自己搞不定就来问人！问人之前，把这个问题写下来阐述清楚，然后发 link 给大家，比如直接写到这个 issue 下面，然后发 link 问人就好了。

[zhaochenyang20]

https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/GRPO/SGLang_GRPO.md

1. 任何不理解或者觉得写的不清楚的地方直接写 comments，回头我和 jin fix；
2. 这些图片全部删了，能用文字表达的地方就用文字，不要用莫名其妙的图片；

@ryang-max

[jhinpan]

RoadMap for Supporting SGLang as Inference Backend for TRL-GRPO

1. 背景与目标

GRPO (Generative Reinforcement Learning with Preference Optimization) 是一种高效的强化学习算法，可用于训练大型语言模型。本项目旨在将 SGLang 作为 inference engine 或者说 Server backend 集成到 Hugging Face 的 TRL (Transformer Reinforcement Learning) 库中，以提高 GRPO 的训练效率和性能。整体算法和结构就不再过多赘述，几周前的第一版可以refer这个文档: https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/GRPO/SGLang_GRPO.md
这个应该还是Server-based为主，有兴趣可以找我要另外一版Engine-based的DevLog

主要目标：

• 实现 SGLang 与 TRL 的无缝集成
• 优化跨 GPU 权重更新机制
• 提高推理速度，降低训练时间

2. 系统架构

2.1 整体架构

系统采用分布式架构，使用多 GPU 协同工作，以4张卡为例：

• 训练卡（3张）：执行模型训练
• 推理卡（1张）：使用 SGLang 进行高效推理

2.2 核心组件

• 训练引擎：基于 Hugging Face TRL 库
• 推理引擎：SGLang
• 权重同步机制：基于 tensor 的跨进程权重更新
• 服务协调：采用服务器模式管理训练和推理进程

3. 主要实现步骤

3.1 Initialization and Dedicated GPU Selection

实现专用 GPU 选择机制，确保训练和推理任务分配到合适的 GPU：

Key Points：

• 明确区分训练 GPU 和推理 GPU
• 实现智能资源分配
• 确保不同进程间 GPU 资源不冲突

3.2 权重更新机制: update_weights_from_tensor

替换原有的基于磁盘的权重更新方法 (update_weights_from_disk)，实现基于 tensor 的高效权重更新：

Advantages：

• 避免磁盘 I/O 开销，显著提高权重更新速度
• 适用于单节点多 GPU 环境
• 减少训练中断时间

3.3 SGLang 作为 Inference Engine

集成 SGLang 作为高性能推理引擎，加速 rollout 过程：

参考 vllm 服务器实现方式，开发类似的 SGLang 服务器模块：

• 创建独立的 SGLang 服务进程
• 实现权重同步 RPC 接口
• 优化推理参数配置

4. 服务器架构重构

参考 Hugging Face 最近的 GRPO vLLM 后端重构，采用基于服务器的架构：

关键改进：

• 使用 SGLang 内部的集体通信机制进行权重同步
• 实现类似 vllm.collective_rpc 的接口
• 优化服务器资源管理

5. 测试与验证

5.1 性能测试

设计测试方案验证实现效果：

• 对比 update_weights_from_disk 和 update_weights_from_tensor 的性能差异
• 测量 SGLang 推理加速效果
• 评估整体训练时间改进，可以对比他们当前的vllm server backend

5.2 集成测试

确保各组件协同工作：

• 验证训练与推理进程通信
• 测试权重同步正确性
• 检查资源使用效率

6. Reference

当前实现基于以下工作：

• 原始仓库：https://github.com/jhinpan/trl-jin/tree/main Server在main branch，Engine在engine branch
• 之前提交的一版 PR：Feature: Add SGLang as inference backend for generation in GRPO huggingface/trl#2981
• Hugging Face vllm 重构 PR：🚀 Scaling GRPO to 70B+ Models and Multi-Node Training with vLLM Server & NCCL Communication huggingface/trl#3094
• vllm Serve 实现：https://github.com/huggingface/trl/blob/main/trl/scripts/vllm_serve.py

有任何问题欢迎随时跟我说～

@renxinx @ryang-max @zhaochenyang20

[zhaochenyang20]

@hebiao064

[ryang-max]

上面的文档都看过了，今天开始看vllm server-based的实现，然后先跑一下

[renxinx]

1. https://docs.sglang.ai/backend/send_request.html
2. https://docs.sglang.ai/backend/openai_api_completions.html
3. https://docs.sglang.ai/backend/native_api.html
4. https://docs.sglang.ai/backend/offline_engine_api.html

@renxinx

这些都看过了

[ryang-max]

TODO:

1. run trl[vllm]
2. implement sglang-serve(server-based sglang for trl, similar to vllm) @ryang-max
3. implement grpo_config and grpo_train @renxinx
4. Integration Test @ryang-max @renxinx
5. Compare performance with vllm @ryang-max @renxinx

cc @zhaochenyang20

[zhaochenyang20]

@ryang-max 你直接更新到最上面，不要 follow up

[zhaochenyang20]

1. 什么是 transformer 的 forward 计算；
2. 四个角色谁需要被训练，谁需要被 inference，每个人用的 engine 是什么？为什么需要用特定的 engine（为什么不能用 sglang 计算 reference model 的 inference）；
3. 第六个文档读完了之后我们再对一下；

@renxinx

[zhaochenyang20]

1. 什么是 transformer 的 forward 计算；
2. 四个角色谁需要被训练，谁需要被 inference，每个人用的 engine 是什么？为什么需要用特定的 engine（为什么不能用 sglang 计算 reference model 的 inference）；
3. 第六个文档读完了之后我们再对一下；

@renxinx

1. 什么是 forward，什么是 backward，什么是 prefill，什么是 decode。对于 prefill 和 decode，可以参考此处。
2. 答案在这张图里面；剩下的问题我们下次讨论；这个文章需要反复练习，烂熟于心。
3. 关于 OpenRLHF 的开发，这些东西需要自己记笔记学习：

• Ray 的意义和简单用法；
• Collocate 策略；本质上是一种资源分配方法，可以参考这个文档学习各种资源分配方法，这也是最核心的做法；

这个进展蛮好的。直接问 GPT 是个好办法，不过 GPT 更多是答疑。如果我给你的 link 直接 link 到了一个文档，那个文档通篇都是值得读的，如果是 link 到一个子标题，只读子标题那部分就好了。我不确定现在你是咋读的哈，反正这些文档我自己写的很有感兴趣，很值得细细品味，能学到不少东西 😂 anyway，真的做的挺棒的，学无止境，这些核心概念的帮助还是很大的，咱们下次继续过。我们也写好 learning log，我预估这些东西得花个四五天。

[ryang-max]

update:
写了半天serve，感觉好像又不是很需要，可能sglang native server api就够用了吧。我直接在jin的codebase上边跑边调了，先看看遗留问题多不多。。。

[zhaochenyang20]

update: 写了半天serve，感觉好像又不是很需要，可能sglang native server api就够用了吧。我直接在jin的codebase上边跑边调了，先看看遗留问题多不多。。。

你的理解是正确的，就用 native server，不用 verl server