[worker] refactor: move the implementation of rm to workers.roles and polish (volcengine#3423)

Author: yyDing1

verl/workers/config/reward_model.py

@@ -19,21 +19,9 @@
 from verl.utils.profiler import ProfilerConfig
 
 from .model import HFModelConfig
+from .rollout import SamplingConfig, ServerConfig
 
-__all__ = ["ServerConfig", "SandboxFusionConfig", "RewardModelConfig"]
-
-
-@dataclass
-class ServerConfig(BaseConfig):
-    """
-    Configuration for SGLang server when running in server mode
-    """
-
-    timeout: float = 60.0
-    max_attempts: int = 3
-    retry_delay: float = 2.0
-    max_connections: int = 1000
-    max_start_wait_time: float = 300.0
+__all__ = ["SandboxFusionConfig", "RewardModelConfig"]
 
 
 @dataclass
@@ -53,50 +41,25 @@ class SandboxFusionConfig(BaseConfig):
 
 @dataclass
 class RewardModelConfig(BaseConfig):
-    """Configuration for reward model scoring.
-
-    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.
-
-    Args:
-        enable (bool): Whether to enable reward model.
-        enable_resource_pool (bool): Whether to deploy the model to a separate resource pool.
-        n_gpus_per_node (int): Number of GPUs per node when using resource pool.
-        nnodes (int): Number of nodes when using resource pool.
-        strategy (str): FSDP strategy: "fsdp" or "fsdp2".
-        model (Dict[str, Any]): Model configuration for reward scoring.
-        micro_batch_size (Optional[int]): Global micro batch size (deprecated).
-        micro_batch_size_per_gpu (Optional[int]): Local per-GPU micro batch size.
-        max_length (Optional[int]): Maximum sequence length to process for scoring.
-        use_dynamic_bsz (bool): Whether to dynamically adjust batch size at runtime.
-        forward_max_token_len_per_gpu (int): Maximum number of tokens per GPU in one forward pass.
-        reward_manager (str): Reward manager type (naive or prime).
-        launch_reward_fn_async (bool): Whether to launch custom reward function asynchronously during log_prob.
-        sandbox_fusion (Dict[str, Any]): Cloud/local sandbox fusion configuration for custom reward logic.
-        profiler (Dict[str, Any]): Profiler configuration for reward model.
-    """
-
     _mutable_fields = BaseConfig._mutable_fields
 
     enable: bool = False
+    model_type: str = "discriminative"
+    name: str = "sglang"
     enable_resource_pool: bool = False
     n_gpus_per_node: int = 0
     nnodes: int = 0
-    # strategy: str = MISSING
-    # model: BaseModelConfig = field(default_factory=BaseModelConfig)
-    # micro_batch_size: Optional[int] = None
-    # micro_batch_size_per_gpu: Optional[int] = None
-    # max_length: Optional[int] = None
-    # use_dynamic_bsz: bool = False
-    # forward_max_token_len_per_gpu: int = 32768
     reward_manager: str = "naive"
     launch_reward_fn_async: bool = False
 
-    tensor_model_parallel_size: int = 2
-    engine_kwargs: dict = field(default_factory=dict)
-    max_num_seqs: int = 1024
     dtype: str = "bfloat16"
     gpu_memory_utilization: float = 0.5
     free_cache_engine: bool = True
+    tensor_model_parallel_size: int = 2
+    sampling_config: SamplingConfig = field(default_factory=SamplingConfig)
+
+    engine_kwargs: dict = field(default_factory=dict)
+    max_num_seqs: int = 1024
 
     sandbox_fusion: SandboxFusionConfig = field(default_factory=SandboxFusionConfig)
     profiler: ProfilerConfig = field(default_factory=ProfilerConfig)

verl/workers/config/rollout.py

@@ -26,6 +26,7 @@
     "CustomAsyncServerConfig",
     "AgentLoopConfig",
     "TraceConfig",
+    "ServerConfig",
     "RolloutConfig",
 ]
 

