Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / dcp_checkpointer.py: 65%
278 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates
2# Copyright 2026 Huawei Technologies Co., Ltd
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ============================================================================
16"""Distributed-checkpoint (DCP) backend for :class:`CheckpointerBase`."""
18__all__ = [
19 "STEP_PREFIX",
20 "DistributedCheckpointer",
21 "initialize_optimizer_state",
22]
24import gc
25import logging
26import os
27from functools import partial
28from typing import Any, Dict, List, Optional
30import torch # pylint: disable=forbidden-backend-import
31import torch.distributed as dist # pylint: disable=forbidden-backend-import
33from hyper_parallel import SkipDTensorDispatch
34from hyper_parallel.components.optim.mixed_precision_optimizer import (
35 FP32_MAIN_PARAM_STATE_KEY,
36 MIXED_PRECISION_OPTIMIZER_STATE_KEY,
37 MixedPrecisionOptimizer,
38)
39from hyper_parallel.models.build_options import (
40 IS_CUDA_AVAILABLE,
41 IS_NPU_AVAILABLE,
42 get_torch_device,
43)
44from hyper_parallel.components.checkpoint.base import CheckpointerBase
45from hyper_parallel.core.distributed_checkpoint import (
46 StandardLoadPlanner,
47 async_save as dcp_async_save,
48 load as dcp_load,
49 save as dcp_save,
50)
51from hyper_parallel.core.distributed_checkpoint.utils import flatten_state_dict
52from hyper_parallel.core.optimizer.adamw import AdamW as HyperAdamW
53from hyper_parallel.core.optimizer.muon import Muon as HyperMuon
54from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer
57logger = logging.getLogger(__name__)
60def _empty_cache() -> None:
61 """Collect host and device memory around checkpoint I/O."""
62 gc.collect()
63 if IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE:
64 get_torch_device().empty_cache()
66# Layout of a checkpoint root:
67#
68# checkpoints/
69# ├── latest_checkpoint_iteration.txt # step number of the newest checkpoint
70# └── global_step_42/
71# ├── .metadata # DCP global metadata, written last
72# ├── __0_0.safetensors, ... # DCP shard payloads
73# └── extra_state/
74# └── extra_state_rank_{R}.pt # per-rank training state
75#
76# ``.metadata`` doubles as the completeness signal: ``FileSystemWriter``'s
77# ``finalize_checkpoint`` only writes it after every rank's shard writes have
78# been collected, so a directory carrying it holds a fully persisted checkpoint
79# and one missing it was interrupted mid-save.
80STEP_PREFIX = "global_step_"
81LATEST_CHECKPOINT_FILE = "latest_checkpoint_iteration.txt"
82_EXTRA_STATE_DIR = "extra_state"
83_EXTRA_STATE_FILE = "extra_state_rank_{rank}.pt"
84_DCP_METADATA_FILE = ".metadata"
86# ``extra_state`` must carry this key; the checkpointer stamps a sentinel over
87# it in the skeleton handed to an embedded-bundle load, so a load that read
88# nothing back is distinguishable from a genuine restore. A real checkpoint is
89# never written at a step this low.
90_STEP_KEY = "global_step"
91_UNRESTORED_STEP = -1
94def _get_rank() -> int:
95 """Return this process' global rank, or 0 outside a process group."""
96 if dist.is_available() and dist.is_initialized():
97 return dist.get_rank()
98 return 0
101def _barrier() -> None:
102 """Synchronize all ranks, tolerating single-process (non-distributed) runs."""
103 if dist.is_available() and dist.is_initialized():
104 dist.barrier()
107def _prepare_main_param_state_for_dcp_load(
108 state_dict: Dict[str, Any],
109 available_keys: set[str],
110) -> None:
111 """Select the ordinary or fp32-main optimizer checkpoint branch."""
112 flattened_state, _ = flatten_state_dict(state_dict)
113 optimizer_state = state_dict.get("optimizer")
114 optimizer_state_is_list = isinstance(optimizer_state, list)
115 optimizer_states = (
116 optimizer_state
117 if optimizer_state_is_list
118 else [optimizer_state]
119 )
120 for optimizer_index, optimizer_state in enumerate(optimizer_states):
121 if not isinstance(optimizer_state, dict):
122 continue
123 if MIXED_PRECISION_OPTIMIZER_STATE_KEY not in optimizer_state:
124 continue
126 state_prefix = (
127 f"optimizer.{optimizer_index}"
128 if optimizer_state_is_list
129 else "optimizer"
130 )
131 main_param_prefix = (
132 f"{state_prefix}.{MIXED_PRECISION_OPTIMIZER_STATE_KEY}."
133 f"{FP32_MAIN_PARAM_STATE_KEY}."
134 )
135 expected_keys = {
136 key for key in flattened_state if key.startswith(main_param_prefix)
137 }
138 if not expected_keys:
139 continue
141 checkpoint_keys = {
142 key for key in available_keys if key.startswith(main_param_prefix)
143 }
144 if not checkpoint_keys:
145 optimizer_state.pop(MIXED_PRECISION_OPTIMIZER_STATE_KEY)
146 continue
148 missing_keys = sorted(expected_keys - checkpoint_keys)
149 if missing_keys:
150 preview = ", ".join(missing_keys[:10])
151 raise RuntimeError(
152 "Mixed-precision optimizer checkpoint is incomplete; missing "
153 f"{len(missing_keys)} fp32 main parameter(s): {preview}"
154 )
157class _ModelStrictLoadPlanner(StandardLoadPlanner):
158 """Allow partial optimizer state while still requiring a complete model.
160 A parameter that never received a gradient --- a frozen weight, an MoE expert
161 no batch routed to, or simply a checkpoint written before the first
162 ``step()`` --- has no optimizer state to save. Meanwhile the load side primes
163 the optimizer (see :func:`initialize_optimizer_state`) so DCP has entries to
164 fill, which means it legitimately asks for keys the checkpoint does not
165 carry. Under the planner's default strict mode that is a hard error, so
166 ``allow_partial_load`` lets those entries keep the zero-initialised values
167 the optimizer would have created on its next step anyway.
169 ``allow_partial_load`` is global to the planner and cannot be scoped to the
170 optimizer subtree, so model completeness is re-checked here: a truncated
171 model checkpoint must still fail loudly instead of resuming half-loaded.
172 """
174 def __init__(self, strict_model: bool = True) -> None:
175 """Load partially, optionally demanding every requested model key."""
176 super().__init__(allow_partial_load=True)
177 self.strict_model = strict_model
179 def configure_planner(
180 self,
181 state_dict: Dict[str, Any],
182 metadata: Any,
183 **kwargs: Any,
184 ) -> None:
185 """Normalize optional optimizer state before flattening the load skeleton."""
186 _prepare_main_param_state_for_dcp_load(
187 state_dict,
188 set(metadata.state_dict_metadata),
189 )
190 super().configure_planner(state_dict, metadata, **kwargs)
192 def build_local_plan(self) -> Any:
193 """Validate model coverage, then defer to the standard planner."""
194 if self.state_dict is None or self.metadata is None:
195 return super().build_local_plan()
197 available = self.metadata.state_dict_metadata
198 missing = [fqn for fqn in self.state_dict if fqn not in available]
200 missing_model = sorted(fqn for fqn in missing if fqn.startswith("model."))
201 if self.strict_model and missing_model:
202 preview = ", ".join(missing_model[:10])
203 suffix = " ..." if len(missing_model) > 10 else ""
204 raise RuntimeError(
205 f"Checkpoint is missing {len(missing_model)} model key(s) required "
206 f"for a full-model resume: {preview}{suffix}"
207 )
209 missing_optimizer = sorted(fqn for fqn in missing if fqn.startswith("optimizer."))
210 if missing_optimizer:
211 preview = ", ".join(missing_optimizer[:3])
212 logger.warning(
213 "%s optimizer entries are absent from the checkpoint and keep their "
214 "zero-initialised values (e.g. %s). Expected for parameters that "
215 "never received a gradient, or for a checkpoint written with "
216 "save_optimizer=false.",
217 len(missing_optimizer),
218 preview,
219 )
221 return super().build_local_plan()
224def _optimizer_parameters(optimizer: Any) -> List[torch.Tensor]:
225 """Return trainable parameters managed by an optimizer."""
226 parameters = []
227 for group in optimizer.param_groups:
228 for param in group["params"]:
229 if param.requires_grad:
230 parameters.append(param)
231 return parameters
234def _prime_optimizer_state(
235 optimizer: Any,
236 params: List[torch.Tensor],
237) -> None:
238 """Run a parameter-neutral optimizer step to materialize its state."""
239 for param in params:
240 param.grad = torch.zeros_like(param)
241 for group in optimizer.param_groups:
242 if "lr" in group:
243 group["lr"] = 0.0
244 if "weight_decay" in group:
245 group["weight_decay"] = 0.0
246 # zeros_like must retain DTensor dispatch so newly created state keeps the parameter mesh and placements.
247 with SkipDTensorDispatch(no_skip={torch.zeros_like}):
248 optimizer.step()
251def _materialize_hyper_optimizer_state(
252 optimizer: Any,
253 params: List[torch.Tensor],
254) -> bool:
255 """Create lazy Hyper AdamW/Muon state without executing an update step."""
256 if isinstance(optimizer, HyperMuon):
257 state_keys = ("momentum_buffer",)
258 elif isinstance(optimizer, HyperAdamW):
259 state_keys = ("exp_avg", "exp_avg_sq")
260 else:
261 return False
263 param_ids = {id(param) for param in params}
264 # zeros_like must retain DTensor dispatch so newly created state keeps the parameter mesh and placements.
265 with SkipDTensorDispatch(no_skip={torch.zeros_like}):
266 for group in optimizer.param_groups:
267 group["step"] = group.get("step") or 0
268 for param in group["params"]:
269 if id(param) not in param_ids:
270 continue
271 state = optimizer.state[param]
272 for state_key in state_keys:
273 state.setdefault(
274 state_key,
275 torch.zeros_like(param, memory_format=torch.preserve_format),
276 )
277 if isinstance(optimizer, HyperAdamW) and group.get("amsgrad", False):
278 state.setdefault(
279 "max_exp_avg_sq",
280 torch.zeros_like(param, memory_format=torch.preserve_format),
281 )
282 return True
285def _restore_optimizer_hyperparams(
286 optimizer: Any,
287 saved_hyperparams: List[tuple[Optional[float], Optional[float]]],
288) -> None:
289 """Restore optimizer hyperparameters after state priming."""
290 for group, (lr, weight_decay) in zip(optimizer.param_groups, saved_hyperparams):
291 if lr is not None:
292 group["lr"] = lr
293 if weight_decay is not None:
294 group["weight_decay"] = weight_decay
295 optimizer.zero_grad(set_to_none=True)
298def _leaf_optimizers(optimizer: Any) -> List[Any]:
299 """Return the concrete optimizers that own lazy tensor state."""
300 inner_optimizer = (
301 optimizer.optimizer
302 if isinstance(optimizer, MixedPrecisionOptimizer)
303 else optimizer
304 )
305 if isinstance(inner_optimizer, ChainedOptimizer):
306 return list(inner_optimizer.chained_optimizers)
307 return [inner_optimizer]
310def initialize_optimizer_state(optimizer: Any) -> bool:
311 """Materialize optimizer state so a DCP load has entries to fill.
313 ``Optimizer.state`` is empty until the first ``step()``, and DCP loads
314 strictly into the keys the caller's state_dict already contains. Restoring a
315 freshly built optimizer therefore reads *nothing* and silently resumes with
316 zeroed moments --- the model continues from the checkpoint while Adam starts
317 over. Priming with zero gradients creates the ``exp_avg`` / ``exp_avg_sq``
318 entries the checkpoint will overwrite.
320 Hyper AdamW and Muon states are created directly. This avoids running Muon's
321 HSDP communication and Newton-Schulz update merely to allocate zero momentum.
322 Unknown optimizer implementations retain the parameter-neutral priming-step
323 fallback, with ``lr`` and ``weight_decay`` forced to zero.
325 Args:
326 optimizer: Optimizer whose state is materialized.
328 Returns:
329 Whether the optimizer now carries per-parameter state.
330 """
331 leaf_optimizers = _leaf_optimizers(optimizer)
332 if all(leaf_optimizer.state for leaf_optimizer in leaf_optimizers):
333 return True
335 leaf_parameters = {
336 leaf_optimizer: _optimizer_parameters(leaf_optimizer)
337 for leaf_optimizer in leaf_optimizers
338 if not leaf_optimizer.state
339 }
340 if any(not params for params in leaf_parameters.values()):
341 return False
342 if any(
343 param.grad is not None
344 for params in leaf_parameters.values()
345 for param in params
346 ):
347 logger.warning(
348 "Optimizer parameters already carry gradients; skipping state priming "
349 "to avoid corrupting them. Optimizer state may not be restored."
350 )
351 return False
353 for leaf_optimizer, params in leaf_parameters.items():
354 if _materialize_hyper_optimizer_state(leaf_optimizer, params):
355 continue
356 saved_hyperparams = [
357 (group.get("lr"), group.get("weight_decay"))
358 for group in leaf_optimizer.param_groups
359 ]
360 try:
361 _prime_optimizer_state(leaf_optimizer, params)
362 finally:
363 _restore_optimizer_hyperparams(leaf_optimizer, saved_hyperparams)
365 return all(leaf_optimizer.state for leaf_optimizer in leaf_optimizers)
368class DistributedCheckpointer(CheckpointerBase):
369 """Persist a training payload through ``hyper_parallel``'s DCP.
371 ``extra_state`` can live in either of two layouts, chosen by
372 ``extra_state_per_rank`` at save time:
374 * ``True`` --- one ``torch.save`` file per rank under ``extra_state/``.
375 Always correct.
376 * ``False`` --- embedded in the DCP payload. Fewer files, but DCP
377 deduplicates entries sharing an FQN across ranks (see
378 ``distributed_checkpoint/utils.py::plan_ownership_masks``), so every rank
379 restores whichever rank's copy won. Only valid when the extra state is
380 rank-invariant.
382 :meth:`load` auto-detects which layout a checkpoint on disk used, so the
383 flag may differ from the run that produced it.
384 """
386 def __init__(self, extra_state_per_rank: bool = True) -> None:
387 """Select the ``extra_state`` layout used by :meth:`save`."""
388 self.extra_state_per_rank = extra_state_per_rank
389 self._async_save_response: Optional[Any] = None
391 # ------------------------------------------------------------------
392 # Layout helpers
393 # ------------------------------------------------------------------
395 @staticmethod
396 def _extra_state_path(checkpoint_dir: str, rank: int) -> str:
397 """Return this rank's extra_state file inside ``checkpoint_dir``."""
398 return os.path.join(
399 checkpoint_dir, _EXTRA_STATE_DIR, _EXTRA_STATE_FILE.format(rank=rank)
400 )
402 @staticmethod
403 def _is_complete(checkpoint_path: str) -> bool:
404 """Report whether ``checkpoint_path`` holds a fully persisted checkpoint.
406 Presence of the DCP metadata file is the signal: it is written only once
407 every rank's shards are on disk. ``use_collectives=False`` names it
408 ``{rank}.metadata`` instead of a single shared ``.metadata``, so both
409 spellings count.
410 """
411 if not os.path.isdir(checkpoint_path):
412 return False
413 if os.path.isfile(os.path.join(checkpoint_path, _DCP_METADATA_FILE)):
414 return True
415 return any(
416 name.endswith(_DCP_METADATA_FILE) for name in os.listdir(checkpoint_path)
417 )
419 @staticmethod
420 def _latest_pointer_path(checkpoint_dir: str) -> str:
421 """Return the path of the newest-checkpoint pointer file."""
422 return os.path.join(checkpoint_dir, LATEST_CHECKPOINT_FILE)
424 def _write_latest_pointer(self, checkpoint_dir: str, global_step: int) -> None:
425 """Record which checkpoint is newest, so resume need not scan the tree.
427 Only the step number is stored: the directory name is
428 ``STEP_PREFIX + step``, and that prefix is a module constant rather than
429 anything configurable, so writing it down would just be a second copy of
430 a value the reader already has.
431 """
432 with open(self._latest_pointer_path(checkpoint_dir), "w", encoding="utf-8") as handle:
433 handle.write(f"{global_step}\n")
435 def _read_latest_pointer(self, checkpoint_dir: str) -> Optional[str]:
436 """Resolve the pointer file to a usable checkpoint directory.
438 Returns ``None`` when the file is absent, unreadable, or names a
439 checkpoint that is missing or incomplete --- every one of which is a
440 reason to fall back to scanning rather than to fail.
441 """
442 pointer_path = self._latest_pointer_path(checkpoint_dir)
443 if not os.path.isfile(pointer_path):
444 return None
446 try:
447 with open(pointer_path, encoding="utf-8") as handle:
448 global_step = int(handle.read().strip())
449 except (OSError, ValueError):
450 logger.warning(
451 "Ignoring unreadable %s; falling back to scanning %s.",
452 pointer_path,
453 checkpoint_dir,
454 )
455 return None
457 candidate = os.path.join(checkpoint_dir, f"{STEP_PREFIX}{global_step}")
458 if not self._is_complete(candidate):
459 logger.warning(
460 "%s points at %s, which is missing or incomplete; falling back to "
461 "scanning %s.",
462 pointer_path,
463 candidate,
464 checkpoint_dir,
465 )
466 return None
467 return candidate
469 def _list_checkpoints(self, checkpoint_dir: str) -> List[tuple]:
470 """List ``(step, path)`` for every complete checkpoint, oldest first."""
471 if not os.path.isdir(checkpoint_dir):
472 return []
474 found = []
475 for entry in os.listdir(checkpoint_dir):
476 entry_path = os.path.join(checkpoint_dir, entry)
477 if not os.path.isdir(entry_path) or not entry.startswith(STEP_PREFIX):
478 continue
479 try:
480 step = int(entry[len(STEP_PREFIX):])
481 except ValueError:
482 continue
483 if self._is_complete(entry_path):
484 found.append((step, entry_path))
486 found.sort(key=lambda item: item[0])
487 return found
489 def find_latest_checkpoint(self, checkpoint_dir: str) -> Optional[str]:
490 """Find the newest complete checkpoint under ``checkpoint_dir``.
492 The pointer file written after each save answers this in one read. It is
493 only a cache, though --- checkpoints produced before it existed, or
494 copied in by hand, are still found by scanning the directory. A
495 checkpoint interrupted mid-save has no DCP metadata and is skipped by
496 both paths, so resume never lands on a half-written one.
497 """
498 from_pointer = self._read_latest_pointer(checkpoint_dir)
499 if from_pointer is not None:
500 return from_pointer
502 candidates = self._list_checkpoints(checkpoint_dir)
503 if not candidates:
504 return None
505 return candidates[-1][1]
507 # ------------------------------------------------------------------
508 # Async coordination
509 # ------------------------------------------------------------------
511 def maybe_wait_for_async_save(self) -> None:
512 """Block until an in-flight async save finishes.
514 Safe to call when nothing is pending. This is the single entry point for
515 async-save coordination --- it runs before every new save, before a
516 restore, and at train end, so two saves never overlap and the process
517 never exits with a half-persisted checkpoint. Publishing the checkpoint
518 is not done here: the persist callback has already written the pointer by
519 the time the future resolves, and waiting for this call to come around
520 would leave the checkpoint unnamed until the next save.
521 """
522 if self._async_save_response is None:
523 return
525 rank = _get_rank()
526 try:
527 logger.info("[rank %s] waiting for async DCP save to finish...", rank)
528 self._async_save_response.persist_completion.result()
529 logger.info("[rank %s] async DCP save has been finished...", rank)
530 except Exception: # pylint: disable=broad-except
531 logger.exception("[rank %s] async DCP save failed", rank)
532 raise
533 finally:
534 self._async_save_response = None
536 def raise_for_failed_async_save(self) -> None:
537 """Re-raise an in-flight async save's failure as soon as it is known.
539 Non-blocking: a save still running is left alone, so a training loop can
540 afford to ask every step. That is the point --- the thread that resolves
541 the future cannot propagate anything into the training loop, so a failure
542 otherwise waits for the next :meth:`maybe_wait_for_async_save`, which is
543 the next save, and a long ``save_steps`` interval turns one failed save
544 into thousands of wasted steps.
545 """
546 response = self._async_save_response
547 if response is None or not response.persist_completion.done():
548 return
550 # Dropped before the raise, so one failed save is reported once rather
551 # than at every step that follows it. The drain has nothing left to wait
552 # for either way.
553 self._async_save_response = None
554 response.persist_completion.result()
556 def _publish_latest(self, save_dir: str, step: int) -> None:
557 """Point the latest-checkpoint file at a checkpoint that is fully on disk.
559 Deliberately collective-free: the async path calls this from its persist
560 callback, which runs on a background thread while the training loop is
561 still issuing its own collectives, and a ``dist`` call from there would
562 race them. Rank 0 writing on its own is enough, because the persisting
563 child processes coordinate among themselves and rank 0 hosts their
564 coordinator --- its success already implies every rank's shards and the
565 shared ``.metadata`` have landed.
566 """
567 if _get_rank() == 0:
568 self._write_latest_pointer(os.path.dirname(save_dir), step)
570 logger.info(
571 "Distributed checkpoint saved successfully: global_step=%s, dir=%s",
572 step,
573 save_dir,
574 )
576 def _finalize_checkpoint(self, save_dir: str, step: int) -> None:
577 """Publish a synchronously written checkpoint as the newest one.
579 The leading barrier holds rank 0 until every rank has finished writing,
580 so the pointer is never published over a checkpoint that is still being
581 assembled; the trailing one keeps the ranks in step afterwards. The async
582 path cannot borrow this and publishes from its persist callback instead.
583 """
584 _barrier()
585 self._publish_latest(save_dir, step)
587 # ------------------------------------------------------------------
588 # Save
589 # ------------------------------------------------------------------
591 def save(
592 self,
593 path: str,
594 state: Dict[str, Any],
595 *,
596 global_step: int,
597 save_async: bool = False,
598 ) -> None:
599 """Write ``state`` into ``path``; see :meth:`CheckpointerBase.save`."""
600 # Never let two saves overlap: they would compete for host memory and,
601 # in async mode, race on the checkpoint directory.
602 self.maybe_wait_for_async_save()
604 payload = dict(state)
605 extra_state = payload.pop("extra_state", None)
607 if extra_state is not None and self.extra_state_per_rank:
608 os.makedirs(os.path.join(path, _EXTRA_STATE_DIR), exist_ok=True)
609 _barrier()
610 # Written before the shards so that every persisted model/optimizer
611 # payload is already backed by the progress state to resume from it.
612 torch.save(extra_state, self._extra_state_path(path, _get_rank()))
613 else:
614 os.makedirs(path, exist_ok=True)
615 _barrier()
616 if extra_state is not None:
617 payload["extra_state"] = extra_state
619 # Free residual training-step allocations before DCP allocates its
620 # collective buffers for the gather pass.
621 _empty_cache()
623 if save_async:
624 # Publishing from the persist callback, rather than from the next
625 # ``maybe_wait_for_async_save``, is what keeps the pointer honest:
626 # a run that dies after an async save has landed but before the
627 # next save would otherwise leave a complete checkpoint on disk
628 # that the pointer never names.
629 self._async_save_response = dcp_async_save(
630 payload,
631 checkpoint_id=path,
632 callback=partial(self._publish_latest, path, global_step),
633 )
634 logger.info("Async DCP save dispatched: %s", path)
635 else:
636 dcp_save(payload, checkpoint_id=path)
637 logger.info("Sync DCP save completed: %s", path)
639 _empty_cache()
641 if not save_async:
642 self._finalize_checkpoint(path, global_step)
644 # ------------------------------------------------------------------
645 # Load
646 # ------------------------------------------------------------------
648 def load(
649 self,
650 path: str,
651 state: Dict[str, Any],
652 *,
653 strict_model: bool = True,
654 extra_state_skeleton: Optional[Dict[str, Any]] = None,
655 ) -> Dict[str, Any]:
656 """Fill ``state`` from ``path``; see :meth:`CheckpointerBase.load`."""
657 self.maybe_wait_for_async_save()
659 # Auto-detect the layout the checkpoint was written with, so the
660 # configured ``extra_state_per_rank`` may differ from the producing run.
661 extra_state_file = self._extra_state_path(path, _get_rank())
662 from_file = os.path.isfile(extra_state_file)
664 if extra_state_skeleton is not None and not from_file:
665 skeleton = dict(extra_state_skeleton)
666 skeleton[_STEP_KEY] = _UNRESTORED_STEP
667 state["extra_state"] = skeleton
669 try:
670 dcp_load(
671 state,
672 checkpoint_id=path,
673 planner=_ModelStrictLoadPlanner(strict_model=strict_model),
674 )
675 except Exception:
676 logger.exception("Failed to load checkpoint from %s", path)
677 raise
679 if extra_state_skeleton is not None:
680 state["extra_state"] = self._read_extra_state(
681 path, state, extra_state_file, from_file
682 )
684 # Hold every rank here until all of them have finished reading, so no
685 # rank starts training against a checkpoint others are still loading.
686 _barrier()
687 return state
689 def _read_extra_state(
690 self,
691 path: str,
692 state: Dict[str, Any],
693 extra_state_file: str,
694 from_file: bool,
695 ) -> Dict[str, Any]:
696 """Return the extra_state bundle from wherever the checkpoint kept it."""
697 if from_file:
698 logger.info("Reading extra_state from %s", extra_state_file)
699 return torch.load(extra_state_file, weights_only=False)
701 extra = state.get("extra_state") or {}
702 if extra.get(_STEP_KEY, _UNRESTORED_STEP) == _UNRESTORED_STEP:
703 raise FileNotFoundError(
704 f"Checkpoint {path} carries no training state for rank "
705 f"{_get_rank()}: neither a per-rank file ({extra_state_file}) nor "
706 "an extra_state entry inside the DCP. Resume needs the progress "
707 "state saved beside the model shards; set "
708 "checkpoint.restore_train_state=false to load the weights only."
709 )
710 logger.info("Read extra_state embedded in the DCP state dict")
711 return extra