Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / checkpoint_callback.py: 86%

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

16 

17__all__ = ["CheckpointerCallback"] 

18 

19import math 

20import os 

21import random 

22import shutil 

23import time 

24from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple 

25 

26import torch # pylint: disable=forbidden-backend-import 

27 

28from hyper_parallel.components.checkpoint import build_checkpointer 

29from hyper_parallel.components.checkpoint.dcp_checkpointer import ( 

30 STEP_PREFIX, 

31 initialize_optimizer_state, 

32) 

33from hyper_parallel.components.optim.mixed_precision_optimizer import ( 

34 MixedPrecisionOptimizer, 

35) 

36from hyper_parallel.models._transformers.checkpoint_loader import CheckpointManager 

37from hyper_parallel.core.distributed_checkpoint.utils import flatten_state_dict 

38from hyper_parallel.core.optimizer.muon import Muon 

39from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer 

40from hyper_parallel.models._transformers.model_builder import ( 

41 validate_model_init_dtype, 

42) 

43from hyper_parallel.trainer.runtime.logging import create_logger 

44from hyper_parallel.trainer.runtime.memory import empty_cache 

45from hyper_parallel.trainer.runtime.device import ( 

46 get_device_rng_state, 

47 set_device_rng_state, 

48) 

49from hyper_parallel.models.external_state import CheckpointRuntime, get_model_external_state 

50from hyper_parallel.trainer.runtime.distributed import all_reduce 

51from .base import Callback, TrainerState 

52 

53 

54if TYPE_CHECKING: 

55 from hyper_parallel.trainer.base import BaseTrainer 

56 

57 

58logger = create_logger(__name__) 

59_HF_CHECKPOINT_DIR = "hf_ckpt" 

60_HF_WEIGHT_FILES = ( 

61 "model.safetensors", 

62 "model.safetensors.index.json", 

63 "pytorch_model.bin", 

64 "pytorch_model.bin.index.json", 

65) 

66 

67 

68def _as_list(value: Any) -> List[Any]: 

69 """Normalize an optional single-or-list component into a list.""" 

70 if value is None: 

71 return [] 

72 return list(value) if isinstance(value, list) else [value] 

73 

74 

75def _state_leaf_keys(value: Any) -> frozenset[str]: 

76 """Use DCP's own traversal to list persisted dense optimizer leaves.""" 

77 flat, _ = flatten_state_dict({"optimizer": value}) 

78 return frozenset(key.removeprefix("optimizer.") for key in flat) 

79 

80 

81def _unwrap_single(values: List[Any]) -> Any: 

82 """Collapse a one-element list so single-component runs keep a flat state.""" 

83 if len(values) == 1: 

84 return values[0] 

85 return values 

86 

87 

88def _cleanup_muon_checkpoint_state(optimizer: Any) -> None: 

89 """Release ordinary Muon replicas after checkpoint consumers finish.""" 

90 if isinstance(optimizer, MixedPrecisionOptimizer): 

91 optimizer = optimizer.optimizer 

92 leaves = optimizer.chained_optimizers if isinstance(optimizer, ChainedOptimizer) else (optimizer,) 

93 for leaf in leaves: 

94 # Swap wrappers and subclasses own their checkpoint state lifecycle. 

95 if type(leaf) is Muon: # pylint: disable=unidiomatic-typecheck 

96 leaf.cleanup_synced_state() 

97 

98 

99class CheckpointerCallback(Callback): 

100 """Decide when to checkpoint and what goes in it. 

101 

102 This callback owns *policy* --- the save cadence, duplicate suppression, and 

103 the mapping between trainer objects and the persisted payload. Writing that 

104 payload to disk and reading it back belongs to the 

105 :class:`~hyper_parallel.components.checkpoint.CheckpointerBase` it delegates 

106 to, so the storage format can change without touching this file. 

107 

108 The payload is the model and optimizer state dicts plus an ``extra_state`` 

109 bundle holding what those cannot represent: ``global_step`` / ``epoch``, the 

110 LR scheduler, the dataloader position, and the CPU / device / Python RNG 

111 states. 

112 

113 DCP saving, Hugging Face export, and restoring are independently controlled 

114 by ``save_ckpt``, ``save_hf_weights``, and ``restore_from``. This permits a 

115 restore-only run as well as exporting Hugging Face weights from an existing 

116 DCP checkpoint without writing another DCP checkpoint. 

117 

118 Restore runs in :meth:`on_train_begin`, i.e. after the model, optimizer, 

119 scheduler and dataloader exist but before the first training step. 

120 """ 

121 

122 def __init__(self, trainer: "BaseTrainer") -> None: 

123 """Read the checkpoint configuration and build the checkpointer.""" 

124 super().__init__(trainer) 

125 ckpt_cfg = trainer.config.checkpoint 

126 # ``save_ckpt`` gates the write path only --- this callback is also 

127 # registered for restore-only runs. Folding it into the cadence fields 

128 # makes "saving is off" structural (there is simply no cadence) instead 

129 # of a second enable check inside every save hook. 

130 self._save_ckpt = ckpt_cfg.save_ckpt 

131 self._checkpoint_dir = ckpt_cfg.checkpoint_dir 

132 self._save_steps = ckpt_cfg.save_steps if self._save_ckpt else 0 

133 self._save_epochs = ckpt_cfg.save_epochs if self._save_ckpt else 0 

134 save_time_interval_minutes = float(ckpt_cfg.save_time_interval_minutes) 

135 if not math.isfinite(save_time_interval_minutes): 

136 raise ValueError("checkpoint.save_time_interval_minutes must be finite") 

137 self._save_time_interval_seconds = ( 

138 max(save_time_interval_minutes, 0.0) * 60.0 

139 if self._save_ckpt 

140 else 0.0 

141 ) 

142 self._next_time_save: Optional[float] = None 

143 self._save_total_limit = ckpt_cfg.save_total_limit 

144 self._is_async = ckpt_cfg.is_async 

145 self._is_peft = ckpt_cfg.is_peft 

146 self._save_optimizer = ckpt_cfg.save_optimizer 

147 self._save_train_state = ckpt_cfg.save_train_state 

148 self._save_hf_weights = ckpt_cfg.save_hf_weights 

149 self._save_extra_state_per_rank = ckpt_cfg.save_extra_state_per_rank 

150 

151 if self._save_hf_weights and self._is_peft: 

152 raise ValueError( 

153 "checkpoint.save_hf_weights does not yet support PEFT checkpoints" 

154 ) 

155 

156 self._restore_from = ckpt_cfg.restore_from 

157 self._restore_optimizer = ckpt_cfg.restore_optimizer 

158 self._restore_train_state = ckpt_cfg.restore_train_state 

159 self._restore_dataloader_state = ckpt_cfg.restore_dataloader_state 

160 

161 external = get_model_external_state(getattr(trainer, "model", None)) 

162 if external is not None: 

163 if self._is_async or self._is_peft: 

164 raise ValueError("Host external state requires synchronous non-PEFT checkpointing") 

165 if self._save_ckpt and (not self._save_optimizer or not self._save_train_state): 

166 raise ValueError("Host checkpoint requires optimizer and train-state saving") 

167 if self._restore_optimizer != self._restore_train_state: 

168 raise ValueError("Host restore requires optimizer and train state together") 

169 if trainer.mesh.pp_size > 1 and not self._save_extra_state_per_rank: 

170 raise ValueError("Host PP checkpoint requires per-rank extra state") 

171 

172 self._last_saved_step: int = -1 

173 self._last_hf_saved_step: int = -1 

174 self.checkpointer = build_checkpointer( 

175 extra_state_per_rank=self._save_extra_state_per_rank, 

176 ) 

177 self._hf_checkpoint_manager = ( 

178 CheckpointManager(trainer.model) if self._save_hf_weights else None 

179 ) 

180 

181 # ------------------------------------------------------------------ 

182 # Hook dispatchers 

183 # ------------------------------------------------------------------ 

184 

185 def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None: 

186 """Log the checkpoint configuration and restore any requested state. 

187 

188 Args: 

189 state: Trainer or external model state. 

190 """ 

191 logger.info( 

192 "Checkpoint configuration: " 

193 "checkpoint_dir=%s, save_ckpt=%s, save_hf_weights=%s, " 

194 "save_steps=%s, save_epochs=%s, save_time_interval_minutes=%s, " 

195 "save_total_limit=%s, " 

196 "is_async=%s, is_peft=%s, " 

197 "save_extra_state_per_rank=%s, restore_from=%s", 

198 self._checkpoint_dir, 

199 self._save_ckpt, 

200 self._save_hf_weights, 

201 self._save_steps, 

202 self._save_epochs, 

203 self._save_time_interval_seconds / 60.0, 

204 self._save_total_limit, 

205 self._is_async, 

206 self._is_peft, 

207 self._save_extra_state_per_rank, 

208 self._restore_from, 

209 ) 

210 self._load_checkpoint() 

211 if self._save_time_interval_seconds > 0: 

212 self._next_time_save = time.monotonic() + self._save_time_interval_seconds 

213 

214 def on_step_end( # pylint: disable=arguments-differ 

215 self, state: TrainerState, **kwargs: Any 

216 ) -> None: 

217 """Save on either the configured step or elapsed-time cadence.""" 

218 self.checkpointer.raise_for_failed_async_save() 

219 step_due = self._save_steps > 0 and state.global_step % self._save_steps == 0 

220 time_due = self._time_checkpoint_due() 

221 if not step_due and not time_due: 

222 return 

223 

224 if state.global_step != self._last_saved_step: 

225 self._save_checkpoint(state) 

226 if time_due: 

227 self._advance_time_deadline() 

228 

229 def on_epoch_end(self, state: TrainerState, **kwargs: Any) -> None: 

230 """Save on the configured epoch cadence. 

231 

232 Args: 

233 state: Trainer or external model state. 

234 """ 

235 if self._save_epochs > 0 and (state.epoch + 1) % self._save_epochs == 0: 

236 if state.global_step != self._last_saved_step: 

237 self._save_checkpoint(state) 

238 else: 

239 logger.info( 

240 "Skipping duplicate checkpoint save at epoch_end " 

241 "(global_step %s already saved at step_end).", 

242 state.global_step, 

243 ) 

244 

245 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None: 

246 """Persist the final step, then drain any in-flight async save. 

247 

248 Always saved when saving is on and the step is not already on disk: 

249 losing the last stretch of training to a cadence that happened not to 

250 land on the final step is never what anyone wants. 

251 

252 Args: 

253 state: Trainer or external model state. 

254 """ 

255 if ( 

256 self._save_ckpt 

257 and state.global_step > 0 

258 and state.global_step != self._last_saved_step 

259 ): 

260 # The process is about to exit, so the last checkpoint is written 

261 # synchronously regardless of ``is_async``. 

262 self._save_checkpoint(state, force_sync=True) 

263 self.wait_for_pending_save() 

264 if ( 

265 self._save_hf_weights 

266 and state.global_step > 0 

267 and state.global_step != self._last_hf_saved_step 

268 ): 

269 save_dir = os.path.join( 

270 self._checkpoint_dir, f"{STEP_PREFIX}{state.global_step}" 

271 ) 

272 self._save_hf_checkpoint(save_dir, state.global_step) 

273 self._rotate_checkpoints() 

274 

275 def wait_for_pending_save(self) -> None: 

276 """Block until the checkpointer's in-flight async save is persisted.""" 

277 self.checkpointer.maybe_wait_for_async_save() 

278 self._rotate_checkpoints() 

279 

280 # ------------------------------------------------------------------ 

281 # Payload assembly 

282 # ------------------------------------------------------------------ 

283 

284 def _model_state_dict(self) -> Dict[str, Any]: 

285 """Return the model state to persist, trainable-only under PEFT.""" 

286 model = self.trainer.model 

287 state_dict = model.state_dict() 

288 external = get_model_external_state(model) 

289 if external is not None: 

290 missing = set(external.parameters) - set(state_dict) 

291 if missing: 

292 raise RuntimeError(f"External model FQNs missing from state_dict: {sorted(missing)}") 

293 state_dict = {key: value for key, value in state_dict.items() 

294 if key not in external.parameters} 

295 if not self._is_peft: 

296 return state_dict 

297 

298 trainable = { 

299 name for name, param in model.named_parameters() if param.requires_grad 

300 } 

301 return {name: value for name, value in state_dict.items() if name in trainable} 

302 

303 def _collect_extra_state(self, state: TrainerState) -> Dict[str, Any]: 

304 """Build the extra_state bundle (progress / scheduler / dataloader / RNG).""" 

305 # Prefer the iterator snapshot: with background prefetching the loader has 

306 # already advanced past the batch the training step actually consumed. 

307 dataloader_state: Dict[str, Any] = {} 

308 data_iterator = getattr(self.trainer, "data_iterator", None) 

309 if data_iterator is not None and hasattr(data_iterator, "state_dict"): 

310 dataloader_state = data_iterator.state_dict() 

311 elif self.trainer.train_dataloader is not None and hasattr( 

312 self.trainer.train_dataloader, "state_dict" 

313 ): 

314 dataloader_state = self.trainer.train_dataloader.state_dict() 

315 

316 schedulers = _as_list(self.trainer.lr_scheduler) 

317 lr_scheduler_sd = _unwrap_single([sch.state_dict() for sch in schedulers]) 

318 

319 return { 

320 "global_step": state.global_step, 

321 "epoch": state.epoch, 

322 "lr_scheduler": lr_scheduler_sd, 

323 "train_dataloader": dataloader_state, 

324 "rng_state": { 

325 "torch_cpu": torch.get_rng_state(), 

326 "torch_device": get_device_rng_state(), 

327 "python": random.getstate(), 

328 }, 

329 } 

330 

331 def _resume_position(self) -> Tuple[int, int]: 

332 """Return the ``(epoch, step)`` position the training loop resumes from. 

333 

334 Read off the same two values the loops use rather than re-derived from 

335 ``global_step``: they re-enter ``state.epoch`` and skip the optimizer steps 

336 its predecessors consumed, and ``state.epoch`` is restored from the 

337 checkpoint. Dividing ``global_step`` instead would be a second answer that 

338 can disagree with the loop, and it would have to divide by 

339 ``train_steps`` --- the optimizer steps one epoch holds --- not by the 

340 dataloader's length, which counts micro-batches. 

341 

342 Returns: 

343 The epoch training resumes in, and its offset within that epoch. 

344 """ 

345 state = self.trainer.state 

346 return state.epoch, state.global_step - state.epoch * self.trainer.train_steps 

347 

348 # ------------------------------------------------------------------ 

349 # Save 

350 # ------------------------------------------------------------------ 

351 

352 def _time_checkpoint_due(self) -> bool: 

353 """Return one rank-consistent decision for the elapsed-time cadence.""" 

354 if self._next_time_save is None: 

355 return False 

356 

357 local_due = int(time.monotonic() >= self._next_time_save) 

358 if torch.distributed.is_available() and torch.distributed.is_initialized(): 

359 return bool(all_reduce(local_due, op="max")) 

360 return bool(local_due) 

361 

362 def _advance_time_deadline(self) -> None: 

363 """Move the time deadline forward without accumulating save latency.""" 

364 if self._next_time_save is None: 

365 return 

366 now = time.monotonic() 

367 elapsed = max(now - self._next_time_save, 0.0) 

368 elapsed_intervals = int(elapsed // self._save_time_interval_seconds) 

369 self._next_time_save += (elapsed_intervals + 1) * self._save_time_interval_seconds 

370 

371 def _save_checkpoint(self, state: TrainerState, force_sync: bool = False) -> None: 

372 """Assemble the payload for this step and hand it to the checkpointer.""" 

373 if ( 

374 self._save_total_limit is not None 

375 and self._save_total_limit > 0 

376 and self._is_async 

377 and self._last_saved_step >= 0 

378 ): 

379 # Complete and rotate the preceding async checkpoint before 

380 # starting a new one. This temporarily allows limit + 1 

381 # directories instead of deleting the last known-good checkpoint 

382 # while its replacement is still being written. 

383 self.wait_for_pending_save() 

384 

385 save_dir = os.path.join(self._checkpoint_dir, f"{STEP_PREFIX}{state.global_step}") 

386 save_async = self._is_async and not force_sync 

387 

388 logger.info( 

389 "Saving checkpoint: global_step=%s, epoch=%s, dir=%s, is_async=%s, " 

390 "extra_state_per_rank=%s, optimizer=%s, train_state=%s", 

391 state.global_step, 

392 state.epoch, 

393 save_dir, 

394 save_async, 

395 self._save_extra_state_per_rank, 

396 self._save_optimizer, 

397 self._save_train_state, 

398 ) 

399 

400 checkpoint_state: Dict[str, Any] = {"model": self._model_state_dict()} 

401 optimizers = _as_list(self.trainer.optimizer) if self._save_optimizer else [] 

402 if self._save_optimizer: 

403 checkpoint_state["optimizer"] = _unwrap_single( 

404 [optimizer.state_dict() for optimizer in optimizers] 

405 ) 

406 if self._save_train_state: 

407 checkpoint_state["extra_state"] = self._collect_extra_state(state) 

408 

409 external = get_model_external_state(self.trainer.model) 

410 if external is not None: 

411 persisted_keys = (_state_leaf_keys(checkpoint_state["optimizer"]) 

412 if "optimizer" in checkpoint_state else frozenset()) 

413 external.before_checkpoint_save(CheckpointRuntime( 

414 step_dir=save_dir, global_step=state.global_step, 

415 mesh_context=self.trainer.mesh, 

416 include_optimizer=self._save_optimizer, 

417 include_train_state=self._save_train_state, 

418 persisted_optimizer_keys=persisted_keys, 

419 )) 

420 

421 model_integration = getattr(self.trainer, "model_integration", None) 

422 if model_integration is not None: 

423 model_integration.capture_checkpoint_payload( 

424 "before_save", 

425 save_dir, 

426 checkpoint_state, 

427 ) 

428 

429 self.checkpointer.save( 

430 save_dir, 

431 checkpoint_state, 

432 global_step=state.global_step, 

433 save_async=save_async, 

434 ) 

435 

436 # Bookkeeping reflects the dispatched step immediately, including in async 

437 # mode: otherwise on_epoch_end would queue the same step again while the 

438 # first save is still in flight. 

439 self._last_saved_step = state.global_step 

440 # save() has finished snapshotting even when disk persistence is async. 

441 del checkpoint_state 

442 cleanup_targets = ( 

443 [external.optimizer_for_dcp()] if external is not None and optimizers else optimizers 

444 ) 

445 for optimizer in cleanup_targets: 

446 _cleanup_muon_checkpoint_state(optimizer) 

447 if self._save_hf_weights: 

448 # The full-weight gather is memory-intensive, so never overlap it 

449 # with an asynchronous DCP persistence job. 

450 self.checkpointer.maybe_wait_for_async_save() 

451 self._save_hf_checkpoint(save_dir, state.global_step) 

452 self._rotate_checkpoints() 

453 elif not save_async: 

454 self._rotate_checkpoints() 

455 

456 def _save_hf_checkpoint(self, save_dir: str, global_step: int) -> None: 

457 """Collect and export one Transformers-compatible model checkpoint.""" 

458 if self._hf_checkpoint_manager is None: 

459 raise RuntimeError("Hugging Face checkpoint manager is not initialized") 

460 

461 hf_dir = os.path.join(save_dir, _HF_CHECKPOINT_DIR) 

462 logger.info( 

463 "Saving Hugging Face weights: global_step=%s, dir=%s", 

464 global_step, 

465 hf_dir, 

466 ) 

467 

468 write_assets = self._hf_checkpoint_manager.save_pretrained(hf_dir) 

469 if write_assets: 

470 self._save_hf_assets(hf_dir) 

471 if torch.distributed.is_available() and torch.distributed.is_initialized(): 

472 torch.distributed.barrier() 

473 

474 self._last_hf_saved_step = global_step 

475 if write_assets: 

476 logger.info( 

477 "Hugging Face checkpoint saved successfully: global_step=%s, dir=%s", 

478 global_step, 

479 hf_dir, 

480 ) 

481 

482 @staticmethod 

483 def _checkpoint_step(name: str) -> Optional[int]: 

484 """Parse an exact ``global_step_<integer>`` directory name.""" 

485 if not name.startswith(STEP_PREFIX): 

486 return None 

487 suffix = name[len(STEP_PREFIX):] 

488 if not suffix or not suffix.isascii() or not suffix.isdigit(): 

489 return None 

490 if len(suffix) > 1 and suffix.startswith("0"): 

491 return None 

492 return int(suffix) 

493 

494 @staticmethod 

495 def _is_complete_checkpoint(path: str) -> bool: 

496 """Return whether a managed step directory contains a complete payload.""" 

497 try: 

498 with os.scandir(path) as entries: 

499 if any( 

500 entry.is_file(follow_symlinks=False) 

501 and ( 

502 entry.name == ".metadata" 

503 or ( 

504 entry.name[:-len(".metadata")].isascii() 

505 and entry.name[:-len(".metadata")].isdigit() 

506 and entry.name.endswith(".metadata") 

507 ) 

508 ) 

509 for entry in entries 

510 ): 

511 return True 

512 except OSError: 

513 return False 

514 

515 hf_dir = os.path.join(path, _HF_CHECKPOINT_DIR) 

516 if os.path.islink(hf_dir) or not os.path.isdir(hf_dir): 

517 return False 

518 return any(os.path.isfile(os.path.join(hf_dir, name)) for name in _HF_WEIGHT_FILES) 

519 

520 def _list_managed_checkpoints(self) -> List[tuple[int, str]]: 

521 """List completed, non-symlink checkpoint directories from oldest to newest.""" 

522 checkpoint_root = os.path.realpath(self._checkpoint_dir) 

523 try: 

524 entries = os.scandir(checkpoint_root) 

525 except OSError: 

526 return [] 

527 

528 checkpoints = [] 

529 with entries: 

530 for entry in entries: 

531 step = self._checkpoint_step(entry.name) 

532 if step is None or entry.is_symlink() or not entry.is_dir(follow_symlinks=False): 

533 continue 

534 resolved_path = os.path.realpath(entry.path) 

535 if os.path.dirname(resolved_path) != checkpoint_root: 

536 logger.warning("Skipping checkpoint outside configured root: %s", entry.path) 

537 continue 

538 if self._is_complete_checkpoint(resolved_path): 

539 checkpoints.append((step, resolved_path)) 

540 

541 checkpoints.sort(key=lambda item: item[0]) 

542 return checkpoints 

543 

544 def _rotate_checkpoints(self) -> None: 

545 """Delete oldest completed checkpoints while preserving active saves.""" 

546 limit = self._save_total_limit 

547 if limit is None or limit <= 0: 

548 return 

549 if torch.distributed.is_available() and torch.distributed.is_initialized(): 

550 if torch.distributed.get_rank() != 0: 

551 return 

552 

553 checkpoints = self._list_managed_checkpoints() 

554 delete_count = len(checkpoints) - limit 

555 if delete_count <= 0: 

556 return 

557 

558 protected_steps = { 

559 step 

560 for step in (self._last_saved_step, self._last_hf_saved_step) 

561 if step >= 0 

562 } 

563 deletable = [item for item in checkpoints if item[0] not in protected_steps] 

564 if len(deletable) < delete_count: 

565 logger.warning( 

566 "Checkpoint limit %s cannot be met without deleting an active checkpoint; " 

567 "keeping %s completed checkpoints.", 

568 limit, 

569 len(checkpoints) - len(deletable), 

570 ) 

571 delete_count = len(deletable) 

572 

573 for step, checkpoint_path in deletable[:delete_count]: 

574 # Revalidate immediately before deletion so a replaced path or 

575 # symlink can never redirect rotation outside checkpoint_dir. 

576 checkpoint_root = os.path.realpath(self._checkpoint_dir) 

577 if ( 

578 os.path.islink(checkpoint_path) 

579 or not os.path.isdir(checkpoint_path) 

580 or os.path.dirname(os.path.realpath(checkpoint_path)) != checkpoint_root 

581 or self._checkpoint_step(os.path.basename(checkpoint_path)) != step 

582 ): 

583 logger.warning("Skipping unsafe checkpoint deletion target: %s", checkpoint_path) 

584 continue 

585 try: 

586 shutil.rmtree(checkpoint_path) 

587 except OSError as exc: 

588 logger.warning("Failed to delete old checkpoint %s: %s", checkpoint_path, exc) 

589 continue 

590 logger.info( 

591 "Deleted old checkpoint due to save_total_limit=%s: %s", 

592 limit, 

593 checkpoint_path, 

594 ) 

595 

596 def _save_hf_assets(self, hf_dir: str) -> None: 

597 """Write the tokenizer or processor assets next to exported weights.""" 

598 processor = getattr(self.trainer, "processor", None) 

599 if processor is not None: 

600 processor.save_pretrained(hf_dir) 

601 return 

602 

603 chat_template = getattr(self.trainer, "chat_template", None) 

604 if chat_template is not None: 

605 chat_template.save_pretrained(hf_dir) 

606 return 

607 

608 tokenizer = getattr(self.trainer, "tokenizer", None) 

609 if tokenizer is not None: 

610 save_tokenizer = getattr(tokenizer, "save_pretrained", None) 

611 if callable(save_tokenizer): 

612 save_tokenizer(hf_dir) 

613 else: 

614 logger.warning( 

615 "Tokenizer %s has no exportable assets; HF weights were saved without tokenizer files.", 

616 type(tokenizer).__name__, 

617 ) 

618 

619 # ------------------------------------------------------------------ 

620 # Load / restore 

621 # ------------------------------------------------------------------ 

622 

623 def _resolve_restore_path(self) -> Optional[str]: 

624 """Resolve ``restore_from`` (including ``LATEST``) to a directory.""" 

625 if self._restore_from is None: 

626 logger.info("No checkpoint to restore (restore_from is None).") 

627 return None 

628 

629 restore_path = self._restore_from 

630 if restore_path.upper() == "LATEST": 

631 logger.info( 

632 "restore_from='LATEST', searching for the latest checkpoint in %s", 

633 self._checkpoint_dir, 

634 ) 

635 resolved = self.checkpointer.find_latest_checkpoint(self._checkpoint_dir) 

636 if resolved is None: 

637 logger.warning( 

638 "restore_from='LATEST' but no checkpoint found in %s; " 

639 "starting from scratch.", 

640 self._checkpoint_dir, 

641 ) 

642 return None 

643 logger.info("Resolved LATEST checkpoint: %s", resolved) 

644 return resolved 

645 

646 if not os.path.isdir(restore_path): 

647 raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}") 

