[rollout] chore: Misc changes for extending internal compatibility (volcengine#3701)

### What does this PR do?

* New config field:
    * rollout: `pipeline_model_parallel_size` for internal compatibility
* ~~legacy_data: `agent_name` for default agent name if not specified in
the rldataset~~
* Registry for `RolloutReplica`
* `VERL_USE_EXTERNAL_MODULES` to import desired modules to trigger
external registration


### Test

Be covered by CI

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [x] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [x] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

Author: pengwu22

verl/__init__.py

@@ -22,6 +22,7 @@
 
 from .protocol import DataProto
 from .utils.device import is_npu_available
+from .utils.import_utils import import_external_libs
 from .utils.logging_utils import set_basic_config
 
 version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))
@@ -35,6 +36,13 @@
 
 __all__ = ["DataProto", "__version__"]
 
+
+modules = os.getenv("VERL_USE_EXTERNAL_MODULES", "")
+if modules:
+    modules = modules.split(",")
+    import_external_libs(modules)
+
+
 if os.getenv("VERL_USE_MODELSCOPE", "False").lower() == "true":
     if importlib.util.find_spec("modelscope") is None:
         raise ImportError("You are using the modelscope hub, please install modelscope by `pip install modelscope -U`")

verl/experimental/agent_loop/agent_loop.py

