Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / config.py: 89%

1# Copyright 2025 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"""parallel dimensions""" 

16 

17import os 

18import copy 

19import json 

20import yaml 

21try: 

22 import tomllib 

23except ModuleNotFoundError: # Python 3.10, where tomllib is not in the standard library 

24 tomllib = None 

25 

26from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger 

27 

28 

29class YamlObject: 

30 """Attributed dictionary""" 

31 

32 def from_dict(self, field): 

33 """init from dictionary""" 

34 for k in field: 

35 if isinstance(field[k], dict): 

36 setattr(self, k, YamlObject(field[k])) 

37 else: 

38 setattr(self, k, field[k]) 

39 

40 def __init__(self, field): 

41 self.from_dict(field) 

42 

43 def __getattr__(self, attr): 

44 if attr not in self.__dict__: 

45 logger.warning( 

46 "Attribute %s does not exist. Value '0' will be assigned.", 

47 str(attr), 

48 ) 

49 return 0 

50 return self.__dict__[attr] 

51 

52 def __copy__(self): 

53 cls = self.__class__ 

54 res = cls.__new__(cls) 

55 res.__dict__.update(self.__dict__) 

56 return res 

57 

58 def __deepcopy__(self, memo): 

59 cls = self.__class__ 

60 res = cls.__new__(cls) 

61 for k, v in self.__dict__.items(): 

62 setattr(res, k, copy.deepcopy(v, memo)) 

63 return res 

64 

65 def __getstate__(self): 

66 return self.to_dict() 

67 

68 def __setstate__(self, field): 

69 return self.from_dict(field) 

70 

71 def to_dict(self): 

72 """Transform from Config to a strict dict class""" 

73 return_dict = {} 

74 for key, val in self.__dict__.items(): 

75 if isinstance(val, YamlObject): 

76 val = val.to_dict() 

77 return_dict[key] = val 

78 return return_dict 

79 

80 def dump(self, file_name, folder=None): 

81 """Dump to given file in given folder""" 

82 if not folder: 

83 folder = os.path.dirname(__file__) 

84 full_file_name = os.path.join(folder, "config_" + file_name) + ".yaml" 

85 with open(full_file_name, "w", encoding="utf-8") as outfile: 

86 yaml.dump(self.to_dict(), outfile, default_flow_style=False) 

87 

88 

89def _load_toml(handle): 

90 """Parse a TOML file opened in text mode. 

91 

92 Only TorchTitan configurations are TOML, so the third-party ``toml`` package 

93 is imported on use rather than at module scope: it is not a declared 

94 dependency, and from Python 3.11 the standard library covers this anyway. 

95 

96 Args: 

97 handle: File object opened in text mode. 

98 

99 Returns: 

100 The parsed document. 

101 """ 

102 if tomllib is not None: 

103 return tomllib.loads(handle.read()) 

104 import toml # pylint: disable=import-outside-toplevel 

105 return toml.load(handle) 

106 

107 

108class Config(YamlObject): 

109 """Yaml config""" 

110 

111 def __init__(self, input_config): 

112 if isinstance(input_config, str): 

113 with open(input_config, encoding="utf-8") as f: 

114 if input_config.endswith("yaml"): 

115 try: 

116 super().__init__(yaml.safe_load(f)) 

117 except yaml.YAMLError as exc: 

118 print(exc) 

119 logger.warning(exc) 

120 elif input_config.endswith("json"): 

121 super().__init__(json.load(f)) 

122 elif input_config.endswith("toml"): 

123 super().__init__(_load_toml(f)) 

124 else: 

125 logger.warning("Current handled file formats: YAML, JSON") 

126 elif isinstance(input_config, Config): 

127 super().__init__(input_config.__dict__) 

128 elif isinstance(input_config, dict): 

129 super().__init__(input_config) 

130 else: 

131 raise TypeError("Expecting path string or Config object") 

132 

133 def __str__(self): 

134 return str( 

135 { 

136 k: (v if not isinstance(v, Config) else vars(v)) 

137 for k, v in vars(self).items() 

138 } 

139 ) 

140 

141 # def global_batch_size(self): 

142 # """Compute the global batch size""" 

143 # pp = Dim.PP.from_config(self) 

144 # dp = Dim.DP.from_config(self) 

145 # mbs = Dim.MBS.from_config(self) 

146 # gbs = dp * mbs 

147 # if pp == 1: 

148 # logger.debug("global_batch_size = DP(%d) * MBS(%d)", dp, mbs) 

149 # # gas = Dim.GAS.from_config(self) 

150 # # gbs *= gas 

151 # else: 

152 # mbn = Dim.MBN.from_config(self) 

153 # gbs *= mbn 

154 # logger.debug( 

155 # "global_batch_size = DP(%d) * MB(%d) * MBS(%d)", dp, mbn, mbs 

156 # ) 

157 # return gbs 

158 

159 def set_par(self, parallel_dimensions): 

160 """Set the given parallelization""" 

161 for dim in parallel_dimensions.dims_val: 

162 dim.to_config(self, parallel_dimensions.val(dim)) 

163 

164 # def layer_num(self): 

165 # """Get the layer number""" 

166 # layers = self.model.model_config.num_layers 

167 # if layers and layers is not None: 

168 # return layers 

169 # return self.model.model_config.num_hidden_layers 

170 

171 # def name(self): 

172 # """Get the model name""" 

173 # return self.trainer.model_name