From 1c01bced5a1a6ad00599ed6054d984f55495c028 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 17 Jul 2026 11:58:20 +0000 Subject: [PATCH] [Feat] Support immediate rollout recovery by export hf asynchronously --- xtuner/v1/rl/rollout/controller.py | 14 +++ xtuner/v1/rl/rollout/health_manager.py | 104 ++++++++++---------- xtuner/v1/rl/rollout/worker.py | 25 +++-- xtuner/v1/rl/trainer/controller.py | 16 ++++ xtuner/v1/rl/trainer/worker.py | 24 +++++ xtuner/v1/train/rl_trainer.py | 125 ++++++++++++++++++++++++- 6 files changed, 244 insertions(+), 64 deletions(-) diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index cee7d7c6e..39ff6c28b 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -161,6 +161,20 @@ async def restart_inactive_workers(self): """Restart inactive groups before a sync-step weight update.""" await asyncio.to_thread(self.health_manager.restart_inactive_workers) + def set_ready_recovery_hf( + self, + *, + model_path: str, + tokenizer_path: str | None = None, + ) -> None: + self.health_manager.set_ready_recovery_hf( + model_path=model_path, + tokenizer_path=tokenizer_path, + ) + + def clear_ready_recovery_hf(self) -> None: + self.health_manager.clear_ready_recovery_hf() + def continue_generation(self): self._broadcast_to_active_workers("continue_generation") self.health_manager.resume() diff --git a/xtuner/v1/rl/rollout/health_manager.py b/xtuner/v1/rl/rollout/health_manager.py index 2026cd655..c9f19bb37 100644 --- a/xtuner/v1/rl/rollout/health_manager.py +++ b/xtuner/v1/rl/rollout/health_manager.py @@ -101,6 +101,12 @@ def mark_failed_ranks(self, worker_health_results: dict[int, bool]) -> set[int]: return failed_ranks +@dataclass(frozen=True) +class _ReadyRecoveryHF: + model_path: str + tokenizer_path: str | None = None + + class RolloutHealthManager: """Own worker health state and recovery after controller startup. @@ -127,10 +133,26 @@ def __init__( self._thread: threading.Thread | None = None self._lifecycle_operation_lock = threading.Lock() self._worker_health_failure_tracker = _WorkerHealthFailureTracker(threshold=self._check_failure_threshold) + self._ready_recovery_hf: _ReadyRecoveryHF | None = None # ------------------------------------------------------------------ # Public lifecycle # ------------------------------------------------------------------ + def set_ready_recovery_hf( + self, + *, + model_path: str, + tokenizer_path: str | None = None, + ) -> None: + self._ready_recovery_hf = _ReadyRecoveryHF( + model_path=model_path, + tokenizer_path=tokenizer_path, + ) + logger.info(f"Ready rollout recovery HF updated: model_path={model_path}, tokenizer_path={tokenizer_path}.") + + def clear_ready_recovery_hf(self) -> None: + self._ready_recovery_hf = None + logger.info("Ready rollout recovery HF cleared.") def start(self) -> None: health_thread_alive = self._thread is not None and self._thread.is_alive() @@ -466,8 +488,7 @@ def _restart_worker_group( self, group: WorkerGroup, ) -> bool: - """Shutdown, restart with empty-init, and health-check one complete - worker group.""" + """Shutdown, restart, and health-check one complete worker group.""" if not group.workers or len(group.workers) != len(group.ranks): logger.error(f"Cannot restart incomplete rollout worker group: ranks={group.ranks}.") return False @@ -481,34 +502,38 @@ def _restart_worker_group( restart_cleanup_needed = True self._checkpoint_not_stopping() - with self._skip_load_weights_during_restart(group): - self._checkpoint_not_stopping() - ray.get( - [ - # reinit() reuses the server launch spec bound during - # controller startup. - worker.actor.reinit.remote() # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) + ready_recovery_hf = self._ready_recovery_hf + if ready_recovery_hf is None: + reinit_kwargs: dict[str, object] = {"skip_load_weights": True} + else: + reinit_kwargs = { + "model_path": ready_recovery_hf.model_path, + "tokenizer_path": ready_recovery_hf.tokenizer_path, + "skip_load_weights": False, + } - self._checkpoint_not_stopping() - health_results = self._check_workers_health(group.workers) - unhealthy_ranks = [ - worker.rank for worker in group.workers if not health_results.get(worker.rank, False) - ] - if unhealthy_ranks: - logger.error( - f"Restarted rollout worker group ranks={group.ranks} has unhealthy ranks={unhealthy_ranks}." - ) - self._shutdown_worker_group(group, wait_server_down=False) - return False + ray.get( + [ + worker.actor.reinit.remote(**reinit_kwargs) # type: ignore[attr-defined] + for worker in group.workers + ], + timeout=ROLLOUT_RAY_GET_TIMEOUT, + ) + + self._checkpoint_not_stopping() + health_results = self._check_workers_health(group.workers) + unhealthy_ranks = [worker.rank for worker in group.workers if not health_results.get(worker.rank, False)] + if unhealthy_ranks: + logger.error( + f"Restarted rollout worker group ranks={group.ranks} has unhealthy ranks={unhealthy_ranks}." + ) + self._shutdown_worker_group(group, wait_server_down=False) + return False + if ready_recovery_hf is None: self._checkpoint_not_stopping() - # Newly restarted workers should return to the same offloaded/sleep - # baseline as the other colocated rollout workers before the sync - # path wakes weights/KV back up. + # Weight-update recovery returns to the offloaded baseline + # before the sync path wakes weights and KV cache back up. ray.get( [worker.actor.offload.remote() for worker in group.workers], # type: ignore[attr-defined] timeout=ROLLOUT_RAY_GET_TIMEOUT, @@ -526,31 +551,6 @@ def _restart_worker_group( self._shutdown_worker_group(group, wait_server_down=False) return False - @contextmanager - def _skip_load_weights_during_restart(self, group: WorkerGroup): - try: - ray.get( - [ - worker.actor.set_skip_load_weights.remote(True) # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) - yield - finally: - try: - ray.get( - [ - worker.actor.restore_skip_load_weights.remote() # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) - except Exception: - logger.exception( - f"Failed to restore rollout worker skip_load_weights after restart: group_ranks={group.ranks}." - ) - def _shutdown_worker_group( self, group: WorkerGroup, diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index f6118fbdf..eab13d3b6 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -536,7 +536,6 @@ def __init__( Defaults to "GPU". """ self.config = config - self._default_skip_load_weights = config.skip_load_weights self.rank = rank self.master_addr = master_addr # ray master self.master_port = master_port @@ -594,9 +593,25 @@ def init(self, server_launch_spec: ServerLaunchSpec) -> RolloutWorkerInitResult: self._bind_server_launch_spec(server_launch_spec) return self._init_server() - def reinit(self) -> RolloutWorkerInitResult: + def reinit( + self, + *, + model_path: str | Path | None = None, + tokenizer_path: str | Path | None = None, + skip_load_weights: bool | None = None, + ) -> RolloutWorkerInitResult: """Reinitialize the rollout server using the previously bound launch spec.""" + config_updates: dict[str, object] = {} + if model_path is not None: + config_updates["model_path"] = str(model_path) + if tokenizer_path is not None: + config_updates["tokenizer_path"] = str(tokenizer_path) + if skip_load_weights is not None: + config_updates["skip_load_weights"] = skip_load_weights + + if config_updates: + self.config = self.config.model_copy(update=config_updates) return self._init_server() def _init_server(self) -> RolloutWorkerInitResult: @@ -616,12 +631,6 @@ def _init_server(self) -> RolloutWorkerInitResult: session_url=self.session_server_url, ) - def set_skip_load_weights(self, skip_load_weights: bool) -> None: - self.config = self.config.model_copy(update={"skip_load_weights": skip_load_weights}) - - def restore_skip_load_weights(self) -> None: - self.config = self.config.model_copy(update={"skip_load_weights": self._default_skip_load_weights}) - def init_dist_port(self) -> tuple[int, str]: """Initialize distributed communication ports. diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index b87da88ce..9fe4a89d8 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -328,6 +328,22 @@ def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16): ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) return + def start_hf_export( + self, + hf_dir: str, + save_dtype: torch.dtype = torch.bfloat16, + ) -> None: + handles = [worker.start_hf_export.remote(hf_dir, save_dtype) for worker in self.workers] + ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) + + def is_hf_export_done(self) -> bool: + handles = [worker.is_hf_export_done.remote() for worker in self.workers] + return all(ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)) + + def wait_hf_export(self) -> str: + handles = [worker.wait_hf_export.remote() for worker in self.workers] + return ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)[0] + def resume(self, load_checkpoint_cfg: LoadCheckpointConfig): """Resume the training workers from the checkpoint.""" handles = [worker.resume.remote(load_checkpoint_cfg) for worker in self.workers] # type: ignore diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index f7e0bc487..c72a723fa 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -3,6 +3,7 @@ import math import os import time +from concurrent.futures import Future from contextlib import contextmanager from pathlib import Path from typing import ( @@ -242,6 +243,7 @@ def __init__( if not worker_cfg.fsdp_cfg.torch_compile: worker_cfg.model_cfg.compile_cfg = False self._engine = self._build_engine(worker_cfg) + self._pending_hf_export: Future[Path] | None = None self._has_ref = False if worker_cfg.loss_cfg.use_kl_loss: @@ -942,6 +944,28 @@ def _reduce_number_across_rank(self, rank_number: int) -> int: def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16): self._engine.save_hf(hf_dir, save_dtype) + @ray_method + def start_hf_export( + self, + hf_dir: str, + save_dtype: torch.dtype = torch.bfloat16, + ) -> None: + self._pending_hf_export = self._engine.async_save_hf(hf_dir, save_dtype) + + @ray_method + def is_hf_export_done(self) -> bool: + pending = cast(Future[Path], self._pending_hf_export) + return pending.done() + + @ray_method + def wait_hf_export(self) -> str: + pending = cast(Future[Path], self._pending_hf_export) + try: + finalized_path = pending.result() + finally: + self._pending_hf_export = None + return str(finalized_path) + @ray_method def get_data_replicate_size(self) -> int: """Get the data replicate size for the training worker.""" diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 9e810ad16..1c4d708aa 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -4,6 +4,7 @@ import random import re import time +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import asdict, dataclass from pathlib import Path from shutil import rmtree @@ -41,7 +42,7 @@ ) from xtuner.v1.rl.rollout.controller import RolloutControllerProxy from xtuner.v1.rl.rollout.worker import RolloutConfig -from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace +from xtuner.v1.rl.trace import TraceConfig, configure_trace from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.rl.trainer.worker import WorkerConfig, WorkerLogItem from xtuner.v1.rl.utils import ( @@ -63,6 +64,7 @@ # TODO: Move DEVICE to `xtuner.utils.device` PG_READY_TIMEOUT = 30 RL_TRAINER_RAY_GET_TIMEOUT = 3600 +HF_EXPORT_POLL_INTERVAL_S = 1.0 DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() @@ -354,6 +356,7 @@ class BaseRLTrainerConfig(BaseModel): checkpoint_maxkeep: int | None = -1 hf_interval: int | None = -1 hf_max_keep: int | None = -1 + enable_immediate_recovery: bool = False checkpoint_no_save_optimizer: bool = False checkpoint_no_save_replay_buffer: bool = False log_dir: Path | str | None = None @@ -439,6 +442,9 @@ class RLColocateTrainerConfig(BaseRLTrainerConfig): Defaults to -1. hf_max_keep (int | None): Maximum number of Hugging Face checkpoints to keep. Defaults to -1. + enable_immediate_recovery (bool): Whether rollout worker recovery may + use fresh ready HF exports before the next weight update. Defaults + to False. checkpoint_no_save_optimizer (bool): Whether to skip optimizer states when saving checkpoints. Defaults to False. checkpoint_no_save_replay_buffer (bool): Whether to skip replay buffer @@ -527,6 +533,9 @@ class RLDisaggregatedTrainerConfig(BaseRLTrainerConfig): Defaults to -1. hf_max_keep (int | None): Maximum number of Hugging Face checkpoints to keep. Defaults to -1. + enable_immediate_recovery (bool): Whether rollout worker recovery may + use fresh ready HF exports before the next weight update. Defaults + to False. checkpoint_no_save_optimizer (bool): Whether to skip optimizer states when saving checkpoints. Defaults to False. checkpoint_no_save_replay_buffer (bool): Whether to skip replay buffer @@ -622,6 +631,14 @@ def _init_load_source(self, cfg: BaseRLTrainerConfig) -> None: def _init_save_config(self, cfg: BaseRLTrainerConfig) -> None: self._hf_max_keep = cfg.hf_max_keep self._hf_interval = cfg.hf_interval + self._enable_immediate_recovery = cfg.enable_immediate_recovery + self._hf_export_executor = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="rl-hf-export") + if self._enable_immediate_recovery + else None + ) + self._pending_hf_export: Future[Path | None] | None = None + self._ready_recovery_hf_path: Path | None = None self._checkpoint_interval = cfg.checkpoint_interval self._checkpoint_maxkeep = cfg.checkpoint_maxkeep @@ -858,6 +875,96 @@ def _maybe_save_hf(self, cur_step: int): if isinstance(self.tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)): self.tokenizer.save_pretrained(str(save_hf_path)) + def _maybe_save_recovery_hf(self, cur_step: int) -> None: + if not self._enable_immediate_recovery: + return + + reuse_regular_hf = ( + self._hf_interval is not None + and self._hf_interval != -1 + and (cur_step % self._hf_interval == 0 or cur_step == self._total_train_steps) + ) + tokenizer_path = self._rollout_config.tokenizer_path or self._rollout_config.model_path + previous_ready_hf_path = self._ready_recovery_hf_path + ray.get( + self.rollout_controller.clear_ready_recovery_hf.remote(), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + self._ready_recovery_hf_path = None + if ( + previous_ready_hf_path is not None + and str(previous_ready_hf_path) not in self._meta.latest_exp.hf_checkpoint_list + ): + rmtree(previous_ready_hf_path, ignore_errors=True) + + save_hf_path = self.exp_dir / self._HF_DIR / f"hf-step-{cur_step}" + if reuse_regular_hf: + try: + ray.get( + self.rollout_controller.set_ready_recovery_hf.remote( + model_path=str(save_hf_path), + tokenizer_path=str(tokenizer_path), + ), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + except Exception: + self.logger.exception(f"Failed to publish recovery HF: path={save_hf_path}.") + return + + self._ready_recovery_hf_path = save_hf_path + return + + save_hf_path.mkdir(parents=True, exist_ok=True) + self.train_controller.start_hf_export(str(save_hf_path)) + executor = cast(ThreadPoolExecutor, self._hf_export_executor) + + def wait_and_publish_recovery_hf() -> Path | None: + try: + while not self.train_controller.is_hf_export_done(): + time.sleep(HF_EXPORT_POLL_INTERVAL_S) + finalized_hf_path = Path(self.train_controller.wait_hf_export()) + ray.get( + self.rollout_controller.set_ready_recovery_hf.remote( + model_path=str(finalized_hf_path), + tokenizer_path=str(tokenizer_path), + ), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + except Exception: + self.logger.exception(f"Async recovery HF export failed: path={save_hf_path}.") + return None + + self._ready_recovery_hf_path = finalized_hf_path + return finalized_hf_path + + self._pending_hf_export = executor.submit(wait_and_publish_recovery_hf) + + def _wait_or_disable_immediate_recovery(self) -> None: + pending = self._pending_hf_export + if pending is None: + return + + disable_immediate_recovery = not pending.done() + if disable_immediate_recovery: + self._enable_immediate_recovery = False + self.logger.warning( + "Disable immediate recovery because the previous recovery HF " + "export did not finish before the next weight sync." + ) + + finalized_hf_path = pending.result() + self._pending_hf_export = None + if not disable_immediate_recovery: + return + + ray.get( + self.rollout_controller.clear_ready_recovery_hf.remote(), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + self._ready_recovery_hf_path = None + if finalized_hf_path is not None: + rmtree(finalized_hf_path, ignore_errors=True) + async def _run_initial_evaluate(self) -> None: try: eval_produce_result = await self.eval_agent_loop_manager.produce_batch( @@ -1639,7 +1746,8 @@ def fit(self): self._fit() finally: self._exp_tracker.close() - close_trace() + if self._hf_export_executor is not None: + self._hf_export_executor.shutdown(wait=True) def _fit(self): self.logger.info("Start RL training") @@ -1748,6 +1856,8 @@ def _get_colocate_rollout_model_step(self, train_step: int) -> int: def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool: """保存后切回共卡 rollout 资源。""" + self._wait_or_disable_immediate_recovery() + should_sync_weights = train_step % self._sync_weights_interval == 0 will_evaluate = self._enable_evaluate and train_step % self._evaluate_step == 0 needs_rollout_ready = train_step < self._total_train_steps or will_evaluate @@ -1781,6 +1891,7 @@ def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool ) self.train_controller.update_weights() self.logger.info("Rollout workers update weights successfully in colocate mode") + self._maybe_save_recovery_hf(train_step) self.train_controller.offload(target="model") else: self.train_controller.offload(target="model") @@ -1870,7 +1981,8 @@ def fit(self): return asyncio_run(self._fit()) finally: self._exp_tracker.close() - close_trace() + if self._hf_export_executor is not None: + self._hf_export_executor.shutdown(wait=True) async def _get_batch_or_raise_producer_failure( self, @@ -2007,12 +2119,16 @@ async def _fit(self): async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): # producer 已暂停；保持 save -> bind -> update 顺序。 + self._wait_or_disable_immediate_recovery() + with timer("save_ckpt", step_timer_dict): await self._maybe_save_checkpoint(model_step) self._maybe_save_hf(model_step) - # TODO: 非共卡需要额外加健康检查恢复worker的逻辑，共卡是在训练之前恢复，但是非共卡不需要在训练之前恢复,挂掉就恢复或者更新权重前恢复，需要评估一下哪种方式更合理。 with timer("sync_weight", step_timer_dict): + # 非共卡在权重更新前恢复 inactive workers；如果没有 ready + # recovery HF，HealthManager 会用空权重启动并等待本次 update。 + await self.rollout_controller.restart_inactive_workers.remote() # type: ignore[attr-defined] bind_train_rollout( train_controller=self.train_controller, rollout_controller=self.rollout_controller, @@ -2022,6 +2138,7 @@ async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): weight_update_port=self._rollout_config.weight_update_port, ) self.update_weights() + self._maybe_save_recovery_hf(model_step) def update_weights(self): # rollout 恢复由 AgentLoopManager 控制。