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

49 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"""cost model parser module""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18 

19import math 

20from abc import ABC 

21from abc import abstractmethod 

22 

23if TYPE_CHECKING: 

24 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import _CostModVar 

25 

26 

27class _CostModelParser(ABC): 

28 """abstract parser class""" 

29 

30 def __init__(self, ccfg: _CostModVar): 

31 self.ccfg = ccfg 

32 self.config = ccfg.config 

33 

34 @abstractmethod 

35 def parse(self): 

36 """Parse the cost model configuration and populate the cost model variables. 

37 

38 Subclasses must implement this method to read framework-specific 

39 configuration values into the shared _CostModVar instance. 

40 """ 

41 

42 def config_optimizer_shard(self, ccfg): 

43 """OP related variables""" 

44 # Non expert params 

45 ccfg.shard_p_os_non_exp_partial = ( 

46 ccfg.os_max_shard if ccfg.has_op else ccfg.t 

47 ) * ccfg.cp 

48 ccfg.shard_p_os_non_exp = ( 

49 (ccfg.d if ccfg.has_op else 1) * ccfg.cp * ccfg.t 

50 ) 

51 ccfg.shard_grad_non_exp = ( 

52 ccfg.shard_p_os_non_exp if ccfg.has_grad_shard else ccfg.t 

53 ) 

54 

55 # Expert params 

56 ccfg.shard_p_os_exp_partial = math.gcd( 

57 ccfg.n_exp, 

58 (ccfg.os_max_shard if ccfg.has_op else 1) * ccfg.t_exp, 

59 ) 

60 ccfg.shard_p_os_exp = ( 

61 (ccfg.d_exp if ccfg.has_op else 1) * ccfg.cp * ccfg.t_exp 

62 ) 

63 ccfg.shard_grad_exp = ( 

64 ccfg.shard_p_os_exp 

65 if ccfg.has_grad_shard 

66 else ccfg.t_exp 

67 ) 

68 ccfg.shard_grad_exp_partial = ( 

69 ccfg.shard_p_os_exp_partial 

70 if ccfg.has_grad_shard 

71 else ccfg.t_exp 

72 ) 

73 

74 # def config_op_level(self, ccfg, strategy): 

75 # def full_partial(): 

76 # return Config({"full":0, "partial":0}) 

77 # def exp_or_not(): 

78 # return Config({ 

79 # "non_exp":full_partial(), 

80 # "exp":full_partial() 

81 # }) 

82 # ccfg.op = Config({ 

83 # "p":exp_or_not(), 

84 # "os":exp_or_not(), 

85 # "grad"exp_or_not() 

86 # }) 

87 # shard_strat = { 

88 # "grad":0, #zero 1 

89 # "os+grad":0, #zero 2 

90 # "p+os+grad":0, # zero 3 

91 # "p+os":0 # zero2 mindspore 

92 # } 

93 # shard_strat[strategy] 

94 

95 def init_hff(self): 

96 """MindFormers format for FFn hidden size""" 

97 # Assuming following 3 variables are already parsed 

98 hidden_size = self.ccfg.h 

99 ffn_dim_multiplier = self.ccfg.fdm 

100 multiple_of = self.ccfg.multiple_of 

101 hff = 4 * hidden_size 

102 if ffn_dim_multiplier: 

103 hff = int((ffn_dim_multiplier + 0.01) * hff) 

104 hff = int(2 * hff / 3) 

105 hff = multiple_of * ((hff + multiple_of - 1) // multiple_of) 

106 return hff 

107 

108 def config_comm_flag(self, ccfg): 

109 """comm flag variables""" 

110 ccfg.comm_d_non_exp = ( 

111 0 

112 if ((ccfg.d == 1) or not ccfg.has_op) 

113 else (2 if not ccfg.has_grad_shard else 3) 

114 ) # data parallel comm factor 

115 ccfg.comm_d_exp = ( 

116 0 

117 if ((ccfg.d_exp == 1) or not ccfg.has_op) 

118 else (2 if not ccfg.has_grad_shard else 3) 

119 ) # data parallel comm factor 

120 ccfg.comm_t = float(ccfg.t > 1) # tensor parallel comm factor 

121 ccfg.comm_ep = float( 

122 ccfg.ep > 1 or ccfg.n_exp > 1 

123 ) # expert parallel comm factor 

124 ccfg.comm_cp = float(ccfg.cp > 1) # context parallel comm factor 

125 

126 def config_dp_tp_exp(self, ccfg): 

127 """MoE strategy variables""" 

128 if ccfg.etp: 

129 ccfg.t_exp = ccfg.etp 

130 # d * t = inner dp * outer dp * etp 

131 # inner dp = EP, outer dp = the rest 

132 ccfg.d_exp = ccfg.d * ccfg.t * ccfg.cp // ccfg.t_exp // ccfg.ep 

133 else: 

134 ccfg.t_exp = ccfg.t 

135 if ccfg.d >= ccfg.ep: 

136 ccfg.d_exp = ccfg.d // ccfg.ep 

137 else: 

138 ccfg.d_exp = ccfg.d * ccfg.t // ccfg.ep 

139 if ccfg.t_exp * ccfg.ep > ccfg.d * ccfg.t: 

140 ccfg.t_exp = 1 

141 

142 exp_group1_invalid = ccfg.d_exp < 1 or ccfg.t_exp < 1 

143 exp_group2_invalid = ccfg.hff_exp < 1 or ccfg.n_exp < 1 

144 if exp_group1_invalid or exp_group2_invalid: 

145 raise TypeError( 

146 f"MoE parsing error: d_exp({ccfg.d_exp})/t_exp({ccfg.t_exp})/" 

147 f"hff_exp({ccfg.hff_exp})/n_exp({ccfg.n_exp})/" 

148 f"DP = {ccfg.d}, TP = {ccfg.t}, EP = {ccfg.ep}/" 

149 )