## File: README.md

# Awesome-ML-SYS-Tutorial

## [English Version](./README.md) | [Chinese Version](./README-cn.md)

My learning notes for ML SYS.

I've been writing this blog series intermittently for over a year now, and it's almost become an RL Infra Learning Note 😂

I often see discussions about whether ML SYS or AI Infra is worth getting into, and how to start. Everyone's choice is different. For me, I simply want to **pursue the truth in algorithms**:

> A large number of RL conclusions derived from papers are based on RL infrastructure in the open-source community that may be extremely flawed. I've been involved in RL infra development for over a year, and I've seen numerous community experts diligently working, but the fact is that RL infra, whether open-source or within major companies, still has many problems. It is absolutely worth questioning whether the high-level conclusions drawn from this flawed infrastructure are correct. When I was reviewing for ICLR this year, I often asked the papers assigned to me, "If the framework you are using has implementation issues itself, can your conclusions still hold?" Although I never deducted points for this reason, no one could provide an answer that resolved my fundamental doubt.
>
> Therefore, some excellent researchers I know are keen to participate in infra development, spending most of their time on foundational work to rigorously ensure that the algorithm they plan to develop next has a correct basis. I greatly admire them and agree with such rigor—they are my role models. The same is true for our SGLang RL community. With so much human power and time, we all hope to provide the most correct and concise RL foundation possible, whether it's for companies training models or researchers developing new algorithms, with the goal of genuinely serving everyone in the community. Thank you for your recognition, and I look forward to hearing from interested friends who wish to contact me and join us!

After a year of going around in circles, this is the resolve that keeps me going in Infra: **to make a contribution to the community by building a correct foundation, thereby helping to ensure correct conclusions.**

