diff --git a/tests/rl/test_rl_trainer_checkpoint.py b/tests/rl/test_rl_trainer_checkpoint.py index cb2977b6c..8c6dea436 100644 --- a/tests/rl/test_rl_trainer_checkpoint.py +++ b/tests/rl/test_rl_trainer_checkpoint.py @@ -135,6 +135,7 @@ def bind_rollout_weight_update( weight_transport_type, weight_update_host=None, weight_update_port=None, + disk_weight_path=None, ): self.rollout_info = { "targets": targets, @@ -143,7 +144,7 @@ def bind_rollout_weight_update( self.weight_transport_type = weight_transport_type self.weight_update_host = weight_update_host self.weight_update_port = weight_update_port - + self.disk_weight_path = disk_weight_path def onload(self, target="all"): return f"onload:{target}" diff --git a/tests/rl/test_update_weight_disk.py b/tests/rl/test_update_weight_disk.py new file mode 100644 index 000000000..5891060ca --- /dev/null +++ b/tests/rl/test_update_weight_disk.py @@ -0,0 +1,326 @@ +import asyncio +import os +import tempfile +import unittest +from pathlib import Path + +import ray + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams +from xtuner.v1.model.compose.qwen3_vl import Qwen3VLDense4BConfig +from xtuner.v1.rl.loss import GRPOLossConfig as LossConfig +from xtuner.v1.rl.rollout.constants import ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY +from xtuner.v1.rl.rollout.controller import RolloutController +from xtuner.v1.rl.rollout.sglang import SGLangWorker +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trainer import ( + TrainingController, + TrainingWorker as BaseTrainingWorker, + WorkerConfig, +) +from xtuner.v1.rl.utils import ( + AcceleratorResourcesConfig, + AutoAcceleratorWorkers, + CPUResourcesConfig, + CPUResourceManager, + clear_cpu_resource_manager, + register_cpu_resources, + set_cpu_resource_manager, +) + + +TEST_TEXT_MESSAGES = [{"role": "user", "content": "Hello!"}] +MODEL_PATH = os.environ.get("QWEN3_VL_DENSE_PATH") + + +class DiskUpdateTestSGLangWorker(SGLangWorker): + """Test-only SGLang worker that keeps recovered workers onloaded.""" + + def offload(self): + # restart_inactive_workers finally calls worker.actor.offload, but disaggregated recovery does not need offload. + # The production recovery path offloads restarted workers to match the + # colocated rollout baseline. This disk-recovery test immediately loads + # weights from disk, so keep the test server alive and probeable. + return {"success": True} + + +class DiskUpdateTestRolloutController(RolloutController): + """Rollout controller with test hooks for targeting one recovered engine.""" + + def _build_remote_worker_cls(self, worker_base_cls): + # Force this E2E to use the test worker above. The passed worker class is + # the production backend class selected from RolloutConfig. + del worker_base_cls + return super()._build_remote_worker_cls(DiskUpdateTestSGLangWorker) + + async def generate_from_weight_update_endpoint( + self, + *, + endpoint_rank: int, + rollout_state: RolloutState, + ) -> RolloutState: + # Production routing is session-based and may choose any active engine. + # This test must verify the exact engine that was killed and restarted. + worker = self.registry.active_entrypoint_by_rank(endpoint_rank) + if worker is None: + raise RuntimeError(f"Rollout endpoint rank={endpoint_rank} is not active.") + + response_ref = worker.actor.generate.remote(rollout_state=rollout_state) # type: ignore[attr-defined] + return await asyncio.wait_for(response_ref, timeout=self.config.rollout_timeout * self.timeout_multiplier) + + def shutdown_weight_update_endpoint(self, endpoint_rank: int | None = None) -> tuple[int, ...]: + + # Tests need a deterministic failure injection point. Mark the lifecycle + # group inactive before shutdown so restart_inactive_workers() can claim it. + targets = [target for target in self.registry.weight_update_targets() if target.is_active] + if not targets: + raise RuntimeError("No active rollout weight-update endpoint can be shut down.") + + target = targets[0] if endpoint_rank is None else next( + (target for target in targets if target.endpoint_rank == endpoint_rank), + None, + ) + if target is None: + raise RuntimeError(f"No active rollout weight-update endpoint rank={endpoint_rank} can be shut down.") + + with self.health_manager._paused_lifecycle_operation(): + groups = self.registry.mark_unhealthy_ranks({target.endpoint_rank}) + shutdown_ranks: list[int] = [] + for group in groups: + if not self.health_manager._shutdown_worker_group(group): + raise RuntimeError(f"Failed to shut down rollout worker group ranks={group.ranks}.") + shutdown_ranks.extend(group.ranks) + return tuple(shutdown_ranks) + +class TestUpdateWeightDisk(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + if MODEL_PATH is None: + raise unittest.SkipTest("QWEN3_VL_DENSE_PATH is not set") + os.environ["XTUNER_USE_FA3"] = "1" + # TODO(shipengcheng): SGLang disaggregated weight update cannot use + # NCCL_CUMEM for now. Remove this after the root cause is fixed. + os.environ["NCCL_CUMEM_ENABLE"] = "0" + + @classmethod + def tearDownClass(cls) -> None: + os.environ.pop("XTUNER_USE_FA3", None) + + def setUp(self): + self._original_pytorch_cuda_alloc_conf = os.environ.pop("PYTORCH_CUDA_ALLOC_CONF", None) + ray.init(num_cpus=128, ignore_reinit_error=True) + self.temp_dir = tempfile.TemporaryDirectory() + self.worker_log_dir = os.path.join(self.temp_dir.name, "work_dirs") + self.init_config() + self.train_pg = AutoAcceleratorWorkers.build_placement_group( + self.train_resources_cfg, + name=f"test_update_weight_disk_train_{id(self)}", + ) + self.rollout_pg = AutoAcceleratorWorkers.build_placement_group( + self.rollout_resources_cfg, + name=f"test_update_weight_disk_rollout_{id(self)}", + ) + set_cpu_resource_manager(CPUResourceManager(accelerator_placement_groups=[self.train_pg, self.rollout_pg])) + + def tearDown(self): + clear_cpu_resource_manager() + ray.shutdown() + self.temp_dir.cleanup() + if self._original_pytorch_cuda_alloc_conf is not None: + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = self._original_pytorch_cuda_alloc_conf + else: + os.environ.pop("PYTORCH_CUDA_ALLOC_CONF", None) + + def init_config(self): + train_num_workers = int(os.environ.get("TRAIN_NUM_WORKERS", "4")) + rollout_tp_size = int(os.environ.get("ROLLOUT_TP_SIZE", "2")) + # Use at least two rollout engines by default so the test exercises + # recovery of one failed engine while another engine remains active. + rollout_num_workers = int(os.environ.get("ROLLOUT_NUM_WORKERS", str(rollout_tp_size * 2))) + if rollout_num_workers < rollout_tp_size * 2: + raise unittest.SkipTest("Disk recovery E2E requires at least two rollout engines.") + if rollout_num_workers % rollout_tp_size != 0: + raise unittest.SkipTest("ROLLOUT_NUM_WORKERS must be divisible by ROLLOUT_TP_SIZE.") + + self.train_resources_cfg = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=train_num_workers, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, + ) + self.rollout_resources_cfg = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=rollout_num_workers, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, + ) + self.rollout_cfg = RolloutConfig( + env="test_rollout_disk", + model_path=MODEL_PATH, + model_name=os.path.basename(MODEL_PATH).lower(), + tokenizer_path=MODEL_PATH, + rollout_cross_node_comm=False, + tensor_parallel_size=rollout_tp_size, + expert_parallel_size=1, + gpus_per_node=int(os.environ.get("GPUS_PER_NODE", "8")), + dtype="bfloat16", + skip_load_weights=False, + context_length=256, + worker_log_dir=self.worker_log_dir, + gpu_memory_utilization=float(os.environ.get("ROLLOUT_GPU_MEMORY_UTILIZATION", "0.5")), + ) + + self.worker_cfg = WorkerConfig( + model_cfg=Qwen3VLDense4BConfig(), + optim_cfg=AdamWConfig(lr=5e-7, foreach=False), + loss_cfg=LossConfig( + policy_loss_cfg=dict( + cliprange_high=0.28, + cliprange_low=0.2, + loss_type="vanilla", + ), + ignore_idx=-100, + use_kl_loss=False, + kl_loss_coef=0.001, + kl_loss_type="low_var_kl", + mode="eager", + ), + lr_cfg=LRConfig(lr_type="constant", warmup_ratio=0, lr_min=5e-7), + fsdp_cfg=FSDPConfig(ep_size=1), + load_from=MODEL_PATH, + sp_size=1, + pack_max_length=1024, + ) + + def _build_training_controller(self) -> TrainingController: + TrainingWorker = ray.remote( + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + } + }, + )(BaseTrainingWorker) + train_workers, _ = AutoAcceleratorWorkers.from_placement_group( + TrainingWorker, + self.worker_cfg, + self.train_pg, + ) + ray.get([worker.test_all_reduce.remote() for worker in train_workers]) + return TrainingController(workers=train_workers) + + def _build_rollout_controller(self): + num_workers = 1 + register_cpu_resources( + name="rollout_controller", + cpu_resources=CPUResourcesConfig(num_workers=num_workers), + ) + test_dir = Path(__file__).resolve().parent + pythonpath = os.pathsep.join( + path for path in (str(test_dir), os.environ.get("PYTHONPATH", "")) if path + ) + return ( + ray.remote( + concurrency_groups={ + "generate": ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, + }, + runtime_env={"env_vars": {"PYTHONPATH": pythonpath}}, + )(DiskUpdateTestRolloutController) + .options(num_cpus=num_workers) + .remote(self.rollout_cfg, self.rollout_pg) + ) + + def _update_weights( + self, + train_controller, + rollout_controller, + weight_transport_type, + target_endpoint_ranks: set[int] | None = None, + **kwargs, + ) -> None: + targets = ray.get(rollout_controller.get_weight_update_targets.remote()) + if target_endpoint_ranks is not None: + targets = tuple(target for target in targets if target.endpoint_rank in target_endpoint_ranks) + self.assertGreater( + len(targets), + 0, + f"No rollout weight-update targets matched endpoint ranks={sorted(target_endpoint_ranks)}.", + ) + train_controller.bind_rollout_weight_update( + targets=targets, + rollout_config=self.rollout_cfg, + weight_transport_type=weight_transport_type, + **kwargs, + ) + train_controller.update_weights() + + @unittest.skipIf(os.environ.get("XTUNER_USE_SGLANG", "0") == "0", "sglang backend is not enabled") + def test_sglang_disaggregated_disk_update_after_engine_recovery(self): + train_controller = self._build_training_controller() + rollout_controller = self._build_rollout_controller() + + try: + sample_params = SampleParams(temperature=0.0, max_tokens=128, top_k=1) + input_state = RolloutState(message=TEST_TEXT_MESSAGES, sample_params=sample_params) + + self._update_weights(train_controller, rollout_controller, "nccl") + + initial_targets = ray.get(rollout_controller.get_weight_update_targets.remote()) + self.assertGreaterEqual( + len(initial_targets), + 2, + "This test requires multiple rollout engines. Set ROLLOUT_NUM_WORKERS >= 2 * ROLLOUT_TP_SIZE.", + ) + failed_endpoint_rank = initial_targets[0].endpoint_rank + + # Warm up every rollout engine before killing one of them. + endpoint_results = { + target.endpoint_rank: ray.get( + rollout_controller.generate_from_weight_update_endpoint.remote( + endpoint_rank=target.endpoint_rank, + rollout_state=input_state.model_copy(deep=True), + ) + ) + for target in initial_targets + } + baseline = endpoint_results[failed_endpoint_rank] + + hf_dir = Path(self.temp_dir.name) / "hf-disk-update" + hf_dir.mkdir(parents=True, exist_ok=True) + train_controller.save_hf(str(hf_dir)) + + shutdown_ranks = ray.get( + rollout_controller.shutdown_weight_update_endpoint.remote(endpoint_rank=failed_endpoint_rank) + ) + self.assertGreater(len(shutdown_ranks), 0) + + targets_after_shutdown = ray.get(rollout_controller.get_weight_update_targets.remote()) + self.assertTrue(any(not target.is_active for target in targets_after_shutdown)) + + ray.get(rollout_controller.restart_inactive_workers.remote()) + targets_after_restart = ray.get(rollout_controller.get_weight_update_targets.remote()) + self.assertTrue(all(target.is_active for target in targets_after_restart)) + self.assertTrue(any(target.endpoint_rank == failed_endpoint_rank for target in targets_after_restart)) + + self._update_weights( + train_controller, + rollout_controller, + "disk", + target_endpoint_ranks={failed_endpoint_rank}, + disk_weight_path=str(hf_dir), + ) + + recovered = ray.get( + rollout_controller.generate_from_weight_update_endpoint.remote( + endpoint_rank=failed_endpoint_rank, + rollout_state=input_state.model_copy(deep=True), + ) + ) + self.assertEqual(recovered.response, baseline.response) + finally: + ray.get(rollout_controller.shutdown.remote(), timeout=60) + + +if __name__ == "__main__": + unittest.main() diff --git a/xtuner/v1/rl/rollout/sglang.py b/xtuner/v1/rl/rollout/sglang.py index 6825fdcfb..0f2f2049d 100644 --- a/xtuner/v1/rl/rollout/sglang.py +++ b/xtuner/v1/rl/rollout/sglang.py @@ -418,6 +418,9 @@ def _transform_rollout_config_to_server_configs(self): if self.config.context_length is not None: init_kwargs["context_length"] = self.config.context_length + if self.config.skip_load_weights and "load_format" not in sglang_config_kwargs: + init_kwargs["load_format"] = "dummy" + sglang_server_args = ServerArgs(**init_kwargs) return sglang_server_args diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index 87638a8ab..cbf542855 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -298,6 +298,7 @@ def bind_rollout_weight_update( weight_transport_type, weight_update_host=None, weight_update_port=None, + disk_weight_path=None, ): ray.get( [ @@ -307,6 +308,7 @@ def bind_rollout_weight_update( weight_transport_type=weight_transport_type, weight_update_host=weight_update_host, weight_update_port=weight_update_port, + disk_weight_path=disk_weight_path, ) for worker in self.workers ] diff --git a/xtuner/v1/rl/utils/ray_utils.py b/xtuner/v1/rl/utils/ray_utils.py index 2e0273980..c16e0fca6 100644 --- a/xtuner/v1/rl/utils/ray_utils.py +++ b/xtuner/v1/rl/utils/ray_utils.py @@ -163,6 +163,7 @@ def bind_train_rollout( weight_transport_type, weight_update_host=None, weight_update_port=None, + disk_weight_path=None, ) -> None: """Bind the training and rollout workers for updating weights. @@ -183,6 +184,7 @@ def bind_train_rollout( weight_transport_type=weight_transport_type, weight_update_host=weight_update_host, weight_update_port=weight_update_port, + disk_weight_path=disk_weight_path, ) for worker in train_workers ] diff --git a/xtuner/v1/rl/weight_update/__init__.py b/xtuner/v1/rl/weight_update/__init__.py index e2268c11f..fff481ce7 100644 --- a/xtuner/v1/rl/weight_update/__init__.py +++ b/xtuner/v1/rl/weight_update/__init__.py @@ -6,11 +6,14 @@ WeightUpdateBatch, ) from .transport import ( + DiskBackendAdapter, + DiskWeightTransport, IPCBackendAdapter, IPCWeightTransport, LMDeployIPCBackendAdapter, NCCLBackendAdapter, NCCLWeightTransport, + SGLangDiskBackendAdapter, SGLangIPCBackendAdapter, SGLangNCCLBackendAdapter, WeightTransport, @@ -21,6 +24,8 @@ __all__ = [ + "DiskBackendAdapter", + "DiskWeightTransport", "IPCBackendAdapter", "IPCWeightTransport", "LMDeployIPCBackendAdapter", @@ -29,6 +34,7 @@ "RolloutBackend", "RolloutWeightUpdateTarget", "RolloutWeightUpdateInfo", + "SGLangDiskBackendAdapter", "SGLangIPCBackendAdapter", "SGLangNCCLBackendAdapter", "UpdateWeighter", diff --git a/xtuner/v1/rl/weight_update/data.py b/xtuner/v1/rl/weight_update/data.py index 08007b5a4..c5a17a81f 100644 --- a/xtuner/v1/rl/weight_update/data.py +++ b/xtuner/v1/rl/weight_update/data.py @@ -12,7 +12,7 @@ RolloutBackend: TypeAlias = Literal["sglang", "vllm", "pytorch", "turbomind"] # Rollout inference backend. -WeightTransportType: TypeAlias = Literal["ipc", "nccl"] # Supported weight transport types. +WeightTransportType: TypeAlias = Literal["ipc", "nccl", "disk"] # Supported weight transport types. def _resolve_rollout_backend(rollout_config: RolloutConfig) -> RolloutBackend: @@ -40,11 +40,15 @@ def _validate_transport_type( assert weight_transport_type is not None, "bind_rollout_weight_update() must set weight_transport_type." transport_type = weight_transport_type.lower() - if transport_type not in ("ipc", "nccl"): - raise ValueError(f"Unsupported weight_transport_type: {weight_transport_type!r}. Expected 'ipc' or 'nccl'.") + if transport_type not in ("ipc", "nccl", "disk"): + raise ValueError( + f"Unsupported weight_transport_type: {weight_transport_type!r}. Expected 'ipc', 'nccl' or 'disk'." + ) transport_type = cast(WeightTransportType, transport_type) if transport_type == "nccl" and backend in ("vllm", "turbomind"): raise NotImplementedError(f"NCCL weight transport is not supported for {backend} backend.") + if transport_type == "disk" and backend != "sglang": + raise ValueError(f"Disk weight transport is not supported for {backend} backend.") return transport_type @@ -86,6 +90,8 @@ class RolloutWeightUpdateInfo: weight_update_host: str | None = None # Optional port used by NCCL external weight update groups. weight_update_port: int | None = None + # Optional disk weight path used by disk weight transport. + disk_weight_path: str | None = None @classmethod def from_targets( @@ -97,6 +103,7 @@ def from_targets( weight_transport_type: WeightTransportType | str, weight_update_host: str | None = None, weight_update_port: int | None = None, + disk_weight_path: str | None = None, ) -> RolloutWeightUpdateInfo: backend = _resolve_rollout_backend(rollout_config) tp = rollout_config.tensor_parallel_size @@ -106,6 +113,8 @@ def from_targets( weight_transport_type=weight_transport_type, backend=backend, ) + if transport_type == "disk" and not disk_weight_path: + raise ValueError("Disk weight transport requires disk_weight_path.") return cls( rollout_config=rollout_config, weight_update_targets=weight_update_targets, @@ -114,6 +123,7 @@ def from_targets( backend=backend, weight_update_host=weight_update_host, weight_update_port=weight_update_port if weight_update_port is not None else 30000, + disk_weight_path=disk_weight_path, ) @property diff --git a/xtuner/v1/rl/weight_update/transport.py b/xtuner/v1/rl/weight_update/transport.py index 637da8648..9ce349a9a 100644 --- a/xtuner/v1/rl/weight_update/transport.py +++ b/xtuner/v1/rl/weight_update/transport.py @@ -838,3 +838,105 @@ def teardown(self) -> None: self.group_name = None self.engine_urls = [] self.external_group_world_size = None + + +class DiskBackendAdapter: + def build_weight_update_payload(self, hf_weight_path: str) -> dict[str, Any]: + raise NotImplementedError + + def build_request( + self, + payload: dict[str, Any], + ) -> WeightUpdateRequest: + raise NotImplementedError + + def update(self, weight_iterator: Any) -> None: + raise NotImplementedError + + def teardown(self) -> None: + return + + +class SGLangDiskBackendAdapter(DiskBackendAdapter): + def __init__(self, *, rank: int, rollout_info: RolloutWeightUpdateInfo): + self.rank = rank + self.rollout_info = rollout_info + self.executor: ThreadPoolExecutor | None = None + + def build_weight_update_payload(self, hf_weight_path: str) -> dict[str, Any]: + # SGLang already owns the disk reload path. XTuner only needs to pass + # the HF checkpoint directory to the rollout server. + return { + "model_path": hf_weight_path, + "load_format": "safetensors", + "abort_all_requests": True, + "flush_cache": True, + } + + def build_request( + self, + payload: dict[str, Any], + ) -> WeightUpdateRequest: + return WeightUpdateRequest(endpoint="update_weights_from_disk", body=payload) + + def update(self, weight_iterator: Any) -> None: + # SGLang consumes the checkpoint path on the rollout server side. + del weight_iterator + + disk_weight_path = self.rollout_info.disk_weight_path + if not disk_weight_path: + raise RuntimeError("Disk weight update requires rollout_info.disk_weight_path from rollout_config.") + + try: + if dist.get_rank() != 0: + dist.barrier() + return + + target_urls = [t.server_url for t in self.rollout_info.active_update_targets] + if not target_urls: + raise RuntimeError("Disk weight update requires at least one rollout server url.") + payload = self.build_weight_update_payload(disk_weight_path) + request = self.build_request(payload) + self.executor = ThreadPoolExecutor(max_workers=max(1, len(target_urls))) + futures = [ + self.executor.submit( + WeightTransport.post_json, + url, + request.endpoint, + request.body, + api_key=self.rollout_info.api_key, + ) + for url in target_urls + ] + for future in futures: + result = future.result() + assert result.get("success", True), f"disk weight update failed: {result.get('message', result)}" + dist.barrier() + finally: + self.teardown() + DEVICE_MODULE.empty_cache() + + def teardown(self) -> None: + if self.executor is not None: + self.executor.shutdown(wait=False, cancel_futures=True) + self.executor = None + + +class DiskWeightTransport(WeightTransport): + _disk_adapter: DiskBackendAdapter + + def __init__(self, *, rank: int, logger: Any, rollout_info: RolloutWeightUpdateInfo, config: Any | None = None): + super().__init__(rank=rank, logger=logger, rollout_info=rollout_info) + self.config = config + self._disk_adapter = self._build_adapter() + + def _build_adapter(self) -> DiskBackendAdapter: + if self.backend == "sglang": + return SGLangDiskBackendAdapter(rank=self.rank, rollout_info=self.rollout_info) + raise ValueError(f"Unsupported disk weight update backend: {self.backend!r}") + + def update(self, weight_iterator: Any) -> None: + self._disk_adapter.update(weight_iterator) + + def send(self, batch: WeightUpdateBatch) -> None: + raise NotImplementedError("DiskWeightTransport bypasses WeightIterator batches.") diff --git a/xtuner/v1/rl/weight_update/update_weighter.py b/xtuner/v1/rl/weight_update/update_weighter.py index db4558845..f8dcaa5bd 100644 --- a/xtuner/v1/rl/weight_update/update_weighter.py +++ b/xtuner/v1/rl/weight_update/update_weighter.py @@ -9,7 +9,7 @@ RolloutWeightUpdateTarget, WeightTransportType, ) -from .transport import IPCWeightTransport, NCCLWeightTransport, WeightTransport +from .transport import DiskWeightTransport, IPCWeightTransport, NCCLWeightTransport, WeightTransport from .weight_iterator import WeightIterator @@ -37,6 +37,7 @@ def bind_rollout_weight_update( weight_transport_type: WeightTransportType, weight_update_host: str | None = None, weight_update_port: int | None = None, + disk_weight_path: str | None = None, ): """Bind this train worker to rollout weight-update targets.""" @@ -47,6 +48,7 @@ def bind_rollout_weight_update( weight_transport_type=weight_transport_type, weight_update_host=weight_update_host, weight_update_port=weight_update_port, + disk_weight_path=disk_weight_path, ) new_transport_signature = self.rollout_info.transport_signature @@ -90,6 +92,8 @@ def _set_transport(self) -> None: ) elif rollout_info.transport_type == "nccl": self._transport = NCCLWeightTransport(rank=self.rank, logger=self.logger, rollout_info=rollout_info) + elif rollout_info.transport_type == "disk": + self._transport = DiskWeightTransport(rank=self.rank, logger=self.logger, rollout_info=rollout_info) else: raise NotImplementedError diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index ad9f4689c..cb5aa935f 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -123,6 +123,7 @@ def bind_train_rollout( weight_transport_type: WeightTransportType | str, weight_update_host: str | None = None, weight_update_port: int | None = None, + disk_weight_path: str | None = None, ) -> None: """Bind the training and rollout workers for update weights.""" targets = ray.get( @@ -135,6 +136,7 @@ def bind_train_rollout( weight_transport_type=weight_transport_type, weight_update_host=weight_update_host, weight_update_port=weight_update_port, + disk_weight_path=disk_weight_path, ) return @@ -470,6 +472,7 @@ class RLColocateTrainerConfig(BaseRLTrainerConfig): agent_loop_manager_cfg: AgentLoopManagerConfig resources: AcceleratorResourcesConfig + weight_transport_type: WeightTransportType = "ipc" def build(self) -> "RLColocateTrainer": return RLColocateTrainer(self) @@ -1549,6 +1552,7 @@ def _log_mini_batch_metrics(self, workers_log_item: List[WorkerLogItem]): class RLColocateTrainer(BaseRLTrainer): _META_PATH = ".xtuner_rl_colocate_trainer" + _DISK_WEIGHT_UPDATE_DIR = "disk_weight_update" agent_loop_manager: AgentLoopManager # 共卡保留资源切换和权重同步流程；通用保存、日志在 BaseRLTrainer。 @@ -1556,6 +1560,12 @@ def __init__(self, cfg: RLColocateTrainerConfig): self._init_common(cfg, meta_path=self._META_PATH, logger_tag="RLTrainer") self._num_workers = float(cfg.resources.num_workers) self._rollout_num_workers = float(cfg.resources.num_workers) + self._weight_transport_type = cfg.weight_transport_type + if self._weight_transport_type not in ("ipc", "disk"): + raise ValueError( + f"RLColocateTrainer only supports weight_transport_type 'ipc' or 'disk', " + f"got {self._weight_transport_type!r}." + ) self._pg = AutoAcceleratorWorkers.build_placement_group(cfg.resources) self._cpu_resource_manager = CPUResourceManager(self._pg) @@ -1598,12 +1608,8 @@ def __init__(self, cfg: RLColocateTrainerConfig): self.train_controller.offload(target="all") self.rollout_controller = self._rollout_config.build(self._pg) - bind_train_rollout( - train_controller=self.train_controller, - rollout_controller=self.rollout_controller, - rollout_config=self._rollout_config, - weight_transport_type="ipc", - ) + if self._weight_transport_type == "ipc": + self._bind_colocate_weight_update() replay_buffer = cfg.replay_buffer_config.build() self._build_agent_loop_components(cfg, replay_buffer) @@ -1615,13 +1621,39 @@ def __init__(self, cfg: RLColocateTrainerConfig): if self._rollout_config.skip_load_weights: self._sync_weights_from_train_workers() + def _bind_colocate_weight_update(self, disk_weight_path: Path | str | None = None) -> None: + bind_train_rollout( + train_controller=self.train_controller, + rollout_controller=self.rollout_controller, + rollout_config=self._rollout_config, + weight_transport_type=self._weight_transport_type, + disk_weight_path=str(disk_weight_path) if disk_weight_path is not None else None, + ) + + def _save_disk_weight_checkpoint(self) -> Path: + save_hf_path = self.exp_dir / self._DISK_WEIGHT_UPDATE_DIR + if save_hf_path.exists(): + rmtree(save_hf_path, ignore_errors=True) + save_hf_path.mkdir(parents=True, exist_ok=True) + self.logger.info(f"Saving disk weight checkpoint to {save_hf_path}") + self.train_controller.save_hf(str(save_hf_path)) + if isinstance(self.tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)): + self.tokenizer.save_pretrained(str(save_hf_path)) + return save_hf_path + def _sync_weights_from_train_workers(self) -> None: self.logger.info("Rollout workers skip load weights, update weights from train workers.") ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) self.train_controller.onload(target="model") + if self._weight_transport_type == "disk": + # disk方式在保存权重到磁盘后可立即offload model，避免模型占用显存。 + disk_weight_path = self._save_disk_weight_checkpoint() + self._bind_colocate_weight_update(disk_weight_path=disk_weight_path) + self.train_controller.offload(target="model") ray.get(self.rollout_controller.onload_weights.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) self.train_controller.update_weights() - self.train_controller.offload(target="model") + if self._weight_transport_type == "ipc": + self.train_controller.offload(target="model") ray.get(self.rollout_controller.onload_kvcache.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) self.logger.info("Rollout workers updated weights from train workers.") @@ -1759,19 +1791,18 @@ def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool self.rollout_controller.restart_inactive_workers.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT, ) - bind_train_rollout( - train_controller=self.train_controller, - rollout_controller=self.rollout_controller, - rollout_config=self._rollout_config, - weight_transport_type="ipc", - ) + if self._weight_transport_type == "disk": + disk_weight_path = self._save_disk_weight_checkpoint() + self.train_controller.offload(target="model") + self._bind_colocate_weight_update(disk_weight_path=disk_weight_path) ray.get( self.rollout_controller.onload_weights.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT, ) self.train_controller.update_weights() self.logger.info("Rollout workers update weights successfully in colocate mode") - self.train_controller.offload(target="model") + if self._weight_transport_type == "ipc": + self.train_controller.offload(target="model") else: self.train_controller.offload(target="model") ray.get(self.rollout_controller.onload_weights.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT)