Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/auto_parallel/_spec

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"""The census's ModelSpec: a model's stack, dimensions and op counts, measured on the layers it builds.

17The resolver (:mod:`hyper_parallel.auto_parallel._hf_model_spec`) states a

18model from its Transformers config's field names, completed by rules a

19family's name selects. The census states it from the model Transformers

20builds of the config instead, the block walk of IR phase 5:

22- every decoder layer is built on the meta device, and its attention and

23 feed-forward are told apart by what they hold: a gated delta rule makes

24 a linear attention, routed experts a MoE feed-forward;

25- one layer of each shape runs a forward on fake tensors, under the

26 kernels the runtime runs (:mod:`~hyper_parallel.auto_parallel._layer_census`),

27 and the kernels' inputs state its heads and their widths, the router's

28 top-k its experts per token;

29- the parameters state the widths, the biases, the norms and the gates.

31A layer's op counts are counted in ND's units: an attention's projections

32by role, the queries', keys' and values' the kernel takes and the output's

33after it; its batched matmuls, softmaxes and score casts per softmax kernel;

34a feed-forward's projections as its weights over one ``h x width``

35projection's; its norms and activation functions as the modules it runs.

36The kinds are the model's family's (``arch``), each layer taking the kind

37whose flavours its own match, and :func:`census_model_spec` states the whole

38as a ModelSpec mapping, the resolver's form, which a train.yaml's

39``model.config_overrides`` takes as it is.

40"""

41from __future__ import annotations

43import dataclasses

44import functools

45import importlib

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

48import torch # pylint: disable=forbidden-backend-import

49from torch.utils._python_dispatch import TorchDispatchMode # pylint: disable=forbidden-backend-import

51from hyper_parallel.auto_parallel._layer_census import _fake_layer, _positions, _run, delta_rule_module

52from hyper_parallel.auto_parallel._model_spec import LayerGroup, ModelSpec, ModelSpecError, OpCounts

53from hyper_parallel.auto_parallel._op_profiles import OpProfile, infer_arch, load_op_profile

55# The modules an activation function's module is defined in.

56_ACTIVATION_MODULES = ("transformers.activations", "torch.nn.modules.activation")

58# The names of an attention's output projection, the projection after its kernel.

59_OUTPUT_PROJECTIONS = ("o_proj", "out_proj", "dense")

62@dataclasses.dataclass(frozen=True)

63class LayerShape:

64 """How one decoder layer is built.

66 Attributes:

67 attention: ``linear`` for an attention running a gated delta rule,

68 ``full`` otherwise.

69 ffn: ``moe`` for a feed-forward holding routed experts, ``dense``

70 otherwise.

71 params: Every parameter's path and shape, which tell two layers of

72 one flavour apart.

73 """

75 attention: str

76 ffn: str

77 params: Tuple[Tuple[str, Tuple[int, ...]], ...]

80@dataclasses.dataclass

81class KindCensus:

82 """What one forward of a layer states of its kind: its op counts and the dimensions its kernels take.

84 Attributes:

85 ops: The op counts, in ND's units.

86 dims: The model fields the layer states, by spec field name.

87 """

89 ops: OpCounts

90 dims: Dict[str, Any]

93def _parts(layer: Any) -> Tuple[Any, Tuple[str, Any], List[Any]]:

94 """A layer's attention, its feed-forward's name and module, and its own norms, as its children are named."""

95 attention = ffn = None

96 norms = []

97 for name, child in layer.named_children():

98 if "norm" in type(child).__name__.lower():

99 norms.append(child)

100 elif "attn" in name or "attention" in name:

101 attention = child

102 elif any(True for _ in child.parameters()):

103 ffn = (name, child)

104 if attention is None or ffn is None:

105 raise ModelSpecError(f"{type(layer).__name__} holds no attention or no feed-forward the census knows")

106 return attention, ffn, norms

107

108

109def _experts(ffn: Any) -> Optional[Any]:

110 """A feed-forward's routed experts, or ``None`` for a dense one."""

111 return getattr(ffn, "experts", None)

112

113

114def layer_shape(layer: Any) -> LayerShape:

115 """The flavours and the parameters of *layer*."""

116 attention, (_, ffn), _ = _parts(layer)

117 linear = any(delta_rule_module(module) for module in attention.modules())

118 return LayerShape("linear" if linear else "full", "moe" if _experts(ffn) is not None else "dense",

119 tuple((name, tuple(param.shape)) for name, param in layer.named_parameters()))

120

121

122def _linears(module: Any) -> List[Tuple[str, Any]]:

123 """The linear projections of *module*, by path."""

124 return [(name, child) for name, child in module.named_modules() if isinstance(child, torch.nn.Linear)]

125

126

127def _down_width(module: Any, hidden: int) -> int:

128 """The width of a feed-forward *module*: the input of its projection back to the hidden width.

129

130 Where every projection maps the hidden width to itself, the

131 feed-forward is as wide as the model.

132 """

133 back = [child.in_features for _, child in _linears(module) if child.out_features == hidden]

134 if not back:

135 raise ModelSpecError(f"{type(module).__name__} holds no projection back to width {hidden}")

136 return next((width for width in back if width != hidden), hidden)

137

138

139def _is_activation(module: Any) -> bool:

140 """Whether *module* is an activation function's module."""

141 return type(module).__module__ in _ACTIVATION_MODULES

142

143

144class _KindOps(TorchDispatchMode):

145 """What a layer's forward runs: its kernels and their inputs, its routers' top-k, its dropouts and projections."""

146

147 def __init__(self, layer: Any) -> None:

148 """Follow which of *layer*'s modules runs."""

149 super().__init__()

150 self.running: List[str] = []

151 self.kernels: List[Tuple[Tuple[int, ...], ...]] = []

152 self.delta_rules = 0

153 self.top_k: Optional[int] = None

154 self.dropouts = 0

155 self.after_kernel = 0

156 self.activations: List[str] = []

157 self.modules = dict(layer.named_modules())

158 for name, module in self.modules.items():

159 if name:

160 module.register_forward_pre_hook(functools.partial(self._enter, name))

161 module.register_forward_hook(self._leave)

162 if hasattr(module, "chunk_gated_delta_rule"):

163 module.chunk_gated_delta_rule = self._counted(module.chunk_gated_delta_rule)

164

165 def _counted(self, rule: Any) -> Any:

166 """*rule*, counting its calls."""

167 def run(*args: Any, **kwargs: Any) -> Any:

168 """The delta rule, its call counted."""

169 self.delta_rules += 1

170 return rule(*args, **kwargs)

171 return run

172

173 def _enter(self, name: str, *_: Any) -> None:

174 """Note that the module at *name* runs: an output projection after the kernel, an activation function."""

175 module = self.modules[name]

176 if (self.kernels or self.delta_rules) and name.rsplit(".", 1)[-1] in _OUTPUT_PROJECTIONS:

177 self.after_kernel += 1

178 if _is_activation(module):

179 self.activations.append(name)

180 self.running.append(name)

181

182 def _leave(self, *_: Any) -> None:

183 """Note that the innermost module running has returned."""

184 self.running.pop()

185

186 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument

187 out = func(*args, **(kwargs or {}))

188 if func is torch.ops.nd_census.flash_attention.default:

189 self.kernels.append(tuple(tuple(tensor.shape) for tensor in args[:3]))

190 elif func is torch.ops.aten.topk.default:

191 self.top_k = int(args[1])

192 elif func is torch.ops.aten.native_dropout.default and args[1] > 0:

193 self.dropouts += 1

194 return out

195

196

197def _biased(module: Any) -> bool:

198 """Whether a linear projection of *module* holds a bias."""

199 return any(child.bias is not None for _, child in _linears(module))

200

201

202def _shared_dims(ffn: Any, hidden: int, width: int) -> Dict[str, Any]:

203 """A MoE feed-forward's shared experts' fields: as many experts of the routed ones' width as they are wide.

204

205 Its gate is a projection of the hidden state to one scale.

206 """

207 shared = [child for name, child in ffn.named_children() if "shared" in name

208 and any(inner.out_features == hidden for _, inner in _linears(child))]

209 dims: Dict[str, Any] = {"shared_expert_gate": any(

210 "shared" in name and "gate" in name and isinstance(child, torch.nn.Linear)

211 for name, child in ffn.named_children())}

212 if shared:

213 shared_width = sum(_down_width(child, hidden) for child in shared)

214 dims.update(shared_expert_intermediate_size=shared_width,

215 num_shared_experts=shared_width // width if shared_width % width == 0 else len(shared),

216 mlp_bias=any(_biased(child) for child in shared))

217 return dims

218

219

220def _ffn_dims(ffn: Any, hidden: int, top_k: Optional[int]) -> Tuple[Dict[str, Any], int]:

221 """A feed-forward's fields, and its projections in ND's units: its weights over an ``h x width`` projection's."""

222 experts = _experts(ffn)

223 if experts is None:

224 width = _down_width(ffn, hidden)

225 weights = sum(child.weight.numel() for _, child in _linears(ffn))

226 return {"dense_width": width, "mlp_bias": _biased(ffn)}, round(weights / (hidden * width))

227 num, width = int(experts.num_experts), int(experts.intermediate_dim)

228 weights = sum(param.numel() for name, param in experts.named_parameters() if "bias" not in name)

229 shared = _shared_dims(ffn, hidden, width)

230 dims = {"num_experts": num, "moe_intermediate_size": width, "num_experts_per_tok": top_k, **shared,

231 "mlp_bias": bool(getattr(experts, "has_bias", False)) or shared.get("mlp_bias", False)}

232 return dims, round(weights / (num * hidden * width))

233

234

235def _attention_dims(attention: Any, counter: _KindOps) -> Dict[str, Any]:

236 """An attention's fields: its heads and widths as its kernel takes them, its projections' shapes."""

237 projections = dict(_linears(attention))

238 inputs = [child for name, child in projections.items() if name.rsplit(".", 1)[-1] not in _OUTPUT_PROJECTIONS]

239 outputs = [child for name, child in projections.items() if name.rsplit(".", 1)[-1] in _OUTPUT_PROJECTIONS]

240 dims: Dict[str, Any] = {

241 "qkv_bias": any(child.bias is not None for child in inputs),

242 "o_bias": any(child.bias is not None for child in outputs),

243 "qk_norm": all(hasattr(attention, name) for name in ("q_norm", "k_norm")),

244 }

245 if counter.delta_rules:

246 for field, attribute in (("linear_num_key_heads", "num_k_heads"), ("linear_key_head_dim", "head_k_dim"),

247 ("linear_num_value_heads", "num_v_heads"), ("linear_value_head_dim", "head_v_dim"),

248 ("linear_conv_kernel_dim", "conv_kernel_size")):

249 dims[field] = int(getattr(attention, attribute))

250 return dims

251 query, key, value = counter.kernels[0]

252 heads, width, value_width = query[1], query[-1], value[-1]

253 dims.update(num_attention_heads=heads, num_key_value_heads=key[1], head_dim=value_width)

254 latent = projections.get("kv_b_proj")

255 if latent is not None:

256 # MLA: its keys and values are up-projected from a latent, beside a

257 # rotary key the heads share.

258 rope = projections["kv_a_proj_with_mqa"].out_features - latent.in_features

259 dims.update(kv_lora_rank=latent.in_features, qk_rope_head_dim=rope, qk_nope_head_dim=width - rope,

260 v_head_dim=value_width)

261 if "q_a_proj" in projections:

262 dims["q_lora_rank"] = projections["q_a_proj"].out_features

263 elif "q_proj" in projections:

264 dims["attn_output_gate"] = projections["q_proj"].out_features == 2 * heads * width

