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: 83%

140 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-07-06 05:41 +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.config import Config 

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

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

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

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.tie_emb_out = self.config.model.model_config.tie_word_embeddings 

36 self.ccfg.cp_algo = ( 

37 self.config.parallel_config.context_parallel_algo 

38 if self.config.parallel_config.context_parallel_algo 

39 else "colossalai_cp" 

40 ) 

41 self.ccfg.d = max( 

42 1, self.config.parallel_config.data_parallel 

43 ) # Data parallel 

44 self.ccfg.t = max( 

45 1, self.config.parallel_config.model_parallel 

46 ) # Tensor parallel 

47 self.ccfg.p = max( 

48 1, self.config.parallel_config.pipeline_stage 

49 ) # Pipeline parallel 

50 self.ccfg.cp = max( 

51 1, self.config.parallel_config.context_parallel 

52 ) # Context parallel 

53 self.ccfg.ep = max( 

54 1, self.config.parallel_config.expert_parallel 

55 ) # Expert parallel 

56 self.ccfg.sp = ( 

57 self.ccfg.t if self.config.parallel_config.use_seq_parallel else 1 

58 ) # Sequence parallel factor 

59 if self.ccfg.cp > 1 and self.ccfg.sp > 1: 

60 logger.warning( 

61 "sequence parallelism and context parallelism are both enabled" 

62 ) 

63 

64 # Interleaving 

65 self.ccfg.n_s_split = 1 # seqpipe 

66 self.ccfg.pp_sched = "1f1b" 

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

68 if self.config.parallel.pipeline_config: 

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

70 self.ccfg.n_s_split = max( 

71 1, self.config.parallel_config.seq_split_num 

72 ) 

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

74 self.ccfg.pp_sched = ( 

75 self.config.parallel.pipeline_config.pipeline_scheduler 

76 ) 

77 

78 def __config_parse_yaml_hyperparameters(self): 

79 """MindFormer format for hyperparams""" 

80 self.ccfg.n_lay = ( 

81 self.config.model.model_config.num_layers 

82 if self.config.model.model_config.num_layers 

83 else self.config.model.model_config.num_hidden_layers 

84 ) 

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

86 self.ccfg.n_lay *= 2 

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

88 self.ccfg.hff = ( 

89 self.config.model.model_config.intermediate_size 

90 if self.config.model.model_config.intermediate_size 

91 else self.init_hff() 

92 ) # Expanded hidden size 

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

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

95 self.ccfg.a = ( 

96 self.config.model.model_config.num_heads 

97 ) # Number of attention (query) heads 

98 if not self.ccfg.a: 

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

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

101 if not self.ccfg.n_kv: 

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

103 if not self.ccfg.n_kv: 

104 self.ccfg.n_kv = self.ccfg.a 

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

106 self.ccfg.dc_kv = ( 

107 self.config.model.model_config.kv_lora_rank 

108 if self.config.model.model_config.kv_lora_rank 

109 else 0 

110 ) # KV compression dimension 

111 self.ccfg.dc_q = ( 

112 self.config.model.model_config.q_lora_rank 

113 if self.config.model.model_config.q_lora_rank 

114 else 0 

115 ) # Q compression dimension 

116 self.ccfg.dhr = ( 

117 self.config.model.model_config.qk_rope_head_dim 

118 if self.config.model.model_config.qk_rope_head_dim 

119 else 0 

120 ) # decoupled QK per head dimension 

121 

122 # Microbatch infos 

123 self.ccfg.b = max( 

124 1, self.config.runner_config.batch_size 

125 ) # Microbatch size 

126 self.ccfg.m = ( 

127 self.config.parallel_config.micro_batch_num 

128 ) # Number of microbatches 

129 if self.ccfg.m <= 0: 

130 logger.warning("num_micro is negative") 

131 

132 def __config_parse_yaml_moe(self): 

133 """MindFormer format for MoE infos""" 

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

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

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

137 if self.config.moe_config: 

138 self.ccfg.n_exp = max( 

139 1, self.config.moe_config.expert_num 

140 ) # Total number of experts 

141 self.ccfg.n_chosen_exp = max( 

142 1, self.config.moe_config.num_experts_chosen 

143 ) # Number of chosen experts 

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

145 if self.config.moe_config.moe_intermediate_size: 

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

147 self.ccfg.cap_fact = max( 

148 1, self.config.moe_config.capacity_factor 

149 ) # Capacity factor 

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

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

152 self.config_dp_tp_exp(self.ccfg) 

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

154 else: 

155 cfg = self.config.model.model_config 

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

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

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

159 if cfg.moe_intermediate_size: 

160 self.ccfg.hff_exp = cfg.moe_intermediate_size 

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

162 self.config_dp_tp_exp(self.ccfg) 

163 self.ccfg.gmm = cfg.moe_grouped_gemm 

164 

165 def __config_parse_yaml_op_recompute(self): 

166 """MindFormer format for select recompute""" 

167 # [HYPOTHESIS] 

168 self.ccfg.rec_op = Config( 

169 {} 

170 ) # recomputed operators (selective recompute only) 

171 self.ccfg.rec_op.attBMM = int( 

172 not ( 

173 self.config.recompute_config.select_recompute 

174 and not self.ccfg.has_fa 

175 and self.ccfg.sp > 1 

176 ) 

177 ) 

178 self.ccfg.rec_op.headCast = int( 

179 not (self.config.recompute_config.select_recompute and self.ccfg.has_fa) 

180 ) 