648 return restore_path 

649 

650 def _build_restore_skeleton(self, restore_path: str, external: Any) -> tuple[Dict[str, Any], List[Any]]: 

651 """Initialize lazy dense moments and create the DCP load skeleton.""" 

652 optimizers = _as_list(self.trainer.optimizer) if self._restore_optimizer else [] 

653 for optimizer in optimizers: 

654 target = external.optimizer_for_dcp() if external is not None else optimizer 

655 if not initialize_optimizer_state(target): 

656 logger.warning( 

657 "Could not materialize optimizer state before loading; " 

658 "optimizer moments may not be restored from %s.", 

659 restore_path, 

660 ) 

661 

662 checkpoint_state: Dict[str, Any] = {"model": self._model_state_dict()} 

663 if optimizers: 

664 checkpoint_state["optimizer"] = _unwrap_single( 

665 [optimizer.state_dict() for optimizer in optimizers] 

666 ) 

667 return checkpoint_state, optimizers 

668 

669 def _load_optimizer_payload(self, checkpoint_state: Dict[str, Any], optimizers: List[Any], 

670 persisted_keys: Optional[frozenset[str]]) -> None: 

671 """Restore dense optimizer leaves or reload dense main parameters.""" 

672 optimizer_sds = _as_list(checkpoint_state.get("optimizer")) 

