diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 1f16cf0b8..bad8565f7 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -154,6 +154,7 @@ def _make_trainer(self, agent_loop_manager): restart_inactive_workers=SimpleNamespace(remote=MagicMock(return_value="rollout_restarted")), pause_generation=SimpleNamespace(remote=MagicMock(return_value="pause")), continue_generation=SimpleNamespace(remote=MagicMock(return_value="continue")), + flush_cache=SimpleNamespace(remote=MagicMock(return_value="flush_cache")), onload_weights=SimpleNamespace(remote=MagicMock(return_value="onload_weights")), onload_kvcache=SimpleNamespace(remote=MagicMock(return_value="onload_kvcache")), validate_registered_workers_to_proxy=SimpleNamespace(remote=AsyncMock(return_value=None)), @@ -266,6 +267,7 @@ def test_fit_rebinds_weight_update_with_rollout_update_address(self): with ( patch("xtuner.v1.train.rl_trainer.asyncio_run", side_effect=asyncio.run), + patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=lambda obj, timeout=None: obj), patch("xtuner.v1.train.rl_trainer.bind_train_rollout") as bind_train_rollout_mock, ): trainer.fit() diff --git a/tests/rl/test_rl_trainer_checkpoint.py b/tests/rl/test_rl_trainer_checkpoint.py index d78afc646..0c4d53fa3 100644 --- a/tests/rl/test_rl_trainer_checkpoint.py +++ b/tests/rl/test_rl_trainer_checkpoint.py @@ -100,6 +100,7 @@ def __init__(self): self.generate = _RemoteMethod(self._generate, async_result=True) self.pause_generation = _RemoteMethod(async_result=True) self.continue_generation = _RemoteMethod(async_result=True) + self.flush_cache = _RemoteMethod(return_value="cache_flushed") self.offload = _RemoteMethod(return_value="rollout_offloaded") self.check_and_shutdown_inactive_workers = _RemoteMethod(return_value="rollout_inactive_workers_shutdown") self.restart_inactive_workers = _RemoteMethod(return_value="rollout_restarted") diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index b2ab50e23..3db2874d0 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -168,6 +168,9 @@ def continue_generation(self): def offload(self): self._broadcast_to_active_workers("offload") + def flush_cache(self): + self._broadcast_to_active_workers("flush_cache") + def onload(self): self._broadcast_to_active_workers("onload_weights") self._broadcast_to_active_workers("onload_kvcache") diff --git a/xtuner/v1/rl/rollout/lmdeploy.py b/xtuner/v1/rl/rollout/lmdeploy.py index 6d368286b..220f60e05 100644 --- a/xtuner/v1/rl/rollout/lmdeploy.py +++ b/xtuner/v1/rl/rollout/lmdeploy.py @@ -213,6 +213,12 @@ def offload(self): """Offloads the model weights and KV cache.""" return self._sleep(level=2) + def flush_cache(self): + """Flushes cache through LMDeploy sleep/wakeup lifecycle.""" + self.offload() + self.onload_weights() + return self.onload_kvcache() + def onload_weights(self): """Onloads the model weights by waking up the model.""" return self._wake_up(tags=["weights"]) diff --git a/xtuner/v1/rl/rollout/vllm.py b/xtuner/v1/rl/rollout/vllm.py index d601988b9..21571c920 100644 --- a/xtuner/v1/rl/rollout/vllm.py +++ b/xtuner/v1/rl/rollout/vllm.py @@ -287,6 +287,12 @@ def offload(self): """Offloads the model weights and KV cache.""" return self.sleep(level=2) + def flush_cache(self): + """Flushes cache through vLLM sleep/wakeup lifecycle.""" + self.offload() + self.onload_weights() + return self.onload_kvcache() + def reset_prefix_cache(self, tags: List[str] | None = None): raise NotImplementedError("The 'reset_prefix_cache' API is not yet implemented in the vLLM server.") diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 61ef1d1db..4d9ab95c8 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -2091,6 +2091,10 @@ async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): # TODO: 非共卡需要额外加健康检查恢复worker的逻辑,共卡是在训练之前恢复,但是非共卡不需要在训练之前恢复,挂掉就恢复或者更新权重前恢复,需要评估一下哪种方式更合理。 with timer("sync_weight", step_timer_dict): + ray.get( + self.rollout_controller.flush_cache.remote(), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) bind_train_rollout( train_controller=self.train_controller, rollout_controller=self.rollout_controller,