flush cache when update weight in disag mode

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()

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")

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")

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"])

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.")
 

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,