SGLang 后端 ============== 最后更新时间：2025 年 5 月 31 日。 **作者：SGLang RL Team，按姓氏字母顺序排列** `Jingyi Chen `_、`Yitong Guan `_、`Zhuobin Huang `_、`Jiajun Li `_、`Ji Li `_、`Shenggui Li `_、`Junrong Lin `_、`Xiang Long `_、`Rui Lu `_、`Jin Pan `_、`Shuai Shi `_、`Yushen Su `_、`Xinyuan Tong `_、`Chendong Wang `_、`Hanchen Zhang `_、`Haoran Wang `_、`Yongan Xiang `_、`Chengxing Xie `_、`Yuhao Yang `_、`Jinwei Yao `_、`Qiaolin Yu `_、`Yuzhen Zhou `_、`Chenyang Zhao `_ 简介 ------------ `SGLang `_ 是一个开源的、最先进的推理服务引擎，已被 xAI 完全采用，以支持 Grok 在研究和部署过程中的所有推理需求。 目前，verl 在部署阶段完全支持使用 SGLang 作为推理引擎。作为部署引擎，SGLang 提供了与 vLLM 相同的功能覆盖，包括内存节省和多节点部署功能。安装 verl 和 SGLang 后，只需在启动脚本中添加 ``actor_rollout_ref.rollout.name=sglang`` 即可在两个推理框架之间无缝切换。 此外，SGLang 团队还在积极支持多轮智能体 RL、VLM RLHF、基于服务器的 RLHF 和部分部署等功能。您可以在 `Roadmap 跟踪 `_ 中跟踪相关开发进展。 安装 ------------ 请务必按照以下命令安装 SGLang 和 verl。 .. code-block:: bash pip install --upgrade pip # 当前版本为 0.4.8，可能会随时更新，请参考 `setup.py` 中指定的最新版本 pip install -e ".[sglang]" 您可以检查您的环境中是否已安装以下依赖项： .. note:: - **PyTorch**: 2.6.0+cu124 - **CUDA**: 12.4 - **flashinfer-python**: 0.2.5+cu124torch2.6 - **SGLang**: 0.4.6.post5 - **sgl-kernel**: 0.1.4 在单台机器上使用 SGLang 作为 PPO 训练的推理后端 ------------------------------------------------------------------------- 我们使用 Qwen/Qwen2-7B-Instruct 在 gsm8k 数据集上进行简单测试。 1. 运行以下命令准备 gsm8k 数据集： .. code-block:: bash python3 examples/data_preprocess/gsm8k.py 2. 运行以下脚本在具有 4 个 GPU 的单台机器上进行 PPO 实验： .. code-block:: bash export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK=True PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ data.train_files=$HOME/data/gsm8k/train.parquet \ data.val_files=$HOME/data/gsm8k/test.parquet \ data.train_batch_size=4096 \ data.max_prompt_length=4096 \ data.max_response_length=4096 \ actor_rollout_ref.rollout.name=sglang \ actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \ actor_rollout_ref.actor.optim.lr=1e-6 \ actor_rollout_ref.actor.ppo_mini_batch_size=64 \ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \ actor_rollout_ref.model.enable_gradient_checkpointing=True \ actor_rollout_ref.actor.fsdp_config.param_offload=True \ actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \ critic.optim.lr=1e-5 \ critic.model.path=Qwen/Qwen2-7B-Instruct \ critic.ppo_micro_batch_size_per_gpu=4 \ critic.model.fsdp_config.param_offload=True \ critic.model.fsdp_config.optimizer_offload=True \ algorithm.kl_ctrl.kl_coef=0.001 \ trainer.logger=console \ trainer.val_before_train=False \ trainer.n_gpus_per_node=4 \ trainer.nnodes=1 \ trainer.save_freq=-1 \ trainer.test_freq=10 \ trainer.total_epochs=15 2>&1 | tee verl_demo.log 为什么要导出 SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK？ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 1. ``verl`` 在 rollout 期间初始化 ``SGLangRollout`` 模块，用于评估/生成样本。 2. ``SGLangRollout`` 会初始化 ``Engine``，并进一步初始化 ``torch.distributed.DeviceMesh``，用于支持张量并行（TP）。 3. ``DeviceMesh.init()`` 内部会检查所有参与设备的可用 GPU 内存。如果差异过大（超过约 10%），它会直接报告错误，以避免初始化失败或死锁。 为什么会出现 GPU 内存不一致？ """""""""""""""""""""""""""""""""" **1. Ray 分布式 Actor 加载模型的时机不同** ``verl`` 使用基于 Ray 的多进程、多 GPU 并发训练。每个 ``WorkerDict`` 可能在不同时间被调用： .. code-block:: python self.rollout = SGLangRollout(...) 不同的 worker 在不同时间初始化模型 → 导致内存使用量不同。 **2. 延迟初始化导致内存偏差** 有些 worker 比其他 worker 更早开始加载模型/进行推理（例如，``generate_sequences()``、``compute_log_prob()``）。 早开始的 worker 已经占用了 GPU 内存 → 后开始的 worker 仍有空闲内存 → 出现内存差异。 **3. SGLang 的 TP 初始化使用“所有设备广播”，但没有统一的释放时机** 尽管 ``SGLangRollout`` 可能只涉及部分 GPU，但其 ``Engine`` 初始化会调用 ``torch.distributed.init_process_group()`` 并广播权重，因此： - 非 rollout 的 GPU 也会加入通信。 - 之后，``DeviceMesh`` 初始化会因为“内存不一致”而失败。 **4. 不同的 FSDP/TP 加载行为也会导致不匹配** 如果使用： .. code-block:: bash actor.fsdp_config.param_offload=True ref.fsdp_config.param_offload=True 那么一些 worker 将参数保留在 CPU 上，而其他 worker 则已将其分片到 GPU → 导致内存布局不对称。 在多台机器上使用 SGLang 作为 PPO 训练的推理后端 ------------------------------------------------------------------------------ SGLang 也支持在 IPv4 和 IPv6 场景下运行 verl 基于 RAY 的跨机器推理。在下面的脚本中，我们使用 TP=16 进行跨机器推理。假设我们有两台互联的机器：node0 的 IP 地址是 10.94.16.4，node1 的 IP 地址是 10.94.16.5。 1. 在 node0 上启动 Ray： .. code-block:: bash ray start --head --dashboard-host=0.0.0.0 您将看到以下提示： .. code-block:: bash Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details. Local node IP: 10.94.16.4 -------------------- Ray runtime started. -------------------- Next steps To add another node to this Ray cluster, run ray start --address='10.94.16.4:6379' 2. 让 node1 加入 Ray 集群： 在 node1 上运行以下命令： .. code-block:: bash ray start --address='10.94.16.4:6379' 运行以下命令确认 Ray 集群现在有两个节点： .. code-block:: bash ray status 您可以看到集群有两个节点，共 16 个 GPU： .. code-block:: bash ======== Autoscaler status: 2025-04-09 09:25:37.694016 ======== Node status --------------------------------------------------------------- Active: 1 node_ef382ffd687d8f6b060c1b68e63ada7341b936fe5b1901dd04de1027 1 node_1eb4d7d07e793114c23a89d1a41f1f76acf6ef5b35af844a4ee8e4ba Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/360.0 CPU 0.0/16.0 GPU 0B/3.39TiB memory 0B/372.53GiB object_store_memory 3. 运行以下脚本，使用 TP=16 在 2 台机器上（共 16 个 GPU）训练 meta-llama/Llama-3.1-8B-Instruct： .. code-block:: bash DATA_DIR=$HOME/data/gsm8k python3 -m verl.trainer.main_ppo \ actor_rollout_ref.rollout.name=sglang \ data.train_files=$DATA_DIR/train.parquet \ data.val_files=$DATA_DIR/test.parquet \ data.train_batch_size=4096 \ data.max_prompt_length=4096 \ data.max_response_length=4096 \ actor_rollout_ref.model.path=meta-llama/Llama-3.1-8B-Instruct \ actor_rollout_ref.actor.optim.lr=1e-6 \ actor_rollout_ref.model.use_remove_padding=True \ actor_rollout_ref.actor.ppo_mini_batch_size=64 \ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ actor_rollout_ref.model.enable_gradient_checkpointing=True \ actor_rollout_ref.actor.fsdp_config.param_offload=True \ actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ actor_rollout_ref.rollout.tensor_model_parallel_size=16 \ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ actor_rollout_ref.rollout.free_cache_engine=True \ actor_rollout_ref.ref.log_prob_micro_batch_size=16 \ actor_rollout_ref.ref.fsdp_config.param_offload=True \ critic.optim.lr=1e-5 \ critic.model.use_remove_padding=True \ critic.model.path=meta-llama/Llama-3.1-8B-Instruct \ critic.model.enable_gradient_checkpointing=True \ critic.ppo_micro_batch_size=16 \ critic.model.fsdp_config.param_offload=True \ critic.model.fsdp_config.optimizer_offload=True \ algorithm.kl_ctrl.kl_coef=0.001 \ trainer.critic_warmup=0 \ trainer.logger=console \ trainer.val_before_train=True \ trainer.n_gpus_per_node=8 \ trainer.nnodes=2 \ trainer.save_freq=-1 \ trainer.test_freq=10 \ trainer.total_epochs=15 2>&1 | tee verl_demo.log