verl/workers/roles/__init__.py

@@ -16,7 +16,7 @@
 from .critic import CriticWorker
 
 try:
-    from .reward import RewardModelWorker
+    from .reward_model import RewardModelWorker
 except ImportError:
     RewardModelWorker = None
 

verl/workers/roles/reward.py renamed to verl/workers/roles/reward_model.py

@@ -34,7 +34,7 @@
 from verl.utils.model import compute_position_id_with_mask
 from verl.utils.profiler import DistProfiler, DistProfilerExtension, log_gpu_memory_usage
 from verl.workers.config import HFModelConfig, RewardModelConfig
-from verl.workers.reward_model.sglang_reward_model import SGLangRewardModel
+from verl.workers.roles.reward_model_engine import get_reward_model_class
 
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
@@ -92,7 +92,7 @@ def _build_reward_model(self):
 
         # 4. build reward model
         log_gpu_memory_usage("Before building sglang reward model", logger=logger)
-        self.reward_model = SGLangRewardModel(
+        self.reward_model = get_reward_model_class(reward_model_config.name)(
             config=reward_model_config, model_config=model_config, device_mesh=reward_model_device_mesh
         )
         log_gpu_memory_usage("After building sglang reward model", logger=logger)

verl/workers/roles/reward_model_engine/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+
+from .base import get_reward_model_class
+
+__all__ = ["get_reward_model_class"]

verl/workers/roles/reward_model_engine/base.py

@@ -0,0 +1,94 @@
+# 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.
+"""
+The base class for reward model
+"""
+
+import importlib
+from abc import ABC, abstractmethod
+
+from torch.distributed.device_mesh import DeviceMesh
+
+from verl import DataProto
+from verl.workers.config import HFModelConfig, RewardModelConfig
+
+__all__ = ["BaseRewardModel"]
+
+
+class BaseRewardModel(ABC):
+    """base class for reward model"""
+
+    def __init__(
+        self,
+        config: RewardModelConfig,
+        model_config: HFModelConfig,
+        device_mesh: DeviceMesh,
+    ):
+        self.config = config
+        self.model_config = model_config
+        self.device_mesh = device_mesh
+
+    @abstractmethod
+    async def resume(self, tags: list[str]):
+        """Resume reward model weights or kv cache in GPU memory.
+
+        Args:
+            tags: weights or kv_cache.
+        """
+        pass
+
+    @abstractmethod
+    async def release(self):
+        """Release weights and kv cache in GPU memory."""
+        pass
+
+    @abstractmethod
+    def compute_reward(self, data: DataProto) -> DataProto:
+        """Computing reward given input_ids. The transformers should output a tensor with shape
+           [batch_size, sequence_length], and the value at [EOS] mask should be gathered.
+
+        Args:
+            data: must contain keys "input_ids", "attention_mask" and "position_ids".
+                - input_ids: [batch_size, sequence_length]
+                - attention_mask: [batch_size, sequence_length]
+                - position_ids: [batch_size, sequence_length]
+
+        Returns: a data pass protocol containing "reward". Only the [EOS] position contains the reward.
+            Other position should have zero reward. Note that this may change in the future if we use
+            dense reward. So, we leave the interface for general case.
+            - reward: [batch_size, sequence_length].
+
+        """
+        pass
+
+
+_REWARD_MODEL_REGISTRY = {
+    "sglang": "verl.workers.roles.reward_model_engine.sglang_reward_model.SGLangRewardModel",
+}
+
+
+def get_reward_model_class(reward_model_name: str) -> type[BaseRewardModel]:
+    """Get the reward model class by name.
+
+    Args:
+        reward_model_name: The name of the reward model.
+
+    Returns:
+        The reward model class.
+    """
+    assert reward_model_name in _REWARD_MODEL_REGISTRY, f"Reward Model {reward_model_name} with mode not found"
+    fqdn = _REWARD_MODEL_REGISTRY[reward_model_name]
+    module_name, class_name = fqdn.rsplit(".", 1)
+    reward_model_module = importlib.import_module(module_name)
+    return getattr(reward_model_module, class_name)

verl/workers/reward_model/sglang_reward_model.py renamed to verl/workers/roles/reward_model_engine/sglang_reward_model.py

@@ -15,27 +15,72 @@
 
 import asyncio
 import logging
+import multiprocessing as mp
 import os
 
+import sglang.srt.entrypoints.engine
 import torch
 import torch.distributed as dist
+from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
+    assert_pkg_version,
     get_ip,
     get_open_port,
+    is_cuda,
+    set_prometheus_multiproc_dir,
+    set_ulimit,
 )
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 
 from verl import DataProto
 from verl.utils.net_utils import is_ipv6
 from verl.workers.config import HFModelConfig, RewardModelConfig
