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

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"""A model's family, applied to its config, and its layer kinds. 

16 

17A family is data: its op profile in ``auto_parallel/op_profiles``, or the op 

18counts its model spec declares, chosen by ``ccfg.arch``, which the parser 

19settles, never by matching the model name. :func:`check_and_apply_custom_hook` 

20applies it where the estimators price a config: what the family decides 

21where the run states nothing (:func:`derive_family`), then the op counts of 

22its default kind. 

23 

24The layer stack is data too (``ccfg.layer_stack``). The estimators read each 

25layer's kind from :func:`layer_groups`, and :func:`apply_layer_kind` gives a 

26layer its kind from the fields :func:`bind_layer_stack` recorded when the 

27family was applied. 

28""" 

29from typing import Any, Dict, List, Optional, Tuple 

30from hyper_parallel.auto_parallel._layer_stack import LayerStack, LinearAttentionDims 

31from hyper_parallel.auto_parallel._model_spec import ModelSpecError, OpCounts 

32from hyper_parallel.auto_parallel._op_profiles import ( 

33 VISION_ARCH, LayerKind, family_profile, known_archs, load_op_profile, 

34) 

35from hyper_parallel.auto_parallel.sapp_nd.nd.common.apply_exec import apply_layer_strategy 

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

37from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import derive_family 

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

39 

40 

41class CWrap: 

42 """Temporary evaluator-like instance""" 

43 

44 def __init__(self, e) -> None: 

45 self.ccfg = e 

46 

47 def set_ccfg(self, hook): 

48 """Apply a hook to the wrapped config object.""" 

49 return hook(self.ccfg) 

50 

51 def get_model_name(self): 

52 """Return model name from the wrapped config.""" 

53 return self.ccfg.model_name 

54 

55 def reset(self, e): 

56 """Replace the wrapped config object.""" 

57 self.ccfg = e 

58 

59 def __getattr__(self, attr): 

60 if attr not in self.__dict__: 

61 return lambda *args, **kwargs: None 

62 return self.__dict__[attr] 

63 

64 def set_strategy(self, **kwargs): 

65 """Forward strategy updates to the wrapped config.""" 

66 self.ccfg.set_strategy(**kwargs) 

67 

68 def get_strategy(self): 

69 """Return strategy from the wrapped config.""" 

70 return self.ccfg.get_strategy() 

71 

72 

73def layer_op_counts(ccfg: Any, arch: str, kind: str) -> OpCounts: 

74 """Return the op counts of one layer kind. 

75 

76 The counts the parser recorded from the model spec when there are any, 

77 else the family's own profile, which is what a config assembled without 

78 a parser gets. 

79 

80 Raises: 

81 ModelSpecError: If the recorded counts have no such kind. 

82 """ 

83 table = getattr(ccfg, "op_counts", None) 

84 if not table: 

85 return load_op_profile(arch).counts(kind) 

86 if kind not in table: 

87 raise ModelSpecError( 

88 f"{ccfg.model_name}: no op counts for layer kind {kind!r}, " 

89 f"the spec declares {sorted(table)}" 

90 ) 

91 return table[kind] 

92 

93 

94def apply_op_counts(ccfg: Any, counts: OpCounts) -> None: 

95 """Set one layer kind's op counts on *ccfg*, each as ``n_<op>``.""" 

96 for name, count in counts.to_dict().items(): 

97 setattr(ccfg, "n_" + name, count) 

98 # A layer keeps a cast beside each matmul where the run says it does, 

99 # derive's keeps_param_casts; a config derive has not seen, where the 

100 # optimizer does not shard. 

101 casts = getattr(ccfg, "keeps_param_casts", None) 

102 if casts is None: 

103 casts = not ccfg.has_op 

104 ccfg.n_attParamCast = ccfg.n_attMM if casts else 0 

105 ccfg.n_ffParamCast = ccfg.n_ffMM if casts else 0 

106 

107 

108# The fields an attention flavour assigns. Every kind of a stack whose kinds 

109# differ in attention writes all of them, so applying kinds in place, one 

110# layer after another and in any order, leaves each layer the same config. 

111_ATTENTION_FIELDS = ( 

112 "attn_kind", "a", "dh", "n_kv", "attn_output_gate", "attn_extra_p", "n_qknorm", 

113 "lin_n_k", "lin_d_k", "lin_n_v", "lin_d_v", "lin_conv", 

114) 

