Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / checkpoint_callback.py: 84%
233 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"""CheckpointerCallback --- save/restore policy on top of a Checkpointer."""
17__all__ = ["CheckpointerCallback"]
19import os
20import random
21from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
23import torch # pylint: disable=forbidden-backend-import
25from hyper_parallel.components.checkpoint import build_checkpointer
26from hyper_parallel.components.checkpoint.dcp_checkpointer import (
27 STEP_PREFIX,
28 initialize_optimizer_state,
29)
30from hyper_parallel.components.optim.mixed_precision_optimizer import (
31 MixedPrecisionOptimizer,
32)
33from hyper_parallel.core.distributed_checkpoint.utils import flatten_state_dict
34from hyper_parallel.models._transformers.model_builder import (
35 validate_model_init_dtype,
36)
37from hyper_parallel.trainer.runtime.logging import create_logger
38from hyper_parallel.trainer.runtime.memory import empty_cache
39from hyper_parallel.trainer.runtime.device import (
40 get_device_rng_state,
41 set_device_rng_state,
42)
43from hyper_parallel.models.external_state import CheckpointRuntime, get_model_external_state
44from .base import Callback, TrainerState
47if TYPE_CHECKING:
48 from hyper_parallel.trainer.base import BaseTrainer
51logger = create_logger(__name__)
54def _as_list(value: Any) -> List[Any]:
55 """Normalize an optional single-or-list component into a list."""
56 if value is None:
57 return []
58 return list(value) if isinstance(value, list) else [value]
61def _state_leaf_keys(value: Any) -> frozenset[str]:
62 """Use DCP's own traversal to list persisted dense optimizer leaves."""
63 flat, _ = flatten_state_dict({"optimizer": value})
64 return frozenset(key.removeprefix("optimizer.") for key in flat)
67def _unwrap_single(values: List[Any]) -> Any:
68 """Collapse a one-element list so single-component runs keep a flat state."""
69 if len(values) == 1:
70 return values[0]
71 return values
74class CheckpointerCallback(Callback):
75 """Decide when to checkpoint and what goes in it.
77 This callback owns *policy* --- the save cadence, duplicate suppression, and
78 the mapping between trainer objects and the persisted payload. Writing that
79 payload to disk and reading it back belongs to the
80 :class:`~hyper_parallel.components.checkpoint.CheckpointerBase` it delegates
81 to, so the storage format can change without touching this file.
83 The payload is the model and optimizer state dicts plus an ``extra_state``
84 bundle holding what those cannot represent: ``global_step`` / ``epoch``, the
85 LR scheduler, the dataloader position, and the CPU / device / Python RNG
86 states.
88 Saving and restoring are independent: ``save_ckpt`` gates the write path,
89 ``restore_from`` the read path. Turning saving off while pointing at a
90 checkpoint is therefore a supported combination --- start from these weights
91 and write nothing further.
93 Restore runs in :meth:`on_train_begin`, i.e. after the model, optimizer,
94 scheduler and dataloader exist but before the first training step.
95 """
97 def __init__(self, trainer: "BaseTrainer") -> None:
98 """Read the checkpoint configuration and build the checkpointer."""
99 super().__init__(trainer)
100 ckpt_cfg = trainer.config.checkpoint
101 # ``save_ckpt`` gates the write path only --- this callback is also
102 # registered for restore-only runs. Folding it into the cadence fields
103 # makes "saving is off" structural (there is simply no cadence) instead
104 # of a second enable check inside every save hook.
105 self._save_ckpt = ckpt_cfg.save_ckpt
106 self._checkpoint_dir = ckpt_cfg.checkpoint_dir
107 self._save_steps = ckpt_cfg.save_steps if self._save_ckpt else 0
108 self._save_epochs = ckpt_cfg.save_epochs if self._save_ckpt else 0
109 self._is_async = ckpt_cfg.is_async
110 self._is_peft = ckpt_cfg.is_peft
111 self._save_optimizer = ckpt_cfg.save_optimizer
112 self._save_train_state = ckpt_cfg.save_train_state
113 self._save_extra_state_per_rank = ckpt_cfg.save_extra_state_per_rank
115 self._restore_from = ckpt_cfg.restore_from
116 self._restore_optimizer = ckpt_cfg.restore_optimizer
117 self._restore_train_state = ckpt_cfg.restore_train_state
118 self._restore_dataloader_state = ckpt_cfg.restore_dataloader_state
120 external = get_model_external_state(getattr(trainer, "model", None))
121 if external is not None:
122 if self._is_async or self._is_peft:
123 raise ValueError("Host external state requires synchronous non-PEFT checkpointing")
124 if self._save_ckpt and (not self._save_optimizer or not self._save_train_state):
125 raise ValueError("Host checkpoint requires optimizer and train-state saving")
126 if self._restore_optimizer != self._restore_train_state:
127 raise ValueError("Host restore requires optimizer and train state together")
128 if trainer.mesh.pp_size > 1 and not self._save_extra_state_per_rank:
129 raise ValueError("Host PP checkpoint requires per-rank extra state")
131 self._last_saved_step: int = -1
132 self.checkpointer = build_checkpointer(
133 extra_state_per_rank=self._save_extra_state_per_rank,
134 )
136 # ------------------------------------------------------------------
137 # Hook dispatchers
138 # ------------------------------------------------------------------
140 def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None:
141 """Log the checkpoint configuration and restore any requested state.
143 Args:
144 state: Trainer or external model state.
145 """
146 logger.info(
147 "Checkpoint configuration: "
148 "checkpoint_dir=%s, save_ckpt=%s, save_steps=%s, save_epochs=%s, "
149 "is_async=%s, is_peft=%s, "
150 "save_extra_state_per_rank=%s, restore_from=%s",
151 self._checkpoint_dir,
152 self._save_ckpt,
153 self._save_steps,
154 self._save_epochs,
155 self._is_async,
156 self._is_peft,
157 self._save_extra_state_per_rank,
158 self._restore_from,
159 )
160 self._load_checkpoint()
162 def on_step_end(self, state: TrainerState, **kwargs: Any) -> None: # pylint: disable=arguments-differ
163 """Surface a failed async save, then save on the configured step cadence.
165 Args:
166 state: Trainer or external model state.
167 """
168 # Asked every step on purpose. An async save that failed can only be
169 # reported on a thread that cannot raise into this loop, so waiting for the
170 # next save to notice would throw away every step in between.
171 self.checkpointer.raise_for_failed_async_save()
172 if self._save_steps > 0 and state.global_step % self._save_steps == 0:
173 if state.global_step == self._last_saved_step:
174 return
175 self._save_checkpoint(state)
177 def on_epoch_end(self, state: TrainerState, **kwargs: Any) -> None:
178 """Save on the configured epoch cadence.
180 Args:
181 state: Trainer or external model state.
182 """
183 if self._save_epochs > 0 and (state.epoch + 1) % self._save_epochs == 0:
184 if state.global_step != self._last_saved_step:
185 self._save_checkpoint(state)
186 else:
187 logger.info(
188 "Skipping duplicate checkpoint save at epoch_end "
189 "(global_step %s already saved at step_end).",
190 state.global_step,
191 )
193 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
194 """Persist the final step, then drain any in-flight async save.
196 Always saved when saving is on and the step is not already on disk:
197 losing the last stretch of training to a cadence that happened not to
198 land on the final step is never what anyone wants.
200 Args:
201 state: Trainer or external model state.
202 """
203 if (
204 self._save_ckpt
205 and state.global_step > 0
206 and state.global_step != self._last_saved_step
207 ):
208 # The process is about to exit, so the last checkpoint is written
209 # synchronously regardless of ``is_async``.
210 self._save_checkpoint(state, force_sync=True)
211 self.wait_for_pending_save()
213 def wait_for_pending_save(self) -> None:
214 """Block until the checkpointer's in-flight async save is persisted."""
215 self.checkpointer.maybe_wait_for_async_save()
217 # ------------------------------------------------------------------
218 # Payload assembly
219 # ------------------------------------------------------------------
221 def _model_state_dict(self) -> Dict[str, Any]:
222 """Return the model state to persist, trainable-only under PEFT."""
223 model = self.trainer.model
224 state_dict = model.state_dict()
225 external = get_model_external_state(model)
226 if external is not None:
227 missing = set(external.parameters) - set(state_dict)
228 if missing:
229 raise RuntimeError(f"External model FQNs missing from state_dict: {sorted(missing)}")
230 state_dict = {key: value for key, value in state_dict.items()
231 if key not in external.parameters}
232 if not self._is_peft:
233 return state_dict
235 trainable = {
236 name for name, param in model.named_parameters() if param.requires_grad
237 }
238 return {name: value for name, value in state_dict.items() if name in trainable}
240 def _collect_extra_state(self, state: TrainerState) -> Dict[str, Any]:
241 """Build the extra_state bundle (progress / scheduler / dataloader / RNG)."""
242 # Prefer the iterator snapshot: with background prefetching the loader has
243 # already advanced past the batch the training step actually consumed.
244 dataloader_state: Dict[str, Any] = {}
245 data_iterator = getattr(self.trainer, "data_iterator", None)
246 if data_iterator is not None and hasattr(data_iterator, "state_dict"):
247 dataloader_state = data_iterator.state_dict()
248 elif self.trainer.train_dataloader is not None and hasattr(
249 self.trainer.train_dataloader, "state_dict"
250 ):
251 dataloader_state = self.trainer.train_dataloader.state_dict()
253 schedulers = _as_list(self.trainer.lr_scheduler)
254 lr_scheduler_sd = _unwrap_single([sch.state_dict() for sch in schedulers])
256 return {
257 "global_step": state.global_step,
258 "epoch": state.epoch,
259 "lr_scheduler": lr_scheduler_sd,
260 "train_dataloader": dataloader_state,
261 "rng_state": {
262 "torch_cpu": torch.get_rng_state(),
263 "torch_device": get_device_rng_state(),
264 "python": random.getstate(),
265 },
266 }
268 def _resume_position(self) -> Tuple[int, int]:
269 """Return the ``(epoch, step)`` position the training loop resumes from.
271 Read off the same two values the loops use rather than re-derived from
272 ``global_step``: they re-enter ``state.epoch`` and skip the optimizer steps
273 its predecessors consumed, and ``state.epoch`` is restored from the
274 checkpoint. Dividing ``global_step`` instead would be a second answer that
275 can disagree with the loop, and it would have to divide by
276 ``train_steps`` --- the optimizer steps one epoch holds --- not by the
277 dataloader's length, which counts micro-batches.
279 Returns:
280 The epoch training resumes in, and its offset within that epoch.
281 """
282 state = self.trainer.state
283 return state.epoch, state.global_step - state.epoch * self.trainer.train_steps
285 # ------------------------------------------------------------------
286 # Save
287 # ------------------------------------------------------------------
289 def _save_checkpoint(self, state: TrainerState, force_sync: bool = False) -> None:
290 """Assemble the payload for this step and hand it to the checkpointer."""
291 save_dir = os.path.join(self._checkpoint_dir, f"{STEP_PREFIX}{state.global_step}")
292 save_async = self._is_async and not force_sync
294 logger.info(
295 "Saving checkpoint: global_step=%s, epoch=%s, dir=%s, is_async=%s, "
296 "extra_state_per_rank=%s, optimizer=%s, train_state=%s",
297 state.global_step,
298 state.epoch,
299 save_dir,
300 save_async,
301 self._save_extra_state_per_rank,
302 self._save_optimizer,
303 self._save_train_state,
304 )
306 checkpoint_state: Dict[str, Any] = {"model": self._model_state_dict()}
307 if self._save_optimizer:
308 checkpoint_state["optimizer"] = _unwrap_single(
309 [optimizer.state_dict() for optimizer in _as_list(self.trainer.optimizer)]
310 )
311 if self._save_train_state:
312 checkpoint_state["extra_state"] = self._collect_extra_state(state)
314 external = get_model_external_state(self.trainer.model)
315 if external is not None:
316 persisted_keys = (_state_leaf_keys(checkpoint_state["optimizer"])
317 if "optimizer" in checkpoint_state else frozenset())
318 external.before_checkpoint_save(CheckpointRuntime(
319 step_dir=save_dir, global_step=state.global_step,
320 mesh_context=self.trainer.mesh,
321 include_optimizer=self._save_optimizer,
322 include_train_state=self._save_train_state,
323 persisted_optimizer_keys=persisted_keys,
324 ))
326 model_integration = getattr(self.trainer, "model_integration", None)
327 if model_integration is not None:
328 model_integration.capture_checkpoint_payload(
329 "before_save",
330 save_dir,
331 checkpoint_state,
332 )
334 self.checkpointer.save(
335 save_dir,
336 checkpoint_state,
337 global_step=state.global_step,
338 save_async=save_async,
339 )
341 # Bookkeeping reflects the dispatched step immediately, including in async
342 # mode: otherwise on_epoch_end would queue the same step again while the
343 # first save is still in flight.
344 self._last_saved_step = state.global_step
346 # ------------------------------------------------------------------
347 # Load / restore
348 # ------------------------------------------------------------------
350 def _resolve_restore_path(self) -> Optional[str]:
351 """Resolve ``restore_from`` (including ``LATEST``) to a directory."""
352 if self._restore_from is None:
353 logger.info("No checkpoint to restore (restore_from is None).")
354 return None
356 restore_path = self._restore_from
357 if restore_path.upper() == "LATEST":
358 logger.info(
359 "restore_from='LATEST', searching for the latest checkpoint in %s",
360 self._checkpoint_dir,
361 )
362 resolved = self.checkpointer.find_latest_checkpoint(self._checkpoint_dir)
363 if resolved is None:
364 logger.warning(
365 "restore_from='LATEST' but no checkpoint found in %s; "
366 "starting from scratch.",
367 self._checkpoint_dir,
368 )
369 return None
370 logger.info("Resolved LATEST checkpoint: %s", resolved)
371 return resolved
373 if not os.path.isdir(restore_path):
374 raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}")
375 return restore_path
377 def _build_restore_skeleton(self, restore_path: str, external: Any) -> tuple[Dict[str, Any], List[Any]]:
378 """Initialize lazy dense moments and create the DCP load skeleton."""
379 optimizers = _as_list(self.trainer.optimizer) if self._restore_optimizer else []
380 for optimizer in optimizers:
381 target = external.optimizer_for_dcp() if external is not None else optimizer
382 if not initialize_optimizer_state(target):
383 logger.warning(
384 "Could not materialize optimizer state before loading; "
385 "optimizer moments may not be restored from %s.",
386 restore_path,
387 )
389 checkpoint_state: Dict[str, Any] = {"model": self._model_state_dict()}
390 if optimizers:
391 checkpoint_state["optimizer"] = _unwrap_single(
392 [optimizer.state_dict() for optimizer in optimizers]
393 )
394 return checkpoint_state, optimizers
396 def _load_optimizer_payload(self, checkpoint_state: Dict[str, Any], optimizers: List[Any],
397 persisted_keys: Optional[frozenset[str]]) -> None:
398 """Restore dense optimizer leaves or reload dense main parameters."""
399 optimizer_sds = _as_list(checkpoint_state.get("optimizer"))
400 for optimizer, optimizer_sd in zip(optimizers, optimizer_sds):
401 optimizer.load_state_dict(optimizer_sd)
402 drop_primed = getattr(optimizer, "drop_unpersisted_dense_state", None)
403 if drop_primed is not None and persisted_keys is not None:
404 drop_primed(persisted_keys)
405 if not optimizers:
406 for optimizer in _as_list(self.trainer.optimizer):
407 if isinstance(optimizer, MixedPrecisionOptimizer):
408 optimizer.reload_model_params()
410 def _apply_restored_payload(self, checkpoint_state: Dict[str, Any], optimizers: List[Any],
411 external: Any, runtime: Any, restore_path: str,
412 load_plan: Any) -> None:
413 """Install the loaded dense and external model, optimizer, and train state."""
414 if external is not None:
415 full_model_keys = set(self.trainer.model.state_dict())
416 dense_model_keys = set(checkpoint_state["model"])
417 if full_model_keys - dense_model_keys != set(external.parameters):
418 raise RuntimeError("Dense checkpoint model keys differ from exact external exclusions")
419 load_result = self.trainer.model.load_state_dict(
420 checkpoint_state["model"], strict=not self._is_peft and external is None
421 )
422 if external is not None and load_result is not None and (
423 set(load_result.missing_keys) != set(external.parameters)
424 or load_result.unexpected_keys):
425 raise RuntimeError("Dense checkpoint model keys differ from exact external exclusions")
426 validate_model_init_dtype(
427 self.trainer.model,
428 self.trainer.config.model_init_dtype,
429 )
430 persisted_keys = load_plan.persisted_optimizer_keys if load_plan is not None else None
431 self._load_optimizer_payload(checkpoint_state, optimizers, persisted_keys)
433 if self._restore_train_state:
434 self._apply_extra_state(checkpoint_state["extra_state"])
435 else:
436 logger.info(
437 "restore_train_state=False: loaded weights only from %s "
438 "(step, scheduler, dataloader and RNG start fresh).",
439 restore_path,
440 )
442 if external is not None:
443 external.after_checkpoint_load(runtime, load_plan)
444 if not self._restore_optimizer:
445 dense = external.optimizer_for_dcp()
446 if isinstance(dense, MixedPrecisionOptimizer):
447 dense.reload_model_params()
448 external.after_weights_only_load()
450 def _load_checkpoint(self) -> None:
451 """Restore a checkpoint into the trainer's live objects."""
452 restore_path = self._resolve_restore_path()
453 if restore_path is None:
454 return
456 logger.info("Loading checkpoint from %s", restore_path)
457 external = get_model_external_state(self.trainer.model)
458 runtime = None
459 load_plan = None
460 if external is not None:
461 runtime = CheckpointRuntime(
462 step_dir=restore_path, global_step=None,
463 mesh_context=self.trainer.mesh,
464 include_optimizer=self._restore_optimizer,
465 include_train_state=self._restore_train_state,
466 )
467 load_plan = external.before_checkpoint_load(runtime)
469 checkpoint_state, optimizers = self._build_restore_skeleton(restore_path, external)
471 # The skeleton gives an embedded extra_state bundle keys to be read into;
472 # the checkpointer decides whether it is actually needed for this layout.
473 extra_state_skeleton = (
474 self._collect_extra_state(self.trainer.state)
475 if self._restore_train_state
476 else None
477 )
479 self.checkpointer.load(
480 restore_path,
481 checkpoint_state,
482 strict_model=not self._is_peft,
483 extra_state_skeleton=extra_state_skeleton,
484 )
485 if load_plan is not None and self._restore_optimizer:
486 missing_optimizer_keys = (load_plan.persisted_optimizer_keys
487 - _state_leaf_keys(checkpoint_state["optimizer"]))
488 if missing_optimizer_keys:
489 raise RuntimeError(f"Dense optimizer skeleton lost persisted keys: {sorted(missing_optimizer_keys)}")
490 model_integration = getattr(self.trainer, "model_integration", None)
491 if model_integration is not None:
492 model_integration.capture_checkpoint_payload(
493 "after_load",
494 restore_path,
495 checkpoint_state,
496 )
498 self._apply_restored_payload(
499 checkpoint_state, optimizers, external, runtime, restore_path,
500 load_plan,
501 )
503 empty_cache()
504 # Computed here because this log is the only consumer: the training loops
505 # derive their own starting position from ``state`` instead, so carrying it
506 # on the trainer would be state nobody reads.
507 start_epoch, start_step = self._resume_position()
508 logger.info(
509 "Checkpoint loaded successfully: path=%s, global_step=%s, "
510 "start_epoch=%s, start_step=%s",
511 restore_path,
512 self.trainer.state.global_step,
513 start_epoch,
514 start_step,
515 )
517 def _apply_extra_state(self, extra: Dict[str, Any]) -> None:
518 """Restore progress, scheduler, dataloader position and RNG state."""
519 trainer = self.trainer
520 trainer.state.global_step = extra["global_step"]
521 persisted_epoch = extra.get("epoch", 0)
522 trainer.state.epoch = persisted_epoch
524 # The restored step is already on disk. Without this, resuming a run that
525 # had nothing left to do would have ``on_train_end`` rewrite the very
526 # checkpoint it just loaded.
527 self._last_saved_step = trainer.state.global_step
529 lr_scheduler_sd = extra.get("lr_scheduler")
530 schedulers = _as_list(trainer.lr_scheduler)
531 if lr_scheduler_sd and schedulers:
532 scheduler_sds = _as_list(lr_scheduler_sd)
533 if len(scheduler_sds) != len(schedulers):
534 logger.warning(
535 "Checkpoint carries %s LR scheduler state dict(s) but this "
536 "run has %s scheduler(s); only the first %s pair(s) are "
537 "restored and any extra scheduler(s) keep their freshly "
538 "initialized state.",
539 len(scheduler_sds),
540 len(schedulers),
541 min(len(scheduler_sds), len(schedulers)),
542 )
543 for scheduler, scheduler_sd in zip(schedulers, scheduler_sds):
544 scheduler.load_state_dict(scheduler_sd)
546 dataloader_sd = extra.get("train_dataloader")
547 if dataloader_sd and not self._restore_dataloader_state:
548 logger.info(
549 "restore_dataloader_state=False: retaining the configured "
550 "dataloader start instead of restoring its checkpoint cursor."
551 )
552 elif dataloader_sd and hasattr(trainer.train_dataloader, "load_state_dict"):
553 trainer.train_dataloader.load_state_dict(dataloader_sd)
554 elif dataloader_sd:
555 logger.warning(
556 "Checkpoint carries a dataloader position but %s is not stateful; "
557 "the resumed epoch replays samples from its start.",
558 type(trainer.train_dataloader).__name__,
559 )
561 rng_state = extra.get("rng_state") or {}
562 torch_cpu_rng = rng_state.get("torch_cpu")
563 if torch_cpu_rng is not None:
564 torch.set_rng_state(torch_cpu_rng)
565 set_device_rng_state(rng_state.get("torch_device"))
566 python_rng = rng_state.get("python")
567 if python_rng is not None:
568 random.setstate(python_rng)