Simplify RL Colocate Trainer initialization by using cfg.build (#1520)

jayhenry · web-flow · commit 36184b7cdf9c · 2026-03-03T20:21:34.000+08:00
* Build XtunerMeta and TrainController by cfg.build

* Build RolloutController by cfg.build

* simplify rl colocate trainer init

* fix some lint errors

* fix some bugs
diff --git a/examples/v1/config/rl_grpo_gsm8k_judge.py b/examples/v1/config/rl_grpo_gsm8k_judge.py
@@ -14,7 +14,8 @@
 from xtuner.v1.model import get_model_config_from_hf
 from xtuner.v1.ray.base import AcceleratorResourcesConfig
 from xtuner.v1.ray.config.worker import RolloutConfig
-from xtuner.v1.ray.judger.gsm8k import GSM8KJudgerConfig
+from xtuner.v1.ray.judger.gsm8k import GSM8KRouterJudgerConfig
+from xtuner.v1.rl.base.replay_buffer import SyncReplayBufferConfig
 from xtuner.v1.rl.base import WorkerConfig
 from xtuner.v1.rl.base.agent_loop import SingleTurnAgentLoopConfig
 from xtuner.v1.rl.base.agent_loop_manager import AgentLoopManagerConfig
@@ -67,7 +68,7 @@
 )
 
 # 3. judger
-judger_config = GSM8KJudgerConfig(judger_name="openai/gsm8k")
+judger_config = GSM8KRouterJudgerConfig(judger_name="openai/gsm8k")
 
 # 4. train worker
 lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6)
@@ -181,7 +182,7 @@
     rollout_config=rollout_config,
     judger_config=judger_config,
     tokenizer_path=model_path,
-    replay_buffer_config=dict(),
+    replay_buffer_config=SyncReplayBufferConfig(),
     agent_loop_manager_cfg=agent_loop_manager_cfg,
     eval_agent_loop_manager_cfg=eval_agent_loop_manager_cfg,
     evaluator_config=evaluator_config,
diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py
@@ -10,13 +10,12 @@
 # ====================================
 # ====== DataFlow 数据流 ==============
 # ====================================
+from xtuner.v1.utils.cache import CacheObj
 from xtuner.v1.utils.logger import get_logger
 
 
 if TYPE_CHECKING:
-    import ray
-
-    RayObjectRef = ray.ObjectRef
+    from ray import ObjectRef as RayObjectRef
 else:
     RayObjectRef: TypeAlias = Any
 
@@ -60,12 +59,12 @@ class Status(Enum):
 
 class MultimodalInfo(TypedDict):
     # 使用TypedDict给出pixel_values的类型提示
-    pixel_values: NotRequired[torch.Tensor | RayObjectRef | None]  # type: ignore[valid-type]
+    pixel_values: NotRequired[torch.Tensor | RayObjectRef | None]
     image_grid_thw: NotRequired[torch.Tensor]
     position_ids: NotRequired[torch.Tensor]
 
 
-class RolloutState(BaseModel):
+class RolloutState(CacheObj, BaseModel):
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
     # --- 数据 ---
@@ -88,22 +87,22 @@ class RolloutState(BaseModel):
     response: str | None = None
     response_ids: list[int] | None = None
     logprobs: list[float] | None = None
-    routed_experts: list[int] | RayObjectRef | None = None  # type: ignore[valid-type]
+    routed_experts: list[int] | RayObjectRef | None = None
     finish_reason: str | None = None
 
-    @field_serializer('routed_experts')
+    @field_serializer("routed_experts")
     def _serialize_routed_experts(self, value: list[int] | RayObjectRef | None) -> list[int] | None:
         """Dump 时跳过 ray.ObjectRef，序列化为 None，避免 PydanticSerializationError。"""
         if value is None:
             return None
         try:
             import ray
+
             if isinstance(value, ray.ObjectRef):
                 return None
         except ImportError:
             pass
-        if type(value).__name__ == 'ObjectRef' and 'ray' in getattr(
-                type(value), '__module__', ''):
+        if type(value).__name__ == "ObjectRef" and "ray" in getattr(type(value), "__module__", ""):
             return None
         return value  # list[int]
 
diff --git a/xtuner/v1/datasets/__init__.py b/xtuner/v1/datasets/__init__.py
@@ -22,7 +22,7 @@
 from .rl_tokenize_fn import RLTextTokenizeFnConfig
 from .sampler import LengthGroupedSampler, ParallelSampler
 from .sft_tokenize_fn import OpenaiTokenizeFunction, OpenaiTokenizeFunctionConfig