181 self.ccfg.rec_op.dropout = 1 

182 self.ccfg.rec_op.softmax = int( 

183 not ( 

184 self.config.recompute_config.select_recompute 

185 and not self.ccfg.has_fa 

186 ) 

187 ) 

188 self.ccfg.rec_op.normOp = int( 

189 not (self.config.recompute_config.select_recompute and self.ccfg.sp > 1) 

190 ) 

191 self.ccfg.rec_op.gather = int( 

192 not ( 

193 self.config.recompute_config.select_comm_recompute 

194 and self.ccfg.sp > 1 

195 ) 

196 ) 

197 self.ccfg.rec_op.ffAct = int( 

198 not (self.config.recompute_config.select_recompute and self.ccfg.sp > 1) 

199 ) 

200 

201 def config_shard_emb(self): 

202 """Configure embedding and output activation sharding.""" 

203 self.ccfg.shard_embed = ( 

204 self.ccfg.d 

205 if (self.ccfg.vocab_emb_dp and self.ccfg.p == 1) 

206 else (self.ccfg.t * self.ccfg.d) 

207 ) 

208 

209 def __config_parse_yaml(self): 

210 """MindFormer format for unimodal""" 

211 self.ccfg.config_format = "yaml" 

212 self.ccfg.model_name = self.config.trainer.model_name 

213 self.ccfg.device_capacity = Memory.from_string( 

214 self.config.context.max_device_memory 

215 ) 

216 # ( 

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

218 # * 1024 

219 # * 1024 

220 # * 1024 

221 # ) 

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

223 # self.ccfg.vp_less_mem = False 

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

225 self.ccfg.gmm = False 

226 self.ccfg.freeze = False 

227 op_cfg = self.config.parallel.parallel_optimizer_config 

228 if op_cfg: 

229 self.ccfg.op_weight_shard = op_cfg.optimizer_weight_shard_size 

230 self.ccfg.optimizer = self.config.optimizer.type 

231 self.ccfg.multiple_of = ( 

232 self.config.model.model_config.multiple_of 

233 if self.config.model.model_config.multiple_of 

234 else 1 

235 ) 

236 self.ccfg.fdm = ( 

237 self.config.model.model_config.ffn_dim_multiplier 

238 if self.config.model.model_config.ffn_dim_multiplier 

239 else 1 

240 ) 

241 self.__config_parse_yaml_parallelism() 

242 self.__config_parse_yaml_hyperparameters() 

243 self.__config_parse_yaml_moe() 

244 

245 # FP byte storages 

246 self.ccfg.bytes_p = self.ccfg.fp_bytes( 

247 self.config.model.model_config.param_init_type, 

248 ) # parameters 

249 if not self.ccfg.bytes_p: 

250 self.ccfg.bytes_p = self.ccfg.fp_bytes( 

251 self.config.model.model_config.params_dtype 

252 ) 

253 self.ccfg.bytes_compute = self.ccfg.fp_bytes( 

254 self.config.model.model_config.compute_dtype 

255 ) # activations 

256 self.ccfg.bytes_softmax = self.ccfg.fp_bytes( 

257 self.config.model.model_config.softmax_compute_type 

258 ) # softmax output 

259 if not self.ccfg.bytes_softmax: 

260 self.ccfg.bytes_softmax = self.ccfg.fp_bytes( 

261 self.config.model.model_config.softmax_compute_dtype 

262 ) 

263 if not self.ccfg.bytes_p: 

264 raise AttributeError("bytes_p not positive") 

265 if not self.ccfg.bytes_compute: 

266 raise AttributeError("bytes_compute not positive") 

267 

268 # Optimizer parallel factors 

269 if self.ccfg.op_weight_shard: 

270 self.ccfg.os_max_shard = self.ccfg.op_weight_shard 

271 elif self.ccfg.has_op: 

272 self.ccfg.os_max_shard = self.ccfg.d * self.ccfg.t 

273 else: 

274 self.ccfg.os_max_shard = 1 

275 self.config_optimizer_shard(self.ccfg) 

276 

277 # Other factors 

278 self.config_shard_emb() 

279 self.ccfg.shard_output_activ = 1 

280 self.ccfg.shard_recompute_input = ( 

281 self.ccfg.t 

282 if self.config.recompute_config.recompute_slice_activation 

283 else 1 

284 ) 

285 self.ccfg.s_fa = ( 

286 self.ccfg.s 

287 if not self.config.model.model_config.use_flash_attention 

288 else self.ccfg.s / self.ccfg.a 

289 ) # flash attention factor [HYPOTHESIS] 

290 self.config_comm_flag(self.ccfg) 

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

292 self.ccfg.n_mtp = ( 

293 self.config.model.model_config.mtp_depth 

294 if self.config.model.model_config.mtp_depth 

295 else 0 

296 ) 

297 if not self.ccfg.n_mtp: 

298 self.ccfg.n_mtp = ( 

299 self.config.model.model_config.num_nextn_predict_layers 

300 ) 

301 self.ccfg.is_mtp_in_offset = False 

302 

303 # Layer custom config 

304 # [(num layers selected, layer custom config to apply)] 

305 self.ccfg.layer_custom_config = [(self.ccfg.n_lay + self.ccfg.n_mtp, None)] 

306 

307 self.__config_parse_yaml_op_recompute() 

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

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

310 self.ccfg.sel_rec = self.config.recompute_config.select_recompute 

311 self.ccfg.full_rec = self.config.recompute_config.recompute 

312 self.ccfg.overwrite_eval_functions = {}