From f8c6bd726218cbb9c4610a754b3ff2d069465500 Mon Sep 17 00:00:00 2001 From: Vincent Date: Sun, 31 May 2026 13:23:39 +0800 Subject: [PATCH] [Feature] Improve async DCP checkpoint reliability and early failure detection --- xtuner/v1/engine/train_engine.py | 87 ++++++++++++++++++------------- xtuner/v1/patch/xtuner_storage.py | 65 ++++++++++++++++++++--- xtuner/v1/train/trainer.py | 19 ++++++- 3 files changed, 126 insertions(+), 45 deletions(-) diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 9032332b80..34dfe19150 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -373,15 +373,6 @@ def _get_async_checkpoint_pg(self) -> dist.ProcessGroup: self._async_checkpoint_pg = dist.new_group(backend="gloo") return self._async_checkpoint_pg - @staticmethod - def _is_async_checkpoint_daemon_init_error(exc: BaseException) -> bool: - message = str(exc) - return ( - "EADDRINUSE" in message - or "address already in use" in message - or "Checkpoint background process is dead" in message - ) - def async_save_dcp( self, weights_dir: Path, @@ -432,34 +423,13 @@ def start_async_save() -> Future: dcp_future = start_async_save() def commit_async_save() -> None: - nonlocal dcp_future - # Retry only PyTorch DCP daemon init port races, such as - # EADDRINUSE from TCPStore. Other checkpoint failures still raise. - max_daemon_init_attempts = 3 - for attempt in range(1, max_daemon_init_attempts + 1): - try: - dcp_future.result() - break - except BaseException as exc: - if attempt == max_daemon_init_attempts or not self._is_async_checkpoint_daemon_init_error(exc): - elapsed = time.time() - t0 - logger.error(f"[DCP async_save for {weights_dir}] failed after {elapsed:.2f}s: {exc}") - logger.error(traceback.format_exc()) - raise - - if dist.get_rank() == 0: - logger.warning( - "[DCP async_save for %s] checkpoint daemon init failed on attempt %s/%s, retrying: %s", - weights_dir, - attempt, - max_daemon_init_attempts, - exc, - ) - if incomplete_dir.exists(): - shutil.rmtree(incomplete_dir) - incomplete_dir.mkdir(parents=True, exist_ok=True) - dist.barrier(group=async_checkpoint_pg) - dcp_future = start_async_save() + try: + dcp_future.result() + except BaseException as exc: + elapsed = time.time() - t0 + logger.error(f"[DCP async_save for {weights_dir}] failed after {elapsed:.2f}s: {exc}") + logger.error(traceback.format_exc()) + raise dist.barrier(group=async_checkpoint_pg) if dist.get_rank() == 0: @@ -493,6 +463,49 @@ def _build_async_storage_writer(self, weights_dir: Path, *, save_optimizer: bool storage_writer.state_dict_cache = self._async_state_dict_cache return storage_writer + def warmup_async_save_dcp(self, timeout: float = 1800.0) -> None: + """Pre-initialize the async DCP checkpoint daemon before training. + + Runs one tiny dummy ``async_save`` so process-group creation / TCPStore binding happens now and errors like + EADDRINUSE surface before any training step is wasted. A rank that fails raises immediately; the launcher then + tears the job down -- the standard path for init-time failures. + + Args: + timeout (float): Seconds to wait for the warmup save before declaring a hang. Defaults to 1800 to match + PyTorch's own daemon-init wait. + """ + async_checkpoint_pg = self._get_async_checkpoint_pg() + # warmup_dir is per-rank and lives in node-local /dev/shm, so each rank + # prepares and cleans up its OWN directory (a rank0-only guard would + # leak every non-zero rank's directory). + warmup_dir = Path(f"/dev/shm/xtuner_dcp_warmup_{dist.get_rank()}") + if warmup_dir.exists(): + shutil.rmtree(warmup_dir, ignore_errors=True) + warmup_dir.mkdir(parents=True, exist_ok=True) + + async_save_kwargs: dict[str, Any] = {} + state_dict_saver = importlib.import_module("torch.distributed.checkpoint.state_dict_saver") + async_checkpointer_type = getattr(state_dict_saver, "AsyncCheckpointerType", None) + if async_checkpointer_type is not None: + async_save_kwargs["async_checkpointer_type"] = async_checkpointer_type.PROCESS + + try: + future = cast(Any, dcp.async_save)( + {"_warmup": torch.zeros(1)}, + checkpoint_id=warmup_dir, + process_group=async_checkpoint_pg, + **async_save_kwargs, + ) + future.result(timeout=timeout) + finally: + shutil.rmtree(warmup_dir, ignore_errors=True) + + # Ensure warmup does not leave stale state that could be mistakenly + # reused by the first real checkpoint (different state_dict structure). + self._async_state_dict_cache = None + + log_rank0.info("[DCP] Async save warmup completed successfully") + def destroy_async_checkpoint_pg(self) -> None: """Destroy the dedicated gloo process group used for async checkpoint.""" diff --git a/xtuner/v1/patch/xtuner_storage.py b/xtuner/v1/patch/xtuner_storage.py index f06426f674..e6e23d9b1c 100644 --- a/xtuner/v1/patch/xtuner_storage.py +++ b/xtuner/v1/patch/xtuner_storage.py @@ -15,7 +15,7 @@ StreamTransformExtension, ) from torch.distributed.checkpoint.filesystem import FileSystem -from torch.distributed.checkpoint.staging import _copy_state_dict, _create_cpu_state_dict +from torch.distributed.checkpoint.staging import _copy_state_dict from torch.distributed.checkpoint.storage import ( WriteResult, ) @@ -132,6 +132,60 @@ def create_stream(self, path: Union[str, os.PathLike], mode: str): _release_write_lock(lock_fd) +def _create_coalesced_shm_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]: + """Create a CPU state dict backed by coalesced shared-memory buffers. + + Instead of creating one shared-memory file per tensor (which leads to + thousands of fds and triggers ``received 0 items of ancdata`` when the + daemon subprocess tries to receive them all), this function groups tensors + by dtype, allocates a single large shared-memory tensor per dtype, and + returns views into that buffer. + + Args: + state_dict (dict[str, Any]): The source state dict (tensors can be on + any device). + + Returns: + dict[str, Any]: A new state dict with the same keys, where every tensor + is a view into a dtype-coalesced shared-memory buffer. + """ + # Collect tensor metadata grouped by dtype + dtype_groups: dict[torch.dtype, list[tuple[str, torch.Size]]] = {} + for key, val in state_dict.items(): + if isinstance(val, torch.Tensor) and val.numel() > 0: + dtype_groups.setdefault(val.dtype, []).append((key, val.size())) + + # Allocate one coalesced buffer per dtype in shared memory + dtype_buffers: dict[torch.dtype, torch.Tensor] = {} + dtype_offsets: dict[torch.dtype, int] = {} + for dtype, items in dtype_groups.items(): + total_numel = sum(size.numel() for _, size in items) + buf = torch.empty(total_numel, dtype=dtype) + buf.share_memory_() + dtype_buffers[dtype] = buf + dtype_offsets[dtype] = 0 + + # Build the output state dict with views into coalesced buffers + result: dict[str, Any] = {} + for key, val in state_dict.items(): + if isinstance(val, torch.Tensor) and val.numel() > 0: + dtype = val.dtype + offset = dtype_offsets[dtype] + numel = val.numel() + view = dtype_buffers[dtype][offset : offset + numel].view(val.size()) + dtype_offsets[dtype] = offset + numel + result[key] = view + elif isinstance(val, torch.Tensor): + # Zero-numel tensors: just create a shared empty tensor + t = torch.zeros_like(val, device="cpu") + t.share_memory_() + result[key] = t + else: + result[key] = val + + return result + + class XtunerCacheWriter(FileSystemWriter): # Save write results for the current rank as computed by `write_data` API # Cached on the local rank. @@ -194,16 +248,13 @@ def stage(self, state_dict: dict[str, Any]) -> dict[str, Any]: self.per_thread_copy_ahead = 0 if not self.cache_staged_state_dict: - staged_state_dict = _create_cpu_state_dict(state_dict, share_memory=True) + staged_state_dict = _create_coalesced_shm_state_dict(state_dict) return _copy_state_dict(state_dict, staged_state_dict, type_check=self.type_check) if self.state_dict_cache is None: if not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0: - logger.info("[DCP async_save] creating shared-memory staged cache") - self.state_dict_cache = _create_cpu_state_dict( - state_dict, - share_memory=True, - ) + logger.info("[DCP async_save] creating shared-memory staged cache (coalesced)") + self.state_dict_cache = _create_coalesced_shm_state_dict(state_dict) return _copy_state_dict(state_dict, self.state_dict_cache, type_check=self.type_check) diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 4dc9efbaab..045e519ea8 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -806,6 +806,10 @@ def fit(self): handles data loading, forward pass, backward pass, optimization, logging, and checkpointing. """ train_begin = time.time() + + if self._async_checkpoint: + self._engine.warmup_async_save_dcp() + time_before_get_data = time.time() for data_batch in self._data_iter(): time_before_train_step = time.time() @@ -861,6 +865,7 @@ def fit(self): self._lr_scheduler.step() self._maybe_check_health() + self._check_async_save_health() self._maybe_save_hf() ckpt_saved = self._maybe_save(is_snapshot=False) if not ckpt_saved: @@ -1177,6 +1182,18 @@ def _maybe_check_health(self): raise RuntimeError("Health check failed, exit training") log_rank0.info(f"Health check passed at step {self.cur_step}") + def _check_async_save_health(self) -> None: + """Non-blocking check for async save failures. + + Called every training step to detect async checkpoint or HF export failures as soon as they happen, rather than + waiting until the next save interval. + """ + if self._pending_checkpoint is not None and self._pending_checkpoint.done(): + exc = self._pending_checkpoint.exception() + if exc is not None: + self._pending_checkpoint = None + raise RuntimeError(f"Async DCP checkpoint failed: {exc}") from exc + def _wait_for_pending_checkpoint(self, timeout: int = 3000) -> None: if self._pending_checkpoint is None: return @@ -1226,7 +1243,7 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool: # Save model and optimizer future: Future | None = None - if self._async_checkpoint and not is_snapshot: + if self._async_checkpoint: future = self._engine.async_save_dcp(weights_dir=weights_path) else: self._engine.save_dcp(weights_dir=weights_path)