673 for optimizer, optimizer_sd in zip(optimizers, optimizer_sds): 

674 optimizer.load_state_dict(optimizer_sd) 

675 drop_primed = getattr(optimizer, "drop_unpersisted_dense_state", None) 

676 if drop_primed is not None and persisted_keys is not None: 

677 drop_primed(persisted_keys) 

678 if not optimizers: 

679 for optimizer in _as_list(self.trainer.optimizer): 

680 if isinstance(optimizer, MixedPrecisionOptimizer): 

681 optimizer.reload_model_params() 

682 

683 def _apply_restored_payload(self, checkpoint_state: Dict[str, Any], optimizers: List[Any], 

684 external: Any, runtime: Any, restore_path: str, 

685 load_plan: Any) -> None: 

686 """Install the loaded dense and external model, optimizer, and train state.""" 

687 if external is not None: 

688 full_model_keys = set(self.trainer.model.state_dict()) 

689 dense_model_keys = set(checkpoint_state["model"]) 

690 if full_model_keys - dense_model_keys != set(external.parameters): 

691 raise RuntimeError("Dense checkpoint model keys differ from exact external exclusions") 

692 load_result = self.trainer.model.load_state_dict( 

693 checkpoint_state["model"], strict=not self._is_peft and external is None 

694 ) 

695 if external is not None and load_result is not None and ( 

696 set(load_result.missing_keys) != set(external.parameters) 

697 or load_result.unexpected_keys): 

698 raise RuntimeError("Dense checkpoint model keys differ from exact external exclusions") 

699 validate_model_init_dtype( 

700 self.trainer.model, 

701 self.trainer.config.model_init_dtype, 

702 ) 

703 persisted_keys = load_plan.persisted_optimizer_keys if load_plan is not None else None 

704 self._load_optimizer_payload(checkpoint_state, optimizers, persisted_keys) 

705 

706 if self._restore_train_state: 

707 self._apply_extra_state(checkpoint_state["extra_state"]) 

