Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / __init__.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved. 

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

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

4# You may obtain a copy of the License at 

5# 

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

7# 

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

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

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

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

12# limitations under the License. 

13# ============================================================================ 

14""" 

15Hyper Parallel Checkpoint Package. 

16 

17This package provides distributed checkpoint saving and loading capabilities, 

18including support for tensor sharding, resharding, and layout management. 

19""" 

20 

21__all__ = [ 

22 # Main API 

23 "save", 

24 "async_save", 

25 "AsyncSaveResponse", 

26 "load", 

27 # Metadata 

28 "Metadata", 

29 "MetadataIndex", 

30 "TensorStorageMetadata", 

31 "BytesStorageMetadata", 

32 "ChunkStorageMetadata", 

33 "TensorProperties", 

34 "CHUNK_INFO", 

35 "ChunkInfo", 

36 # Planner interfaces 

37 "SavePlanner", 

38 "LoadPlanner", 

39 "SavePlan", 

40 "LoadPlan", 

41 "WriteItem", 

42 "ReadItem", 

43 "WriteItemType", 

44 "LoadItemType", 

45 # Standard planners 

46 "StandardSavePlanner", 

47 "StandardLoadPlanner", 

48 # Loads of checkpoints laid out differently from the state dict 

49 "DeferredRead", 

50 "RemapBlock", 

51 "RemapLoadPlanner", 

52 # Storage interfaces 

53 "StorageWriter", 

54 "StorageReader", 

55 "StorageInfo", 

56 "WriteResult", 

57 # File system storage 

58 "FileSystemWriter", 

59 "FileSystemReader", 

60 # Safetensors checkpoints written outside DCP 

61 "HuggingFaceStorageReader", 

62 "TorchShardedSafetensorsReader", 

63 # Optimizer state dict 

64 "get_optim_state_dict", 

65 "set_optim_state_dict", 

66] 

67 

68# Main API 

69from hyper_parallel.core.distributed_checkpoint.api import ( 

70 AsyncSaveResponse, 

71 async_save, 

72 load, 

73 save, 

74) 

75 

76# Metadata structures 

77from hyper_parallel.core.distributed_checkpoint.metadata import ( 

78 BytesStorageMetadata, 

79 ChunkStorageMetadata, 

80 Metadata, 

81 MetadataIndex, 

82 TensorProperties, 

83 TensorStorageMetadata, 

84 CHUNK_INFO, 

85 ChunkInfo 

86) 

87 

88# Planner interfaces and data structures 

89from hyper_parallel.core.distributed_checkpoint.planner import ( 

90 LoadItemType, 

91 LoadPlan, 

92 LoadPlanner, 

93 ReadItem, 

94 SavePlan, 

95 SavePlanner, 

96 WriteItem, 

97 WriteItemType, 

98) 

99 

100# Standard planner implementations 

101from hyper_parallel.core.distributed_checkpoint.standard_planner import ( 

102 StandardLoadPlanner, 

103 StandardSavePlanner, 

104) 

105 

106# Loads of checkpoints laid out differently from the state dict 

107from hyper_parallel.core.distributed_checkpoint.remap_planner import ( 

108 DeferredRead, 

109 RemapBlock, 

110 RemapLoadPlanner, 

111) 

112 

113# Storage interfaces and data structures 

114from hyper_parallel.core.distributed_checkpoint.storage import ( 

115 StorageInfo, 

116 StorageReader, 

117 StorageWriter, 

118 WriteResult, 

119) 

120 

121# File system storage implementations 

122from hyper_parallel.core.distributed_checkpoint.filesystem_storage import ( 

123 FileSystemReader, 

124 FileSystemWriter, 

125) 

126 

127# Safetensors checkpoints written outside DCP 

128from hyper_parallel.core.distributed_checkpoint.hf_storage import ( 

129 HuggingFaceStorageReader, 

130 TorchShardedSafetensorsReader, 

131) 

132 

133# Optimizer state dict 

134from hyper_parallel.core.distributed_checkpoint.state_dict import ( 

135 get_optim_state_dict, 

136 set_optim_state_dict, 

137)