Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / config.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""CheckpointingConfig typed config for checkpoint save/load. 

16 

17Following design doc 04_checkpoint.md. 

18""" 

19 

20from dataclasses import dataclass 

21from typing import Literal, Optional 

22 

23 

24@dataclass 

25class CheckpointingConfig: 

26 """Checkpointing typed config. 

27 

28 Following design doc 04_checkpoint.md. 

29 

30 ``save_hf_weights`` additionally exports a Transformers-compatible model 

31 under the current global-step directory. The remaining consolidated-export 

32 fields are declared by the design doc but are not yet consumed by 

33 :class:`~hyper_parallel.trainer.callbacks.checkpoint_callback.CheckpointerCallback`. 

34 """ 

35 

36 # Master switch for the *write* path only. Restoring is governed solely by 

37 # ``restore_from``, so ``save_ckpt=false`` with a ``restore_from`` set means 

38 # "start from this checkpoint but do not write any more" rather than 

39 # silently doing nothing. 

40 save_ckpt: bool = True 

41 checkpoint_dir: str = "./checkpoints" 

42 save_steps: int = 0 # save every N optimizer steps (0 = disabled) 

43 save_epochs: int = 1 # save every N epochs 

44 # Save after this many elapsed training minutes, at the next completed 

45 # optimizer step. A non-positive value disables time-based saving. 

46 save_time_interval_minutes: float = 0.0 

47 # Maximum number of completed ``global_step_*`` directories to retain. 

48 # None or a non-positive value disables checkpoint rotation. 

49 save_total_limit: Optional[int] = None 

50 is_async: bool = False 

51 is_peft: bool = False # persist trainable (adapter) weights only 

52 save_optimizer: bool = True 

53 save_train_state: bool = True 

54 save_hf_weights: bool = False 

55 # Where the per-step training state (step/epoch, lr_scheduler, dataloader 

56 # position, RNG) is written: 

57 # True --- one torch.save file per rank under ``extra_state/``. Always 

58 # correct, and the only safe choice when any of that state 

59 # differs across ranks. 

60 # False --- embedded in the DCP state dict. Fewer files, but DCP 

61 # deduplicates entries that share an FQN across ranks, so every 

62 # rank restores whichever rank's copy won. Only use this when the 

63 # extra state is rank-invariant (deterministic training, uniform 

64 # input shapes). 

65 # The load side auto-detects which layout a checkpoint used, so this flag can 

66 # be flipped between runs. 

67 save_extra_state_per_rank: bool = True 

68 

69 restore_from: Optional[str] = None # "LATEST" or specific path 

70 restore_optimizer: bool = True 

71 restore_train_state: bool = True # step/epoch, lr_scheduler, dataloader, RNG 

72 # Restore progress/scheduler/RNG while deliberately replaying data from the 

73 # loader's configured start. This is useful for a validation warm start in 

74 # which every topology must consume the same fixed replay, including when 

75 # the loader cannot reshard a saved cursor across DP world sizes. Keep this 

76 # enabled for ordinary K-to-K+N training continuation. 

77 restore_dataloader_state: bool = True 

78 

79 model_save_format: str = "safetensors" 

80 # Whether to additionally export merged HF weights: "none" (never) | 

81 # "final" (only at the end of training) | "every" (on every save). 

82 # The disabled setting is uniformly "none" (a YAML-safe Literal value), 

83 # replacing the old design's "false" enum value ("false" would be parsed 

84 # as a bool by PyYAML; see the decision in 04_checkpoint.md section 4.2). 

85 save_consolidated: Literal["none", "final", "every"] = "final" 

86 staging_dir: Optional[str] = None 

87 best_metric_key: str = "default"