115 

116# The fields a feed-forward flavour assigns, likewise. 

117_FFN_FIELDS = ("hff", "n_chosen_exp", "n_exp", "n_shared_exp", "ep") 

118 

119# Strategy a kind still sets. The guard refuses a strategy write through 

120# set_ccfg, so the applier gives them through apply_layer_strategy. 

121_STRATEGY_FIELDS = ("ep",) 

122 

123 

124def _linear_attention(snapshot: Any, linear: LinearAttentionDims) -> Dict[str, Any]: 

125 """The attention fields of a gated-DeltaNet layer. 

126 

127 The flavour maps onto the q/k/v/o formula: the value heads carry the 

128 q-side width, the key heads the kv-side, and the output gate is a second 

129 q-wide tensor. What the formula does not describe is stated apart: the 

130 short convolution over the projected stream, the two per-head gates' 

131 projections, each head's decay and time-step bias and the gated output 

132 norm's weight as extra parameters, the recurrent state update as the 

133 ``linrec`` op. 

134 """ 

135 n_k, d_k = linear.num_key_heads, linear.key_head_dim 

136 n_v, d_v = linear.num_value_heads, linear.value_head_dim 

137 qkv_width = 2 * n_k * d_k + n_v * d_v 

138 return { 

139 "attn_kind": "linear", 

140 "a": n_v, 

141 "dh": d_v, 

142 "n_kv": n_k * d_k / d_v, 

143 "attn_output_gate": True, 

144 "attn_extra_p": linear.conv_kernel_dim * qkv_width + 2 * snapshot.h * n_v + 2 * n_v + d_v, 

145 # The kernel normalizes its queries and keys itself, with no weights. 

146 "n_qknorm": 0, 

147 "lin_n_k": n_k, 

148 "lin_d_k": d_k, 

149 "lin_n_v": n_v, 

150 "lin_d_v": d_v, 

151 "lin_conv": linear.conv_kernel_dim, 

152 } 

153 

154 

155def _feed_forward(snapshot: Any, flavour: Optional[str]) -> Dict[str, Any]: 

156 """The feed-forward fields of a layer of *flavour*, from the bound config. 

157 

158 A MoE layer runs the routed experts at their width. A dense layer runs 

159 one expert at the model's feed-forward width, the parser's ``hff``, with 

160 no shared expert and no expert parallelism. A kind without a flavour 

161 keeps the model's own. 

162 """ 

163 if flavour == "dense": 

164 return {"hff": snapshot.hff, "n_chosen_exp": 1, "n_exp": 1, "n_shared_exp": 0, "ep": 1} 

165 fields = {name: getattr(snapshot, name) for name in _FFN_FIELDS} 

166 if flavour == "moe": 

167 fields["hff"] = snapshot.hff_exp 

168 return fields 

169 

170 

171def _kind_fields(snapshot: Any, stack: LayerStack, kind: LayerKind) -> Dict[str, Any]: 

172 """The fields *kind* assigns beyond its op counts, from the bound config. 

173 

174 For every flavour some kind of the stack states, the fields it assigns, 

175 with the kind's values or the model's; then the kind's census record, 

176 None where the spec states none. 

177 """ 

178 kinds = stack.distinct_kinds() 

179 fields: Dict[str, Any] = {} 

180 if any(other.attention != "full" for other in kinds): 

181 if kind.attention == "linear": 

182 fields.update(_linear_attention(snapshot, stack.linear)) 

183 else: 

184 fields.update({name: getattr(snapshot, name) for name in _ATTENTION_FIELDS}) 

185 if any(other.ffn is not None for other in kinds): 

186 fields.update(_feed_forward(snapshot, kind.ffn)) 

187 fields["kind_activations"] = (getattr(snapshot, "census", None) or {}).get(kind.name) 

188 return fields 

189 

190 

191def bind_layer_stack(ccfg: Any) -> None: 

192 """Record, per kind of the config's stack, the fields it assigns. 

193 

194 Called where the family is applied, per search candidate and per 

195 estimate, so the values a kind restores are the model's own, as 

196 applying its family left them. A config without a stack gets no 

197 binding. 

198 """ 

199 stack = getattr(ccfg, "layer_stack", None) 

