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

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. 

16 

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. 

25 

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. 

32 

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`). 

38 

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`). 

44 

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 

55 

56import copy 

57import os 

58import tempfile 

59from types import SimpleNamespace 

60from typing import Any, Dict, List, Mapping, NamedTuple, Tuple 

61 

62import yaml 

63 

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 

94 

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

97 

98# The census's parts of a feed-forward, which the time model prices as one. 

99_FFN_PARTS = ("ffn", "routed", "shared", "router") 

100 

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

104 

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

109 

110 

111class VerifyRow(NamedTuple): 

112 """One part's parameters: ND's count, the census's, and how many layers hold it.""" 

113 

114 where: str 

115 part: str 

116 nd: float 

117 census: float 

118 count: int = 1 

119 

120 

121class SpecRow(NamedTuple): 

122 """One field of the model spec, as ND prices it: the resolver's value and the census's.""" 

123 

124 field: str 

125 resolved: Any 

126 census: Any 

127 

128 

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

131 

132 where: str 

133 part: str 

134 moved: float 

135 parameter: float 

136 launches: int 

137 flops: float 

138 count: int = 1 

139 

140 

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 

151 

152 

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. 

155 

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} 

165 

166 

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. 

169 

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} 

177 

178 

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]] 

184 

185 

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. 

188 

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. 

191 

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 

199 

200 

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 

210 

211 

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. 

214 

215 Args: 

216 yaml_path: An AutoModels train.yaml, whose model names a Transformers 

217 checkpoint. 

218 

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. 

224 

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 

244 

245 

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. 

248 

249 Args: 

250 yaml_path: An AutoModels train.yaml, whose model names a Transformers 

251 checkpoint. 

252 

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. 

259 

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 

271 

272 

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. 

275 

276 Args: 

277 yaml_path: An AutoModels train.yaml, whose model names a Transformers 

278 checkpoint. 

279 

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. 

284 

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 

298 

299 

300def spec_as_priced(data: Mapping[str, Any]) -> Dict[str, Any]: 

301 """The fields of the model spec *data*, each as ND prices it. 

302 

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. 

309 

310 Args: 

311 data: A model spec, as a resolver states it. 

312 

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 

329 

330 

331def _train_yaml(yaml_path: str) -> Dict[str, Any]: 

332 """The AutoModels train.yaml at *yaml_path*, whose model names a Transformers checkpoint. 

333 

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 

342 

343 

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. 

346 

347 Args: 

348 yaml_path: An AutoModels train.yaml, whose model names a Transformers 

349 checkpoint. 

350 

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. 

354 

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))] 

362 

363 

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. 

366 

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. 

371 

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 

385 

386 

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])) 

392 

393 

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. 

396 

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. 

401 

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. 

405 

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])] 

419 

420 

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) 

429 

430 

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 

444 

445 

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 

455 

456 

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. 

459 

460 Args: 

461 yaml_path: An AutoModels train.yaml, whose model names a Transformers 

462 checkpoint. 

463 

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'. 

471 

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 

493 

494 

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. 

497 

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