-from verl.workers.reward_model import BasePPORewardModel
+from verl.workers.roles.reward_model_engine.base import BaseRewardModel
 from verl.workers.rollout.sglang_rollout.http_server_engine import AsyncHttpServerAdapter
 from verl.workers.rollout.sglang_rollout.utils import broadcast_pyobj
 
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
 
 
+# patch to avoid issue https://github.com/sgl-project/sglang/issues/6723
+def _set_envs_and_config(server_args: ServerArgs):
+    # Set global environments
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+    os.environ["NCCL_CUMEM_ENABLE"] = "0"
+    os.environ["NCCL_NVLS_ENABLE"] = str(int(server_args.enable_nccl_nvls))
+    os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+    os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "4"
+    os.environ["CUDA_MODULE_LOADING"] = "AUTO"
+
+    # Set prometheus env vars
+    if server_args.enable_metrics:
+        set_prometheus_multiproc_dir()
+
+    # Set ulimit
+    set_ulimit()
+
+    # Check flashinfer version
+    if server_args.attention_backend == "flashinfer":
+        assert_pkg_version(
+            "flashinfer_python",
+            "0.2.5",
+            "Please uninstall the old version and reinstall the latest version by following the instructions at https://docs.flashinfer.ai/installation.html.",
+        )
+    if is_cuda():
+        assert_pkg_version(
+            "sgl-kernel",
+            "0.1.1",
+            "Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`",
+        )
+
+    # Set mp start method
+    mp.set_start_method("spawn", force=True)
+
+
+sglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config
+
+
 def _pre_process_inputs(
     attention_mask: torch.Tensor,
     prompt_token_ids: torch.Tensor,
@@ -54,7 +99,7 @@ def _map_each_output(output):
     return scores
 
 
-class SGLangRewardModel(BasePPORewardModel):
+class SGLangRewardModel(BaseRewardModel):
     def __init__(
         self,
         config: RewardModelConfig,
@@ -66,14 +111,13 @@ def __init__(
         actor_module = model_config.local_path
         trust_remote_code = model_config.trust_remote_code
         port = None
-        kwargs = {}
 
         os.environ.setdefault("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK", "true")
 
-        self._init_distributed_env(device_mesh_cpu=None, **kwargs)
+        self._init_distributed_env(device_mesh_cpu=None)
         self._init_inference_engine(trust_remote_code, actor_module, port)
 
-    def _init_distributed_env(self, device_mesh_cpu, **kwargs):
+    def _init_distributed_env(self, device_mesh_cpu):
         self._device_mesh_cpu = device_mesh_cpu
         os.environ.setdefault("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK", "true")
         self.tensor_parallel_size = self.config.get("tensor_model_parallel_size", 1)
@@ -211,7 +255,7 @@ def compute_reward(self, data: DataProto):
         return reward_score
 
     async def resume(self, tags: list[str]):
-        """Resume rollout weights or kv cache in GPU memory.
+        """Resume reward model weights or kv cache in GPU memory.
 
         Args:
             tag: weights or kv_cache.