[SGLang] Add support between mcore0.11 and sglang (#1055)

BearBiscuit05 · web-flow · commit 312a8cbcebf9 · 2025-05-07T08:57:03.000-07:00
Based on the ongoing alignment between mcore and vllm #851 , I believe we can simultaneously advance the alignment between mcore and sglang, as their interfaces are similar. In the end, we will only need to obtain a generator parameter. [link](sgl-project/sglang#5345)
diff --git a/examples/grpo_trainer/run_qwen2-7b_sgl_megatron.sh b/examples/grpo_trainer/run_qwen2-7b_sgl_megatron.sh
@@ -0,0 +1,48 @@
+set -x
+
+gsm8k_train_path=$HOME/data/gsm8k/train.parquet
+gsm8k_test_path=$HOME/data/gsm8k/test.parquet
+math_train_path=$HOME/data/math/train.parquet
+math_test_path=$HOME/data/math/test.parquet
+
+train_files="['$gsm8k_train_path', '$math_train_path']"
+test_files="['$gsm8k_test_path', '$math_test_path']"
+
+python3 -m verl.trainer.main_ppo --config-path=config \
+    --config-name='ppo_megatron_trainer.yaml'\
+    algorithm.adv_estimator=grpo \
+    data.train_files="$train_files" \
+    data.val_files="$test_files" \
+    data.train_batch_size=1024 \
+    data.max_prompt_length=1024 \
+    data.max_response_length=1024 \
+    data.filter_overlong_prompts=True \
+    data.truncation='error' \
+    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=256 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
+    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \
+    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=2 \
+    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=4 \
+    actor_rollout_ref.actor.use_kl_loss=True \
+    actor_rollout_ref.actor.kl_loss_coef=0.001 \
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+    actor_rollout_ref.rollout.name=sglang \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+    actor_rollout_ref.rollout.n=5 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
+    algorithm.use_kl_in_reward=False \
+    trainer.critic_warmup=0 \
+    trainer.logger=['console','wandb'] \
+    trainer.project_name='verl_grpo_example_gsm8k' \
+    trainer.experiment_name='qwen2_7b_function_rm_megatron' \
+    trainer.n_gpus_per_node=8 \
+    trainer.nnodes=1 \
+    trainer.save_freq=-1 \
+    trainer.test_freq=5 \
+    trainer.total_epochs=15 $@
diff --git a/verl/utils/megatron_utils.py b/verl/utils/megatron_utils.py
@@ -32,6 +32,8 @@
 from transformers import PretrainedConfig
 
 from verl.utils.torch_dtypes import PrecisionType
+from verl.utils.model import normalize_model_name
+import verl.utils.megatron.tensor_parallel as tp_utils
 
 
 def get_model_config(model):
@@ -619,3 +621,140 @@ def broadcast_str_from_megatron_pp(obj: Any):
     torch.distributed.broadcast_object_list(object_list=obj_output, src=global_rank, group=mpu.get_pipeline_model_parallel_group())
 
     return obj_output[0]
+
+def default_tp_concat_fn(layer_name_mapping, name, train_params, infer_params, model_config, convert_qkv_gate_up_by_simple_split=False):
+    """
+    name: name of the parameter
+    train_params: training parameters
+    infer_params (Iterable[torch.Tensor]): a iterator towards list of parameters all-gathered from micro_dp_group
+    model_config: huggingface model_config
+    TODO(zhangchi.usc1992): currently, the implementation is adhoc. We can move this function to the model
+    definition so that it is model-agnostic. If the model doesn't implement this function, 
+    we can throw an error to force user disable TP HybridEngine.
+    """
+    from megatron.core import mpu
+
+    if layer_name_mapping.get("qkv_layer_name") in name and "layer_norm" not in name:
+        # if the tensor is qkv, for each param on tp, split into q, k, v
+        # concat q, k, v separately.
+        q_lst = []
+        k_lst = []
+        v_lst = []
+        assert model_config.num_attention_heads % model_config.num_key_value_heads == 0
+        num_q_per_kv = model_config.num_attention_heads // model_config.num_key_value_heads
+        assert infer_params[0].shape[0] % (num_q_per_kv + 2) == 0
+        kv_size_per_tp = infer_params[0].shape[0] // (num_q_per_kv + 2)
+        split_size = [kv_size_per_tp * num_q_per_kv, kv_size_per_tp, kv_size_per_tp]
+        for infer_param in infer_params:
+            num_query_groups_per_partition = model_config.num_key_value_heads // mpu.get_tensor_model_parallel_world_size(
+            )
+            for chunk in infer_param.chunk(num_query_groups_per_partition):
+                split_size = [
+                    kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition,
+                    kv_size_per_tp // num_query_groups_per_partition,
+                    kv_size_per_tp // num_query_groups_per_partition
+                ]
+                q, k, v = chunk.split(split_size)
+                q_lst.append(q)
+                k_lst.append(k)
+                v_lst.append(v)
+        q = torch.cat(q_lst, dim=0)
+        k = torch.cat(k_lst, dim=0)
+        v = torch.cat(v_lst, dim=0)
+        if not convert_qkv_gate_up_by_simple_split:
+            infer_params = torch.cat((q, k, v), dim=0)
+        else:
+            infer_params = [q, k, v]
+
+    elif layer_name_mapping.get("gate_proj_layer_name") in name:
+        # if the tensor is gate and proj
+        gate_lst = []
+        up_lst = []
+        for infer_param in infer_params:
+            gate, up = infer_param.chunk(2)
+            gate_lst.append(gate)
+            up_lst.append(up)
+        gate = torch.cat(gate_lst, dim=0)
+        up = torch.cat(up_lst, dim=0)
+        if not convert_qkv_gate_up_by_simple_split:
+            infer_params = torch.cat((gate, up), dim=0)
+        else:
+            infer_params = [gate, up]
+
+    else:
+        # concat tensor
+        infer_params = torch.cat(infer_params, dim=tp_utils.get_tensor_parallel_partition_dim(train_params))
+
+    return infer_params
+
+
+def per_tensor_generator(actor_module, model_config, weight_converter, layer_name_mapping, convert_qkv_gate_up_by_simple_split=True):
+    from megatron.core import parallel_state as mpu
+    pp_rank = mpu.get_pipeline_model_parallel_rank()
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    vpp_size = len(actor_module)
+    all_gather_group = mpu.get_tensor_model_parallel_group()
+    all_gather_group_size = torch.distributed.get_world_size(group=all_gather_group)
+
+    def tensor_generator():
+        for scan_vpp_idx in range(vpp_size):
+            yield from actor_module[scan_vpp_idx].named_parameters()
+
+    # we need first make all rank get full model information
+    meta_info = []
+    for scan_vpp_idx in range(vpp_size):
+        for idx, (name, _) in enumerate(actor_module[scan_vpp_idx].named_parameters()):
+            meta_info.append((pp_rank, scan_vpp_idx, idx, name))
+
+    obj_spec_output = [None] * mpu.get_pipeline_model_parallel_world_size()
+    torch.distributed.all_gather_object(
+        object_list=obj_spec_output, obj=meta_info, group=mpu.get_pipeline_model_parallel_group()
+    )
+    layer_list_meta = [item for sublist in obj_spec_output for item in sublist]
+
+    gen_func = tensor_generator()
+
+    # lazy load tensor for full model
+    for cur_pp_rank, scan_vpp_idx, idx, name in layer_list_meta:
+        if cur_pp_rank == pp_rank:
+            try:
+                cur_name, cur_tensor = next(gen_func)
+            except StopIteration:
+                cur_name, cur_tensor = None, None
+            cur_name = normalize_model_name(
+                name, cur_pp_rank, scan_vpp_idx, pp_size, vpp_size, model_config.num_hidden_layers
+            )
+        else:
+            cur_tensor, cur_name = None, None
+
+        # pp broadcast model tensor and name
+        cur_name = broadcast_str_from_megatron_pp(cur_name)
+        broad_pp_tensor = broadcast_from_megatron_pp(cur_tensor)
+
+        # (xya): this is a hack to fix the name of the parameters
+        while cur_name.startswith("module."):
+            cur_name = cur_name[len("module.") :]
+
+        # tp all gather
+        if tp_utils.is_tensor_parallel_param(broad_pp_tensor):
+            # allocate a new tensor with proper size
+            if all_gather_group_size <= 1:
+                infer_params = [broad_pp_tensor]
+            else:
+                infer_params = [torch.empty_like(broad_pp_tensor) for _ in range(all_gather_group_size)]
+                torch.distributed.all_gather(
+                    infer_params, broad_pp_tensor, group=mpu.get_tensor_model_parallel_group()
+                )
+            infer_params = default_tp_concat_fn(
+                layer_name_mapping, cur_name, broad_pp_tensor, infer_params, model_config, convert_qkv_gate_up_by_simple_split
+            )
+        else:
+            infer_params = broad_pp_tensor
+
+
+        if not isinstance(infer_params, list):
+            infer_params = [infer_params]
+        converted_names, converted_params = weight_converter.convert_param(cur_name, infer_params)
+
+        yield from zip(converted_names, converted_params)
+
diff --git a/verl/workers/megatron_workers.py b/verl/workers/megatron_workers.py
@@ -209,6 +209,10 @@ def megatron_actor_model_provider(pre_process, post_process):
         return actor_module, actor_optimizer, self.hf_config, optim_config
 
     def _build_rollout(self, trust_remote_code=False):
