Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / verify.py: 99%
204 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 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"""Verify mode: the parameters and FLOPs ND prices of a model, beside those of the layers Transformers builds of it.
17``run_nd -f hyper_v2 -y <train.yaml> -V`` builds the first layer of each
18kind of the model the train.yaml trains, on fake tensors, and counts its
19parameters part by part
20(:func:`~hyper_parallel.auto_parallel._layer_census.census_parameters`).
21Beside each part it sets what ND's formulas price of it, and the same for
22the embedding and the output layer. It reports each part's two counts and
23their difference, rather than a pass or a fail: a fact the parser misread
24shows as the part it feeds.
26It then counts each kind's forward FLOPs on one sequence of the run's
27length (:func:`~hyper_parallel.auto_parallel._layer_census.census_flops`)
28beside what the time model's op table prices of the same forward: a third
29of each matmul entry, which prices the backward as twice the forward. The
30table's feed-forward entry covers the whole feed-forward, its experts,
31shared expert and router included.
33It sets what each kind keeps for its backward for each op, per token of
34the whole layer, as the op records price it (shared decision S1,
35:class:`~hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block.EvalRecords`)
36beside what the census measures each op saving
37(:func:`~hyper_parallel.auto_parallel._layer_census.census_saved_ops`).
39Last, it sets the model spec the resolver reads from the checkpoint's
40config beside the one the census measures on the layers Transformers
41builds of it (:func:`~hyper_parallel.auto_parallel._spec_census.census_model_spec`),
42field by field as ND prices each (:func:`verify_spec`), and prices the run
43with each (:func:`verify_estimate`).
45It also sets the bytes each kind's forward moves
46(:func:`~hyper_parallel.auto_parallel._layer_census.census_traffic`) beside
47the FLOPs the op table prices of the same forward. The time model counts
48the FLOPs and nothing else, and the cluster's kernel tables put those at
496.3 to 7.1% of a profiled step of Qwen3.5-35B-A3B at 8192 tokens, so this
50section reports the other quantity a step could be priced by rather than a
51difference: the two are in different units, and what a round compares is
52each against the time it took.
53"""
54from __future__ import annotations
56import copy
57import os
58import tempfile
59from types import SimpleNamespace
60from typing import Any, Dict, List, Mapping, NamedTuple, Tuple
62import yaml
64from hyper_parallel.auto_parallel._hf_model_spec import (
65 checkpoint_configs,
66 is_auto_models_schema,
67 resolve_hf_model_spec,
68)
69from hyper_parallel.auto_parallel._layer_census import (
70 census_final_norm,
71 census_flops,
72 census_parameters,
73 census_saved_ops,
74 census_traffic,
75)
76from hyper_parallel.auto_parallel._model_spec import ModelSpec
77from hyper_parallel.auto_parallel._op_profiles import family_profile, resolve_ops
78from hyper_parallel.auto_parallel._op_records import load_op_records
79from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
80from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalRecords
81from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.head import EvalHead
82from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.tail import EvalTail
83from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import check_and_apply_custom_hook, layer_groups
84from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
85from hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware import device_map
86from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
87from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.comm_time import prepare_context
88from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.estimate import (
89 BACKWARD_RATIO,
90 _flavour_tables,
91 estimate_performance,
92)
93from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.getters import get_layer_group_configs
95# The time model's matmul entries, and the part of a layer each prices.
96_FLOP_PARTS = (("n_attMM", "attention"), ("n_attBMM", "scores"), ("n_linrec", "linrec"), ("n_ffMM", "ffn"))
98# The census's parts of a feed-forward, which the time model prices as one.
99_FFN_PARTS = ("ffn", "routed", "shared", "router")
101# The model facts a spec states as true or false.
102_FLAGS = ("attn_output_gate", "tie_word_embeddings", "qk_norm", "qkv_bias", "o_bias", "mlp_bias", "norm_bias",
103 "shared_expert_gate")
105# What a spec states that verify mode does not compare: its name, what the
106# AutoModels parser drops (MTP layers, which Transformers does not build),
107# the census records, and the vision tower, which the census does not walk.
108_NOT_COMPARED = ("name", "mtp_depth", "layer_types", "activations", "output_activations", "vision", "ops", "layers")
111class VerifyRow(NamedTuple):
112 """One part's parameters: ND's count, the census's, and how many layers hold it."""
114 where: str
115 part: str
116 nd: float
117 census: float
118 count: int = 1
121class SpecRow(NamedTuple):
122 """One field of the model spec, as ND prices it: the resolver's value and the census's."""
124 field: str
125 resolved: Any
126 census: Any
129class TrafficRow(NamedTuple):
130 """One part of a forward: the bytes it moves, the parameters among them, its ops, and the FLOPs ND prices of it."""
132 where: str
133 part: str
134 moved: float
135 parameter: float
136 launches: int
137 flops: float
138 count: int = 1
141def nd_parameters(lccfg: Any, ctx: Any) -> Dict[str, float]:
142 """The parameters ND prices of a layer of *lccfg*, by part, with the formulas *ctx* names."""
143 parts = {"attention": ctx.attn_num_p(lccfg, ctx), "norm": ctx.norm_num_p(lccfg, ctx)}
144 if lccfg.n_exp == 1:
145 parts["ffn"] = ctx.ffn_num_p(lccfg, ctx)
146 else:
147 parts["routed"] = ctx.ffn_routed_num_p(lccfg, ctx)
148 parts["shared"] = ctx.ffn_shared_num_p(lccfg, ctx)
149 parts["router"] = ctx.ffn_router_num_p(lccfg, ctx)
150 return parts
153def nd_flops(ccfg: Any, lccfg: Any) -> Dict[str, float]:
154 """The forward FLOPs the time model prices one sequence of a layer of *lccfg* at, by part.
156 A third of each matmul entry of the layer's op table, for one sequence
157 of the whole layer: the table prices the backward as twice the
158 forward, a micro-batch of ``b`` sequences, and one TP and CP rank's
159 share.
160 """
161 base, experts = _flavour_tables(ccfg, lccfg)
162 table = experts if lccfg.n_exp > 1 else base
163 scale = ccfg.b * (1 + BACKWARD_RATIO) * ccfg.bytes_p / ccfg.t / ccfg.cp
164 return {part: getattr(lccfg, op) * table[op] / scale for op, part in _FLOP_PARTS if op in table}
167def nd_saved_ops(lccfg: Any) -> Dict[str, float]:
168 """What the op records price a layer of *lccfg* keeping for each op, per token of the whole layer.
170 The records' slots evaluated on a copy of the layer's config at TP, CP
171 and sequence parallelism 1 and micro-batch 1, divided by its tokens.
172 """
173 whole = copy.copy(lccfg)
174 whole.t = whole.cp = whole.sp = whole.b = 1
175 ctx = SimpleNamespace(current_node=LayerType.NOT_REC_LAYER, micro_factor=1, dropless_tok_factor=1)
176 return {op: sum(EvalRecords.op_bytes(whole, ctx, op).values()) / whole.s for op in load_op_records().ops}
179def _text_model(ccfg: Any) -> Any:
180 """The config of *ccfg*'s language model: the config itself, or a multimodal config's main submodule's."""
181 if not getattr(ccfg, "multimodal", False):
182 return ccfg
183 return ccfg.mm_ccfgs[getattr(ccfg, "mm_main", None) or ccfg.mm_order[-1]]
186def _priced(yaml_path: str) -> Tuple[Any, Any, Any]:
187 """The Transformers configs a train.yaml trains, its model's and its language model's, and ND's config of it.
189 ND's config is its language model's, with the family's op counts and
190 fields applied as the estimates apply them, on a copy.
192 Raises:
193 ValueError: The train.yaml names no Transformers checkpoint.
194 """
195 config, text = checkpoint_configs(_train_yaml(yaml_path)["model"])
196 ccfg = copy.deepcopy(_text_model(CostModelConfig(yaml_path, framework="hyper_v2")))
197 check_and_apply_custom_hook(ccfg)
198 return config, text, ccfg
201def _kinds(ccfg: Any) -> Dict[str, List[Any]]:
202 """Each layer kind of *ccfg*'s stack: its first layer, how many layers it has, and their config."""
203 kinds: Dict[str, List[Any]] = {}
204 first = 0
205 for (kind, count), (lccfg, _) in zip(layer_groups(ccfg), get_layer_group_configs(ccfg)):
206 name = kind.name if kind is not None else "decoder"
207 kinds.setdefault(name, [first, 0, lccfg])[1] += count
208 first += count
209 return kinds
212def verify_parameters(yaml_path: str) -> List[VerifyRow]:
213 """The parameters ND prices of the model a train.yaml trains, part by part, beside the census's.
215 Args:
216 yaml_path: An AutoModels train.yaml, whose model names a Transformers
217 checkpoint.
219 Returns:
220 A row per part of the first layer of each kind of its language
221 model, then the embedding's, the output layer's, and the model's
222 whole. A table the embedding and the output layer share on one
223 stage is counted at the output, as ND prices it.
225 Raises:
226 ValueError: The train.yaml names no Transformers checkpoint.
227 """
228 config, text, ccfg = _priced(yaml_path)
229 ctx = prepare_context()
230 rows = []
231 for name, (index, count, lccfg) in _kinds(ccfg).items():
232 nd, census = nd_parameters(lccfg, ctx), census_parameters(text, index)
233 rows += [VerifyRow(f"{name} x{count}", part, nd.get(part, 0.0), census.get(part, 0), count)
234 for part in sorted(set(nd) | set(census))]
235 table = text.vocab_size * text.hidden_size
236 shared = bool(getattr(config, "tie_word_embeddings", False)) and ccfg.p <= 1
237 rows.append(VerifyRow("embedding", "table", 0.0 if EvalHead.shares_output_table(ccfg)
238 else EvalHead.num_params_embed(ccfg, ctx), 0 if shared else table))
239 rows.append(VerifyRow("output", "table, norm", EvalTail.num_params_output(ccfg, ctx),
240 table + census_final_norm(text)))
241 rows.append(VerifyRow("model", "total", sum(row.count * row.nd for row in rows),
242 sum(row.count * row.census for row in rows)))
243 return rows
246def verify_activations(yaml_path: str) -> List[VerifyRow]:
247 """What each kind of the model a train.yaml trains keeps for each op, the records' beside the census's.
249 Args:
250 yaml_path: An AutoModels train.yaml, whose model names a Transformers
251 checkpoint.
253 Returns:
254 A row per op either prices or measures of the first layer of each
255 kind of its language model, in bytes per token of the whole layer
256 at the run's length (:func:`nd_saved_ops`,
257 :func:`~hyper_parallel.auto_parallel._layer_census.census_saved_ops`),
258 the census's ``other`` its own code's, then the kind's whole.
260 Raises:
261 ValueError: The train.yaml names no Transformers checkpoint.
262 """
263 _, text, ccfg = _priced(yaml_path)
264 rows = []
265 for name, (index, count, lccfg) in _kinds(ccfg).items():
266 nd, census = nd_saved_ops(lccfg), census_saved_ops(text, index, int(ccfg.s))
267 rows += [VerifyRow(f"{name} x{count}", op, nd.get(op, 0.0), census.get(op, 0.0), count)
268 for op in [*nd, "other"] if nd.get(op) or census.get(op)]
269 rows.append(VerifyRow(f"{name} x{count}", "total", sum(nd.values()), sum(census.values()), count))
270 return rows
273def verify_flops(yaml_path: str) -> List[VerifyRow]:
274 """The forward FLOPs ND prices of the model a train.yaml trains, part by part, beside the census's.
276 Args:
277 yaml_path: An AutoModels train.yaml, whose model names a Transformers
278 checkpoint.
280 Returns:
281 A row per part of the first layer of each kind of its language
282 model on one sequence of the run's length, the census's feed-forward
283 parts summed as the time model prices them, then the layers' whole.
285 Raises:
286 ValueError: The train.yaml names no Transformers checkpoint.
287 """
288 _, text, ccfg = _priced(yaml_path)
289 rows = []
290 for name, (index, count, lccfg) in _kinds(ccfg).items():
291 nd, census = nd_flops(ccfg, lccfg), census_flops(text, index, int(ccfg.s))
292 census["ffn"] = sum(census.pop(part, 0) for part in _FFN_PARTS)
293 rows += [VerifyRow(f"{name} x{count}", part, nd.get(part, 0.0), census.get(part, 0), count)
294 for _, part in _FLOP_PARTS if nd.get(part) or census.get(part)]
295 rows.append(VerifyRow("layers", "total", sum(row.count * row.nd for row in rows),
296 sum(row.count * row.census for row in rows)))
297 return rows
300def spec_as_priced(data: Mapping[str, Any]) -> Dict[str, Any]:
301 """The fields of the model spec *data*, each as ND prices it.
303 A field left unstated is compared as the value ND prices for it: a head
304 as wide as derive prices it, an MLA family's at its value heads' width;
305 a flag as false; the routed and shared experts' widths as the spec
306 resolves them; each kind's op counts as its profile states them, one
307 field per op. The layer stack is its body, the groups in model order:
308 the AutoModels parser prices no MTP layer.
310 Args:
311 data: A model spec, as a resolver states it.
313 Returns:
314 ``{field: value}``, ``ops.<kind>.<op>`` for each op count.
315 """
316 spec = ModelSpec.from_dict(data)
317 out = {key: value for key, value in spec.to_dict().items() if key not in _NOT_COMPARED}
318 value_width = family_profile(spec.arch).model.get("v_head_dim")
319 out["head_dim"] = (spec.v_head_dim or value_width) if value_width is not None else spec.effective_head_dim
320 out.update({flag: bool(getattr(spec, flag)) for flag in _FLAGS})
321 if spec.is_moe:
322 out["moe_intermediate_size"] = spec.routed_expert_width
323 if spec.num_shared_experts:
324 out["shared_expert_intermediate_size"] = spec.shared_expert_width
325 for kind, counts in resolve_ops(spec.arch, spec.ops).items():
326 out.update({f"ops.{kind}.{op}": count for op, count in counts.to_dict().items()})
327 out["layers"] = tuple((group.kind, group.count) for group in spec.layers or () if not group.mtp)
328 return out
331def _train_yaml(yaml_path: str) -> Dict[str, Any]:
332 """The AutoModels train.yaml at *yaml_path*, whose model names a Transformers checkpoint.
334 Raises:
335 ValueError: The yaml is not one.
336 """
337 with open(yaml_path, encoding="utf-8") as handle:
338 raw = yaml.safe_load(handle)
339 if not isinstance(raw, dict) or not is_auto_models_schema(raw) or not isinstance(raw.get("model"), dict):
340 raise ValueError(f"{yaml_path}: verify mode needs an AutoModels train.yaml naming a Transformers checkpoint")
341 return raw
344def verify_spec(yaml_path: str) -> List[SpecRow]:
345 """The model spec the resolver reads of the model a train.yaml trains, beside the census's, field by field.
347 Args:
348 yaml_path: An AutoModels train.yaml, whose model names a Transformers
349 checkpoint.
351 Returns:
352 A row per field either spec states, each as ND prices it
353 (:func:`spec_as_priced`), in the order of the fields' names.
355 Raises:
356 ValueError: The train.yaml names no Transformers checkpoint.
357 """
358 model = _train_yaml(yaml_path)["model"]
359 resolved = spec_as_priced(resolve_hf_model_spec(model))
360 census = spec_as_priced(resolve_hf_model_spec(model, census_spec=True))
361 return [SpecRow(field, resolved.get(field), census.get(field)) for field in sorted(set(resolved) | set(census))]
364def census_spec_yaml(yaml_path: str, folder: str) -> str:
365 """Write the census's spec of the model a train.yaml trains to *folder*, as the run would price it.
367 The spec carries the census records of the layers' activations where
368 the train.yaml states ``context.census``, measured at its dataset's
369 length, else at 4096 tokens; a train.yaml's ``model.config_overrides``
370 takes it as it is.
372 Returns:
373 The path written, ``model_spec.yaml`` in *folder*.
374 """
375 raw = _train_yaml(yaml_path)
376 seq_len = 0
377 if (raw.get("context") or {}).get("census"):
378 seq_len = int(((raw.get("dataset") or {}).get("data_transform") or {}).get("max_seq_len") or 4096)
379 spec = resolve_hf_model_spec(raw["model"], census_seq_len=seq_len, census_spec=True)
380 os.makedirs(folder, exist_ok=True)
381 path = os.path.join(folder, "model_spec.yaml")
382 with open(path, "w", encoding="utf-8") as handle:
383 yaml.safe_dump(spec, handle, sort_keys=False)
384 return path
387def _estimate(yaml_path: str, device_type: str) -> Tuple[float, float]:
388 """The peak memory, in MB, and the step's time score ND prices the run of a train.yaml at."""
389 ccfg = CostModelConfig(yaml_path, framework="hyper_v2")
390 peak = EvaluatorV2(yaml_path, framework="hyper_v2", log_level=0).estimate_peak()
391 return float(peak), float(estimate_performance(_text_model(ccfg), device_type=device_map[device_type]))
394def verify_estimate(yaml_path: str, device_type: str = "A2") -> List[VerifyRow]:
395 """The run a train.yaml states, priced with the resolver's model spec and with the census's.
397 Args:
398 yaml_path: An AutoModels train.yaml, whose model names a Transformers
399 checkpoint, and whose run is priced as it states it.
400 device_type: The device to price the step's time on.
402 Returns:
403 The peak memory and the step's time score, with the resolver's spec
404 in the ND column and the census's in the census column.
406 Raises:
407 ValueError: The train.yaml names no Transformers checkpoint.
408 """
409 raw = _train_yaml(yaml_path)
410 resolved = _estimate(yaml_path, device_type)
411 raw.setdefault("context", {})["census_spec"] = True
412 with tempfile.TemporaryDirectory() as folder:
413 path = os.path.join(folder, os.path.basename(yaml_path))
414 with open(path, "w", encoding="utf-8") as handle:
415 yaml.safe_dump(raw, handle)
416 census = _estimate(path, device_type)
417 return [VerifyRow("estimate", "peak MB", resolved[0], census[0]),
418 VerifyRow("estimate", "time score", resolved[1], census[1])]
421def _spec_value(value: Any) -> str:
422 """One field's value as a spec report shows it: a layer stack by its kinds' layers and groups."""
423 if isinstance(value, tuple):
424 layers: Dict[str, int] = {}
425 for kind, count in value:
426 layers[kind] = layers.get(kind, 0) + count
427 return f"{len(value)} groups: " + ", ".join(f"{kind} x{count}" for kind, count in layers.items())
428 return "unstated" if value is None else str(value)
431def report_spec(rows: List[SpecRow]) -> List[str]:
432 """The lines of a spec report: each field's two values, the fields that differ marked."""
433 shown = [(row.field, _spec_value(row.resolved), _spec_value(row.census), row.resolved != row.census)
434 for row in rows]
435 widths = [max([len(title)] + [len(line[column]) for line in shown])
436 for column, title in enumerate(("field", "resolver", "census"))]
437 lines = [f"{'field':{widths[0]}s} {'resolver':>{widths[1]}s} {'census':>{widths[2]}s}"]
438 for field, resolved, census, differs in shown:
439 mark = " <- differs" if differs else ""
440 lines.append(f"{field:{widths[0]}s} {resolved:>{widths[1]}s} {census:>{widths[2]}s}{mark}")
441 differ = sum(line[3] for line in shown)
442 lines.append(f"{len(rows) - differ} fields agree, {differ} differ")
443 return lines
446def report(rows: List[VerifyRow]) -> List[str]:
447 """The lines of a verify report: each part's two counts, their difference and its share of the census's."""
448 lines = [f"{'where':26s} {'part':12s} {'ND':>17s} {'census':>17s} {'ND - census':>15s} {'%':>9s}"]
449 for row in rows:
450 difference = row.nd - row.census
451 share = f"{100 * difference / row.census:+8.3f}%" if row.census else ""
452 lines.append(f"{row.where:26s} {row.part:12s} {row.nd:17,.0f} {row.census:17,.0f} "
453 f"{difference:+15,.0f} {share}")
454 return lines
457def verify_traffic(yaml_path: str) -> List[TrafficRow]:
458 """The bytes each kind of the model a train.yaml trains moves in a forward, beside the FLOPs ND prices.
460 Args:
461 yaml_path: An AutoModels train.yaml, whose model names a Transformers
462 checkpoint.
464 Returns:
465 A row per part of the first layer of each kind of its language model
466 on one sequence of the run's length, the census's feed-forward parts
467 summed as the time model prices them, each kind's whole, then the
468 layers' whole. A part's parameters are the share of its bytes that
469 a strategy shards: on a layer that routes, the feed-forward row's
470 are its experts'.
472 Raises:
473 ValueError: The train.yaml names no Transformers checkpoint.
474 """
475 _, text, ccfg = _priced(yaml_path)
476 rows = []
477 for name, (index, count, lccfg) in _kinds(ccfg).items():
478 flops, traffic = nd_flops(ccfg, lccfg), census_traffic(text, index, int(ccfg.s))
479 moved, parameter, launches = (dict(field) for field in traffic)
480 for part in _FFN_PARTS[1:]:
481 for field in (moved, parameter, launches):
482 field["ffn"] = field.get("ffn", 0) + field.pop(part, 0)
483 where = f"{name} x{count}"
484 rows += [TrafficRow(where, part, moved.get(part, 0.0), parameter.get(part, 0.0),
485 launches.get(part, 0), flops.get(part, 0.0), count)
486 for part in sorted(set(moved) | set(flops)) if moved.get(part) or flops.get(part)]
487 rows.append(TrafficRow(where, "total", sum(moved.values()), sum(parameter.values()),
488 sum(launches.values()), sum(flops.values()), count))
489 rows.append(TrafficRow("layers", "total", *(sum(row.count * getattr(row, field) for row in rows
490 if row.part == "total")
491 for field in ("moved", "parameter", "launches", "flops"))))
492 return rows
495def traffic_report(rows: List[TrafficRow]) -> List[str]:
496 """The lines of a traffic report: each part's bytes, the parameters among them, its ops and ND's FLOPs.
498 The last column is the bytes a part moves for each FLOP ND prices of it,
499 and the spread down it is the reading: a part that moves much and
500 multiplies little is priced near nothing. The FLOPs are the op table's
501 matmul entries, as the FLOPs section reports them, so a part with no
502 matmul of its own shows none.
503 """
504 lines = [f"{'where':26s} {'part':12s} {'bytes moved':>17s} {'parameters':>15s} "
505 f"{'ops':>8s} {'ND FLOPs':>17s} {'bytes/FLOP':>11s}"]
506 for row in rows:
507 ratio = f"{row.moved / row.flops:11.4f}" if row.flops else f"{'':>11s}"
508 lines.append(f"{row.where:26s} {row.part:12s} {row.moved:17,.0f} {row.parameter:15,.0f} "
509 f"{row.launches:8,d} {row.flops:17,.0f} {ratio}")
510 return lines