708 else: 

709 logger.info( 

710 "restore_train_state=False: loaded weights only from %s " 

711 "(step, scheduler, dataloader and RNG start fresh).", 

712 restore_path, 

713 ) 

714 

715 if external is not None: 

716 external.after_checkpoint_load(runtime, load_plan) 

717 if not self._restore_optimizer: 

718 dense = external.optimizer_for_dcp() 

719 if isinstance(dense, MixedPrecisionOptimizer): 

720 dense.reload_model_params() 

721 external.after_weights_only_load() 

722 

723 def _load_checkpoint(self) -> None: 

724 """Restore a checkpoint into the trainer's live objects.""" 

725 restore_path = self._resolve_restore_path() 

726 if restore_path is None: 

727 return 

728 

729 logger.info("Loading checkpoint from %s", restore_path) 

730 external = get_model_external_state(self.trainer.model) 

731 runtime = None 

732 load_plan = None 

733 if external is not None: 

734 runtime = CheckpointRuntime( 

735 step_dir=restore_path, global_step=None, 

736 mesh_context=self.trainer.mesh, 

737 include_optimizer=self._restore_optimizer, 

738 include_train_state=self._restore_train_state, 

739 ) 

740 load_plan = external.before_checkpoint_load(runtime) 

741 

742 checkpoint_state, optimizers = self._build_restore_skeleton(restore_path, external) 

743 

744 # The skeleton gives an embedded extra_state bundle keys to be read into; 

745 # the checkpointer decides whether it is actually needed for this layout. 

746 extra_state_skeleton = ( 

747 self._collect_extra_state(self.trainer.state) 

748 if self._restore_train_state 

749 else None 

750 ) 

751 if extra_state_skeleton is not None and self._restore_dataloader_state: 

752 loader_template = getattr(self.trainer.train_dataloader, "state_dict_for_load", None) 

753 if callable(loader_template): 

754 extra_state_skeleton["train_dataloader"] = loader_template() 

755 

756 self.checkpointer.load( 

757 restore_path, 

758 checkpoint_state, 

759 strict_model=not self._is_peft, 

760 extra_state_skeleton=extra_state_skeleton, 

761 ) 

762 if load_plan is not None and self._restore_optimizer: 

763 missing_optimizer_keys = (load_plan.persisted_optimizer_keys 

764 - _state_leaf_keys(checkpoint_state["optimizer"])) 

765 if missing_optimizer_keys: 

766 raise RuntimeError(f"Dense optimizer skeleton lost persisted keys: {sorted(missing_optimizer_keys)}") 

767 model_integration = getattr(self.trainer, "model_integration", None) 

768 if model_integration is not None: 

769 model_integration.capture_checkpoint_payload( 

770 "after_load", 

771 restore_path, 

772 checkpoint_state, 

773 ) 

774 

775 self._apply_restored_payload( 

776 checkpoint_state, optimizers, external, runtime, restore_path, 

777 load_plan, 

778 ) 

779 

780 del checkpoint_state 

781 cleanup_targets = ( 

782 [external.optimizer_for_dcp()] if external is not None and optimizers else optimizers 

783 ) 

784 for optimizer in cleanup_targets: 

785 _cleanup_muon_checkpoint_state(optimizer) 

786 empty_cache() 

787 # Computed here because this log is the only consumer: the training loops 

788 # derive their own starting position from ``state`` instead, so carrying it 

789 # on the trainer would be state nobody reads. 

790 start_epoch, start_step = self._resume_position() 

791 logger.info( 

792 "Checkpoint loaded successfully: path=%s, global_step=%s, " 

793 "start_epoch=%s, start_step=%s", 

794 restore_path, 

795 self.trainer.state.global_step, 

796 start_epoch, 

797 start_step, 

798 ) 

799 

800 def _apply_extra_state(self, extra: Dict[str, Any]) -> None: 

801 """Restore progress, scheduler, dataloader position and RNG state.""" 

802 trainer = self.trainer 

803 trainer.state.global_step = extra["global_step"] 

804 persisted_epoch = extra.get("epoch", 0) 

805 trainer.state.epoch = persisted_epoch 

806 

807 # The restored step is already on disk. Without this, resuming a run that 

808 # had nothing left to do would have ``on_train_end`` rewrite the very 

809 # checkpoint it just loaded. 

810 self._last_saved_step = trainer.state.global_step 

811