Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / perf_estimation / getters.py: 100%

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"""Getters""" 

16 

17from copy import copy 

18from typing import Any, Dict, List, Mapping, Optional, Tuple 

19from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import CWrap, apply_layer_kind, layer_groups 

20from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import runs_hyper_selective 

21from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_order import get_model_order, layer_switches 

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

23from hyper_parallel.auto_parallel.sapp_nd.nd.logger import perf_logger as logger 

24 

25 

26def get_layer_group_configs(cfg: Any) -> List[Tuple[Any, int]]: 

27 """Return each group of layers' config and layer count, in model order. 

28 

29 A config whose layers run on it as it stands is one group, priced with 

30 ``cfg`` itself. Otherwise each kind is applied once, to a copy of 

31 ``cfg`` that every group of that kind shares. A kind sets top-level 

32 fields only (``apply_layer_kind``), so a shallow copy leaves ``cfg`` 

33 intact, as ``comm_time``'s walk copies it for the same reason; a deep 

34 copy cost 0.36 ms, twice a candidate and once a kind. 

35 """ 

36 groups = layer_groups(cfg) 

37 if all(kind is None for kind, _ in groups): 

38 return [(cfg, cfg.n_lay)] 

39 configs = {} 

40 for kind, _ in groups: 

41 if kind.name not in configs: 

42 configs[kind.name] = copy(cfg) 

43 apply_layer_kind(CWrap(configs[kind.name]), kind) 

44 return [(configs[kind.name], count) for kind, count in groups] 

45 

46 

47def get_layer_configs_by_position(cfg: Any, stages: List) -> Dict[Tuple[int, int, int], Any]: 

48 """Map each regular layer's position to the config it is priced with. 

49 

50 Each group of the layer stack covers exactly its own count of layers in 

51 model order; layers past the stack's count keep the last group, as a 

52 config whose layers need no kind keeps ``cfg``. 

53 """ 

54 lccfgs = get_layer_group_configs(cfg) 

55 per_layer = [lccfg for lccfg, count in lccfgs for _ in range(count)] 

56 last = lccfgs[-1][0] 

57 return { 

58 position: per_layer[idx] if idx < len(per_layer) else last 

59 for idx, position in enumerate(get_model_order(cfg, stages)) 

60 } 

61 

62 

63def get_layer_switches_by_position(cfg: Any, stages: List) -> Dict[Tuple[int, int, int], Any]: 

64 """Map each regular layer's position to its own recompute switches, where its config's ranges state several. 

65 

66 Empty where the config's ``rec_op`` holds the one setting every 

67 selective layer runs (:func:`layer_switches`). 

68 """ 

69 switches = layer_switches(cfg) 

70 if switches is None: 

71 return {} 

72 return {position: switches[idx] for idx, position in enumerate(get_model_order(cfg, stages)) 

73 if idx < len(switches)} 

74 

75 

76# The switch an op answers to where it has none of its own: a QK-norm is a norm. 

77_SWITCH_OF = {"qknorm": "normOp"} 

78 

79# The ops whose share of FLOPs a selective layer of HyperParallel's policy 

80# runs again a census states, and the fields that state them. 

81_CENSUS_SHARES = {"attMM": "selective_attention_mm", "ffMM": "selective_ffn_mm"} 

82 

83 

84def mla_weights(cfg: Any) -> Tuple[Any, Any]: 

85 """An MLA layer's projection weights, all of them and its up-projections', as the time model prices them. 

86 

87 The queries' latent and its up-projection, or one projection to every 

88 head; the keys' and values' shared down-projection beside the rotary 

89 key; their up-projections, a head's key at its own width and its value 

90 at the value heads'; the output projection. The up-projections build 

91 the queries' heads from their latent, where they have one, and the 

92 keys' and values' from theirs. 

93 

94 Returns: 

95 ``(weights, up)``. 

96 """ 

97 d_nope = getattr(cfg, "qk_nope_head_dim", None) or cfg.dh 

98 heads = cfg.a * (d_nope + cfg.dhr) 

99 query = cfg.dc_q * (cfg.h + heads) if cfg.dc_q else cfg.h * heads 

100 weights = (query + cfg.h * (cfg.dc_kv + cfg.dhr) + cfg.dc_kv * cfg.n_kv * (d_nope + cfg.dh) 

101 + cfg.a * cfg.dh * cfg.h) 

102 up = (cfg.dc_q * heads if cfg.dc_q else 0) + cfg.dc_kv * cfg.n_kv * (d_nope + cfg.dh) 

103 return weights, up 

104 

105 

106def selective_shares(lccfg: Any, layer: Any, switches: Optional[Mapping[str, int]] = None) -> Dict[str, float]: 

107 """The share of each matmul op's FLOPs a selective layer runs again, where no switch prices it. 

108 

109 A layer of HyperParallel's selective policy, whose switches it has 

110 (:func:`runs_hyper_selective`), *switches* where the layer has its own, 

111 runs again the shares its kind's census measured: the policy recomputes 

112 every other matmul, which no switch covers. An MLA layer whose 

113 ``attUp`` switch is 0 runs its up-projections again, their share of its 

114 projections' FLOPs. Empty otherwise, and the switches price every op. 

115 """ 

116 if layer != LayerType.SEL_REC_LAYER: 

117 return {} 

118 census = getattr(lccfg, "kind_activations", None) 

119 if census is not None and runs_hyper_selective(lccfg, switches): 

120 return {op: getattr(census, name) for op, name in _CENSUS_SHARES.items() 

121 if getattr(census, name, None) is not None} 

122 if not getattr(lccfg, "dc_kv", 0): 

123 return {} 

124 stated = switches if switches is not None else vars(lccfg.rec_op) if lccfg.rec_op is not None else {} 

125 if stated.get("attUp", 1): 

126 return {} 

127 weights, up = mla_weights(lccfg) 

128 return {"attMM": up / weights} 

129 

130 

131def get_recomp_factor(lccfg, layer, op_name, switches=None): 

132 """Whether a layer of this type runs the op again in its backward pass. 

133 

134 A selective layer recomputes exactly the ops whose switch in 

135 ``lccfg.rec_op`` is 0. A switch at 1 keeps the op's activation, which is 

136 how the memory model's ``EvalUtils.rec_coeff`` reads it, and an op with no 

137 switch is kept. The switches are read from ``vars`` because a ``Config`` 

138 answers 0 for any attribute it lacks. *switches*, a layer's own where 

139 it has them, stand for the config's. 

140 """ 

141 if layer == LayerType.FULL_REC_LAYER: 

142 return 1 

143 if layer == LayerType.NOT_REC_LAYER: 

144 return 0 

145 if layer == LayerType.SEL_REC_LAYER: 

146 if switches is None: 

147 switches = vars(lccfg.rec_op) if lccfg.rec_op is not None else {} 

148 return int(not switches.get(_SWITCH_OF.get(op_name, op_name), 1)) 

149 logger.warning("Unrecognized recompute type %s", layer) 

150 return 0 

151 

152 

153def get_table_quantity(lccfg, table, layer, with_recomp, shares=None, switches=None): 

154 """op compute load from given table; *shares* sets the recompute factor of the ops it names, *switches* the rest""" 

155 shares = shares or {} 

156 qt_layer = 0 

157 for op, quantity in table.items(): 

158 op_name = op[2:] 

159 factor = shares[op_name] if op_name in shares else get_recomp_factor(lccfg, layer, op_name, switches) 

160 

161 qt_layer += ( 

162 (1 + with_recomp * factor) 

163 * getattr(lccfg, op) 

164 * quantity 

165 ) 

166 return qt_layer