265 return dims

266

267

268def census_kind(config: Any, layer_index: int, seq_length: int = 64) -> KindCensus:

269 """The op counts and the fields one forward of layer *layer_index* of *config* states.

270

271 Args:

272 config: The language model's Transformers config.

273 layer_index: The layer to build.

274 seq_length: Tokens of the sequence it runs; the counts and the

275 fields do not depend on it.

276

277 Returns:

278 The layer's :class:`KindCensus`.

279 """

280 with _fake_layer(config, layer_index) as (layer, rotary):

281 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)

282 counter = _KindOps(layer)

283 with counter:

284 _run(layer, hidden, _positions(rotary, hidden, seq_length))

285 attention, (ffn_name, ffn), norms = _parts(layer)

286 dims = _attention_dims(attention, counter)

287 ffn_dims, projections = _ffn_dims(ffn, config.hidden_size, counter.top_k)

288 dims.update(ffn_dims)

289 dims.update(layer_norms=len(norms), norm_bias=any(norm.bias is not None for norm in norms

290 if hasattr(norm, "bias")))

291 # The activation functions of one expert, or of the dense feed-forward.

292 runs = f"{ffn_name}.experts." if _experts(ffn) is not None else f"{ffn_name}."

293 softmaxes = len(counter.kernels)

294 routed = [name for name in counter.activations if name.startswith(runs)]

295 ops = OpCounts(

296 attMM=(3 + counter.after_kernel) if softmaxes or counter.delta_rules else len(_linears(attention)),

297 attBMM=2 * softmaxes, ffMM=projections, softmax=softmaxes, dropout=counter.dropouts,

298 normOp=len(norms), gather=4, headCast=softmaxes, ffAct=len(routed), linrec=counter.delta_rules)

299 return KindCensus(ops, dims)

300

301

302def _kind_name(profile: OpProfile, shape: LayerShape) -> str:

303 """The kind of *profile* a layer of *shape* is: the one whose flavours match its own."""

304 for ffn in (shape.ffn, None):

305 names = [kind.name for kind in profile.layer_kinds.values()

306 if kind.attention == shape.attention and kind.ffn == ffn]

307 if len(names) == 1:

308 return names[0]

309 raise ModelSpecError(

310 f"op profile {profile.arch!r} has no single kind for a layer of {shape.attention} attention "

311 f"and a {shape.ffn} feed-forward")

312

313

314def _language_model(config: Any) -> Any:

315 """The Transformers causal LM of *config*, built on the meta device."""

316 transformers = importlib.import_module("transformers")

317 with torch.device("meta"):

318 return transformers.AutoModelForCausalLM.from_config(config)

319

320

321def _stack(profile: OpProfile, shapes: List[LayerShape]) -> Tuple[Tuple[LayerGroup, ...], Dict[str, int]]:

322 """The layer stack *shapes* make, each layer of its kind, and each kind's first layer.

323

324 Raises:

325 ModelSpecError: If the profile gives two differently built layers

326 one kind.

327 """

328 kinds = [_kind_name(profile, shape) for shape in shapes]

329 firsts: Dict[str, int] = {}

330 for index, (kind, shape) in enumerate(zip(kinds, shapes)):

331 first = firsts.setdefault(kind, index)

332 if shapes[first] != shape:

333 raise ModelSpecError(

334 f"op profile {profile.arch!r} prices layers {first} and {index} as one kind {kind!r}, "

335 f"but they are built differently: {shapes[first].ffn} and {shape.ffn} feed-forwards")

336 groups: List[List[Any]] = []

337 for kind in kinds:

338 if groups and groups[-1][0] == kind:

339 groups[-1][1] += 1

340 else:

341 groups.append([kind, 1])

342 return tuple(LayerGroup(kind, count) for kind, count in groups), firsts

343

344

345def _model_fields(model: Any) -> Dict[str, Any]:

346 """The fields of the model around its layers: its width, vocabulary, tied table and positions."""