-from .utils import CachableTokenizeFunction, CacheObj, calculate_file_sha256, calculate_xxhash, tokenizer_hash
+from .utils import CachableTokenizeFunction, CacheDict, calculate_file_sha256, calculate_xxhash, tokenizer_hash
 from .vlm_jsonl import VLMJsonlDataset
 
 
@@ -32,7 +32,7 @@
 __all__ = [
     "JsonlDataset",
     "CachableTokenizeFunction",
-    "CacheObj",
+    "CacheDict",
     "calculate_file_sha256",
     "calculate_xxhash",
     "tokenizer_hash",
diff --git a/xtuner/v1/datasets/_hardcode_patch.py b/xtuner/v1/datasets/_hardcode_patch.py
@@ -26,10 +26,10 @@
 from xtuner.v1.utils import get_logger
 
 from .ftdp import FtdpTokenizeFunction
-from .mllm_tokenize_fn import Qwen3VLTokenizeFunction
-from .pt_tokenize_fn import PretrainTokenizeFunction
+
 # from .rl_tokenize_fn.rl_tokenize_fn import InternS1VLTokenizeFunction
-from .mllm_tokenize_fn import InternS1VLTokenizeFunction
+from .mllm_tokenize_fn import InternS1VLTokenizeFunction, Qwen3VLTokenizeFunction
+from .pt_tokenize_fn import PretrainTokenizeFunction
 from .sft_tokenize_fn import OpenaiTokenizeFunction
 
 
diff --git a/xtuner/v1/datasets/jsonl.py b/xtuner/v1/datasets/jsonl.py
@@ -25,7 +25,7 @@
 from xtuner.v1.datasets.data_item import CacheItem
 from xtuner.v1.utils import SharedMemory, get_logger
 
-from .utils import CachableTokenizeFunction, CacheObj, calculate_xxhash
+from .utils import CachableTokenizeFunction, CacheDict, CacheObj, calculate_xxhash
 
 
 T = TypeVar("T")
@@ -439,11 +439,11 @@ def count_offsets(self, cache_dir=None):
     @staticmethod
     def _tokenize_by_offset(
         data: bytes,
-        tokenize_fn: Callable[[dict], CacheObj],
+        tokenize_fn: Callable[[dict], CacheDict | CacheObj],
     ) -> dict:
         line = data.decode()
         tokenized = tokenize_fn(json.loads(line))
-        if hasattr(tokenized, "num_tokens"):
+        if isinstance(tokenized, CacheObj):
             num_tokens = tokenized.num_tokens
         else:
             num_tokens = tokenized["num_tokens"]
diff --git a/xtuner/v1/datasets/utils.py b/xtuner/v1/datasets/utils.py
@@ -9,7 +9,7 @@
 import numpy as np
 import xxhash
 from PIL import Image
-from typing_extensions import TypedDict
+from xtuner.v1.utils.cache import CacheDict, CacheObj
 
 from .data_item import CacheItem
 
@@ -20,10 +20,6 @@
     from transformers import PreTrainedTokenizer
 
 
-class CacheObj(TypedDict, total=False):
-    num_tokens: int
-
-
 class CachableTokenizeFunction(ABC, Generic[T]):
     def __init__(self, tokenizer, *args, **kwargs):
         self.tokenizer = tokenizer
diff --git a/xtuner/v1/ray/config/worker.py b/xtuner/v1/ray/config/worker.py
@@ -2,7 +2,11 @@
 import os
 import socket
 from pathlib import Path
-from typing import Any, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union
+
+
+if TYPE_CHECKING:
+    from ray.util.placement_group import PlacementGroup
 
 from cyclopts import Group, Parameter
 from pydantic import BaseModel, ConfigDict, PrivateAttr
@@ -313,6 +317,25 @@ def model_post_init(self, __context: Any) -> None:
 
         self.worker_log_dir.mkdir(parents=True, exist_ok=True)
 
+    def build(self, placement_group: "PlacementGroup"):
+        """Build and return a Ray remote RolloutController from this config.
+
+        Args:
+            placement_group: The placement group for scheduling RolloutWorker actors.
+
+        Returns:
+            A Ray actor handle (proxy) of RolloutController.
+        """
+        import ray
+
+        from xtuner.v1.ray.rollout.controller import RolloutController
+
+        return (
+            ray.remote(RolloutController)
+            .options(max_concurrency=int(os.environ.get("RAY_MAX_CONCURRENCY", 1000)))
+            .remote(self, placement_group)
+        )
+
 
 if __name__ == "__main__":
     from cyclopts import App, Group, Parameter
diff --git a/xtuner/v1/ray/judger/native.py b/xtuner/v1/ray/judger/native.py
@@ -15,6 +15,7 @@
 
 class Judger(ABC):
     @abstractmethod
+    @ray_method
     async def judge(self, rollout_state: RolloutState) -> RolloutState: ...
 
 
diff --git a/xtuner/v1/ray/rollout/controller.py b/xtuner/v1/ray/rollout/controller.py
@@ -10,9 +10,9 @@
 from uuid import uuid4
 
 import ray
-from ray.actor import ActorProxy
 import uvicorn
 from fastapi import FastAPI
+from ray.actor import ActorProxy
 from ray.util.placement_group import PlacementGroup
 
 from transformers import AutoTokenizer
@@ -468,5 +468,6 @@ def _init_workers(self):
         self.logger.info(f"Rollout worker server URLs: {list(self.workers_info.keys())}")
         return engine_rank_mesh_array, worker_server_urls_map
 
+
 RayRolloutController = ray.remote(RolloutController)
-RolloutControllerProxy = ActorProxy[RayRolloutController]
+RolloutControllerProxy = ActorProxy[RayRolloutController]
diff --git a/xtuner/v1/ray/rollout/lmdeploy.py b/xtuner/v1/ray/rollout/lmdeploy.py
@@ -107,14 +107,14 @@ def _get_request_payload(self, rollout_state: RolloutState) -> dict:
                 prompt_token_ids = self.tokenizer(text_prompt, add_special_tokens=False)["input_ids"]
                 payload["input_ids"] = prompt_token_ids
             sample_params.return_routed_experts = True if self.enable_return_routed_experts else False
-            lmdeploy_sample_params = self._transform_sample_params(sample_params)   
+            lmdeploy_sample_params = self._transform_sample_params(sample_params)
             payload.update(sample_params)
         else:
             payload = {
                 "model": self.model_name,
                 "messages": rollout_state.message,
             }
-            lmdeploy_sample_params = self._transform_sample_params(sample_params)   
+            lmdeploy_sample_params = self._transform_sample_params(sample_params)
             lmdeploy_sample_params.pop("no_stop_trim", None)
             lmdeploy_sample_params.pop("return_logprob", None)
             lmdeploy_sample_params.pop("stop_token_ids", None)
diff --git a/xtuner/v1/rl/base/agent_loop.py b/xtuner/v1/rl/base/agent_loop.py
@@ -47,6 +47,11 @@ def __init__(
     @abstractmethod
     async def generate_sample(self, rollout_state: RolloutState) -> RolloutState: ...
 
+    async def pause(self) -> None:
+        """Pause the agent loop if supported by the implementation."""
+        # Default implementation is a no-op to keep behavior unchanged.
+        return None
+
     async def generate_group(self, rollout_state: list[RolloutState]) -> list[RolloutState]:
         pending_tasks = []
         for state in rollout_state:
diff --git a/xtuner/v1/rl/base/agent_loop_manager.py b/xtuner/v1/rl/base/agent_loop_manager.py
@@ -1,17 +1,15 @@
 from pydantic import BaseModel, ConfigDict
 
-from xtuner.v1.data_proto import Status
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+from xtuner.v1.data_proto import RolloutState, Status
+from xtuner.v1.ray.judger import Judger
+from xtuner.v1.ray.rollout import RolloutController
 from xtuner.v1.rl.base.producer import ProduceStrategy, Sampler
 from xtuner.v1.rl.base.replay_buffer import ReplayBuffer
-from xtuner.v1.data_proto import RolloutState
 
-from .agent_loop import AgentLoop
-from .agent_loop import AgentLoopConfig
+from .agent_loop import AgentLoop, AgentLoopConfig
 from .producer import ProduceStrategyConfig, SyncProduceStrategyConfig
 from .sampler import SamplerConfig
-from xtuner.v1.ray.rollout import RolloutController
-from xtuner.v1.ray.judger import Judger
-from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
 
 class AgentLoopManagerConfig(BaseModel):
@@ -62,7 +60,9 @@ async def produce_batch(self, batch_size: int) -> list[list[RolloutState]]:
         await self._scheduler.produce_batch(
             self._agent_loop, self._data_sampler, self._replay_buffer, batch_size, self.task_name
         )
-        batch_rollout_states: list[list[RolloutState]] = await self._replay_buffer.get(batch_size, self.task_name, Status.COMPLETED)
+        batch_rollout_states: list[list[RolloutState]] = await self._replay_buffer.get(
+            batch_size, self.task_name, Status.COMPLETED
+        )
         return batch_rollout_states
 
     # # 非共卡
diff --git a/xtuner/v1/rl/base/producer.py b/xtuner/v1/rl/base/producer.py
@@ -35,7 +35,7 @@ class ProduceStrategyConfig(ABC, BaseModel):
     should_continue_fn: ShouldContinueFn = default_should_continue_fn
 
     @abstractmethod
-    def build(self) -> "ProduceStrategyConfig": ...
+    def build(self) -> "ProduceStrategy": ...
 
 
 class SyncProduceStrategyConfig(ProduceStrategyConfig):
diff --git a/xtuner/v1/rl/base/worker.py b/xtuner/v1/rl/base/worker.py
@@ -4,7 +4,13 @@
 import time
 from itertools import chain
 from pathlib import Path
-from typing import Dict, Iterable, List, Sequence, TypeAlias, TypedDict, cast
+from typing import TYPE_CHECKING, Dict, Iterable, List, Sequence, TypeAlias, TypedDict, cast
+
+
+if TYPE_CHECKING:
+    from ray.util.placement_group import PlacementGroup
+
+    from .controller import TrainingControllerProxy
 
 import ray
 import requests
@@ -169,6 +175,26 @@ class WorkerConfig(BaseModel):
     rollout_steps_per_sft: int = 1
     sft_loss_cfg: CELossConfig = CELossConfig()
 
+    def build(self, placement_group: "PlacementGroup") -> "TrainingControllerProxy":
+        """Build training workers and controller from this config and placement
+        group."""
+        # import here to avoid circular import
+        from xtuner.v1.ray.base import AutoAcceleratorWorkers
+        from xtuner.v1.rl.base.controller import TrainingController
+
+        TrainingWorkerCls = ray.remote(
+            runtime_env={
+                "env_vars": {
+                    "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+                    "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+                    "HCCL_NPU_SOCKET_PORT_RANGE": "auto",
+                }
+            }
+        )(TrainingWorker)
+        train_workers, _ = AutoAcceleratorWorkers.from_placement_group(TrainingWorkerCls, self, placement_group)
+        ray.wait([w.ready.remote() for w in train_workers])
+        return TrainingController.remote(workers=train_workers)
+
 
 class WorkerInputItem(TypedDict):
     seq_ctx: SequenceContext
diff --git a/xtuner/v1/rl/evaluator.py b/xtuner/v1/rl/evaluator.py
@@ -1,6 +1,8 @@
-from pydantic import BaseModel, ConfigDict, Field
+from collections.abc import Mapping
 from typing import Annotated, Protocol, cast, runtime_checkable
+
 from cyclopts import Parameter
+from pydantic import BaseModel, ConfigDict, Field
 
 from xtuner.v1.data_proto import RolloutState
 
@@ -11,7 +13,17 @@ def __call__(self, samples: list[RolloutState]) -> dict[str, float]: ...
 
 
 def default_compute_metric_func(samples: list[RolloutState]) -> dict[str, float]:
-    return {"accuracy": sum(s.reward["score"] > 0 for s in samples) / len(samples)}
+    if not samples:
+        return {"accuracy": 0.0}
+
+    positive = 0
+    for s in samples:
+        reward = s.reward
+        assert isinstance(reward, Mapping)
+        score = reward["score"]
+        if score > 0:
+            positive += 1
+    return {"accuracy": positive / len(samples)}
 
 
 class Evaluator:
@@ -54,7 +66,9 @@ def build(self, total_eval_samples: int = 0) -> "Evaluator":
         if self.eval_sample_num > 0:
             eval_batch_size = self.eval_sample_num
         else:
-            assert total_eval_samples > 0, "Total eval samples must be greater than 0 if eval sample num is not provided"
+            assert total_eval_samples > 0, (
+                "Total eval samples must be greater than 0 if eval sample num is not provided"
+            )
             if self.eval_sample_ratio > 0:
                 eval_batch_size = int(total_eval_samples * self.eval_sample_ratio)
             else:
@@ -63,4 +77,4 @@ def build(self, total_eval_samples: int = 0) -> "Evaluator":
         return Evaluator(
             compute_metric_func=self.compute_metric_func,
             eval_batch_size=eval_batch_size,
-        )
+        )
diff --git a/xtuner/v1/train/rl_colocate_trainer.py b/xtuner/v1/train/rl_colocate_trainer.py
diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py
diff --git a/xtuner/v1/utils/type_helper.py b/xtuner/v1/utils/type_helper.py
