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

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"""Head submodule""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18 

19if TYPE_CHECKING: 

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

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

22 

23 

24class EvalHead: 

25 """Head layer formulas class""" 

26 

27 @staticmethod 

28 def num_params_embed(ccfg: CostModelConfig, _) -> float: 

29 """Parameter count""" 

30 return ccfg.h * ccfg.v 

31 

32 @staticmethod 

33 def stat_embed_p(ccfg: CostModelConfig, ctx: Context) -> float: 

34 """model param""" 

35 if ccfg.tie_emb_out: 

36 return 0 

37 param_size = ctx.eval.num_p(ccfg, ctx) 

38 param_size /= ccfg.shard_embed 

39 b_p = ccfg.bytes_p 

40 b_p /= ccfg.cp 

41 return param_size * b_p 

42 

43 @staticmethod 

44 def stat_embed_os(ccfg: CostModelConfig, ctx: Context) -> float: 

45 """optim state""" 

46 if ctx.swap_os or ccfg.tie_emb_out: 

47 return 0 

48 param_size = ctx.eval.num_p(ccfg, ctx) 

49 param_size /= ccfg.shard_embed 

50 b_os = 2 * ccfg.bytes_os 

51 b_os /= ccfg.cp 

52 return param_size * b_os 

53 

54 @staticmethod 

55 def stat_embed_grad(ccfg: CostModelConfig, ctx: Context) -> float: 

56 """gradient""" 

57 if ccfg.tie_emb_out: 

58 return 0 

59 param_size = ctx.eval.num_p(ccfg, ctx) 

60 param_size /= ccfg.shard_embed 

61 b_grad = ccfg.bytes_grad 

62 b_grad /= ccfg.cp 

63 return param_size * b_grad 

64 

65 @staticmethod 

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

67 """DP Communication size""" 

68 return ( 

69 ccfg.comm_d_non_exp 

70 * ctx.eval.num_p(ccfg, ctx) 

71 / (ccfg.shard_embed * ccfg.cp) 

72 ) 

73 

74 @staticmethod 

75 def tp_comm_embed(ccfg: CostModelConfig, _) -> float: 

76 """TP Communication size""" 

77 return ( 

78 ccfg.rec_op.gather 

79 * ccfg.comm_t 

80 * ccfg.s 

81 * ccfg.h 

82 * ccfg.b 

83 * (ccfg.t - 1) 

84 / (ccfg.t * ccfg.cp) 

85 ) 

86 

87 @staticmethod 

88 def activ_embed(ccfg: CostModelConfig, ctx: Context) -> float: 

89 """activations""" 

90 micro_factor = ctx.micro_factor 

91 activ_size = micro_factor * ccfg.bytes_compute / (ccfg.t * ccfg.cp) 

92 activ_size *= ccfg.s * ccfg.b * ccfg.h 

93 return activ_size