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 ``model_save_format``, ``save_consolidated``, ``staging_dir`` and 

31 ``best_metric_key`` are declared by the design doc but not yet consumed by 

32 :class:`~hyper_parallel.trainer.callbacks.checkpoint_callback.CheckpointerCallback`; 

33 they are reserved for the consolidated HuggingFace-format export. 

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 is_async: bool = False 

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

46 save_optimizer: bool = True 

47 save_train_state: bool = True 

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

49 # position, RNG) is written: 

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

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

52 # differs across ranks. 

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

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

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

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

57 # input shapes). 

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

59 # be flipped between runs. 

60 save_extra_state_per_rank: bool = True 

61 

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

63 restore_optimizer: bool = True 

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

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

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

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

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

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

70 restore_dataloader_state: bool = True 

71 

72 model_save_format: str = "safetensors" 

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

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

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

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

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

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

79 staging_dir: Optional[str] = None 

80 best_metric_key: str = "default"