@@ -789,6 +789,7 @@ def _initialize_llm_servers(self):
         rollout_world_size = (
             self.config.actor_rollout_ref.rollout.tensor_model_parallel_size
             * self.config.actor_rollout_ref.rollout.data_parallel_size
+            * self.config.actor_rollout_ref.rollout.pipeline_model_parallel_size
         )
         world_size = (
             self.worker_group.world_size

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -190,6 +190,7 @@ actor_rollout_ref:
     tensor_model_parallel_size: 2
     data_parallel_size: 1
     expert_parallel_size: 1
+    pipeline_model_parallel_size: 1
     max_num_batched_tokens: 8192
     max_model_len: null
     max_num_seqs: 1024

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -177,6 +177,7 @@ actor_rollout_ref:
     tensor_model_parallel_size: 2
     data_parallel_size: 1
     expert_parallel_size: 1
+    pipeline_model_parallel_size: 1
     max_num_batched_tokens: 8192
     max_model_len: null
     max_num_seqs: 1024

verl/trainer/config/rollout/rollout.yaml

@@ -55,6 +55,9 @@ data_parallel_size: 1
 # EP size for rollout
 expert_parallel_size: 1
 
+# PP size for rollout.
+pipeline_model_parallel_size: 1
+
 # max number of tokens in a batch
 max_num_batched_tokens: 8192
 

verl/utils/memory_utils.py

@@ -24,7 +24,8 @@
 
 from verl.utils.device import get_torch_device, is_cuda_available
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
 
 
 def aggressive_empty_cache(force_sync: bool = True, max_retries: int = 3) -> None:

verl/workers/config/rollout.py

@@ -121,6 +121,7 @@ class RolloutConfig(BaseConfig):
     data_parallel_size: int = 1
     expert_parallel_size: int = 1
     tensor_model_parallel_size: int = 2
+    pipeline_model_parallel_size: int = 1
     max_num_batched_tokens: int = 8192
 
     # TODO: enable train_kwargs
@@ -183,3 +184,9 @@ def __post_init__(self):
             assert self.expert_parallel_size == (self.tensor_model_parallel_size * self.data_parallel_size), (
                 "expert_parallel_size must be equal to tensor_model_parallel_size * data_parallel_size"
             )
+
+        if self.pipeline_model_parallel_size > 1:
+            if self.name == "vllm" or self.name == "sglang":
+                raise NotImplementedError(
+                    f"Current rollout {self.name=} not implemented pipeline_model_parallel_size > 1 yet."
+                )

verl/workers/fsdp_workers.py

@@ -571,19 +571,24 @@ def _build_rollout(self, trust_remote_code=False):
 
         # 2. build rollout device mesh
         infer_tp = self.config.rollout.tensor_model_parallel_size * self.config.rollout.data_parallel_size
-        dp = self.world_size // infer_tp
-        assert self.world_size % infer_tp == 0, (
-            f"rollout world_size: {self.world_size} is not divisible by infer_tp: {infer_tp}"
+        infer_pp = self.config.rollout.pipeline_model_parallel_size
+        infer_world_size = infer_tp * infer_pp
+        dp = self.world_size // infer_world_size
+        assert self.world_size % infer_world_size == 0, (
+            f"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}"
         )
         rollout_device_mesh = init_device_mesh(
-            device_name, mesh_shape=(dp, infer_tp), mesh_dim_names=["dp", "infer_tp"]
+            device_name, mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=["dp", "infer_tp", "infer_pp"]
         )
         rollout_name = self.config.rollout.name
 
         if rollout_name == "hf":
             self._register_dispatch_collect_info("rollout", dp_rank=self.rank, is_collect=True)
         else:
-            is_collect = rollout_device_mesh["infer_tp"].get_local_rank() == 0
+            is_collect = (
+                rollout_device_mesh["infer_tp"].get_local_rank() == 0
+                and rollout_device_mesh["infer_pp"].get_local_rank() == 0
+            )
             self._register_dispatch_collect_info(
                 "rollout", dp_rank=rollout_device_mesh["dp"].get_local_rank(), is_collect=is_collect
             )

verl/workers/megatron_workers.py

@@ -397,15 +397,20 @@ def _build_rollout(self, trust_remote_code=False):
 
         # 2. build rollout device mesh
         infer_tp = self.config.rollout.tensor_model_parallel_size * self.config.rollout.data_parallel_size
-        dp = self.world_size // infer_tp
-        assert self.world_size % infer_tp == 0, (
-            f"rollout world_size: {self.world_size} is not divisible by infer_tp: {infer_tp}"
+        infer_pp = self.config.rollout.pipeline_model_parallel_size
+        infer_world_size = infer_tp * infer_pp
+        dp = self.world_size // infer_world_size
+        assert self.world_size % infer_world_size == 0, (
+            f"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}"
         )
         rollout_device_mesh = init_device_mesh(
-            get_device_name(), mesh_shape=(dp, infer_tp), mesh_dim_names=["dp", "infer_tp"]
+            get_device_name(), mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=["dp", "infer_tp", "infer_pp"]
         )
 
-        is_collect = rollout_device_mesh["infer_tp"].get_local_rank() == 0
+        is_collect = (
+            rollout_device_mesh["infer_tp"].get_local_rank() == 0
+            and rollout_device_mesh["infer_pp"].get_local_rank() == 0
+        )
         self._register_dispatch_collect_info(
             "rollout", dp_rank=rollout_device_mesh["dp"].get_local_rank(), is_collect=is_collect
         )

verl/workers/rollout/replica.py

@@ -16,7 +16,7 @@
 import os
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Optional
+from typing import Callable, Optional
 
 from pydantic import BaseModel
 from ray.actor import ActorHandle
@@ -85,7 +85,11 @@ def __init__(
         self.config = omega_conf_to_dataclass(config)
         self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)
 
-        self.world_size = self.config.tensor_model_parallel_size * self.config.data_parallel_size
+        self.world_size = (
+            self.config.tensor_model_parallel_size
+            * self.config.data_parallel_size
+            * self.config.pipeline_model_parallel_size
+        )
         self.gpus_per_node = min(gpus_per_node, self.world_size)
         assert self.world_size % self.gpus_per_node == 0, (
             f"world_size {self.world_size} must be divisible by gpus_per_node {self.gpus_per_node}"
@@ -171,32 +175,57 @@ async def sleep(self):
         await asyncio.gather(*[server.sleep.remote() for server in self.servers])
 
 
+class RolloutReplicaRegistry:
+    """Factory for managing rollout replica implementations."""
+
+    _registry: dict[str, Callable[[], type[RolloutReplica]]] = {}
+
+    @classmethod
+    def register(cls, name: str, loader: Callable[[], type[RolloutReplica]]) -> None:
+        """Register a new rollout replica type."""
+        cls._registry[name] = loader
+
+    @classmethod
+    def get(cls, name: str) -> type[RolloutReplica]:
+        """Get a rollout replica class by name."""
+        if name not in cls._registry:
+            raise ValueError(f"Unknown rollout mode: {name}. Available: {list(cls._registry.keys())}")
+        return cls._registry[name]()
+
+
+# Loader functions for built-in types
+def _load_vllm():
+    from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMReplica
+
+    return vLLMReplica
+
+
+def _load_sglang():
+    os.environ["SGLANG_USE_CPU_ENGINE"] = "1"
+
+    try:
+        import vllm  # noqa: F401
+    except ImportError:
+        import sys
+        from unittest.mock import Mock
+
+        mock_vllm = Mock()
+        mock_vllm._custom_ops = Mock()
+        mock_vllm._custom_ops.scaled_fp8_quant = Mock()
+        sys.modules["vllm"] = mock_vllm
+        sys.modules["vllm._custom_ops"] = mock_vllm._custom_ops
+
+    from verl.workers.rollout.sglang_rollout.async_sglang_server import SGLangReplica
+
+    del os.environ["SGLANG_USE_CPU_ENGINE"]
+    return SGLangReplica
+
+
+# Register built-in types
+RolloutReplicaRegistry.register("vllm", _load_vllm)
+RolloutReplicaRegistry.register("sglang", _load_sglang)
+
+
+# Original function for backward compatibility
 def get_rollout_replica_class(rollout: str) -> type[RolloutReplica]:
-    if rollout == "vllm":
-        from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMReplica
-
-        return vLLMReplica
-    elif rollout == "sglang":
-        # NOTE: verl driver is cpu only, avoid sglang fp8 quantization import error.
-        os.environ["SGLANG_USE_CPU_ENGINE"] = "1"
-
-        # TODO: remove this once we bump to sglang>=0.5.1
-        try:
-            import vllm  # noqa: F401
-        except ImportError:
-            import sys
-            from unittest.mock import Mock
-
-            mock_vllm = Mock()
-            mock_vllm._custom_ops = Mock()
-            mock_vllm._custom_ops.scaled_fp8_quant = Mock()
-
-            sys.modules["vllm"] = mock_vllm
-            sys.modules["vllm._custom_ops"] = mock_vllm._custom_ops
-
-        from verl.workers.rollout.sglang_rollout.async_sglang_server import SGLangReplica
-
-        del os.environ["SGLANG_USE_CPU_ENGINE"]
-        return SGLangReplica
-    else:
-        raise ValueError(f"Unknown rollout mode: {rollout}")
+    return RolloutReplicaRegistry.get(rollout)