Coming back to the topic, this series of podcasts started in August 2024, when I began learning ML SYS notes following the opportunity to use [SGLang](https://github.com/sgl-project/sglang) during my research. It's largely written by me, with content focusing on **RL infra, online/offline inference systems, and some fundamentals of AI Infra**. Over the past year, starting from two or three articles and thirty to fifty Github Stars, to now exceeding 4.5K Stars, I have become a minor technical influencer. I am deeply honored and grateful for the support.

**I would like to thank my advisors, Professor Quanquan Gu, Dr. Ying Sheng, and Dr. Linmin Zheng**, for the immense help and guidance they gave me in my study of AI Infra, career choices, and life path. Although I am no longer pursuing a Ph.D. at UCLA due to personal reasons, this journey after my undergraduate graduation has been an incredibly valuable experience. I have now joined RadixArk full-time, continuing my research in RL Infra. We will continue to share AI Infra-related technology and thoughts through my blog, via unofficial channels. **I also hope readers interested in AI Infra reach out to us, join the SGLang open-source community, and together build open-source AI Infra that changes the world and is worth being proud of for a lifetime!**

## Omni Model Inference Notes

- [Codec、RVQ、Dual AR、Thinker-Talker——A Deep Dive into Omni Model Inference for Qwen3-Omni and S2 Pro](./transformers/omni/readme-en.md): available in [Chinese version](./transformers/omni/readme.md) and [Zhihu](https://zhuanlan.zhihu.com/p/2023413564229042799).
- [When SGLang OOMs, What Exactly Runs Out of Memory?](./sglang/kvcache-code-walk-through/mem-fraction-static-en.md): available in [Chinese version](./sglang/kvcache-code-walk-through/mem-fraction-static.md) and [Zhihu](https://zhuanlan.zhihu.com/p/2041811635405336677).
- [SGLang Omni: Redesigning the Inference Framework for Multi-Stage Generative Models from Decoding Computation Characteristics](./sglang/sglang-omni/why-sglang-omni-en.md): why multi-stage decoding is the right axis to slice the model landscape, what computational characteristics multi-stage decoding models share, and how SGLang Omni's scheduling decoupling / layered communication / memory isolation map to these characteristics. Available in [Chinese version](./sglang/sglang-omni/why-sglang-omni.md) and [Zhihu](https://zhuanlan.zhihu.com/p/2043972741326451200).
- [Optimizing TTS Inference: Engineering Lessons from Profiling to Streaming in SGLang Omni](./sglang/sglang-omni/tts-optimization.md): how we made Higgs and MOSS-TTS-Local TTS serving 1.9–3.4× faster on SGLang-Omni — encoder LRU caching, delay-pattern / backbone+inner codec scheduling, CUDA Graph + async CPU–GPU decode, and vocoder batching / windowed streaming. Available in [Chinese version](./sglang/sglang-omni/tts-optimization-zh.md) and [Zhihu](https://zhuanlan.zhihu.com/p/2054821717529343214).
- [Lifecycle Management and Framework Abstractions: Refactoring TTS Serving in SGLang Omni](./sglang/sglang-omni/tts-refactor.md): how six architecturally different TTS backends were refactored around shared contracts for engine startup, state transport, caching, vocoder lifecycles, and scheduling; and how the refactor removed a net 2,840 non-test lines while simplifying new-model onboarding. Available in [Chinese version](./sglang/sglang-omni/tts-refactor-zh.md).

## RLHF System Development Notes

### slime Framework

- [One Rollout to Rule Them All: Seamless Multi-Turn RL for LLM and VLM](./rlhf/slime/vlm-multi-turn/readme-en.md): Introduces the seamless multi-turn RL for LLM and VLM, including the technical details and the implementation of the seamless multi-turn RL for LLM and VLM. Also available in [Chinese version](./rlhf/slime/vlm-multi-turn/readme.md) and [Zhihu](https://zhuanlan.zhihu.com/p/1999486437280745217).
- [Squeezing 1TB Model Rollout into a Single H200: INT4 QAT RL End-to-End Practice](./rlhf/slime/int4/readme-en.md): Introduces the INT4 QAT RL end-to-end practice, including the technical details and the implementation of the INT4 QAT RL end-to-end practice. Also available in [Chinese version](./rlhf/slime/int4/readme.md) and [Zhihu](https://zhuanlan.zhihu.com/p/1999403679393014924).
- 【Not finished】[Let Speed Be With Stability: All-In-One Solution to Training-Inference Mismatch with Miles](./rlhf/slime/mismatch/blog-en.md): Introduces two solutions provided by the Miles framework for the train-inference mismatch problem: achieving perfect True On-Policy training through kernel-level alignment, and mitigating the mismatch using algorithms like TIS/MIS. Also available in [Chinese version](./rlhf/slime/mismatch/blog-cn.md).
- [Support FSDP2 as A Training Backend for slime](./rlhf/slime/fsdp/readme_en.md): Added FSDP as a training backend to slime, and aligned it with Megatron. FSDP is more flexible in supporting models with architectural innovations like Qwen3-Next/gpt-oss and helps us further support VLM RL. Also available in [Chinese version](./rlhf/slime/fsdp/readme.md) and on [Zhihu](https://zhuanlan.zhihu.com/p/1979141713449742500).
- [Unified FP8: Moving Beyond Mixed Precision for Stable and Accelerated MoE RL](./rlhf/slime/fp8/readme_en.md): Fully utilizing FP8 for both sampling (Rollout) and training (Training) in RL. Also available in [Chinese version](./rlhf/slime/fp8/readme.md) and on [Zhihu](https://zhuanlan.zhihu.com/p/1974681194017865986).
- [Power Up Speculative Decoding In Reinforcement Learning](./rlhf/slime/spec/readme-en.md): Introduces speculative decoding into the RL sampling process, significantly boosting sampling speed when the batch size is appropriate; moreover, the draft model is updated during training. Compared to freezing the draft model, the accepted length remains consistently high, yielding long-term stable positive returns. Also available in [Chinese version](./rlhf/slime/spec/readme.md).
- [An In-Depth Look at the Elegant Design and Source Code of the slime RL Framework](./rlhf/slime/code-walk-through/readme_en.md): slime source code appreciation. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1946402397409740613) and in [Chinese version](./rlhf/slime/code-walk-through/readme.md).
- [Pending Review] [slime FSDP Setup Guide](./rlhf/slime/fsdp/release_log/setup_fsdp.md): Records how to test FSDP on slime, including H-cards and B-cards, and both Colocate and Disaggregated placement methods.
- [Pending Review] [Chunked Parallel Computation of GAE in PPO (slime Implementation)](./rlhf/slime/batch-GAE/ppo-gae-chunk.md): Rewrites the standard backward recurrence of GAE into chunk-based parallel prefix scanning, significantly mitigating the GAE computation bottleneck in long sequence scenarios, achieving about $100\times–300\times$ acceleration in slime. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1975237289425798560).

### AReal Framework

- [AReal Code Walk Through](./rlhf/areal/code-walk-through_EN.md) AReal source code appreciation. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1983417813080236770) and in [Chinese version](./rlhf/areal/code-walk-through_CN.md).


### System Design and Optimization

- [Deep Dive into DeepSeek MoE with Classic Secondary Development of EP on FSDP](./rlhf/sys-design/readme-4-en.md): Deep dive into DeepSeek MoE with classic secondary development of EP on FSDP. Also available in [Chinese version](./rlhf/sys-design/readme-4.md) and [zhihu](https://zhuanlan.zhihu.com/p/1990790333823481023).
- [Deep Thoughts on RL Systems: In-Depth Understanding of Weight Update Mechanism](./rlhf/sys-design/readme-1-EN.md): Summary of half a year's work, in-depth understanding of the weight update mechanism. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1925210722704531547) and in [Chinese version](./rlhf/sys-design/readme-1.md).
- [Deep Thoughts on RL Systems: FSDP Training Backend](./rlhf/sys-design/readme-2-en.md): Discusses the principles and implementation of FSDP, and analyzes verl's use of FSDP. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1929115059113693341) and in [Chinese version](./rlhf/sys-design/readme-2.md).
- [Pending Review] [Deep Thoughts on RL Systems: Megatron](./rlhf/sys-design/readme-3.md): Brief analysis of Megatron's basic features, focusing on its use in the RL framework.
- [Extending the OpenRLHF Inference Engine](./rlhf/OpenRLHF/develop-log.md): Development notes on integrating SGLang into OpenRLHF. The entire process was very painful, and there's still an nccl hang error that a DeepSpeed core contributor is currently fixing.
- [Pending Review] [SGLang as rollout engine of GRPO trainer](./rlhf/GRPO/SGLang_GRPO.md): Introduction on how to use SGLang as the inference backend for the GRPO Trainer in TRL. GRPO is a PPO variant that optimizes PPO's memory usage while improving mathematical reasoning capabilities.
- [Pending Review] [Light-DuoAttention: Achieves efficient Long-Context Inference using CuTeDSL and runs within SGLang](sglang/light-duoattention/light-duoattention.md): Introduction DuoAttention, a solution designed for Long-Context Inference, implemented with CuteDSL and supported in SGLang. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1984978732436431358).

### verl Framework

- [Analyzing VLM RL Training Memory Leaks via Torch Memory Snapshot](./torch/mem-snapshot/readme-en.md): Analysis of SGLang memory leak issues and solutions. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1943202817247519535) and in [Chinese version](./torch/mem-snapshot/readme.md).
- [Latency optimization for weight updates](./sglang/latency-accelerate-for-weight-updates/readme.md): A debug process for efficiency. Also available on [Zhihu: A record of optimizing SGLang weight update latency](https://zhuanlan.zhihu.com/p/9908228168).
- [In-Depth Understanding of verl Source Code (Initialization)](./rlhf/verl/multi-turn/code-walk-through/readme_EN.md): Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1920751852749849692) and in [Chinese version](./rlhf/verl/multi-turn/code-walk-through/readme.md).
- [In-Depth Understanding of verl Source Code (Rollout)](./rlhf/verl/multi-turn/code-walk-through/readme-2-EN.md): Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1923349757566388159) and in [Chinese version](./rlhf/verl/multi-turn/code-walk-through/readme-2.md).
- [Pending Review] [In-Depth Understanding of verl Source Code (Make Experience)](./rlhf/verl/multi-turn/code-walk-through/readme-3.md): Analysis of the logic for the make experience part in verl.
- [AgentLoop Source Code Analysis](./rlhf/verl/multi-turn/code-walk-through/readme-6.md): Analysis of the multi-turn RL implementation based on AgentLoop in verl.
- [verl Parameter Quick Reference](./rlhf/verl/multi-turn/code-walk-through/readme-5-EN.md): Quick reference for verl parameters. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1925041836998783250) and in [Chinese version](./rlhf/verl/multi-turn/code-walk-through/readme-5.md).
- [Analyzing the Complexity of Agentic Multi-Turn Training from a Tokenizer Perspective](./rlhf/verl/multi-turn/fast_tokenization/multiturn_tokenization_and_masking.md): Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1917126584806139373) and in [Chinese version](./rlhf/verl/multi-turn/fast_tokenization/multiturn_tokenization_and_masking_ZH.md).
- [Pending Review] [DAPO Dynamic Filtering Implementation and Batch Size Analysis](./rlhf/verl/multi-turn/code-walk-through/dapo.md): Exploring how to achieve higher parallelism by padding prompts to a smaller batch size.
- [Systematic Analysis of Time Consumption in verl Multi-Turn Training](./rlhf/verl/multi-turn/tool_examples/profile_en.md): verl multi-turn interaction and tool call profile analysis. Also available in [Chinese version](./rlhf/verl/multi-turn/tool_examples/profile.md) and on [Zhihu](https://zhuanlan.zhihu.com/p/1929748460212552414).
- [SGLang, verl, OpenBMB, and Tsinghua University Team Jointly Open Source: First Support for Multi-Turn Interaction and Tool Calling in Mainstream RLHF Frameworks](./rlhf/verl/multi-turn/release_log/verl-multiturn-rollout-Release_ZH.md): First support for multi-turn interaction and tool calling in mainstream RLHF frameworks. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1906007821889283171).
- [Search-R1 & veRL-SGLang: Train LLMs with Multi-Turn RL to Reason and Call a Search Engine](./rlhf/verl/multi-turn/tool_examples/verl-multiturn-searchR1-like_ZH.md): Integrating the Search-R1 framework into the verl-sglang ecosystem. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1912156329751081620).
- [SGLang-veRL Server: From Engine to Server, We Need More Flexible RLHF Rollout Interfaces](./rlhf/verl/server-based/veRL-server-based-rollout.md): To implement more complex RLHF systems, we are gradually replacing the rollout engine in veRL with a rollout server. Also available on [Zhihu: SGLang-veRL Server](https://zhuanlan.zhihu.com/p/1890631652486665464).
- [HybridFlow veRL Original Paper Analysis](./rlhf/verl/readme.md): Principles and implementation of SGLang's hybrid engine. Also available on [Zhihu: HybridFlow veRL Original Paper Analysis](https://zhuanlan.zhihu.com/p/24682036412).

### OpenRLHF Framework

- [Illustrated Series on LLM RLHF: PPO Principles and Source Code Interpretation for Everyone](https://zhuanlan.zhihu.com/p/677607581) and [Illustrated Distributed Training Process based on Ray in OpenRLHF](https://zhuanlan.zhihu.com/p/12871616401): Excellent RLHF introductory resources by Ms. Mengyuan. After reading, you will have a good understanding of RLHF's computational flow and the OpenRLHF PPO framework. I have also added my own understanding in [RLHF Computational Flow](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/rlhf/OpenRLHF#rlhf-%E7%9A%84%E8%AE%A1%E7%AE%97%E6%B5%81).
- [Brief Analysis of the Computational Flow of Post-Training Systems Represented by OpenRLHF](./rlhf/OpenRLHF/readme.md): Further complement to Ms. Mengyuan's article. The Github native rendering is terrible; you might as well look at [Zhihu](https://zhuanlan.zhihu.com/p/16370000391).


### Algorithms and Theory

- [Kimi K1.5: Successful Practice of Long Context RL](./rlhf/partial-rollout/readme.md): Industrial implementation of Long Context RLHF. I have always liked the technical reports from the Kimi team. Also available on [Zhihu: Kimi K1.5: Successful Practice of Long Context RL](https://zhuanlan.zhihu.com/p/1894282607325344277).
- [Rule-based Reward](https://zhuanlan.zhihu.com/p/13211508979): Only on Zhihu, a brief write-up. Honestly, I didn't particularly like the original paper, but determined reward is indeed charming.
- [SWE-Bench: How to Construct an Excellent Benchmark in the LLM Era](https://zhuanlan.zhihu.com/p/16292266518): Reading notes on the SWE-Bench paper. How to construct a good benchmark to provide fine-grained reward for post-training is an eternal and beautiful topic.
- [Brief Analysis of Mainstream Alignment Algorithms and the NeMo-Aligner Framework](https://zhuanlan.zhihu.com/p/5220718268)


## SGLang Learning Notes

### SGLang Diffusion Learning Notes

- [Power Up Diffusion LLMs: Day‑0 Support for LLaDA 2.0](./sglang/diffusion-llm/readme-en.md): Introduction to the implementation of LLaDA2.0-flash-CAP in SGLang. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1985516215326749534) and in [Chinese version](./sglang/diffusion-llm/readme.md).
- [SGLang Diffusion Code Walk Through](./sglang/code-walk-through/sgl_diffusion_en.md): Basic principles of the diffusion model, and the entire process of a request being handled by SGLang-Diffusion. Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1982441236066480797) and in [Chinese version](./sglang/code-walk-through/sgl_diffusion.md).

### Core Architecture and Optimization

- [From KV Cache to Zero Overhead Scheduling, Understanding SGLang's Scheduling Ingenuity](./sglang/scheduler/readme-en.md): Also available on [Zhihu](https://zhuanlan.zhihu.com/p/1992587332189197731) and in [Chinese version](./sglang/scheduler/readme.md).
- [SGLang Code Walk Through](./sglang/code-walk-through/readme.md): The entire process of a request being handled by the SGLang Engine. Some parts are unfinished, but most are okay and have served as a starting point for many SGLang beginners. [Chinese version is here](./sglang/code-walk-through/readme-CN.md).

- [Walk Through SGLang Scheduler](./sglang/sglang-scheduler/readme-CN.md)
- [Pending Review] [KV Cache Code Walkthrough](./sglang/kvcache-code-walk-through/readme.md): Overview of KV cache management implementation, starting from the Scheduler component, detailing the update process of KV cache and memory pool during prefill and decode stages.
- [Pending Review] [SGLang Multimodal Request Lifecycle: A Deep Architectural Analysis with Qwen2.5-VL as an Example](./sglang/code-walk-through/multimodal_request_lifecycle.md): Provides a detailed analysis of the multimodal request processing flow within the SGLang framework, using Qwen2.5-VL as a reference model.
- [Pending Review] [How A Model is Loaded in Hugging Face and SGLang](./sglang/how-model-is-loaded/readme.md): Documents the process of loading models in Hugging Face and SGLang to help understand the weight loading mechanism.
- [Pending Review] [Speculative Decoding](./sglang/speculative-decoding/speculative-decoding.md): Introduces the speculative decoding optimization technique, which uses a smaller draft model to predict the next $K$ tokens, achieving up to $K$-fold acceleration.
- [Pending Review] [Zero-Overhead Batch Scheduler](./sglang/zero-overhead-scheduler/zero-overhead-batch-scheduler.md): Introduces the zero-overhead batch scheduler, which solves the GPU Bubble problem caused by serial execution of CPU scheduling and GPU computation in traditional inference systems.
- [Pending Review] [Data Parallelism Attention](./sglang/dp-attention/readme.md): Detailed introduction to the principles and implementation of DP Attention, specifically for models like DeepSeek that use MLA and only have one KV head, to avoid KV cache duplication caused by tensor parallelism.
- [Brief Analysis of SGLang Framework's Quantization Design and Ideas](./sglang/quantization/quantization_architecture_en.md): Also available on [Zhihu: Brief Analysis of SGLang Framework's Quantization Design and Ideas](https://zhuanlan.zhihu.com/p/1971183020338832111) and in [Chinese version](./sglang/quantization/quantization_architecture.md).
- [Constraint Decoding: Concepts, Methods, and Optimization](./sglang/constraint-decoding/readme.md): Also available on [Zhihu: Understanding Constraint Decoding: Concepts, Methods, and Optimization in one article](https://zhuanlan.zhihu.com/p/18336995950).
- [Pending Review] [Online Update Weights](./sglang/online-update-weights/readme.md): Introduction to the implementation of the `online_update_weights` interface in SGLang. Unlike `update_weights` which reads weights from the disk, this interface broadcasts new weights directly from the training engine via NCCL.
- [Pending Review] [SGLang Verl Engine Optimization Analysis](./sglang/sglang-verl-engine/readme.md): Analysis of optimizations in the SGLang verl engine, including the implementation of interfaces like `update_weights_from_tensor`.
- [Latency Accelerate For Weight Updates](./sglang/latency-accelerate-for-weight-updates/readme-CN.md)

### Usage and Practice

- [Pending Review] [Qwen3-Coder Usage](./sglang/qwen/coder.md): Introduction to using Qwen3-coder in SGLang, including the use of tool-parser.
- [Pending Review] [NVIDIA Dynamo](./sglang/nvidia-dynamo/dynamo.md): Introduction to NVIDIA Dynamo, a high-throughput, low-latency inference framework designed for generative AI and inference model serving in multi-node distributed environments.
- [Viewing HuggingFace Model Structure](https://zhuanlan.zhihu.com/p/9912733791)
- [SGLang Backend Original Paper Analysis](https://zhuanlan.zhihu.com/p/716543182)
- [Brief Analysis of the Status Quo of Reward / Embed Model Server Engine](https://zhuanlan.zhihu.com/p/4148050391)
- [Newbie Perspective: Experience and Gains from Migrating vllm to SGLang](https://zhuanlan.zhihu.com/p/714833359)
- [Newbie Perspective: Using SGL to Serve Embedding Model](https://zhuanlan.zhihu.com/p/715805386)
- [Newbie Perspective: Using vllm to serve a new Embedding Model](https://zhuanlan.zhihu.com/p/715857723)

## Scheduling and Routing

- [Mooncake: Carrying the P/D Separation to the End](https://zhuanlan.zhihu.com/p/1711346141)
- [Should Prefill and Decode be Separated onto Different Cards?](https://zhuanlan.zhihu.com/p/1280567902)
- [Understanding Prefill and Decode Computation Characteristics Based on Chunked Prefill](https://zhuanlan.zhihu.com/p/718715866)
- [ModelServer: A Frontend Distribution System Based on SGLang](https://zhuanlan.zhihu.com/p/718015016)

## ML System Fundamentals

### Transformers & Model Architecture

- [Pending Review] [Cross-Attention Mechanism in Transformer](./transformers/attention/cross_attention_en.md): Introduction to the cross-attention mechanism in Transformers, allowing the decoder to access and use relevant information from the encoder. Also available in [Chinese version](./transformers/attention/cross_attention.md).
- [Understanding Special Tokens and Chat Templates in One Article](./transformers/special_tokens/special_tokens.md): Also recorded on Zhihu [Understanding Special Tokens and Chat Templates in One Article](https://zhuanlan.zhihu.com/p/17052593700).

### CUDA & GPU

- [Revisiting CUDA Graph: Core Mechanisms, Multi-Graph Memory Sharing, and Unified Coverage for Dual AR Models](./torch/cuda-graph/readme-2-en.md): Also available on [Zhihu](https://zhuanlan.zhihu.com/p/2017950447520980998) and in [Chinese version](./torch/cuda-graph/readme-2.md).
- [Brief Analysis of CUDA Graph Based on torch-memory-savor](./torch/cuda-graph/readme_en.md): Also available on [Zhihu: Brief Analysis of CUDA Graph Based on torch-memory-savor](https://zhuanlan.zhihu.com/p/1921726788574360686) and in [Chinese version](./torch/cuda-graph/readme.md).

### Distributed Training & Communication

- [Pending Review] [Implementing Tensor Parallelism From Scratch](./torch/tensor-parallelism/readme.md): Implementation and practice of Tensor Parallelism.
- [Pending Review] [Expert Parallelism](./rlhf/sys-design/readme-4.md)
- [NCCL and NVIDIA TOPO](./torch/nccl/readme.md): Introduction to NCCL and NVIDIA GPU detection. Also available on [NCCL and NVIDIA TOPO](https://zhuanlan.zhihu.com/p/6160835906).
- [NCCL and SGLang](./torch/nccl/readme_en.md): Application of NCCL in SGLang. This is very similar to the Chinese content but includes some additional notes on parallel strategies. I probably won't complete this note and will write a separate one to record parallel strategies.
- [PyTorch Distributed](./torch/torch-distributed/readme.md): Communication practice with `torch.distributed`, details on GIL and `all_reduce`. This part is also available on [Zhihu: PyTorch Communication Practice](https://zhuanlan.zhihu.com/p/5853094319).
- [[Original][In-Depth][PyTorch] DDP Series Part 1: Introductory Tutorial](https://zhuanlan.zhihu.com/p/178402798): Although I didn't fully grasp the DDP content, I used this to learn about GIL and ring all reduce. This step is recorded in the [Postscript of torch-distributed](./torch/torch-distributed/readme.md#gil).
- [Detailed Explanation of nvidia-smi Command and Some Advanced Tips](https://www.yourmetaverse.cn/deep_learning/199/): Mainly about network topology; my local results are recorded in the [NCCL section](./torch/nccl/readme.md#nvlink-查询).

### Quantization

- [Give me BF16 or Give Me Death: Comprehensive Evaluation of Current Quantization Methods](https://zhuanlan.zhihu.com/p/5485556270)
- [AWQ: Model Quantization Should Focus on Activation Values](https://zhuanlan.zhihu.com/p/942485319)

## Developer Guide

- [How to use docker](./engineer/how-to-use-docker/readme_en.md): How to use Docker to manage development environments. Please note that to collectively foster a good research environment and prevent others from being annoyed by the baseline "it runs on my machine," learning Docker is essential for everyone. We also have a [Chinese version](./engineer/how-to-use-docker/readme.md) and [Zhihu](https://zhuanlan.zhihu.com/p/1916764175230801287).
- [Setting up a Clean Development Environment](./engineer/uv/readme.md): Setting up a clean development environment. Also available on [Zhihu: Setting up a Clean Development Environment](https://zhuanlan.zhihu.com/p/23440683394).
- [Compiling and Deploying Jupyter Notebooks as Documentation on CI](https://zhuanlan.zhihu.com/p/2382351079)

---

## File: engineer/uv/readme.md

# 如何配置一台爽快的开发机器

前天，我的一位好朋友批评我：“你和世界上 98% 的男人都一样，见到什么美好的事物都想着占为己有，而不是欣赏这份美好”。

我觉得她说的很对，但是，我就想问，有几个男人见了这个，能把持得住？




总之，我是把持不住了，终于有幸摸到了金子做成的 H200。这个笔记简单记录下自己好几个月以来为了开发 SGLang for RLHF，踩坑无数次后的配环境经验。天下苦 conda 久已，我尝试完全基于原生的 python 虚拟环境管理，搭配 uv 包管理，希望能帮助读者零帧起手，配置一台爽快的开发机器。

## 配置 bash/zsh

任何一台集群起步，我都会建议先配置好 bash/zsh，最大的好处是在进行任何安装前，我们就可以确定好所有的数据路径，避免数据被写到 `/root` 或者 `/home` 等集群共用的目录下。据我所知和，大部分开发集群的数据路径都不是 `/home` 或者 `/root`，**如果向这两个路径下写入大量内容而占据所有的磁盘，会导致 `/root/tmp` 或者 `/home/tmp` 也无法写入，而 ssh 登录集群需要向这两个关键的 `tmp` 目录写入数据，所以一直往这两个目录写入数据，会导致 ssh 登录失败，集群得返厂重修。**

这里分享下我自己喜欢的一套配置，可以参考下：


我喜欢用的 .bashrc 文件，zsh 同理

```
/* Detailed source-code truncated for AI context efficiency. */
```



## 安装 uv

uv 是更加现代的 python 包管理器，可以完全替代 conda。轻量级，方便，快速且强大，是未来趋势。

首先，登录崭新的集群账号，这一刻会使用全集群共用的 python 环境。这一环境在比较安全的集群管理上是不可修改的，所以我们需要开辟自己的虚拟环境来安装 uv。

```bash
# 创建虚拟环境
python3 -m venv ~/.python/sglang

# 激活虚拟环境
source ~/.python/sglang/bin/activate

# 安装 uv
pip install uv
```



## 配置 ssh

```bash
ssh-keygen
```

把公钥扔到 github 就好了。然后本地配置 git config：

```bash
git config --global user.name "zhaochenyang20"
git config --global user.email "zhaochenyang20@gmail.com"
```

## 配置 oh-my-zsh

其实我一般不自己折腾换 shell，但是如果默认就是 zsh 的话，我还是喜欢用 oh-my-zsh：

**注意 oh-my-zsh 会覆盖之前所有的 .zshrc 文件，所以需要先备份！！！**


```bash
sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
```

---

## File: engineer/how-to-use-docker/readme.md

# Docker 零帧起手

## Docker 的用途

我们经常诟病深度学习代码的不可复现性，某次在小红书见人锐评——学术圈的哥们要是真在乎自己工作的 real impact，都应该学会如何用 docker，而不是留下一堆没法复现的代码让几年后的 reviewer 拿出来恶心人。从我上述讨论中，想必大家已经感受到了，docker 是一种将开发环境连同代码复制并分发的强大软件。在我的日常工作中，docker 起到了如下作用：

1. 提供高度可复现且隔离的开发环境：Docker 不仅确保了代码的可复现性，还提供了一个隔离的沙箱环境。在开发过程中，我们可以在容器内自由安装依赖、修改系统配置，而不会影响宿主机的环境。这种隔离性特别适合在共享集群上工作，因为即使没有 sudo 权限，我们也能在容器内完成所有必要的环境配置。同时，通过使用统一的 Docker 镜像，我们确保了所有开发者都使用完全相同的环境，**避免了"在我机器上能跑"的问题**。Talk is cheap, show me your code.
2. 便于我们在集群上 share 磁盘空间：这是一个非常现实的问题，在我们的开发集群上，磁盘空间往往是不够用的，比如 8 * H100 集群可能磁盘只有 3T，倘若大量开发者共用集群，各自建立各自的用户。每个人都会在自己的用户路径下存放自己的 huggingface cache，即便强行指定 huggingface cache 为一个统一路径，其实 huggingface 也会根据登录用户的不同来鉴权，每个用户还是独立存了自己的 cache。模型的磁盘大小都不低，而且几十号人可能都得用 Llama 3.1 8B 这种模型，磁盘当然是吃不消的。因此，我们的集群强制每个人必须使用统一的账号登录，然后建立自己的 docker，并且将外部的 huggingface cache 路径统一映射到 docker 内部，避免了 huggingface 的 cache 被反复存放，显著节省了磁盘空间。当然，给我们的开发者带来了一定不便。不过，也促进了每个人都得搞明白 docker 如何使用。当然，坏处也明显，总有人隔三差五会把别人的 docker 误删了，反过来促进了大家随时 commit and push 代码。
3. ~~便于绕开 sudo 用户直接 kill 僵尸进程。~~

总之，为了共同塑造良好的科研环境，避免有人用 baseline "在我的机器上能跑"来恶心别人，学习 docker 对任何人都是必不可少的。

## 安装 docker

- 通常来讲，服务器管理员需要预先安装 docker 软件，无需你来手动安装。如果他拒绝安装 docker，建议和他据理力争。毕竟，连 docker 管理都不做，这管理员不如不要当了。
- 倘若服务器上找不到 docker，也可以使用 nerdctl 作为替代。nerdctl 完全兼容 docker 的命令。把指令替换成 `nerdctl xxx` 即可。
- 如果你是管理员，可参考 [Install Docker](https://docs.docker.com/engine/install/)。

## 下载 docker 镜像

绝大多数 docker 镜像都被发布在 [Docker Hub](https://hub.docker.com/)。为了开发 sglang，我通常使用 SGLang 的官方镜像 [lmsysorg/sglang Tags | Docker Hub](https://hub.docker.com/r/lmsysorg/sglang/tags)。当然，如果你关注 SGLang RL 小组的工作，其实我们也有专门为了 verl-SGLang 搭建的镜像，比如 jurong 的 `ocss884/verl-sglang:ngc-th2.6.0-cu126-sglang0.4.5.post3`。你当然发现了，这个版本带有 SGLang 的版本号，并不是最新的 SGLang。因此，这个 docker 我更多用于 verl—SGLang 的 CI。实际开发 verl—SGLang，我还是会用 SGLang 本身的 docker，然后在其中安装 verl 和最新的 SGLang。

```bash
# 下载镜像
# docker pull 
docker pull lmsysorg/sglang:latest
```

## 在容器内运行 docker 镜像

下载的**镜像**相当于压缩包，我们要把镜像解成**容器**才可以运行。

运行容器的指令格式是：

```bash
docker run [OPTIONS] IMAGE [COMMAND]
```

- OPTIONS: 运行容器时附加的参数
- IMAGE: 镜像名
- COMMAND: 容器启动时运行什么指令

## 常用 OPTIONS

1. `-it` 交互式终端：有该参数你才能在 docker 里用交互式终端

```bash
# 使用 -it：可以进入容器并执行命令
docker run -it ubuntu bash
# 此时你可以输入命令，比如 ls、cd 等

# 不使用 -it：容器会立即退出
docker run ubuntu bash
# 容器会立即结束，因为无法接收输入
```

2. `--name <container-name>` 容器名，方便下次重启。
  - 标记 docker 的用途或者拥有者。
  - 命名规则请参考服务器准则。在 SGLang 的开发机器上，随意命名的容器会被直接删除。

- `--shm-size <shared-memory-size>` 共享 CPU 内存大小。一般我们做 RL 需要较大的内存，默认的 64MB 会导致崩溃，建议设置为 16g 及以上。

- `--gpus all`  允许容器 access 哪些 GPU。如无特殊需求，设置成 all 即可。

- `-v <host-path>:<container-path>` 目录挂载：这可能是最重要的功能。举个例子，前文就提到了，我们将开发集群上的统一登录用户的 huggingface cache 挂载到了每个 docker 下，避免了每个人一个 cache 满天飞。更具体的来说，`-v` 可将宿主机目录 `<host-path>` 的全部内容挂载到容器目录 `<container-path>` 。宿主机和 docker 容器会共享目录下所有的文件和文件夹。容器对内容修改对宿主机可见，反之亦然。该参数可以多次添加，通常用于映射代码工作区，数据集，模型文件，配置文件等。

我们将这些常见参数组合起来，得到如下指令：

```bash
docker run -it --name <container-name> --shm-size 16g --gpus all -v <host-path>:<container-path> IMAGE
```

## 可选 OPTIONS

1. `-p <host-port>:<container-port>` 端口映射：将容器端口 `<container-port>` 映射到宿主机端口 `<host-port>` 。使得外部可以通过宿主机端口来访问容器内运行的服务。

2. `--network host` 网络共享：使容器直接使用宿主机的网络。共享 ip，端口，网络资源等。
  - 部分服务器在国内，添加 `--network host` 以共享网络代理。
  - 使用 `--network host` 时，`-p` 参数会被忽略。

3. `-e <cv-name>=<cv-value>` 环境变量：设置容器内环境变量`<cv-name>` 的值为 `<cv-value>` ，可多次添加。

4. `--ipc=host` 进程间通信的命名空间共享：允许容器内的进程与宿主机上的进程进行通信，共享 IPC 命名空间。

5. `-d` 在后台运行容器，输入 exit 时容器不关闭。

6. `--rm` 容器关闭后自动删除。

## COMMAND

在启动 container 的同时，我们可以指定 container 立即执行的指令。譬如：

1. 在容器内启动 `bash` 终端，可以输入指令，管理文件等。

```bash
docker run -it [other OPTIONS]  bash
```

2. 在容器内启动 sglang server 服务（运行在容器的 30000 端口上），并映射到宿主机的 30000 端口，以供外部访问。同时定义环境变量 HF_TOKEN 来鉴权。

```bash
docker run -p 30000:30000 --env "HF_TOKEN=hf_xxx" [other OPTIONS]  python3 -m sglang.launch_server [other paras]
```

## 容器管理

### 容器生命周期

在 docker 中，容器的生命周期管理是最基础的操作。这里是几个核心命令的区别：

1. `docker run`：创建并启动新容器：
   - 容器不存在时：创建并启动一个新的容器，相当于"买一台新电脑并开机"
   - 容器已存在时：会报错 `Error: Conflict. The container name is already in use`，因为不能创建同名容器

2. `docker start`：启动已停止的容器：
   - 容器不存在时：报错 `Error: No such container`；
   - 容器已存在时：启动容器，保持原有配置不变，相当于"把已经关机的电脑重新开机"；

3. `docker restart`：重启容器
   - 容器不存在时：报错 `Error: No such container`；
   - 容器已存在时：先停止再启动，相当于"重启电脑"，用于 docker 崩溃时重启；

4. `docker exec`：在运行中的容器执行命令
   - 容器不存在时：报错 `Error: No such container`；
   - 容器已存在但未运行：报错 `Error: Container is not running`；
   - 容器正在运行：在容器中执行命令，相当于"在已经开机的电脑上打开一个新的终端窗口"。

### 容器操作

1. 查看容器：
   - `docker ps` 查看正在运行的容器
   - `docker ps -a` 查看所有容器（包括已停止的）
   - `docker ps -a -s` 查看所有容器及其大小

2. 关闭容器：
   - 退出会话并关闭容器：输入 `exit` 或按 Ctrl + D
   - 退出会话但保持容器运行：
     - 按 Ctrl + P, 再 Ctrl + Q（在部分 IDE 中可能失效）
     - 使用 `-d` 参数启动容器
     - 直接关闭终端

3. 删除容器：
   - 需要先停止容器：`docker stop <container-name>`
   - 然后删除容器：`docker rm <container-name>`
   - 或使用 `--rm` 参数：容器关闭后自动删除

### 使用场景举例

```bash
# 1. 首次创建并运行容器
docker run -it --name my_container ubuntu bash

# 2. 容器停止后，重新启动
docker start my_container

# 3. 容器运行中，需要重启
docker restart my_container

# 4. 容器运行中，需要打开新的终端
docker exec -it my_container bash

# 5. 如果容器已存在，想创建新容器
docker run -it --name my_container_new ubuntu bash
```

### 特别说明

- 使用 `--rm` 参数时，容器停止后会自动删除，此时 `start` 和 `restart` 都会失败
- `exec` 只能用于运行中的容器，不能用于已停止的容器
- `run` 是创建新容器，其他命令都是操作已存在的容器

## 镜像构建

在我们的日常开发中，为 CI 测试机器发布稳定的 docker 镜像必不可少。具体而言，细心的读者可能早已发现，verl-sglang 的不少 docker 都是我们自己 build 的，比如 `ocss884/verl-sglang:ngc-th2.6.0-cu126-sglang0.4.5.post3`。自己 build docker 可以要安装特定版本的 CUDA、PyTorch 驱动，或者加入一些自己写的脚本、代码，配置环境变量等。比如说，我们希望：

1. 配置 Ubuntu 20.04 系统
2. 配置 tmux, git, python 3.10, cuda 12.4.1
3. 把项目代码拷到 `/workspace`
4. 设置一些必要的环境变量

Docker 构建镜像的过程就基于 `Dockerfile` 忠实地执行这些操作。

### Dockerfile 语法

`Dockerfile` 的常见指令：

1. `FROM :<tag>`：指定基础镜像，告诉 Docker 要基于哪个现有的镜像开始构建；
2. `WORKDIR /path/to/workdir`：设置工作目录，后续指令都会在这个目录下执行。如果目录不存在，`WORKDIR` 会自动创建；
3. `RUN <command>`：在镜像内部执行命令，每条 `RUN` 指令都会在当前镜像的基础上创建一个新的层。为了减少镜像层数和体积，加快构建速度，通常把多个 `apt-get install` 或者 `pip install` 命令用 `&&` 连在一条 `RUN` 指令里；
4. `COPY <host-path> <container-path>`：将宿主机的文件或目录复制到镜像内的指定路径；
5. `CMD ["command", "param1", "param2"]`：指定容器启动时默认执行的命令。还可写成 `CMD command param1 param2` (shell form)。注意，一个 Dockerfile 里只能有一条 `CMD` 指令，如果有多条，只有最后一条生效。如果在 `docker run` 时指定了命令，那么 `CMD` 的命令会被覆盖；
6. `ENV <key>=<value>`：设置容器**运行**时的环境变量，镜像构建时不生效；
7. `ARG <key>=<value>`：设置镜像**构建**时的临时变量，容器运行时不生效；

### 示例代码

```dockerfile
# 导入 Nvidia 官方镜像
# 该镜像把 CUDA 12.1.1 和 cuDNN 8.9.0 都装好了，不需要再处理版本兼容问题
FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

# 设置工作目录为 /sgl-workspace
WORKDIR /sgl-workspace

# 设置环境变量，如 huggingface token 用于下载模型
ENV HF_TOKEN=hf_xxxyyyzzz

# 安装开发工具
RUN apt-get update && apt-get install -y \
    vim \
    tmux \
    wget \
    git \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get clean

# 将当前目录下的 app.py 文件复制到容器内的工作目录
COPY app.py /sgl-workspace/app.py

# 默认执行命令，启动 Python 脚本
CMD ["python3", "app.py"]
```

### COMMAND

在命令行构建 docker 的指令：

```
docker build -t :<tag> -f <docker-file> 
```

1. `-t :<tag>`：镜像的名字和标签；
2. `-f <docker-file>`：`Dockerfile` 的路径；
3. ``：指定 Docker 构建镜像时的构建上下文路径，即 Docker 可以访问的文件和目录的根路径。它决定了 Dockerfile 中文件引用（如 `COPY、ADD`）的查找范围。通常设为 `.` 表示当前目录，但也可以是其他本地目录或 Git 仓库 URL；

## 镜像上传

镜像构建好以后，我们可以把本地镜像 `docker push`到远程镜像仓库，如 Docker Hub 等。

### 给镜像打标签

对于 Docker Hub，镜像名通常是 `<dockerhub-username>/<repo-name>:<tag>`。 如果你在 `docker build` 时已经用了这个格式，那这步可以跳过。如果没用，或者你想推送到不同的仓库/用户下，就需要重新打标签。

```
docker tag :<tag> <target-image>:<tag>
```

### 登陆到镜像仓库

在 docker hub 注册账号后，使用 docker login 在本地登录

```
docker login
```

### 推送镜像

确保 `:<tag>` 是你上一步打好标签的、符合仓库规范的完整名称。然后 push 即可

```
docker push :<tag>
```

---

## File: rlhf/verl/multi-turn/tool_examples/weave_trace_readme.md

# verl 中的 Wandb Weave 功能

在 Agentic RL 中，为了帮助我们更好的分析 trajectory 中的的多轮对话和工具调用对优化训练过程，verl 提供了 Trace 功能，可记录指定函数的输入、输出及时间戳，并支持在可视化界面中查看，目前支持 `wandb weave`。我们在这篇文档中解析如何使用 weave。

感谢 chengxi li @CMU 和 Chenyang Zhao @Amazon 的贡献。

## Quick Start

在 `config.yaml` 或命令行中添加以下参数即可开启 Trace：

```yaml
actor_rollout_ref:
  rollout:
    trace:
      backend: weave       # 目前仅支持 weave
      token2text: true     # 是否在 Trace 中展示解码后的文本
```

或者在 bash 中追加：

```bash
actor_rollout_ref.rollout.trace.backend=weave \
actor_rollout_ref.rollout.trace.token2text=True \
```


还需要以下配置作为前置条件:

| 场景 | 必要参数 | 备注 |
| ---- | -------- | ---- |
| 使用 Weave 并记录日志到 wandb | `trainer.logger=["console","wandb"]` | 建议同时开启 wandb 日志，实现一处查看所有信息 |
| 启用async rollout | `actor_rollout_ref.rollout.mode=async` 且 `actor_rollout_ref.rollout.multi_turn.enable=true` | Trace 现在只在 `agent_loop` 启用，sglang本身不需要设置 `mode=async` 开启异步，但是需要此设置使 Trace 生效  |

不过，直接在本地的 fork 修改这些参数大概率会 fail，因为 weave 依赖于 verl 新的 AgentLoop 特性。我们接下来提供一套逐步的复现过程。

## 环境

1. 设置环境变量 `WANDB_API_KEY`

```bash
export WANDB_API_KEY=your_wandb_api_key
```

## 数据集要求

数据集需新增一列 `agent_name`，在 `map_fn` 中补充即可:

[example](https://github.com/volcengine/verl/blob/ada82bb719e4d15ed4974f118bc86ec4d78c871d/recipe/retool/retool.py#L96)

```python
# python
data = {
    ...,
    "agent_name": "tool_agent",  # 新增列
}
```

## 如何使用

请严格按照以下步骤，否则会非常痛苦。

### 创建新的 docker

使用前需要配置好 `WANDB_API_KEY`，参考[这个过程](https://community.wandb.ai/t/where-can-i-find-the-api-token-for-my-project/7914)。

```bash
# 如果你的系统没有配置过 HF_TOKEN 和 WANDB_API_KEY，请先配置好
docker run -it --name h100_verl_{your_name} --gpus all \
    --shm-size 32g \
    -v {your_cache_path}:/root/.cache \
    --env "HF_TOKEN=$HF_TOKEN" \
    --env "WANDB_API_KEY=$WANDB_API_KEY" \
    --ipc=host \
    lmsysorg/sglang:latest \
    /bin/bash
```

进入 docker 后，可以查看被映射的环境变量：

```bash
echo $HF_TOKEN
echo $WANDB_API_KEY
```

以后每次从 docker 里面 exit 出来，再用这个指令可以重启：

```bash
docker start -i h100_verl_{your_name}
```

### 基于源码安装 SGLang

配置 python 环境

```bash
mkdir -p /tmp
chmod 1777 /tmp
sudo apt update
sudo apt install -y python3.10 python3.10-venv
sudo python3 -m ensurepip --upgrade
sudo python3 -m venv ~/.python/verl-sglang
source ~/.python/verl-sglang/bin/activate
sudo python3 -m pip install --upgrade pip
sudo python3 -m pip install --upgrade uv
```

先安装 veRL，再安装 SGLang。

```bash
cd ~
git clone https://github.com/volcengine/verl.git
cd verl
python3 -m uv pip install -e ".[sglang,geo]"
python3 -m uv pip install -r ./requirements.txt
```

会遇到这个报错：

```bash
Resolved 130 packages in 1.96s
  × Failed to build `flash-attn==2.8.1`
  ├─▶ The build backend returned an error
  ╰─▶ Call to `setuptools.build_meta:__legacy__.build_wheel` failed (exit status: 1)
```

按照下面的步骤 fix：

```bash
python3 -m uv pip install wheel
python3 -m uv pip install -r ./requirements.txt --no-build-isolation
```

然后安装 SGLang upstream：

```bash
cd ~
git clone https://github.com/sgl-project/sglang.git
cd sglang
python3 -m uv pip install --upgrade pip
python3 -m uv pip install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python
```

额外安装 vllm 和 weave 的依赖：

```bash
python3 -m uv pip install vllm==0.9.1
python3 -m uv pip install weave
```

### 修改并运行

我们可以通过对现有脚本进行简单修改，在运行脚本中启用 `multi_turn` 和 `async rollout`，在数据集处理脚本中的 `def make_map_fn(split)` 增加一列 `agent_name`。

打开你 docker 里面的 `~/verl/examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh` 文件，去掉结尾一行的 `$@`，更改如下参数：

``` bash
    # 注意去掉原本 total_epochs 这行结尾的 $@
    # 不要把这些两行注释也写进去，否则会报错
    trainer.total_epochs=15 \
    actor_rollout_ref.rollout.trace.backend=weave \
    actor_rollout_ref.rollout.trace.token2text=True \
    actor_rollout_ref.rollout.mode=async \
    actor_rollout_ref.rollout.multi_turn.enable=true
```

在 `~/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py` 中追加 `"agent_name": "tool_agent"`

```python
def make_map_fn(split):
        def process_fn(example, idx):
            question_raw = example.pop("question")

            question = question_raw + " " + instruction_following

            answer_raw = example.pop("answer")
            solution = extract_solution(answer_raw)
            data = {
                "data_source": data_source,
                # new column for weave trace
                "agent_name": "tool_agent",
                "prompt": [
                    {
                        #...
                    }
                ]
            }
            return data

        return process_fn
```

接下来测试即可：

```bash
cd ~/verl
python3 -m uv pip install .
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

# 拉取并预处理 gsm8k 数据集
python examples/data_preprocess/gsm8k_multiturn_w_tool.py
```

启动 8 卡训练即可。

```bash
bash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh
```

## Debug

如果你在启动 bash 后发现了这个错误：

```bash
raise ValueError(f"Feature type '{_type}' not found. Available feature types: {list(_FEATURE_TYPES.keys())}")
ValueError: Feature type 'List' not found. Available feature types: ['Value', 'ClassLabel', 'Translation', 'TranslationVariableLanguages', 'LargeList', 'Sequence', 'Array2D', 'Array3D', 'Array4D', 'Array5D', 'Audio', 'Image', 'Video', 'Pdf']
```

其实这不是实际的报错，这个报错让我费解了非常非常久，我仔细看了 log 才发现问题，其实可以向上看几行报错。在报错栈最开始的地方，报错的 python 环境是 `/root/.python/verl-sglang/lib/python3.10`，结果到了栈底部成了 `/usr/local/lib/python3.10`。毫无疑问，是 python 环境错位了；

1. 主进程使用虚拟环境：`/root/.python/verl-sglang/lib/python3.10/site-packages/`
2. Ray worker 进程使用系统 Python：`/usr/local/lib/python3.10/dist-packages/`

最后对这个问题的解决方式是修改 `verl/trainer/constants_ppo.py` 文件，直接改为：

```python
import os
import sys

# 获取当前Python解释器路径和虚拟环境路径
python_executable = sys.executable
virtual_env = os.environ.get("VIRTUAL_ENV", "")
python_path = os.environ.get("PYTHONPATH", "")

# 如果当前在虚拟环境中，确保包含虚拟环境的site-packages
if virtual_env:
    site_packages = os.path.join(virtual_env, "lib", "python3.10", "site-packages")
    if site_packages not in python_path:
        python_path = f"{site_packages}:{python_path}" if python_path else site_packages

PPO_RAY_RUNTIME_ENV = {
    "env_vars": {
        "TOKENIZERS_PARALLELISM": "true",
        "NCCL_DEBUG": "WARN",
        "VLLM_LOGGING_LEVEL": "WARN",
        "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "true",
        # 添加Python环境配置
        "PYTHONPATH": python_path,
        "VIRTUAL_ENV": virtual_env,
    },
    # 指定Python解释器
    "python": python_executable,
}
```

## 查看 Trace

登录 `$WANDB_API_KEY` 对应的账号，在 project 里找到 `gsm8k_async_rl`，侧边栏选择 `Trace`，即可看到多轮对话和工具调用的信息。

---

## File: rlhf/verl/multi-turn/code-walk-through/readme.md

# verl Multi-turn Code Walk Through（Part 1）

承蒙社区厚爱，Agentic RL 如火如荼，我们 SGLang RL 小组的工作也在夜以继日。考虑到领域令人恐惧的发展速度，社区巨大的二次开发需求，我们选择以 verl 出发，分析其 end to end mutli-turn RL 训练的全过程。整体上，我们希望覆盖所有重要的 class 以及函数，更细粒度的代码不再展开。我们的写作风格希望能够 follow SGLang 的 code-walk-through：

[SGLang Code Walk Through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/sglang/code-walk-through/readme-CN.md)

为了前后内容的一致性，我们基于 [76f63cffa5](https://github.com/volcengine/verl/commit/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39) 的 commit 进行分析。

感谢来自 Amazon，LinkedIn，阿里等公司和 SGLang RL 小组的朋友们的贡献。虽然本文以分析 verl 的代码为主，写完之后我们才意识到，系统设计问题是非常通用的。诸如“log probs 重计算”，“Rollout Engine 显存管理”等等系统设计，是各大 RL 框架都需要考虑的核心问题。我们希望本文对于开源社区理解 RL 框架系统设计能提供可迁移的经验 😂

如果您对我们的工作感兴趣，欢迎来联系我们参与一些工作～

特别致谢：zhuoran yin @ CMU，changyi yang @ CMU，zhuohao li @ 阿里，ji li @待业在家（🤣），biao he @ Linkedin 和 xinpeng wei & chenyang zhao @ Amazon。

--------------------------------

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

```mermaid
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](https://huggingface.co/datasets/openai/gsm8k) 为例，预处理脚本是 `examples/data_preprocess/gsm8k_multiturn_w_tool.py`。整个脚本只做了经典的 huggingface datasets mapping，核心逻辑如下：

1. 加载 openai/gsm8k 原始数据集（train/test）。
2. 对每条原始数据，生成带有工具调用要求的 prompt（比如在 user turn 强调模型可以调用 `calc_gsm8k_reward` 工具，每个qa至少调用一次）。
3. 同样对于每条原始数据，解析答案；将 ground truth 写入 extra_info 字段。
4. 存储为 parquet 文件，分别保留为 train.parquet 和 test.parquet，默认路径为 `~/data/gsm8k/`。

## 启动训练

一个典型的启动命令如下：

```bash
# now 用于生成实验启动的时间尾缀，避免重复启动实验时覆盖已有 wandb log

function now() {
    date '+%Y-%m-%d-%H-%M'
}

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
nohup bash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh \
    trainer.experiment_name=qwen2.5-3b_rm-gsm8k-sgl-multiturn-$now \
    > logs/gsm8k-$now.log 2>&1 &

```

## 脚本配置

verl 的各项参数实属复杂，我们会单独编写文档来分享对 verl 各类参数的理解。在这篇文档中，我们想要格外强调的是 verl 各类 config 的覆盖关系。verl 的配置文件利用 hydra 进行了**分层覆盖**的设计模式。



Hydra 简介

[**Hydra**](https://github.com/facebookresearch/hydra) 是一个由 Facebook Research 开发的 Python 框架，旨在**优雅地配置复杂的应用程序**。它特别适用于需要管理大量参数和进行多组实验的场景，例如机器学习项目。Hydra 的核心特点在于其**动态、分层和可组合的配置管理能力**。Hydra 的核心优势：

* **分层配置 (Hierarchical Configuration)**：可以将配置分解成多个小型、模块化的 YAML 文件，并以目录结构进行组织。这使得配置更加清晰、易于管理和复用。
* **配置组合 (Configuration Composition)**：Hydra 能够将这些独立的配置模块动态地组合起来，形成一个完整的配置对象。你可以通过在主配置文件中指定 `defaults` 列表来选择和组合不同的配置组件。
* **命令行覆盖 (Command-line Overrides)**：这是 Hydra 最强大的功能之一。你可以在运行应用程序时，直接通过命令行参数来覆盖配置中的任何值。这使得进行实验和快速迭代变得非常方便，无需修改配置文件本身。
* **多运行模式 (Multi-run)**：Hydra 允许你通过一个命令运行多个具有不同配置的实验。这对于超参数搜索和模型比较非常有用。
* **动态工作目录 (Dynamic Working Directory)**：每次运行应用程序时，Hydra 都会自动创建一个独立的工作目录，并将当前运行的配置和输出保存到该目录中，确保实验的可复现性。
* **对象实例化 (Object Instantiation)**：Hydra 可以直接从配置中实例化 Python 对象（类或函数），这大大简化了代码，使配置更具声明性。


Hydra 实现分层覆盖的主要机制是**组合 (Composition)** 和 **命令行覆盖 (Command-line Overrides)**。

1.  **分层配置的组织**：

通常会创建一个 `conf` 目录，并在其中组织配置。例如：

```yaml
.
├── my_app.py
└── conf
    ├── config.yaml
    ├── model
    │   ├── cnn.yaml
    │   └── rnn.yaml
    └── dataset
        ├── cifar10.yaml
        └── imagenet.yaml
```

`config.yaml` 是你的主配置文件。在 `model` 目录下，你可以定义不同的模型配置（如 `cnn.yaml`、`rnn.yaml`），在 `dataset` 目录下定义不同的数据集配置（如 `cifar10.yaml`、`imagenet.yaml`）。

2.  **`defaults` 列表进行组合**：

在 `config.yaml` 中，你可以使用特殊的 `defaults` 列表来指定默认加载哪些配置组件。

**`conf/config.yaml` 示例：**

```yaml
defaults:
    - model: cnn       # 默认加载 conf/model/cnn.yaml
    - dataset: cifar10 # 默认加载 conf/dataset/cifar10.yaml
    - _self_          # 确保当前文件中的其他配置项也被加载

# 其他应用级别的默认配置
learning_rate: 0.001
epochs: 10
```

当 Hydra 加载 `config.yaml` 时，它会根据 `defaults` 列表中的指示，自动将 `conf/model/cnn.yaml` 和 `conf/dataset/cifar10.yaml` 的内容合并到最终的配置对象中。

3.  **命令行覆盖**：

这是实现灵活覆盖的关键。你可以通过命令行参数来覆盖任何已加载的配置值，包括在 `defaults` 列表中指定的组件或其内部的任何参数。

* **覆盖整个配置组**：
要切换模型从 `cnn` 到 `rnn`，你可以在命令行中这样运行：

```bash
python my_app.py model=rnn
```

这将指示 Hydra 加载 `conf/model/rnn.yaml`，并用它来替换默认的 `cnn` 配置。

* **覆盖特定参数**：
你可以深入到配置的任何层级来覆盖特定的参数。例如，如果你想修改学习率或数据集的某个参数：

```bash
python my_app.py learning_rate=0.01 dataset.batch_size=64
```

这里，`learning_rate` 直接覆盖了 `config.yaml` 中的值，而 `dataset.batch_size` 则覆盖了 `conf/dataset/cifar10.yaml`（或者你通过 `dataset=imagenet` 指定的其他数据集配置文件）中的 `batch_size` 参数。

* **添加新参数 (使用 `+`)**：
如果你想添加一个在默认配置中不存在的新参数，可以使用 `+` 前缀：

```bash
python my_app.py +optimizer.name=AdamW
```

* **动态覆盖 (使用 `++`)**：
如果你希望修改一个已有字段，或者在原配置中没有该字段时自动创建它，可以使用 ++。这种方式适用于需要动态添加或覆盖配置项的场景，确保字段总是被设置为你指定的值，无论它是否已存在。

```bash
python my_app.py ++model.num_layers=10
```

Hydra 内部使用 [OmegaConf](https://www.google.com/search?q=https://omegaconf.readthedocs.io/en/2.3_latest/) 库来处理这些配置对象，它提供了强大的合并和解析功能，使得分层覆盖和值插值（例如，引用其他配置值或环境变量）变得非常容易。




回到 verl multi turn，在我们启动的 `run_qwen2.5-3b_gsm8k_multiturn.sh` 中，设置了：

```bash
PROJECT_DIR="$(pwd)"
CONFIG_PATH="$PROJECT_DIR/examples/sglang_multiturn/config"

python3 -m verl.trainer.main_ppo \
    --config-path="$CONFIG_PATH" \
    --config-name='gsm8k_multiturn_grpo' \
```

这意味着这次任务的默认 config 是 `CONFIG_PATH` 下的 `gsm8k_multiturn_grpo.yaml`，且接下来的参数会覆盖 `gsm8k_multiturn_grpo.yaml` 中的默认值。更进一步，我们来观察 `gsm8k_multiturn_grpo.yaml` 的内容：

```yaml
hydra:
  searchpath:
    - file://verl/trainer/config

defaults:
  - ppo_trainer
  - _self_

data:
  max_prompt_length: 1024
  max_response_length: 1024
  train_batch_size: 256
  return_raw_chat: True

actor_rollout_ref:
  hybrid_engine: True
  rollout:
    name: sglang
    multi_turn:
      enable: True
      max_turns: 5
      # tool_config_path: "./config/tool_config/gsm8k_tool_config.yaml"
```

这里 hydra 语法，会去 `verl/trainer/config` 目录下寻找 `ppo_trainer.yaml` 作为基础配置，并且覆盖。因此，启动 `run_qwen2.5-3b_gsm8k_multiturn.sh` 时，先加载 `gsm8k_multiturn_grpo.yaml` 作为基础配置并覆盖，然后加载 `ppo_trainer.yaml` 并覆盖。最终合并这三级配置，得到最终的 config。

最后，注意到在 `run_qwen2.5-3b_gsm8k_multiturn.sh` 的最后，我们，我们设置了 `actor_rollout_ref.rollout.multi_turn.tool_config_path="$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml"`，这里指定 multi_turn 的 tool_config_path 为 `examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml`。这一文件仅仅配置了 gsm8k 的 tool 调用，并不会覆盖之前训练的 config。

## 训练主入口与初始化

### Ray Actor，Ray Task 和 Ray Worker

在介绍 verl 的训练主入口之前，我们先介绍 Ray 的一些核心概念。Ray 是一个统一计算框架，旨在实现简单地从单机到大型分布式集群的扩展，提供构建和运行分布式应用的底层基础设施和一组核心原语。Ray 通过以下功能实现这一目标：

1. **统一 API**：Ray 提供了一套简单易用的 Python API，将普通函数转换为分布式任务，将 Python 类转换为分布式服务，也即 Ray Actor。Ray Actor 内部持久存储的数据称为状态，可以在 Actor 的整个生命周期内被多次访问、修改和维护，而不会在每次方法调用结束后消失。
2. **弹性伸缩**：Ray 可以将应用从单个机器无缝扩展到拥有数千个节点的集群，并能根据需求自动扩缩容。
3. **容错性**：Ray 内置了容错机制，可以处理节点故障和任务失败，确保应用的健壮性。
4. **性能优化**：Ray 优化了分布式任务调度、内存管理和数据传输，以实现高效的并行计算。

Ray Task 和 Ray Actor 都是用于分布式计算的核心原语，但它们各自服务于不同的目的，主要区别在于**是否维护状态**。

Ray Task 是 Ray 中最基本的计算单元，代表一个无状态的远程函数。Ray Task 的每次执行都是独立的，不保留之前的任何信息。就像调用一个普通函数，执行完后就清除内部状态。我们调用一个 Ray Task 后，会立即返回得到一个 Ray ObjectRef，而不是实际的结果。主程序可以继续执行其他操作，而 Ray Task 则在后台并行运行。我们需要使用 `ray.get()` 来获取 Task 的实际结果。 Ray Task 非常适合并行执行大量独立、一次性的计算任务，譬如数据批处理、独立的模型推理等场景。

Ray Actor 是一种特殊的 Ray Task，正如前文所述，它是一个持续运行的、有自己的状态和方法的远程对象。当我们创建一个 Ray Actor 后，Ray 会在集群中的某个 **Ray Worker** 上启动一个专门的进程来托管这个对象。该进程会一直运行，直到被销毁。Actor 可以维护内部变量，并且这些变量在 Actor 的生命周期内是持久存在的。每次调用 Actor 的方法，都可以访问和修改这些状态。这与普通的 Ray Task 不同，普通 Task 执行完会清除内部状态。Ray Actor 支持并发请求，Ray 会负责将这些请求序列化执行，保证 Actor 内部状态的一致性和线程安全。我们可以通过 `@ray.remote` 装饰器将一个 Python 类转换为一个 Ray Actor 类，然后通过 `.remote()` 方法实例化一个远程 Actor。

最后，Ray Worker 是 Ray 集群中真正执行代码的工作单元。一个 Ray 集群通常由一个 Head Node 和多个 Worker Nodes 组成。每个节点上都会运行一个或多个 Ray Worker 进程。无论是普通的 Ray Task 还是 Ray Actor 的方法，最终都是由 Ray Worker 进程来执行的。每个 Ray Worker 都会被分配一定的计算资源（如 CPU、GPU）。当你提交一个 Ray Task 或创建一个 Ray Actor 时，Ray 的调度器会找到一个有足够资源的 Worker 来运行它。Worker 进程之间以及 Worker 进程与头节点之间会进行通信，以协调任务执行、传输数据和管理状态。一个 Ray Worker 通常就是一个独立的 Python 进程。对于普通的 Ray Task，Ray Worker 相当于函数解释器，执行完任务后可能会被复用去执行其他任务。而对于 Ray Actor，Ray 会启动一个专门的 Worker 进程来托管这个 Actor，这个 Worker 进程的生命周期与 Actor 的生命周期绑定。

### `run_ppo()` 和 `TaskRunner.run()`

有了 ray 的概念，我们回到整个 RL 训练流程的起点：`verl.trainer.main_ppo.py` 中的 [`run_ppo()`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/main_ppo.py#L35)，它负责初始化 Ray 集群，配置 CPU 资源和运行时环境变量，并创建远程 TaskRunner 实例。

```python
def run_ppo(config) -> None:
    # 初始化 Ray 集群，配置 CPU 资源和运行时环境变量
    ray.init(
        runtime_env={"env_vars": {...}},
        num_cpus=config.ray_init.num_cpus,
    )

    # 创建远程 TaskRunner 实例
    # TaskRunner 是 Ray 中的一个远程 actor，它将在 Ray 集群上异步执行主要的训练任务
    runner = TaskRunner.remote()
    # 异步执行远程任务 runner.run()，并等待其完成
    # 通过 ray.get() 阻塞直到远程任务执行完毕，确保整个初始化流程的顺序性
    ray.get(runner.run.remote(config))
```

### ActorRolloutRefWorker 和 RayWorkerGroup 的相互关系

[TaskRunner](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/main_ppo.py#L64) 是 verl 中实现 PPO/GRPO 训练的核心组件，它通过将整个 RL 训练流程封装在一个独立的 Ray Actor 中，实现了任务的封装、资源隔离和分布式协调。为了解释清楚 `TaskRunner`，我们将 verl 当中最让人费解且最复杂的 `ActorRolloutRefWorker` 和 `RayWorkerGroup` 这两个类提前解释清楚。

我们先不讨论这两个类及其基类的具体意义，先讨论清楚其实例对象的创建过程。我们注意到这段 `TaskRunner` 的初始化中引入 `ActorRolloutRefWorker` 和 `RayWorkerGroup` 的相关代码：


TaskRunner 中引入 ActorRolloutRefWorker

```python
# Define worker classes based on the actor strategy.
if config.actor_rollout_ref.actor.strategy in ["fsdp", "fsdp2"]:
    assert config.critic.strategy in ["fsdp", "fsdp2"]
    from verl.single_controller.ray import RayWorkerGroup
    from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker

    actor_rollout_cls = AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
    ray_worker_group_cls = RayWorkerGroup

elif config.actor_rollout_ref.actor.strategy == "megatron":
    assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
    from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup
    from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker

    actor_rollout_cls = AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
    ray_worker_group_cls = NVMegatronRayWorkerGroup

else:
    raise NotImplementedError

from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role

# Map roles to their corresponding remote worker classes.
role_worker_mapping = {
    Role.ActorRollout: ray.remote(actor_rollout_cls),
    Role.Critic: ray.remote(CriticWorker),
}

# Define the resource pool specification.
# Map roles to the resource pool.
global_pool_id = "global_pool"
resource_pool_spec = {
    global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
}
mapping = {
    Role.ActorRollout: global_pool_id,
    Role.Critic: global_pool_id,
}
```



可以观察到，在 `TaskRunner` 的初始化中，会根据各类配置引入对应的 `ActorRolloutRefWorker / AsyncActorRolloutRefWorker` 类以及 `RayWorkerGroup / NVMegatronRayWorkerGroup` 类。对于 SGLang 而言，不存在 `AsyncActorRolloutRefWorker`。`ActorRolloutRefWorker` 类直接通过 `ray.remote(ActorRolloutRefWorker)` 创建一个远程的 Ray Actor，将其包装成一个 Ray Actor 类。此时还还没有创建任何实例，也没有分配资源。那么，`ActorRolloutRefWorker` 类到底在哪儿实例化并分配资源的呢？

实际上，在 `main_ppo.py` 的 [172 行](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/main_ppo.py#L172)，构造了 `RayPPOTrainer` 类，随后调用了 `RayPPOTrainer.init_workers()` 方法，我们进一步查看 `RayPPOTrainer.init_workers()` 方法的[相关代码](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/ppo/ray_trainer.py#L715)，我们观察到，每一个 RL worker 类（比如 ActorRolloutRefWorker）都会创造一个 work group（verl 中的各种 wg 变量），随后调用每个 worker group 的 `init_model()` 方法，而这些 worker group 实际上都是 `RayWorkerGroup` 的实例。`RayWorkerGroup` 的核心作用是资源调度的核心中间层，统一了各种 RL worker（比如 ActorRolloutRefWorker、CriticWorker）的接口，进行统一管理：

```python

# RayWorkerGroup 实例，指定资源池 并规定角色和对应的类
wg_dict = self.ray_worker_group_cls(
    resource_pool=resource_pool,  # 只需要指定资源池
    ray_cls_with_init=worker_dict_cls,  # 一个包含数个worker的类 （e.g. actor_roll， critic, ref）
    device_name=self.device_name,
)

#通过.spawn()获取角色对Ray Actor实例的映射
wg_dict.spawn(prefix_set=class_dict.keys())


# 所有 worker 都通过相同的模式创建，我这里进行简化，实际上的代码比较繁琐
actor_rollout_wg = RayWorkerGroup(resource_pool, actor_rollout_cls)
critic_wg = RayWorkerGroup(resource_pool, critic_cls)
ref_policy_wg = RayWorkerGroup(resource_pool, ref_policy_cls)
```


各种 worker group 实际上的初始化

这部分代码在 [`ray_trainer.py`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/ppo/ray_trainer.py#L771) 中：

```
/* Detailed source-code truncated for AI context efficiency. */
```

注意到 `ray_worker_group_cls` 就是 `RayWorkerGroup` 类，而 `worker_dict_cls` 就是 `ActorRolloutRefWorker` 类，所以我的简化是很合理的。



如此以来，`ActorRolloutRefWorker` 委托给 `RayWorkerGroup` 进行初始化。`RayWorkerGroup` 这个类就是专门用于资源调度的。通过其统一的 `_init_with_resource_pool` [方法](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/single_controller/ray/base.py#L313)，为每个 GPU 创建一个 worker，最终实例化每种 RL worker 并分配资源。

```python
def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, ...):
    # 从 Ray 申请 Placement Groups
    pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)
    
    # 为每个 GPU 创建一个 worker
    for local_rank in range(local_world_size):
        worker = ray_cls_with_init(placement_group=pg, placement_group_bundle_idx=local_rank, ...)
        self._workers.append(worker)
```

读到这里，我们基本对 verl 有了一些感觉。注意到，在 verl 当中有两个带有 Worker 的 base class，一个就叫做 [`Worker`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/single_controller/base/worker.py#L77)，另一个叫做 [`WorkerGroup`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/single_controller/base/worker_group.py#L121)。`Worker` 是 RL 里面的逻辑类（比如 actor 和 critic）,实际管理 RL 的数据流，而 `WorkerGroup` 只用于分布式系统的资源调度。


此外，从 `actor_rollout_wg` 和 `ref_policy_wg` 的实例化当中，也能看出一些学问。在 `ActorRolloutRefWorker` 的设计当中，Actor Training，Actor Rollout 和 Reference model 是用同一个 worker class 进行管理的。但是，之后委托给 `RayWorkerGroup` 创建 worker group 并且调用资源的时候，Actor Training 和 Actor Rollout 是由同一组 `RayWorkerGroup` 进行资源管理的（这二者本来就要被放在同一个资源组上做 hybird engine），而 Reference Model 是由另一组 `RayWorkerGroup` 管理资源的。

最后，我去问了相关开发者，他们也认为把 Actor Rollout，Actor Training 和 Reference Model 放在同一个 worker 里是 bad design 😂，不用纠结这种设计是否有什么高瞻远瞩，完全没有。

### [`ActorRolloutRefWorker.__init__()`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/workers/fsdp_workers.py#L101)

如前文所说，`ActorRolloutRefWorker` 是 verl 中用于管理 Actor Training，Actor Rollout 和 Reference Model 的 worker class。我们具体来分析其逻辑上实现的功能。注意，本文档只分析 FSDP backend 下的实现，megatron 留作后文。

1. 调用 Worker 基类的构造函数，并保存配置。
2. 如果 PyTorch 分布式环境尚未初始化，则进行初始化，包括设置通信后端和进程组。
3. 为 FSDP 创建设备网格，用于模型参数的分片。
4. 如果启用 Ulysses 序列并行，则初始化其设备网格。
5. 根据传入的 `role` 参数设置 Worker 的具体角色（actor, rollout, ref）。
6. 根据 Worker 角色配置 profiler，用于性能分析。
7. 配置 parameter offload 和 optimizer offload。
8. 为 Actor，Rollout 和 Reference 分别 normalize batch size。

第 8 步中配置了非常多的 batch size；verl 的 batch size 参数满天飞，虽然我个人认为名字基本是准确的，但是由于名字太像了，一定要做出一些区分。事实上，参数分析我们有单独的文档，我先把一部分内容提前公布了。

1. `data.train_batch_size`：在一次完整的 PPO 迭代（从 rollout 到 train）中，从数据集中采样并用于生成 experience 的总样本数量，决定了每次 policy 更新所依据的数据量。
2. `actor_rollout_ref.actor.ppo_mini_batch_size`：这个参数的名字其实是准确的，因为 mini batch SGD 就是数据到达了一个 mini batch 就更新一次模型参数。在 verl 中，模型会在数据累积到一个 mini batch 后更新一次参数。
3. `actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu`：这里其实是 gradient accumulation 的参数。由于一个 mini batch 的数据量可能仍然太大，无法一次性前向和反向传播，因此需要将其进一步拆分为 micro batch。每个 micro batch 会计算一次梯度并且累计，但是不会立刻更新模型参数。处理完整个 mini batch 后，才用累积的梯度进行一次参数更新。

此外，在 verl 中，由于 verl 强调 SPMD 策略，可以理解为每个 RL worker 所占据的每个 GPU 上希望进行完全一致的操作，所以 verl 会要求每个 GPU 的 micro batch size 相同。因此，verl 会检查 train batch size / gpu 是否整除 [(ref)](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/trainer/ppo/ray_trainer.py#L363)，如果不整除，则报错。这个设定其实完全没必要；对于 rollout 而言，SGLang 完全不需要发送的请求数量整除 DP 或者 TP size，更何况直接要整除 gpu 数量呢？但是，因为 verl 会用 all gather 从 rollout 的每个 worker 里收集数据，这就要求 rollout 的每个 worker 上分到的数据一致。更进一步，为了 SPMD，又要求 rollout 的每个 gpu 上分到的数据一致。最终，这就导致 verl 的 train batch size 必须整除 gpu 数量；在 GRPO 下是 real train batch size 需要整除 n gpus，等于 train batch size * sampling params 中的 n。

区分好 mini batch 和 micro batch 后，我也是最近才明白 PPO 中是如何维护 on policy 的。我之前一直以为我们都是在做严格 on policy 的训练，但是一个 train batch size 下有好几个 mini batch，似乎第一个 mini batch 结束之后，目标策略（target policy，被训练的 policy）和行为策略（behavior policy，用于在环境中采样的 policy）就不一致了。一次采样会训练很多个 mini batch，从第一个 mini batch 结束就不是 on policy 了。事实也是如此，我们注意到 PPO 的 loss function：

$$ L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min(r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t) \right] $$

其中的 $r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)}$，这是一个对优势函数的矫正比例，而 $\hat{A}_t$ 就是 advantage。对于 LLM 的 PPO 而言，$\pi_{\theta_{old}}(a_t | s_t)$ 代表着采样时 behavior policy 在给定 $s_t$ 时，选择 $a_t$ 的概率，而 $\pi_\theta(a_t | s_t)$ 就是 target policy 在训练中的每一步给定 $s_t$ 时，选择 $a_t$ 的概率。对 LLM 而言，`s_t` 是 prompt 前缀，而 `a_t` 仅仅是 prompt 后的那一个 token。这一概率其实就是 inference 得到的 log probs；我们将收集得到的 (prompt, action) 分别经过 target policy 和 behaviour policy 得到 log probs，然后二者 log probs 相减再取对数，就是矫正项的值。从而，即便第一个 mini batch 之后 target policy 就已经和 behaviour policy 不一致了，仍然可以通过 log probs 进行矫正，也即 importance sampling。

这样一来，又有了两个问题：log probs 应该如何得到？实际上每次采样时都是发送给 rollout 固定数量的 requests，如果每个 (prompt, action) 对都会计算一次 loss 的话，岂不是更长的 sequence 会计算更多次？

对于第一个问题，这又是经典的[精度问题](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/readme.md#introduction)。如同我在链接到的文章中所说的，rollout engine 目前只有采样得到的 token 能用，而得到的 log probs 以及 reward 精度都不够，不能用于训练。behaviour policy 和 target policy 为了做 importance sampling 所需的 log probs 都得用 training engine 重算。不过要算起来也不麻烦，在第一个 mini batch 启动前，这时候 target behaviour 是一致的，重算 log probs 并且存下来即可。

对于第二个问题，的确如此。一条很长的 prompt + answer 序列确实会产生非常多的 (prompt, action) 对，其中每个对都可以看作一个 (state, action) 对。而且理论上每个这样的 (prompt, action) 对都会参与 Loss 的计算。这确实可能导致长序列中的 token 会在 Loss 计算中占据更大的比例，让模型过度关注长序列的优化，而对短序列的优化不足。不过，verl 的 rollout engine 会自动对每个 (prompt, action) 对进行加权，从而让长序列和短序列的 token 在 Loss 计算中占据相同的权重。为了缓解这种情况，有很多相关方法：


样本加权方法

序列级别加权： 一种直接的方法是在计算 Loss 时，给来自不同序列的样本赋予不同的权重。例如，给每个完整序列一个固定的权重（比如 1），然后将这个权重均匀分配给该序列中的每个 (prompt, action) 对。这样，无论序列多长，它对总 Loss 的贡献都相同。如果一个序列有 N 个 token，那么每个 (prompt, action) 对的权重就是 1/N。

按长度分桶： 在数据收集后，可以根据序列长度对样本进行排序，并尝试将相似长度的序列放入同一个 mini-batch。这有助于提高计算效率，因为可以减少 padding，但对于解决 Loss 贡献不均衡的作用有限。

固定 Token 数量的批次： 最常见且有效的方法是构建批次时，不固定样本数量，而是固定批次中的总 token 数量。这样，一个 mini-batch 可能包含 4 条长序列，也可能包含 40 条短序列，确保每次更新时处理的总计算量和梯度来源的总 token 数是恒定的，从而缓解长短序列的不均衡问题。

Loss 归一化：在计算每个 mini-batch 的 Loss 时，可以将其除以该 mini-batch 中实际的 token 数量。这确保了 Loss 值不会仅仅因为批次中包含了更多 token 而增大，从而为不同大小的 mini-batches（如果不是按固定 token 数构建）提供一个公平的比较基础。

截断：设定一个 max_length 参数，限制模型生成的最大 token 数量。虽然这不直接解决已有长序列的权重问题，但可以防止生成过长的序列，从而限制极端不均衡的发生。



whatever，解释了这么多，顺着理解 verl 的框架进一步学习了 RL 算法和系统，这里其实和 multi-turn 都还没有关系，我们还是回到 `ActorRolloutRefWorker` 的源码上。


 ActorRolloutRefWorker.__init__ 源码 

```
/* Detailed source-code truncated for AI context efficiency. */
```



### [`ActorRolloutRefWorker._build_model_optimizer()`](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/workers/fsdp_workers.py#L177)

这部分源码和类写的还是很直白的，不用太多解释：

1. 初始化 Hugging Face 配置，获取 Generation Config，并设置模型的数据类型（Actor 使用 fp32，Reference 使用 bf16）。
2. 使用 Hugging Face 的 `AutoModelForCausalLM` 或 `AutoModelForVision2Seq` 从预训练模型加载基础模型。
3. 应用各种优化技术，包括 Liger kernel、融合 kernel、梯度检查点、LoRA 等。
4. 根据配置选择 FSDP 或 FSDP2 策略，将模型封装到分布式训练框架中，支持参数分片和混合精度训练。
5. 如果当前 Worker 是 Actor 角色，则初始化 AdamW 优化器和学习率调度器。


 ActorRolloutRefWorker._build_model_optimizer 源码 

```
/* Detailed source-code truncated for AI context efficiency. */
```



这里代码很直白。有一个点值得单独拎出来讲一下：仔细观察 `actor_module` 的 dtype，直觉告诉我，`actor_module` 的 dtype 应该是 bf16 的，而 gradient 和 optimizer 的 dtype 是 fp32 的。可是 `actor_module` 的 default dtype 被设为了 fp32，然后从 fp32 load 了模型。实际上这是因为 pytorch 的各种 optimizer 都是直接和 parameter 绑定的，用 bf16 的 parameter 初始化的 optimizer 也是 bf16。所以 model 先 load 了 fp32，然后初始化 optimizer 作为混合精度，最后把 model 转成 bf16。

### [`ActorRolloutRefWorker._build_rollout()`](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/workers/fsdp_workers.py#L394)

这是对我而言最清晰的地方，实际上也是最熟悉的。在这里终于引入了 SGLang：

1.  **设备网格创建**：为 Rollout 创建推理张量并行（`infer_tp`）设备网格。
2.  **SGLang Rollout 构建**：导入并实例化 `SGLangRollout` 和 `FSDPSGLangShardingManager`。`FSDPSGLangShardingManager` 负责在 FSDP 训练格式和 SGLang 推理格式之间转换模型权重。


ActorRolloutRefWorker._build_rollout 部分源码

```
/* Detailed source-code truncated for AI context efficiency. */
```



### [`SGLangRollout.__init__()`](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/workers/rollout/sglang_rollout/sglang_rollout.py#L208)

事已至此，再往下看一层 SGLang 具体的初始化：

1. 调用父类构造函数并设置配置和设备网格。
2. 通过 `_initialize_tools()` 初始化工具 schemas、map 和解析器，支持 Multi-turn 对话中的工具使用。
3. 初始化 SGLang 推理所需的分布式环境。
4. 通过 `_verify_config()` 验证模型配置。
5. 通过 `_init_inference_engine()` 初始化 SGLang 推理引擎。
6. 通过 `_init_sampling_params()` 初始化生成序列的采样参数。
7. 设置 Tokenizer 和 padding token ID。


SGLangRollout.__init__ 部分源码

```
/* Detailed source-code truncated for AI context efficiency. */
```



### [`SGLangRollout.AsyncEngine`](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/workers/rollout/sglang_rollout/sglang_rollout.py#L124)

关于 `SGLangRollout` 调用 tool 的部分，我们在下文的训练循环中再展开，这里先讨论完 SGLang 的初始化。为了调用 SGLang engine 的接口，verl 进行了一层封装，实现了我们对 SGLang 除开 rollout 之外的所有接口：

1. release and resume memory occupation：在训练时释放掉显存占用并在训练后恢复。
2. update weights from tensor：训练结束后更新模型权重。
3. flush cache：模型参数更新后刷新 KV cache，因为之前的 KV cache 已经失效了。

这里涉及到了非常深入的内存管理问题，读者对 SGLang engine 在 verl 里的显存管理感兴趣，欢迎阅读标哥的博客 [optimizing Memory Usage in verl](https://hebiao064.github.io/rl-memory-management)，写的非常深入浅出。


SGLangRollout 何时需要 flush cache

这一部分内容需要单独拎出来讲讲。SGLang engine 的 release 和 resume 需要保留 CUDA Graph，否则 rollout 效率会大幅降低。因此，我们基于 tom 的 [torch_memory_saver](https://github.com/fzyzcjy/torch_memory_saver) 实现了独立的显存管理。简单来说，我们有：

1. `pause`；保留 mem savor 作用域内指定 tensor 的 virtual address，但是将其 physical memory 释放回显存管理器。
2. `resume`；将先前 `pause` 的 tensor 重新申请一块 physical memory，并将其 virtual address 映射到新的 physical memory。

注意，整个 pause 和 resume 的过程中，tensor 的 virtual address 不会发生变化，只是这块 virtual address 映射到的 physical memory 改变了。因此，CUDA Graph 并没有失效，不变的 virtual address 让计算流仍旧可以正常执行。

verl 内的 `release_memory_occupation` 和 `resume_memory_occupation` 就是基于 `pause` 和 `resume` 实现的。听上去是个完美的故事，我们甚至实现了 [mutli-stage 的显存管理](https://github.com/fzyzcjy/torch_memory_saver/pull/20)，能够独立 release 和 resume kv cache 和 model weights。

不过，对于 kv cache 而言，在 kv cache 被 release 掉之后，实际上 kv cache 的 tensor 仍旧保留，只是其 virtual address 映射到的 physical memory 被释放了。与此同时，radix tree 仍旧索引着整个 kv cache。当 kv cache 被 resume 之后，一方面之前物理内存上之前的 kv cache 已经不复存在了，另一方面模型的参数也被更新。出于这两点，我们一定要使用 flush cache 接口来刷新 kv cache 的索引（radix tree）。

这里又有个非常有趣的设计。乍一想 kv cache 的管理这么麻烦，还要 flush，为什么不直接 delete kv cache 以及 delete model weights 再重新初始化呢？显然，这样没法利用已有的 cuda graph，非常消耗时间。保留 virtual address 不变但是更换 physical memory 的方案，让 verl 能够持续利用已建好的 cuda graph。

最后一个问题，一共要几次 flush cache 呢？我个人理解，在一整个 training engine 被 pause，resume 然后 update weights 的过程中，必须要有一次 flush cache 来刷新 kv cache 的索引，只是 verl 当中为了保险，刷新了很多次罢了。



SGLangRollout.AsyncEngine 源码

```python
class AsyncEngine(sglang.srt.entrypoints.engine.Engine):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # default to use dummy load format, which need to reload weights in first time
        self._need_reload = True

    async def release_memory_occupation(self):
        """Release GPU occupation temporarily."""
        obj = ReleaseMemoryOccupationReqInput()
        return await self.tokenizer_manager.release_memory_occupation(obj, None)

    async def resume_memory_occupation(self):
        return await self.tokenizer_manager.resume_memory_occupation(obj, None)

    async def update_weights_from_tensor(
        self,
        named_tensors: List[Tuple[str, torch.Tensor]],  # noqa: UP006
        load_format: Optional[str] = None,
        flush_cache: bool = True,
    ):
        """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false
        to avoid duplicated cache cleaning operation."""
        obj = UpdateWeightsFromTensorReqInput(
            serialized_named_tensors=[MultiprocessingSerializer.serialize(named_tensors) for _ in range(self.server_args.tp_size)],
            load_format=load_format,
            flush_cache=flush_cache,
        )
        return await self.tokenizer_manager.update_weights_from_tensor(obj, None)

    async def flush_cache(self):
        return await self.tokenizer_manager.flush_cache()
```



### [`SGLangRollout._init_inference_engine()`](https://github.com/volcengine/verl/blob/e67ee86f8b94bfa141da95402a254966733cba08/verl/workers/rollout/sglang_rollout/sglang_rollout.py#L325)

`SGLangRollout._init_inference_engine()` 初始化了封装的 `AsyncEngine`。


SGLangRollout._init_inference_engine 源码

```
/* Detailed source-code truncated for AI context efficiency. */
```



这里最值得一提的是，SGLang engine 并没有严格实现 verl 所希望的 SPMD 模式（每个 GPU 上的进程完全一样），而是采用了 mock 的 SPMD。举例来说，假设 tp size = 4，按照 verl 的设计，应该要 4 张 GPU 上每个都运行一个相同的 SGLang engine。实际上的实现是在 GPU 0 上启动一个进程占据全部 GPU，而 GPU 1 2 3 上仅仅保留一个空进程 `None`。虽然 verl team 起初设定中认为严格的 SPMD 意义巨大，但实际使用中，我们认为 mock 的 SPMD 已经足够满足性能需求。

【TODO】 这么描述可能不严谨。

### [`TaskRunner.run()`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/main_ppo.py#L64)

往下走了这么多层，我们终于能够继续回到 `TaskRunner` 类。😭

【TODO】上文其实主要是 Actor Rollout，还没有具体说 Actor 的 training forward and backward。以及 Reference，reward 和 critic 的 training forward and backward。

1. 加载、解析和验证训练任务的配置（使用 `OmegaConf`），确保所有参数的正确性和一致性。
2. 将模型文件从远程路径复制到本地，确保所有 Worker 都可以访问。
3. 组件初始化：
    * 初始化 Tokenizer 和 Processor，用于文本和多模态数据的处理。
    * 根据配置中指定的 Actor 策略（如 `fsdp` 或 `megatron`），动态选择相应的 Worker 类（例如 `ActorRolloutRefWorker` 和 `CriticWorker`），并确定使用的 `RayWorkerGroup` 类型。
    * 定义 Ray 资源池的规格和角色到资源池的映射，用于 GPU 资源的分配和管理。
    * 加载用于训练和验证的奖励模型。
    * 创建训练和验证数据集，以及训练数据采样器。
4. 创建 `RayPPOTrainer` 实例，它是管理所有计算资源和训练流程的中央协调器。
5. 调用 `RayPPOTrainer` 的 `init_workers()` 方法，将配置的 Worker 类实例化到 Ray 集群的 GPU 上，为实际计算做准备。
6. 调用 `RayPPOTrainer` 的 `fit()` 方法，启动 PPO 训练循环。


TaskRunner.run 源码

```
/* Detailed source-code truncated for AI context efficiency. */
```



### [`RayPPOTrainer.__init__()`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/ppo/ray_trainer.py#L277)

1. 保存传入的配置对象、tokenizer、processor、角色到 Worker 的映射、资源池管理器以及 WorkerGroup 类。
2. 根据配置启用或禁用 Critic、Reference Policy、Reward Model 和 Hybrid Engine 等功能组件。
3. 调用 `_validate_config()` 方法验证配置的合理性。
4. 存储训练和验证数据集、collate 函数和训练数据采样器。


RayPPOTrainer 源码

```
/* Detailed source-code truncated for AI context efficiency. */
```



### [`RayPPOTrainer.init_workers()`](https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/trainer/ppo/ray_trainer.py#L715)

`init_workers()` 函数负责在 Ray 集群上实例化和初始化 ActorRollout、Critic、Reference Policy 和 Reward Model Workers。

1.  **创建资源池**：通过 `ResourcePoolManager` 创建 Ray 资源池。
2.  **初始化资源池到类的映射**：为每个资源池创建一个字典，用于存储不同角色 Worker 的 `RayClassWithInitArgs` 包装器。`RayClassWithInitArgs` 用于延迟初始化 Worker，存储了 Worker 的类和初始化参数。
3.  **创建不同角色的 Worker 的 `RayClassWithInitArgs` 实例**：根据配置启用情况，为 ActorRollout、Critic、Reference Policy 和 Reward Model 创建对应的 `RayClassWithInitArgs` 实例。
4.  **初始化 WorkerGroup**：遍历所有资源池，将同一资源池中的多个 Worker 类通过 `create_colocated_worker_cls` 组合成一个共置类，然后实例化 `RayWorkerGroup`。`RayWorkerGroup` 负责在多个 GPU 上启动多个 Worker 实例。最后调用 `spawn()` 方法在 Ray 中实际创建 Worker 实例。
5.  **初始化各个 Worker**：根据角色从创建的 WorkerGroup 字典中获取对应的 WorkerGroup，并调用其 `init_model()` 方法，按照依赖关系依次初始化不同的 Worker 模块。ActorRollout Worker 通常最后初始化以优化内存使用。


RayPPOTrainer.init_workers 源码

```
/* Detailed source-code truncated for AI context efficiency. */
```

---

## File: rlhf/verl/readme.md

# HybridFlow veRL 原文浅析

众所周知，我一直在 SGLang team 负责端茶倒水 + RLHF。对于后者，这段时间我一直在学习 veRL 的整体框架，hybrid engine 的想法真令人眼前一亮。其实 SGLang 团队一直有成员在负责这块工作的推动，但是由于 DeepSeek 模型支持的工作强度巨大，加之团队成员的精力有限，还没能将这系列的改动 up stream 到 main branch 上。不过，随着 grok 的发布，我们会将 SGLang 对于 hybrid engine 的支持，快速 up stream 到 main branch 上，这里浅浅预告一番。

回到这篇文章本身，这是 SGLang-veRL 系列工作的开端笔记，也督促自己更加全面地学习 training and inference co-design。

PS：之前也有基于 nemo-aligner 和 OpenRLHF 做一些解析，欢迎大家参考，也感谢这些框架作者的精彩贡献。大家都是写框架的人，其中的酸甜苦辣，自不必多说。

- [浅析主流 Alignment 算法与 NeMo-Aligner 框架](https://zhuanlan.zhihu.com/p/5220718268)
- [浅析以 OpenRLHF 为代表的 post-training 系统的计算流程](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/OpenRLHF/readme.md)

## Single-Controller vs Multi-Controller

在梳理 introduction 之前，我们先来引入 veRL 所依赖的一大重要概念——single controller 和 multi controller。

**在一个复杂的工作流程中，single controller 只有一个程序负责管理，而其他的子模块只负责执行。所有控制逻辑都写在唯一 controller 上，实现简单，便于调试。然而，single controller 所承担的控制压力巨大，一来，通讯强度大而效率堪忧，二来，倘若 single controller 崩溃，整个系统将彻底失效。反过来，multi controller 则有多个控制程序来管理不同的子模块，每个子模块仍旧只负责执行自己的功能。如此以来，单个控制程序的管理压力降低，系统更加鲁棒可扩展。然而，控制逻辑分散在多个程序中，实现复杂，难以调试。**

有了这个直观的理解，我们回顾下 [PPO 的粗略工作流程](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/OpenRLHF/readme.md#%E6%9B%B4%E6%96%B0%E6%B5%81%E7%A8%8B)。注意，actor 在 RLHF 会进行 auto-regressive decoding，而 critic, reward 和 reference 则只会 prefill，不会 decode。所以，我们将 actor 的推理特定称为 rollout，而其他模型的推理称为 inference。

1. 准备一个 batch 的 prompts；
2. 将这个 batch 的 prompts 输入给 Actor，rollout 得到 responses；
3. 将 prompt + responses 输入给 Critic/Reward/Reference，进行 inference，分别计算得得到 values、reward 和 log probs，将这些整合称为 experiences；
4. 根据 experiences 多轮计算 actor loss 和 critic loss 并更新 Actor 和 Critic。

如此以来，一个很自然的想法是用 single controller 来管理整个工作流程，然后每个子模块（Actor、Critic、Reward、Reference）再各自由一个 single controller 来管理。然而，这种 naive 的实现里，两层控制结构都是 single controller，因此系统内的通讯 overhead 非常大。一个可能不太直观的点是——最高层负责总调度的 single controller 承担的通讯压力，反而比每个子模块的 single controller 要小。

可以简单理解，rollout engine 只需要把 prompt + response return 给上层 controller 就好了，上下两层的通讯量最多 M 级别；而 rollout engine 内部的通讯就非常大了，不然怎么会把 NV Link 接近 T 为单位的通讯速度都拉满呢？

因此，每个子部件内的 controller 承担的通讯压力反而是更大的，如果这一层还都是 single controller 的，可以想见效率堪忧。事实也是如此，主流的训练引擎都是 multi controller 的，譬如 FSDP、Megatron 和 DeepSpeed。既然 Actor 的 Training Engine 是 multi controller 的（具体来说是 SPMD 的），那么 Actor 的 rollout engine，是否也该是 multi controller 的？从理论上来讲，是的，直觉告诉我们，SPMD 的 training engine 和 SPMD 的 rollout engine 相互通讯，会比和 single controller 的 rollout engine 通讯要高效得多。这个直觉也禁得起推敲。SPMD 的 training engine 与 rollout engine 都采用了分布式、多控制器的设计，各节点之间可以直接点对点通信，这样可以充分利用高速互联（例如 NVLink 或 InfiniBand），实现数据并行传输和计算。反过来，training engine 采用 SPMD 而 rollout engine 采用 single controller 的话，所有训练节点的数据都必须汇总到某个控制节点上，才可以完成 training engine到 rollout engine 间的参数更新，这就不可避免地引入了通信瓶颈和单点故障风险。

因此，SPMD 的 training engine 和 SPMD 的 rollout engine 简直是绝配。然而，由于历史原因，目前主流的 rollout engine 还是以 single controller 为主的。所以，将 SGLang 由 single controller 改为 SPMD 是我们的一个重要工作目标。事实上，我们已经有了成熟的 PR，可以[参考这个 branch](https://github.com/fzyzcjy/sglang/tree/feat/overall_verl)。

注：上面描述的 single controller 和 SPMD 模式只是经验上的 practice，但实际上他们描述的是：single controller 主要关注于控制流是否单点；而 SPMD 模式更关注的是分布式场景下数据执行流。即使是 single controller，也是可以使用 SPMD 模式进行数据流执行的。

总之，这篇文章的开篇花费了巨大的篇幅来简述 single controller 和 multi controller，以及为什么 SGLang 需要从 single controller 改为 multi controller。理解了这些概念后，我们可以正式进入 veRL 的 introduction 了。

## Introduction

正如前文所提到的，multi-contoller 能够有效的降低通讯压力，提升系统鲁棒性。然而，如果最顶层的 controller 也是 multi-controller 的，其实对用户会非常复杂。在一个 controller 内代码的修改，需要将所有 dependency 都修改一遍。很难想象读 ML researcher 会愿意接受这一点。

因此，veRL 在上层暴露出 single controller 的接口，并进行完善的封装。用户能够基于算法设计，自由组合并行策略（3D 并行、ZeRO 还有 FSDP），直接对子模块进行拼装；而在每个子模块内部，采用 multi-controller，提供强劲的效率。当然，可能更改子模块就会相对麻烦。

**有了 single controller 和 multi-controller 的概念后，这里引入veRL 的第二个核心概念：hybrid engine。在 RLHF 流程中，actor model 的 generation 和 rollout 占据了绝大多数运行时间（在 veRL 是 58.9%）。并且，由于 PPO 是 on-policy 算法，经验（experiences）必须来自于被 train 的模型本身，因此，rollout 和 training 是必须串行的。如果这两者使用不同的资源组，比如 rollout 用 2 张卡，而 training 用 4 张卡，rollout 的时候 training 的资源闲置，training 的时候 rollout 的资源闲置，无论如何都会浪费大量的计算资源。由此，veRL 将 training 和 rollout engine 放置在同一个资源组中串行执行。training 时，将 rollout engine 的显存回收（offload 到 CPU 上 或者直接析构掉），rollout 时，再将 training engine 的显存释放掉。这种将 actor model 的不同 engine 放置在同一个资源组上的方案，就称为 hybrid engine。**

注意到，除开 hybrid engine 之外，类似共用资源组的方法还有 collocate。在讲述 collocate 策略之前，我们回顾下四个子模块分别需要什么 engine：

1. actor model 需要 training engine 和 rollout engine。前者是用现代 training engine，比如 Megatron 或者 FSDP，后者得用现代推理引擎，比如 SGLang 或者 vllm 作为 rollout engine。这里思考一个小问题，为什么不能拿着 training engine 得到的 logits 做 sampling 然后 decode，貌似也可以用去 rollout？简单来说，太慢了，用训练引擎做 decode 的效果自然不如专用的推理引擎。
2. critic model 需要 training engine 和 inference engine。前者还是是现代的训练引擎，但是后者，可以用现代的推理引擎的高效 prefill 来得到 value 么？其实不能，critic model 的 inference 会直接复用 training engine 的 forward 来得到 value，所以 critic 的 inference engine 和 training engine 其实是同一个。其中的原因在此旧事重提：

> 推理引擎的 kernal fusion 和 training engine 差距不小，batch size 不一样时，推理请求 dispatch 到不同的 kernal 上，然后 numerical 误差逐层累计，到了 log probs 这层就到了不可忽视的程度了。这个问题在 bert 时代就有了，training engine 和 inference engine 的精度差异无法规避，而且全心来搞一两个月可能内都没法修复。所以现在推理引擎在 RLHF 中更多是加速 sampling，reward 和 embedding 还得用训练脚本来算，可能得半年后花好几个月研究研究这个问题。

3. reference model 和 reward model 只需要 inference，因为二者不需要训练，但是如同我之前提到的一样，用现代推理引擎得到的 log probs 和 reward 的精度不如用现代训练引擎得到的精度，所以这里选择用 training engine 的 forward 来做 inference，得到 log probs 和 reward。

有了这些认识，我们再来看 collocate 策略。collocate 策略将 actor 的 training engine 和 reference 的 inference engine 放置在同一个资源组上，将 critic 的 training/inference engine 和 reward 的 inference engine 放置在同一个资源组上，最后单独放置 actor 的 rollout engine。

与此相对的是，hybrid engine 单独强调了将 actor model 的 rollout engine 和 training engine 放置在同一个资源组上，而 collate 则强调的是不同子模块之间的。可以见到，hybrid 共用资源组的 engine 都属于 actor，二者区别显著更大，更容易 OOM。当然，collocate 和 hybrid engine 都可以提高 GPU 利用率，不过速度自然会有所损失。

总之，这两个概念是 veRL 最强大的贡献。实际上 veRL 还提供了一套基于贪心搜索的 placement（资源组）分配算法，不过按照作者描述，这个 feature 效果不太诱人，现在比较冷门了。

## Background

这一部分是一些背景补充：

- 现代的分布式训练框架（Megatron-LM，MegaScale，DeepSpeed）都支持了 3D 并行，也即 DP PP TP。LLM serving 也有对应的策略和概念，不过其中只有模型参数和 KV cache 会被 sharded，不存在优化器和梯度的需求。
- Actor model 训练是 compute bounded 的，通常倾向于更高的 TP 或者 PP size。而采用同样的高 TP 或者 PP size 的 rollout engine 则效率不佳，实际上 rollout engine 通常希望加大 DP size。所以，为了提高两个阶段各自的效率，actor model 的 training engine 和 rollout engine 会采用不同的并行策略。然而，不同的并行策略导致两个阶段之间的参数更新需要 resharding，导致通讯和访存开销显著。

## Hybrid Engine Performance

- 为了提供灵活的并行策略供用户组合，veRL 提供了 `3DParallelWorker`，`FSDPWorker` 和 `ZeROWorker` 三个基类，并用子类支持各种并行策略。为了做到 Training Engine 和 Rollout Engine 之间的 parameter update，veRL 提供了 8 种 transfer protocols，包括但不限于 `3D_PROTO`, `DP_PROTO`, `ONE_TO_ALL`  等等。
- 费力构造好的 Hybrid Engine 自然要大显身手，veRL 鼓励用户采用 Hybrid 的方式来控制 actor 的 rollout 和 training。不过目前 Hybrid 使用的 SPMD rollout engine 相比起单独放置的 single controller rollout engine 会有一定的推理速度下降。
- veRL 采用 mixed precision 训练，模型参数采用 BF16，梯度和 Adam 优化器采用 FP32，actor 的 rollout 和 其他模型的 inference 采用 FP16。

这一部分记录一些 placement 的经验。希望以后我能有实力基于机器的物理属性直接分析得到这些结论。主要对比这四种策略：



1. fully collocate：所有的子模块都放在同一个资源组上，也即 DeepSpeed-Chat。
2. hybrid：actor 的 rollout engine 和 training engine 放在同一个资源组上，其他子模块进行部分 collocate，这是 veRL 所提出的策略，但是 veRL 原文其实给出了一个搜索方法，能够贪心搜索所有的策略，选择出最佳策略。
3. split collocate：actor 的 training engine 和 reference 的 inference engine 放在同一个资源组上，critic 的 training/inference engine 和 reward 的 inference engine 放在同一个资源组上；最后单独放置 actor 的 rollout engine，这是 OpenRLHF 和 NeMo-Aligner 的默认策略。
4. stand alone：所有子模块都单独放置，早期 OpenRLHF 会这么做，现在自然不会了。

这里直接给结论了：
- 16 ~ 64 GPUs 范围内，fully collocate 效果最好；
- 96 ~ 128 GPUs with 34B models 或者 96 GPUs with 13B models，split collocate 效果最佳。
- 暴力搜索总可以得到最佳策略，但是搜索成本过大了 😂
- 每个子模块都可以高强度利用计算资源时，fully collocate 效果最好。
- 在大规模训练时，actor 和 critic 分开放置效果更佳。

最后，veRL 顶层的 single controller 还带了一个好处，方便利用 rule-based reward，这就是一个新的故事了。