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

137 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-10-25 05:45 +0800

1# Copyright 2025-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"""parser child class""" 

16from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser 

17from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory 

18from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger 

19from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import derive 

20 

21 

22class CostModelParserMindformers(_CostModelParser): 

23 """parser class for MindFormers format""" 

24 

25 def parse(self): 

26 self.__config_parse_yaml() 

27 

28 def __config_parse_yaml_parallelism(self): 

29 """MindFormer format for strategy""" 

30 self.ccfg.has_op = self.config.parallel.enable_parallel_optimizer 

31 op_cfg = self.config.parallel.parallel_optimizer_config 

32 if op_cfg: 

33 self.ccfg.has_grad_shard = op_cfg.gradient_accumulation_shard 

34 self.ccfg.vocab_emb_dp = self.config.parallel_config.vocab_emb_dp 

35 self.ccfg.emb_dp_sharded = True 

36 self.ccfg.tie_emb_out = self.config.model.model_config.tie_word_embeddings 

37 self.ccfg.cp_algo = ( 

38 self.config.parallel_config.context_parallel_algo 

39 if self.config.parallel_config.context_parallel_algo 

40 else "colossalai_cp" 

41 ) 

42 self.ccfg.d = max( 

43 1, self.config.parallel_config.data_parallel 

44 ) # Data parallel 

45 self.ccfg.t = max( 

46 1, self.config.parallel_config.model_parallel 

47 ) # Tensor parallel 

48 self.ccfg.p = max( 

49 1, self.config.parallel_config.pipeline_stage 

50 ) # Pipeline parallel 

51 self.ccfg.cp = max( 

52 1, self.config.parallel_config.context_parallel 

53 ) # Context parallel 

54 self.ccfg.ep = max( 

55 1, self.config.parallel_config.expert_parallel 

56 ) # Expert parallel 

57 self.ccfg.sequence_parallel = bool(self.config.parallel_config.use_seq_parallel) 

58 if self.ccfg.cp > 1 and self.ccfg.sequence_parallel and self.ccfg.t > 1: 

59 logger.warning( 

60 "sequence parallelism and context parallelism are both enabled" 

61 ) 

62 

63 # Interleaving 

64 self.ccfg.n_s_split = 1 # seqpipe 

65 self.ccfg.pp_sched = "1f1b" 

66 self.ccfg.vp = max(1, self.config.model.model_config.pp_interleave_num) 

67 if self.config.parallel.pipeline_config: 

68 if self.config.parallel.pipeline_config.pipeline_interleave: 

69 self.ccfg.n_s_split = max( 

70 1, self.config.parallel_config.seq_split_num 

71 ) 

72 if self.config.parallel.pipeline_config.pipeline_scheduler: 

73 self.ccfg.pp_sched = ( 

74 self.config.parallel.pipeline_config.pipeline_scheduler 

75 ) 

76 

77 def __config_parse_yaml_hyperparameters(self): 

78 """MindFormer format for hyperparams""" 

79 self.ccfg.n_lay = ( 

80 self.config.model.model_config.num_layers 

81 if self.config.model.model_config.num_layers 

82 else self.config.model.model_config.num_hidden_layers 

83 ) 

84 if self.config.model.model_config.is_encoder_decoder: 

85 self.ccfg.n_lay *= 2 

86 self.ccfg.h = self.config.model.model_config.hidden_size # Hidden size 

87 self.ccfg.hff = ( 

88 self.config.model.model_config.intermediate_size 

89 if self.config.model.model_config.intermediate_size 

90 else self.init_hff() 

91 ) # Expanded hidden size 

92 self.ccfg.v = self.config.model.model_config.vocab_size # Vocabulary size 

93 self.ccfg.s = self.config.model.model_config.seq_length # Sequence length 

94 self.ccfg.a = ( 

95 self.config.model.model_config.num_heads 

96 ) # Number of attention (query) heads 

97 if not self.ccfg.a: 

98 self.ccfg.a = self.config.model.model_config.num_attention_heads 

99 self.ccfg.n_kv = self.config.model.model_config.n_kv_heads # Number of keys - values heads 

100 if not self.ccfg.n_kv: 

101 self.ccfg.n_kv = self.config.model.model_config.num_key_value_heads 

102 if not self.ccfg.n_kv: 

103 self.ccfg.n_kv = self.ccfg.a 

