From 74f14e262a28f8d5633a96ad2c1bfb3efc420d55 Mon Sep 17 00:00:00 2001 From: Pengcheng Shi Date: Fri, 21 Aug 2026 10:35:52 +0800 Subject: [PATCH 1/2] flush cache when update weight in disag mode --- xtuner/v1/rl/rollout/controller.py | 3 +++ xtuner/v1/rl/rollout/lmdeploy.py | 6 ++++++ xtuner/v1/rl/rollout/vllm.py | 6 ++++++ xtuner/v1/train/rl_trainer.py | 4 ++++ 4 files changed, 19 insertions(+) diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index b2ab50e23c..3db2874d03 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 6d368286b6..220f60e05e 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 d601988b91..21571c920f 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 61ef1d1dbf..4d9ab95c8e 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, From 01076568a6f98fa8f29fa309b207e97eedea4c91 Mon Sep 17 00:00:00 2001 From: Pengcheng Shi Date: Mon, 24 Aug 2026 13:41:08 +0800 Subject: [PATCH 2/2] Fix tests when using flush_cache --- tests/rl/test_rl_disaggregated_trainer.py | 2 ++ tests/rl/test_rl_trainer_checkpoint.py | 1 + 2 files changed, 3 insertions(+) diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 1f16cf0b86..bad8565f74 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 d78afc6463..0c4d53fa3e 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")