+        layer_name_mapping = {
+            "qkv_layer_name": "self_attention.linear_qkv.",
+            "gate_proj_layer_name": "linear_fc1.weight",
+        }
         if self.config.rollout.name == "vllm":
             from torch.distributed.device_mesh import init_device_mesh
 
@@ -217,10 +221,7 @@ def _build_rollout(self, trust_remote_code=False):
 
             # NOTE(sgm): If the QKV and gate_up projection layer are concate together in actor,
             # we will reorganize their weight format when resharding from actor to rollout.
-            layer_name_mapping = {
-                "qkv_layer_name": "self_attention.linear_qkv.",
-                "gate_proj_layer_name": "linear_fc1.weight",
-            }
+            
 
             infer_tp = self.config.rollout.tensor_model_parallel_size
             dp = self.world_size // infer_tp
@@ -259,6 +260,30 @@ def _build_rollout(self, trust_remote_code=False):
                 weight_converter=weight_converter,
             )
             log_gpu_memory_usage("After building sharding manager", logger=logger)
+        elif self.config.rollout.name == 'sglang':
+            from verl.workers.rollout.sglang_rollout import SGLangRollout
+            # NOTE(linjunrong): Due to recent fp8 support in SGLang. Now importing any symbol relate to SGLang's model_runner would check CUDA device capability.
+            # However, due to veRL's setting, the main process of ray can not find any CUDA device, which would potentially lead to:
+            # "RuntimeError: No CUDA GPUs are available".
+            # For this reason, sharding_manager.__init__ should not import FSDPSGLangShardingManager and we import it here use the abs path.
+            # check: https://github.com/sgl-project/sglang/blob/00f42707eaddfc2c0528e5b1e0094025c640b7a0/python/sglang/srt/layers/quantization/fp8_utils.py#L76
+            from verl.workers.sharding_manager.megatron_sglang import MegatronSGLangShardingManager
+            local_path = copy_to_local(self.config.model.path)
+            log_gpu_memory_usage(f'Before building {self.config.rollout.name} rollout', logger=None)
+            rollout = SGLangRollout(actor_module=local_path,
+                                    config=self.config.rollout,
+                                    tokenizer=self.tokenizer,
+                                    model_hf_config=self.actor_model_config)
+            log_gpu_memory_usage(f'After building {self.config.rollout.name} rollout', logger=None)
+
+            from verl.models.mcore import get_mcore_weight_converter
+            weight_converter = get_mcore_weight_converter(self.actor_model_config, self.dtype)
+            sharding_manager = MegatronSGLangShardingManager(actor_module=self.actor.actor_module,
+                                                             inference_engine=rollout.inference_engine,
+                                                             model_config=self.actor_model_config,
+                                                             layer_name_mapping=layer_name_mapping,
+                                                             weight_converter=weight_converter,)
+            log_gpu_memory_usage('After building sharding manager', logger=logger)
         else:
             raise NotImplementedError("Only vllmRollout is supported with Megatron now")
 