104 self.ccfg.dh = self.ccfg.h / self.ccfg.a # Per head dimension 

105 # An MLA model's value-head width, which its family prices dh at. 

106 self.ccfg.v_head_dim = self.config.model.model_config.v_head_dim or None 

107 # Its heads' non-rotary key width, the value heads' unless stated. 

108 self.ccfg.qk_nope_head_dim = self.config.model.model_config.qk_nope_head_dim or None 

109 self.ccfg.dc_kv = ( 

110 self.config.model.model_config.kv_lora_rank 

111 if self.config.model.model_config.kv_lora_rank 

112 else 0 

113 ) # KV compression dimension 

114 self.ccfg.dc_q = ( 

115 self.config.model.model_config.q_lora_rank 

116 if self.config.model.model_config.q_lora_rank 

117 else 0 

118 ) # Q compression dimension 

119 self.ccfg.dhr = ( 

120 self.config.model.model_config.qk_rope_head_dim 

121 if self.config.model.model_config.qk_rope_head_dim 

122 else 0 

123 ) # decoupled QK per head dimension 

124 # Whether each head's queries and keys are normalized, as MindFormers' 

125 # Qwen3 states it. 

126 self.ccfg.qk_norm = bool(self.config.model.model_config.qk_layernorm) 

127 

128 # Microbatch infos 

129 self.ccfg.b = max( 

130 1, self.config.runner_config.batch_size 

131 ) # Microbatch size 

132 self.ccfg.m = ( 

133 self.config.parallel_config.micro_batch_num 

134 ) # Number of microbatches 

135 if self.ccfg.m <= 0: 

136 logger.warning("num_micro is negative") 

137 

138 def __config_parse_yaml_moe(self): 

139 """MindFormer format for MoE infos""" 

140 self.ccfg.n_exp, self.ccfg.n_chosen_exp, self.ccfg.n_shared_exp = 1, 1, 0 

141 self.ccfg.hff_exp, self.ccfg.cap_fact = self.ccfg.hff, 1 

142 self.ccfg.t_exp, self.ccfg.d_exp = self.ccfg.t, self.ccfg.d 

143 if self.config.moe_config: 

144 self.ccfg.n_exp = max( 

145 1, self.config.moe_config.expert_num 

146 ) # Total number of experts 

147 self.ccfg.n_chosen_exp = max( 

148 1, self.config.moe_config.num_experts_chosen 

149 ) # Number of chosen experts 

150 self.ccfg.n_shared_exp = self.config.moe_config.shared_expert_num 

151 if self.config.moe_config.moe_intermediate_size: 

152 self.ccfg.hff_exp = self.config.moe_config.moe_intermediate_size 

153 self.ccfg.cap_fact = max( 

154 1, self.config.moe_config.capacity_factor 

155 ) # Capacity factor 

156 self.ccfg.k_1st_dense = self.config.moe_config.first_k_dense_replace 

157 self.ccfg.etp = self.config.moe_config.expert_model_parallel 

158 self.ccfg.gmm = self.config.moe_config.use_gmm 

159 else: 

160 cfg = self.config.model.model_config 

161 self.ccfg.n_exp = max(self.ccfg.n_exp, cfg.n_routed_experts) 

162 self.ccfg.n_chosen_exp = max(self.ccfg.n_chosen_exp, cfg.num_experts_per_tok) 

163 self.ccfg.n_shared_exp = max(self.ccfg.n_shared_exp, cfg.n_shared_experts) 

164 if cfg.moe_intermediate_size: 

165 self.ccfg.hff_exp = cfg.moe_intermediate_size 

166 self.ccfg.k_1st_dense = max(self.ccfg.k_1st_dense, cfg.first_k_dense_replace) 

167 self.ccfg.gmm = cfg.moe_grouped_gemm 

168 

169 def __config_parse_yaml_recompute(self): 

170 """MindFormer format for recompute""" 

171 rc = self.config.recompute_config 

172 self.ccfg.sel_rec = rc.select_recompute 

173 self.ccfg.sel_comm_rec = rc.select_comm_recompute 

174 self.ccfg.sel_rec_rule = "mindformers" 

175 self.ccfg.full_rec = rc.recompute 

176 self.ccfg.recompute_slice_activation = bool(rc.recompute_slice_activation) 

177 

178 def __config_parse_yaml_fp_bytes(self): 

