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

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"""Body module""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

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

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

20 

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 

24 from typing import Tuple 

25 

26 

27class EvalBody: 

28 """Body layer formulas class""" 

29 

30 @staticmethod 

31 def num_params_layer( 

32 ccfg: CostModelConfig, ctx: Context 

33 ) -> Tuple[float, float, float]: 

34 """Parameters count. 

35 

36 Returns a 3-tuple (non_exp, routed, shared): 

37 - non_exp: attention + norm params (and dense FFN if n_exp==1) 

38 - routed: routed expert params (0 if n_exp==1) 

39 - shared: shared expert params (0 if n_shared_exp==0 or no pointer) 

40 """ 

41 non_exp = ctx.attn_num_p(ccfg, ctx) + ctx.norm_num_p(ccfg, ctx) 

42 routed = 0.0 

43 shared = 0.0 

44 if ccfg.n_exp == 1: 

45 non_exp += ctx.ffn_num_p(ccfg, ctx) 

46 else: 

47 if ctx.ffn_routed_num_p is not None: 

48 routed = ctx.ffn_routed_num_p(ccfg, ctx) 

49 if ctx.ffn_shared_num_p is not None: 

50 shared = ctx.ffn_shared_num_p(ccfg, ctx) 

51 return (non_exp, routed, shared) 

52 

53 @staticmethod 

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

55 """model param""" 

56 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx) 

57 # Routed experts: EP sharding 

58 routed_mem = routed_p / ccfg.ep * ccfg.bytes_p / ccfg.shard_p_os_exp 

59 # Shared experts: partial DP sharding 

60 shared_mem = shared_p * ccfg.bytes_p / ccfg.shard_p_os_exp_partial 

61 # Non expert 

62 non_exp_mem = non_exp_p * ccfg.bytes_p / ccfg.shard_p_os_non_exp_partial 

63 return non_exp_mem + routed_mem + shared_mem 

64 

65 @staticmethod 

66 def stat_os_layer(ccfg: CostModelConfig, ctx: Context) -> float: 

67 """optim state""" 

68 if ctx.swap_os: 

69 return 0 

70 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx) 

71 # Routed experts 

72 routed_mem = routed_p / ccfg.ep * 2 * ccfg.bytes_os / ccfg.shard_p_os_exp 

73 # Shared experts 

74 shared_mem = shared_p * 2 * ccfg.bytes_os / ccfg.shard_p_os_exp_partial 

75 # Non expert 

76 non_exp_mem = non_exp_p * 2 * ccfg.bytes_os / ccfg.shard_p_os_non_exp_partial 

77 return non_exp_mem + routed_mem + shared_mem 

78 

79 @staticmethod 

80 def stat_grad_layer(ccfg: CostModelConfig, ctx: Context) -> float: 

81 """gradients""" 

82 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx) 

83 # Routed experts 

84 routed_mem = routed_p / ccfg.ep * ccfg.bytes_grad / ccfg.shard_grad_exp 

85 # Shared experts: use shard_grad_exp_partial (independent of os sharding) 

86 shared_mem = shared_p * ccfg.bytes_grad / ccfg.shard_grad_exp_partial 

87 # Non expert 

88 non_exp_mem = non_exp_p * ccfg.bytes_grad / ccfg.shard_grad_non_exp 

89 return non_exp_mem + routed_mem + shared_mem 

90 

91 # No recompute and select recompute 

92 

93 @staticmethod 

94 def layer_activ(ccfg: CostModelConfig, ctx: Context) -> float: 

95 """activations""" 

96 attn_size = sum( 

97 [ 

98 ctx.attn_qkv_activ(ccfg, ctx), 

99 ctx.attn_score_activ(ccfg, ctx), 

100 ctx.attn_proj_activ(ccfg, ctx), 

101 ] 

102 ) 

103 if ccfg.n_exp == 1: 

104 ffn_size = ctx.ffn_activ(ccfg, ctx) 

105 else: 

106 ffn_size = ctx.ffn_moe_activ(ccfg, ctx) 

107 norm_size = ctx.norm_activ(ccfg, ctx) 

108 return attn_size + ffn_size + norm_size 

109 

110 # Full recompute 

111 

112 @staticmethod 

113 def fullrec_layer_activ(ccfg: CostModelConfig, ctx: Context) -> float: 

114 """activations""" 

115 micro_factor = ctx.micro_factor 

116 forward_activation = ( 

117 micro_factor * ccfg.bytes_compute * ccfg.s * ccfg.b * ccfg.h 

118 ) 

119 forward_activation /= ccfg.shard_recompute_input 

120 return forward_activation 

121 

122 @staticmethod 

123 def fullrec_layer_activ_gradclip( 

124 ccfg: CostModelConfig, ctx: Context 

125 ) -> float: 

126 """special case with gradient clipping""" 

127 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx) 

128 grad_clip_mem = ( 

129 non_exp_p 

130 + routed_p / ccfg.ep * ccfg.bytes_os / ccfg.shard_p_os_exp 

131 + shared_p * ccfg.bytes_os / ccfg.shard_p_os_exp_partial 

132 ) 

133 grad_clip_mem *= ccfg.bytes_os / ccfg.shard_p_os_non_exp_partial 

134 grad_clip_mem *= int(ccfg.has_clip) 

135 forward_activation = EvalBody.fullrec_layer_activ(ccfg, ctx) 

136 dp_comm_size = ctx.eval.dyn.comm.dp(ccfg, ctx) 

137 if forward_activation + dp_comm_size > grad_clip_mem: 

138 return forward_activation 

139 logger.debug( 

140 "gradient clipping %s > %s", 

141 EvalUtils.mb(grad_clip_mem), 

142 EvalUtils.mb(forward_activation + dp_comm_size), 

143 ) 

144 return grad_clip_mem 

145 

146 @staticmethod 

147 def fullrec_layer_comm_gradclip( 

148 ccfg: CostModelConfig, ctx: Context 

149 ) -> float: 

150 """special case with gradient clipping""" 

151 if EvalBody.fullrec_layer_activ_gradclip(ccfg, ctx) > 0: 

152 return ctx.eval.dyn.comm.dp(ccfg, ctx) 

153 return 0