347 table = model.get_input_embeddings().weight

348 body = model.get_decoder()

349 fields = {"hidden_size": int(table.shape[1]), "vocab_size": int(table.shape[0]),

350 "num_hidden_layers": len(body.layers),

351 "tie_word_embeddings": model.get_output_embeddings().weight is table}

352 positions = getattr(getattr(body, "rotary_emb", None), "original_max_seq_len", None)

353 if positions:

354 fields["max_position_embeddings"] = int(positions)

355 return fields

356

357

358def _widths(spec: Dict[str, Any], stack_dims: Mapping[str, Mapping[str, Any]]) -> None:

359 """Fill *spec*'s widths from its kinds' fields: attention's, linear attention's, feed-forward's."""

360 ordered = sorted(stack_dims.items(), key=lambda item: "linear_num_key_heads" in item[1])

361 for _, dims in ordered:

362 for field, value in dims.items():

363 if field != "dense_width" and value is not None:

364 spec.setdefault(field, value)

365 for field in ("qkv_bias", "o_bias", "qk_norm", "mlp_bias", "norm_bias", "shared_expert_gate",

366 "attn_output_gate"):

367 # A flag no layer states is one no layer has: a dense model's

368 # shared expert gate, a linear attention's output gate.

369 spec[field] = any(dims.get(field, False) for dims in stack_dims.values())

370 dense = [dims["dense_width"] for dims in stack_dims.values() if "dense_width" in dims]

371 # A model whose every layer routes its tokens states the dense width

372 # as the resolver does: its shared expert's, else its experts'.

373 width = dense[0] if dense else spec.get("shared_expert_intermediate_size") or spec.get("moe_intermediate_size")

374 if width:

375 spec["intermediate_size"] = width

376

377

378def _first_dense(shapes: List[LayerShape]) -> Optional[int]:

379 """How many dense layers lead a stack whose other layers route their tokens, or ``None``."""

380 moe = [index for index, shape in enumerate(shapes) if shape.ffn == "moe"]

381 if not moe or not moe[0] or any(shape.ffn != "dense" for shape in shapes[:moe[0]]):

382 return None

383 return moe[0]

384

385

386def census_model_spec(config: Any, arch: Optional[str] = None, name: Optional[str] = None) -> Dict[str, Any]:

387 """The ModelSpec the model Transformers builds of *config* states, as a mapping.

388

389 Args:

390 config: The model's Transformers config; a multimodal config's

391 language model is the one described.

392 arch: The family whose kinds and op profile the spec names; the one

393 the model's name belongs to where omitted.

394 name: The spec's name, the config's model type where omitted.

395

396 Returns:

397 The spec, in the resolver's form: every field the layers state,

398 the layer stack in ``layers``, each kind's op counts in ``ops``.

399

400 Raises:

401 ModelSpecError: If the family has no kind for a layer, or prices

402 two differently built layers as one kind.

403 """

404 name = str(name or config.model_type)

405 text = getattr(config, "text_config", None) or config

406 profile = load_op_profile(arch or infer_arch(name))

407 model = _language_model(text)

408 spec: Dict[str, Any] = {"name": name, "arch": profile.arch, **_model_fields(model)}

409 shapes = [layer_shape(layer) for layer in model.get_decoder().layers]

410 layers, firsts = _stack(profile, shapes)

411 censuses = {kind: census_kind(text, index) for kind, index in firsts.items()}

412 _widths(spec, {kind: record.dims for kind, record in censuses.items()})

413 spec["first_k_dense_replace"] = _first_dense(shapes)

414 # A kind of the family no layer is keeps its profile's counts.

415 spec["ops"] = {kind: counts.to_dict() for kind, counts in profile.kinds.items()}

416 spec["ops"].update({kind: record.ops.to_dict() for kind, record in censuses.items()})

417 spec["layers"] = [group.to_dict() for group in layers]

418 return ModelSpec.from_dict(spec).to_dict()

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _spec_census.py: 98%

209 statements