diff --git a/verl/workers/rollout/sglang_rollout/sglang_rollout.py b/verl/workers/rollout/sglang_rollout/sglang_rollout.py
@@ -370,3 +370,12 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
             self.inference_engine._engine.flush_cache()
 
         return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)
+
+    # this function is left for uniform train-inference resharding
+    def update_weights(self, params_iter):
+        self.inference_engine.resume_memory_occupation()
+        self.inference_engine.update_weights_from_tensor(params_iter, load_format=None)
+
+    # this function is left for uniform train-inference resharding
+    def offload(self):
+        self.inference_engine.release_memory_occupation()
diff --git a/verl/workers/sharding_manager/megatron_sglang.py b/verl/workers/sharding_manager/megatron_sglang.py
@@ -0,0 +1,109 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This file contains a Megatron style Hybrid Engine that shares the weights of the actor with the inference engine.
+"""
+
+import importlib
+import logging
+import os
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from verl.utils.model import normalize_model_name
+from verl.utils.megatron_utils import broadcast_from_megatron_pp, broadcast_str_from_megatron_pp
+
+from verl.utils.megatron_utils import get_model, unwrap_model
+from verl.utils.debug import log_gpu_memory_usage
+from verl.utils.megatron_utils import convert_megatron_model_to_transformers_model
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN'))
+"""
+Megatron Hybrid Engine:
+- During training, only the current pp stage holds the parameters
+- Before inference, broadcast the parameters of the current pp rank to all other pp ranks (all pp ranks holds all the parameters)
+- Bind the parameters to the inference engine
+- Do inference in tp. pp is treated as additional dp
+- After inference, all the parameters that doesn't belong to this pp rank is freed.
+"""
+
+from .base import BaseShardingManager
+
+import torch
+from torch import nn
+import torch.distributed
+from torch.distributed import new_group
+from torch.distributed._tensor import DTensor
+from typing import Dict, Iterable, Union, Tuple
+
+from verl import DataProto
+from verl.protocol import all_gather_data_proto
+from verl.utils.torch_functional import (broadcast_dict_tensor, allgather_dict_tensors)
+from sglang.srt.entrypoints.verl_engine import VerlEngine
+from verl.utils.debug import GPUMemoryLogger
+
+import verl.utils.megatron.tensor_parallel as tp_utils
+from verl.utils.megatron_utils import per_tensor_generator, default_tp_concat_fn
+
+_MICRO_DATA_PARALLEL_GROUP = None
+
+
+class MegatronSGLangShardingManager(BaseShardingManager):
+
+    def __init__(self, actor_module: nn.ModuleList, inference_engine: VerlEngine, model_config, layer_name_mapping, weight_converter):
+        from megatron.core import parallel_state as mpu
+        self.actor_module = actor_module
+        self.inference_engine = inference_engine
+        self.model_config = model_config
+        self.layer_name_mapping = layer_name_mapping
+        self.weight_converter = weight_converter
+        global _MICRO_DATA_PARALLEL_GROUP
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+
+        self.infer_tp_size = self.inference_engine._tp_size
+        self.train_tp_size = mpu.get_tensor_model_parallel_world_size()
+        self.need_tp_reshard = self.infer_tp_size == self.train_tp_size
+
+        assert self.infer_tp_size <= self.train_tp_size, \
+            'Not implemented for infer_tp > train_tp'
+        assert self.train_tp_size % self.infer_tp_size == 0
+
+        micro_dp_size = self.train_tp_size // self.infer_tp_size
+        num_micro_dp_groups = world_size // micro_dp_size
+        assert _MICRO_DATA_PARALLEL_GROUP is None, ("micro data parallel group is already initialized")
+        for i in range(num_micro_dp_groups):
+            ranks = range(i * micro_dp_size, (i + 1) * micro_dp_size)
+            group = new_group(ranks=ranks)
+            if rank in ranks:
+                _MICRO_DATA_PARALLEL_GROUP = group
+
+    @GPUMemoryLogger(role="MegatronSGLangShardingManager enter", logger=logger)
+    def __enter__(self):
+        per_tensor_param = per_tensor_generator(self.actor_module, self.model_config, self.weight_converter, self.layer_name_mapping)
+        self.inference_engine.resume_memory_occupation()
+        self.inference_engine.update_weights_from_tensor(per_tensor_param, load_format=None)
+
+    @GPUMemoryLogger(role="MegatronSGLangShardingManager exit", logger=logger)
+    def __exit__(self, exc_type, exc_value, traceback):
+        log_gpu_memory_usage('Before SGLang offload in sharding manager', logger=logger)
+        self.inference_engine.release_memory_occupation()
+        log_gpu_memory_usage('After SGLang offload in sharding manager', logger=logger)
+
+        for model in self.actor_module:
+            model.train()
+        # add empty cache after each compute
+        torch.cuda.empty_cache()
