Skip to content

[Feat] Support immediate rollout recovery by export hf asynchronously #1966

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YanhuiDua wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Feat] Support immediate rollout recovery by export hf asynchronously #1966

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions xtuner/v1/rl/rollout/controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
104 changes: 52 additions & 52 deletions xtuner/v1/rl/rollout/health_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down
25 changes: 17 additions & 8 deletions xtuner/v1/rl/rollout/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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.

Expand Down
16 changes: 16 additions & 0 deletions xtuner/v1/rl/trainer/controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions xtuner/v1/rl/trainer/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()

Comment on lines +957 to +959

@claude claude Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude: Defensiveness (Nit): is_hf_export_done and wait_hf_export both use cast(Future[Path], self._pending_hf_export) without guarding against None. If either is ever called out of sequence (before start_hf_export), this will produce a confusing AttributeError: 'NoneType' object has no attribute 'done'.

Consider raising an explicit RuntimeError if self._pending_hf_export is None:

if self._pending_hf_export is None:
    raise RuntimeError("No pending HF export — call start_hf_export first.")
return self._pending_hf_export.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."""
Expand Down
Loading
Loading