200 if stack is None: 

201 ccfg.layer_binding = None 

202 return 

203 ccfg.layer_binding = { 

204 kind.name: _kind_fields(ccfg, stack, kind) for kind in stack.distinct_kinds() 

205 } 

206 

207 

208def apply_layer_kind(e: Any, kind: LayerKind) -> None: 

209 """Make the layer about to be priced one of *kind*. 

210 

211 The memory backbone calls it with an evaluator, the performance path 

212 with a bare config. A config without a stack gives the kind its op 

213 counts only. 

214 """ 

215 if isinstance(e, CostModelConfig): 

216 e = CWrap(e) 

217 if getattr(e.ccfg, "layer_binding", None) is None: 

218 bind_layer_stack(e.ccfg) 

219 binding = e.ccfg.layer_binding 

220 fields = binding[kind.name] if binding is not None else {} 

221 

222 def assign(c: Any) -> None: 

223 """Give config *c* the kind's counts, then its model fields.""" 

224 apply_op_counts(c, kind.ops) 

225 for name, value in fields.items(): 

226 if name not in _STRATEGY_FIELDS: 

227 setattr(c, name, value) 

228 

229 e.set_ccfg(assign) 

230 apply_layer_strategy(e.ccfg, {name: fields[name] for name in _STRATEGY_FIELDS if name in fields}) 

231 

232 

233def _needs_kinds(stack: LayerStack) -> bool: 

234 """Whether a stack's layers differ from the config applying its family leaves.""" 

235 kinds = stack.distinct_kinds() 

236 return len(kinds) > 1 or any(kind.attention != "full" or kind.ffn is not None for kind in kinds) 

237 

238 

239def layer_groups(ccfg: Any) -> List[Tuple[Optional[LayerKind], int]]: 

240 """Return the groups of layers the estimators price alike, in model order. 

241 

242 Returns: 

243 ``(kind, count)`` per group of the config's stack. A config without 

244 a stack, or whose stack's one kind is the config applying its family 

245 leaves, is one group of all its layers, MTP included, with no kind: 

246 they are priced on the config as it stands. 

247 """ 

248 stack = getattr(ccfg, "layer_stack", None) 

249 if stack is None or not _needs_kinds(stack): 

250 return [(None, int(ccfg.n_lay + ccfg.n_mtp))] 

251 return [(group.kind, group.count) for group in stack.groups] 

252 

253 

254def layer_kinds(ccfg: Any) -> List[Optional[LayerKind]]: 

255 """Return the kind of every layer in model order, as :func:`layer_groups` groups them.""" 

256 return [kind for kind, count in layer_groups(ccfg) for _ in range(count)] 

257 

258 

259def apply_family(ccfg: Any) -> None: 

260 """Give *ccfg* what its family decides, and the op counts of its default kind. 

261 

262 A vision tower takes the vision profile's encoder counts, whatever 

263 counts it carries from its language model. Each layer of a stack takes 

264 its own kind's counts; the default kind's are the ones the model's 

265 config prices its embedding and output layer with, so t5 names its 

266 encoder. A family that names none gives the model no counts. 

267 """ 

268 derive_family(ccfg) 

269 arch = getattr(ccfg, "arch", None) 

270 if arch == VISION_ARCH: 

271 apply_op_counts(ccfg, load_op_profile(VISION_ARCH).counts("encoder")) 

272 return 

273 profile = family_profile(arch) 

274 if profile.default is not None: 

275 apply_op_counts(ccfg, layer_op_counts(ccfg, profile.arch, profile.default)) 

276 

277 

278def check_and_apply_custom_hook(e: Any) -> None: 

279 """Apply the family the config declares in ``ccfg.arch``, :func:`apply_family`. 

280 

281 A config whose arch names no profile is priced as the default family. 

282 Then bind the config's layer stack, so each kind restores the values 

283 applying the family left. 

284 """ 

285 if isinstance(e, CostModelConfig): 

286 e = CWrap(e) 

287 arch = getattr(e.ccfg, "arch", None) 

288 if arch not in known_archs(): 

289 logger.warning( 

290 "No op profile for: %s (arch %r). The default one is chosen", 

291 e.get_model_name(), arch, 

292 ) 

293 e.set_ccfg(apply_family) 

294 bind_layer_stack(e.ccfg)