Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/auto_parallel/sapp_nd/memory_estimation/evaluators/layer

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

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

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"""Layer's blocks submodule"""

16from __future__ import annotations

17from typing import TYPE_CHECKING

18from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType

19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils

21if TYPE_CHECKING:

22 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig

23 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context

25mb = EvalUtils.mb

28class EvalAttn:

29 """Attention formulas class"""

31 @staticmethod

32 def num_params_mla(ccfg: CostModelConfig, _) -> float:

33 """Parameters count for Multi-Head Latent Attention"""

34 # W_up_q = ccfg.dc_q * ccfg.dh * ccfg.a

35 # W_up_k = ccfg.dc_kv * ccfg.dh * ccfg.n_kv

36 # W_up_v = ccfg.dc_kv * ccfg.dh * ccfg.n_kv

37 # W_down_q = ccfg.dc_q * ccfg.h

38 # W_down_kv = ccfg.dc_kv * ccfg.h

39 # W_q_rope = ccfg.a * ccfg.dhr * ccfg.dc_q

40 # W_k_rope = ccfg.dhr * ccfg.h

41 # Wo = ccfg.h * ccfg.a * ccfg.dh

43 c_kv_fact = ccfg.dc_kv * (ccfg.n_kv * ccfg.dh + ccfg.h)

44 c_q_fact = ccfg.dc_q * (ccfg.a * ccfg.dh + ccfg.h + ccfg.a * ccfg.dhr)

45 rest_fact = (ccfg.h * ccfg.a * ccfg.dh) + (ccfg.h * ccfg.dhr)

46 res = (

47 0.5 * ccfg.n_attMM * c_kv_fact

48 + 0.25 * ccfg.n_attMM * c_q_fact

49 + 0.25 * ccfg.n_attMM * rest_fact

50 )

51 return res

53 @staticmethod

54 def num_params_attn(ccfg: CostModelConfig, ctx: Context) -> float:

55 """Parameters count for Multi-Head/Grouped-Q./Multi-Q. Attention"""

56 if ccfg.dc_kv == 0:

57 # Q,O and K,V have distinct shapes

58 return 0.5 * ccfg.n_attMM * (

59 ccfg.h * ccfg.h + ccfg.h

60 ) + 0.5 * ccfg.n_attMM * (ccfg.h * ccfg.n_kv * ccfg.dh + ccfg.h)

61 return EvalAttn.num_params_mla(ccfg, ctx)

63 @staticmethod

64 def attn_qkv_activations(ccfg: CostModelConfig, ctx: Context) -> float:

65 """QKV linear Activations"""

66 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER

67 att_qkv_size = 0

68 if ccfg.dc_kv == 0:

69 n_op = ccfg.n_attMM + ccfg.n_attParamCast

70 att_qkv_size = (

71 ccfg.s

72 * ccfg.b

73 * ccfg.bytes_compute

74 * (

75 0.25 * n_op * ccfg.h

76 + 0.5 * n_op * ccfg.dh * ccfg.n_kv

77 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.attBMM)

78 * ccfg.n_attBMM

79 * ccfg.dh

80 )

81 )

82 else:

83 q_size = (

84 0.25

85 * (ccfg.n_attMM + ccfg.n_attParamCast)

86 * (ccfg.dc_q + 2 * ccfg.a * (ccfg.dh + ccfg.dhr))

87 )

88 k_size = (

89 0.25

90 * (ccfg.n_attMM + ccfg.n_attParamCast)

91 * (ccfg.dhr + ccfg.n_kv * (2 * ccfg.dh + ccfg.dhr))

92 )

93 v_size = (

94 0.25

95 * (ccfg.n_attMM + ccfg.n_attParamCast)

96 * (ccfg.n_kv * ccfg.dh + ccfg.dc_kv)

97 )

98 att_qkv_size = (

99 ccfg.s

100 * ccfg.b

101 * ccfg.bytes_compute

102 * (

103 q_size

104 + k_size

105 + v_size

106 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.attBMM)

107 * ccfg.n_attBMM

108 * ccfg.dh

109 )

110 )

111 micro_factor = ctx.micro_factor

112 return micro_factor * att_qkv_size / (ccfg.t * ccfg.cp)

113

114 @staticmethod

115 def attn_score_activations(ccfg: CostModelConfig, ctx: Context) -> float:

116 """Score/Softmax Activations"""

117 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER

118 att_score = (

119 ccfg.s_fa

120 * ccfg.b

121 * ccfg.a

122 * ccfg.s

123 * (

124 ccfg.n_softmax

125 * (

126 ccfg.rec_op.softmax * ccfg.bytes_softmax

127 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.dropout)

128 * ccfg.bytes_dropout

129 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.headCast)

130 * ccfg.bytes_compute

131 )

132 )

133 )

134 micro_factor = ctx.micro_factor

135 return micro_factor * att_score / (ccfg.t * ccfg.cp)

136

137 @staticmethod

138 def attn_proj_activations(ccfg: CostModelConfig, ctx: Context) -> float:

139 """Output projection Activations"""

140 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER

141 att_proj = (

142 ccfg.s

143 * ccfg.b

144 * ccfg.h

145 * ccfg.bytes_compute

146 * (

147 0.25 * (ccfg.n_attMM + ccfg.n_attParamCast)

148 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.dropout)

149 * ccfg.n_dropout

150 * ccfg.bytes_dropout

151 )

152 )

153 micro_factor = ctx.micro_factor

154 return micro_factor * att_proj / max(ccfg.sp, ccfg.cp)

155

156

157class EvalFFn:

158 """Feed-forward formulas class"""

159

160 @staticmethod

161 def num_params_ffn(ccfg: CostModelConfig, _) -> float:

162 """Parameters count"""

163 experts_param_size = (

164 (ccfg.n_exp + ccfg.n_shared_exp)

165 * max(ccfg.n_ffMM, ccfg.n_ffBMM)

166 * (ccfg.hff * ccfg.h + ccfg.hff)

167 )

168 return experts_param_size

169

170 @staticmethod

171 def num_params_routed_expert(ccfg: CostModelConfig, _) -> float:

172 """Routed expert parameters count (with ETP correction)"""

173 hff_sliced = ccfg.hff_exp / max(ccfg.etp, 1)

174 return ccfg.n_exp * max(ccfg.n_ffMM, ccfg.n_ffBMM) * (hff_sliced * ccfg.h + hff_sliced)

175

176 @staticmethod

177 def num_params_shared_expert(ccfg: CostModelConfig, _) -> float:

178 """Shared expert parameters count"""

179 return ccfg.n_shared_exp * max(ccfg.n_ffMM, ccfg.n_ffBMM) * (ccfg.hff * ccfg.h + ccfg.hff)

180

181 @staticmethod

182 def ffn_activations(ccfg: CostModelConfig, ctx: Context) -> float:

183 """ "Activations count"""

184 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER

185 tok_size = ccfg.s * ccfg.b

186 n_mm = max(ccfg.n_ffMM, ccfg.n_ffBMM)

187 if n_mm % 2 == 0:

188 matmul = 0.5 * ccfg.h + 0.5 * ccfg.hff

189 else:

190 matmul = 1 / 3 * ccfg.h + 2 / 3 * ccfg.hff

191 matmul *= ccfg.bytes_compute * n_mm

192 activ_fun = ccfg.bytes_compute * ccfg.hff

193 activ_fun *= EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.ffAct)

194 pcast = ccfg.bytes_compute * ccfg.hff * ccfg.n_ffParamCast

195 activ_size = matmul + pcast + activ_fun

196 micro_factor = ctx.micro_factor

197 return micro_factor * tok_size * activ_size / (ccfg.t * ccfg.cp)

198

199 @staticmethod

200 def ffn_router_and_concat_activations(

201 ccfg: CostModelConfig, ctx: Context

202 ) -> float:

203 """MoE router and output activations"""

204 # Router activations (logits, probs, mask)

205 r = ccfg.s * ccfg.b * ccfg.bytes_compute

206 r *= 2 * ccfg.n_exp + ccfg.n_chosen_exp

207 # Concat all exp output

208 c = ccfg.s * ccfg.b * ccfg.bytes_compute * ccfg.h

209 micro_factor = ctx.micro_factor

210 return micro_factor * (r + c) / (ccfg.t * ccfg.cp)

211

212 @staticmethod

213 def shared_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:

214 """Shared expert activations"""

215 return ccfg.n_shared_exp * EvalFFn.ffn_activations(ccfg, ctx)

216

217 @staticmethod

218 def routed_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:

219 """MoE topK activations"""

220 tok_size = ccfg.s * ccfg.b

221 activ_size = EvalFFn.ffn_activations(ccfg, ctx) / tok_size

222 avg_num_toks = tok_size * ccfg.n_chosen_exp / ccfg.n_exp

223 if not ccfg.gmm: # Capacity mode

224 expert_capacity = avg_num_toks * ccfg.cap_fact * ccfg.n_exp

225 routed_activ = activ_size * expert_capacity

226 else: # Dropless mode

227 load = avg_num_toks * ccfg.n_exp * ctx.dropless_tok_factor

228 routed_activ = load * activ_size

229 return routed_activ

230

231 @staticmethod

232 def ffn_moe_activations(ccfg: CostModelConfig, ctx: Context) -> float:

233 """Sum of Activations"""

234 return (

235 EvalFFn.routed_exp_activations(ccfg, ctx)

236 + EvalFFn.shared_exp_activations(ccfg, ctx)

237 + EvalFFn.ffn_router_and_concat_activations(ccfg, ctx)

238 )

239

240

241class EvalNorm:

242 """Normalization formulas class"""

243

244 @staticmethod

245 def num_params_norm(ccfg: CostModelConfig, _) -> float:

246 """Parameters count"""

247 return ccfg.n_normOp * 2 * ccfg.h

248

249 @staticmethod

250 def norm_activations(ccfg: CostModelConfig, ctx: Context) -> float:

251 """Activations"""

252 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER

253 norm = (

254 ccfg.s

255 * ccfg.b

256 * ccfg.bytes_norm

257 * ccfg.h

258 * ccfg.n_normOp

259 * EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.normOp)

260 )

261 micro_factor = ctx.micro_factor

262 return micro_factor * norm / (ccfg.t * ccfg.cp)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / evaluators / layer_block.py: 97%

104 statements