179 """FP byte storages of parameters, activations and softmax outputs.""" 

180 self.ccfg.bytes_p = self.ccfg.fp_bytes( 

181 self.config.model.model_config.param_init_type, 

182 ) # parameters 

183 if not self.ccfg.bytes_p: 

184 self.ccfg.bytes_p = self.ccfg.fp_bytes( 

185 self.config.model.model_config.params_dtype 

186 ) 

187 self.ccfg.bytes_compute = self.ccfg.fp_bytes( 

188 self.config.model.model_config.compute_dtype 

189 ) # activations 

190 self.ccfg.bytes_softmax = self.ccfg.fp_bytes( 

191 self.config.model.model_config.softmax_compute_type 

192 ) # softmax output 

193 if not self.ccfg.bytes_softmax: 

194 self.ccfg.bytes_softmax = self.ccfg.fp_bytes( 

195 self.config.model.model_config.softmax_compute_dtype 

196 ) 

197 if not self.ccfg.bytes_p: 

198 raise AttributeError("bytes_p not positive") 

199 if not self.ccfg.bytes_compute: 

200 raise AttributeError("bytes_compute not positive") 

201 

202 def __config_parse_yaml_optimizer_shard(self): 

203 """Optimizer parallel factors.""" 

204 if self.ccfg.op_weight_shard: 

205 self.ccfg.os_max_shard = self.ccfg.op_weight_shard 

206 elif self.ccfg.has_op: 

207 self.ccfg.os_max_shard = self.ccfg.d 

208 else: 

209 self.ccfg.os_max_shard = 1 

210 

211 def __config_parse_yaml(self): 

212 """MindFormer format for unimodal""" 

213 self.ccfg.config_format = "yaml" 

214 self.ccfg.model_name = self.config.trainer.model_name 

215 self.config_op_counts(self.ccfg) 

216 self.ccfg.device_capacity = Memory.from_string( 

217 self.config.context.max_device_memory 

218 ) 

219 # ( 

220 # float(self.config.context.max_device_memory[:-2]) 

221 # * 1024 

222 # * 1024 

223 # * 1024 

224 # ) 

225 self.ccfg.has_fa = self.config.model.model_config.use_flash_attention 

226 # self.ccfg.vp_less_mem = False 

227 self.ccfg.has_clip = self.config.runner_wrapper.use_clip_grad 

228 self.ccfg.gmm = False 

229 self.ccfg.freeze = False 

230 op_cfg = self.config.parallel.parallel_optimizer_config 

231 if op_cfg: 

232 self.ccfg.op_weight_shard = op_cfg.optimizer_weight_shard_size 

233 # Read from the section's own fields: a missing section reads as 0, 

234 # whose .type raised, where the HyperParallel parser takes AdamW. 

235 optimizer = vars(self.config).get("optimizer") 

236 self.state_optimizer(self.ccfg, getattr(optimizer, "__dict__", {}).get("type")) 

237 self.ccfg.multiple_of = ( 

238 self.config.model.model_config.multiple_of 

239 if self.config.model.model_config.multiple_of 

240 else 1 

241 ) 

242 self.ccfg.fdm = ( 

243 self.config.model.model_config.ffn_dim_multiplier 

244 if self.config.model.model_config.ffn_dim_multiplier 

245 else 1 

246 ) 

247 self.__config_parse_yaml_parallelism() 

248 self.__config_parse_yaml_hyperparameters() 

249 self.__config_parse_yaml_moe() 

250 self.__config_parse_yaml_fp_bytes() 

251 self.__config_parse_yaml_optimizer_shard() 

252 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m 

253 self.ccfg.n_mtp = ( 

254 self.config.model.model_config.mtp_depth 

255 if self.config.model.model_config.mtp_depth 

256 else 0 

257 ) 

258 if not self.ccfg.n_mtp: 

259 self.ccfg.n_mtp = ( 

260 self.config.model.model_config.num_nextn_predict_layers 

261 ) 

262 self.ccfg.is_mtp_in_offset = False 

263 

264 self.config_layer_stack(self.ccfg) 

265 

266 # By default, 100% of layers use a unique custom config (if specified) 

267 self.ccfg.offset = self.config.model.model_config.offset 

268 self.__config_parse_yaml_recompute() 

269 self.ccfg.overwrite_eval_functions = {} 

270 derive(self.ccfg)