[Feature] Improve async DCP checkpoint reliability and early failure …

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

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)
 

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)