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

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"""run parallelization""" 

16 

17import argparse 

18import json 

19import os 

20import sys 

21 

22import yaml 

23 

24from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory 

25from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import set_strict 

26from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger, set_verbose_level 

27import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par 

28from hyper_parallel.auto_parallel.sapp_nd.nd import _pool 

29import hyper_parallel.auto_parallel.sapp_nd.nd.debug as Debug 

30import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim 

31import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard 

32import hyper_parallel.auto_parallel.sapp_nd.nd.ratios as Ratios 

33from hyper_parallel.auto_parallel.sapp_nd.nd.recompute_dimension import read_recompute_modes 

34from hyper_parallel.auto_parallel.sapp_nd.recompute.candidate import trainer_runs_layer_plans 

35from hyper_parallel.auto_parallel.sapp_nd.nd.verify import ( 

36 census_spec_yaml, 

37 report, 

38 report_spec, 

39 traffic_report, 

40 verify_activations, 

41 verify_estimate, 

42 verify_flops, 

43 verify_parameters, 

44 verify_spec, 

45 verify_traffic, 

46) 

47 

48 

49def _link_ratios(cli_parser, cache_file): 

50 """The ratios of a ratios file that price the copies: FORWARD, else COMPUTE, for a copy; COMPUTE for its cost. 

51 

52 FORWARD, where a round measured the forward on its own, is the 

53 milliseconds one unit of the forward the estimate prices stands for: the 

54 window the copies to the host hide in. COMPUTE stands for the forward and 

55 the backward together, which the estimate may split otherwise than the 

56 device runs them, and for the score's units, which the copies' cost to 

57 the step is charged in. 

58 

59 Args: 

60 cli_parser: The parser, to report a file that states neither. 

61 cache_file: The ratios file of -c, or ``None``. 

62 

63 Returns: 

64 ``(copy, cost)``, each ``None`` without a file, the second also 

65 where the file states no COMPUTE ratio. 

66 """ 

67 if cache_file is None: 

68 return None, None 

69 with open(cache_file, encoding="utf-8") as handle: 

70 ratios = json.load(handle) 

71 name = "FORWARD" if "FORWARD" in ratios else "COMPUTE" 

72 ratio = ratios.get(name) 

73 if not isinstance(ratio, (int, float)) or ratio <= 0: 

74 cli_parser.error(f"{cache_file} states no positive {name} ratio to calibrate the host link with") 

75 compute = ratios.get("COMPUTE") 

76 return float(ratio), float(compute) if isinstance(compute, (int, float)) and compute > 0 else None 

77 

78 

79def _link_figures(cli_parser, cli_args): 

80 """The host link figures the CLI states for -ao, ``None`` for each it leaves to the device. 

81 

82 With a ratios file (-c), its FORWARD or COMPUTE ratio turns a copy's 

83 seconds into the estimate's units, in place of the device's sustained 

84 throughput, and its COMPUTE ratio the copies' cost (:func:`_link_ratios`). 

85 

86 Args: 

87 cli_parser: The parser, to report a ratios file with neither ratio. 

88 cli_args: The parsed CLI namespace. 

89 

90 Returns: 

91 ``gib_per_s``, ``sustained_tflops``, ``overlap``, ``ms_per_unit``, 

92 ``score_ms_per_unit`` and ``copy_cost_ms_per_gib``. 

93 """ 

94 copy_ratio, cost_ratio = _link_ratios(cli_parser, getattr(cli_args, "cache_file", None)) 

95 return { 

96 "gib_per_s": cli_args.host_link_gibps, 

97 "sustained_tflops": cli_args.sustained_tflops, 

98 "overlap": getattr(cli_args, "host_link_overlap", None), 

99 "ms_per_unit": copy_ratio, 

100 "score_ms_per_unit": cost_ratio, 

101 "copy_cost_ms_per_gib": getattr(cli_args, "host_link_cost", None), 

102 } 

103 

104 

105def _host_link(cli_parser, cli_args, device): 

106 """The host link -ao offloads over: the device's, with what the CLI states replacing its figures. 

107 

108 Args: 

109 cli_parser: The parser, to report a device with no link to offload over. 

110 cli_args: The parsed CLI namespace. 

111 device: The device type the search prices. 

112 

113 Returns: 

114 The link, or ``None`` without -ao. 

115 """ 

116 if not cli_args.auto_offload: 

117 return None 

118 try: 

119 return Hard.HostLink.of(device.host_link, _link_figures(cli_parser, cli_args)) 

120 except ValueError as error: 

121 cli_parser.error(f"device {device}: {error}") 

122 return None 

123 

124 

125def _non_negative_int(value: str) -> int: 

126 """Parse a non-negative integer command-line value.""" 

127 try: 

128 parsed = int(value) 

129 except ValueError as exc: 

130 raise argparse.ArgumentTypeError("must be a non-negative integer") from exc 

131 if parsed < 0: 

132 raise argparse.ArgumentTypeError("must be a non-negative integer") 

133 return parsed 

134 

135 

136def _apply_cli_overrides(search_cfg, cli_args): 

137 """Override the search config's batch, memory budget and devices from the CLI. 

138 

139 Args: 

140 search_cfg: The search config read from ``-s/--search-config``. 

141 cli_args: The parsed CLI namespace. 

142 

143 Raises: 

144 ValueError: When ``-d`` is not a whole number of nodes. 

145 """ 

146 if cli_args.device_type is not None: 

147 # -A sets the device the search prices, as it does on the CLI path; 

148 # only the search config's device_type reached the search before. 

149 search_cfg.cluster_spec["device_type"] = cli_args.device_type 

150 if cli_args.global_batch_size is not None: 

151 search_cfg.constraint["global_batch_size"] = cli_args.global_batch_size 

152 if getattr(cli_args, "auto_recompute", False): 

153 search_cfg.estimator["recompute_strategy"] = "auto" 

154 if cli_args.max_mem is not None: 

155 # -M sets the device budget the search checks against, exactly as it 

156 # does on the CLI path, instead of being silently ignored here. 

157 search_cfg.cluster_spec["device_memory_gb"] = ( 

158 Memory.from_string(cli_args.max_mem.strip()).to_gb().size 

159 ) 

160 if cli_args.devices is not None: 

161 cards_per_node = search_cfg.cluster_spec.get("cards_per_node") 

162 if not cards_per_node: 

163 # The device type knows its node size (A3: 16); defaulting to 8 

164 # silently halves an A3 node and invalidates every candidate. 

165 device = Hard.device_map.get(cli_args.device_type) 

166 cards_per_node = device.intra_node_num() if device else 8 

167 search_cfg.cluster_spec["cards_per_node"] = cards_per_node 

168 logger.info( 

169 "cluster.cards_per_node not set, using %d from device type %s", 

170 cards_per_node, cli_args.device_type, 

171 ) 

172 cards_per_node = max(1, cards_per_node) 

173 nodes = max(1, cli_args.devices // cards_per_node) 

174 if cli_args.devices % cards_per_node: 

175 # A search config counts whole nodes, so a part of one was rounded 

176 # and another cluster than the one asked for priced. 

177 raise ValueError( 

178 f"-d {cli_args.devices} is not a whole number of nodes of {cards_per_node} cards: a search " 

179 f"config counts whole nodes and would price {nodes * cards_per_node} devices. Give -d a " 

180 f"multiple of {cards_per_node}, or price exactly {cli_args.devices} with -y and no -s" 

181 ) 

182 search_cfg.cluster_spec["num_nodes"] = nodes 

183 if getattr(cli_args, "recompute", None) is not None: 

184 # --recompute states the search config's parallelism.recompute, as the 

185 # other flags state their settings: the modes the search chooses 

186 # among for each strategy, or for each layer where the search config's 

187 # recompute says per_layer, and auto every mode it can price. 

188 modes = read_recompute_modes(cli_args.recompute, "--recompute") 

189 if search_cfg.estimator.get("recompute_strategy") != "per_layer": 

190 search_cfg.estimator["recompute_strategy"] = "auto" 

191 if [str(value).strip().lower() for value in cli_args.recompute] == ["auto"]: 

192 search_cfg.estimator.pop("recompute_modes", None) 

193 else: 

194 search_cfg.estimator["recompute_modes"] = modes 

195 if getattr(cli_args, "cache_file", None) is not None: 

196 # -c names the ratios the search scores with, as on the CLI path; 

197 # the search config path dropped it and ranked on the raw score. 

198 search_cfg.estimator["ratios_file"] = cli_args.cache_file 

199 

200 

201def _recompute_modes(cli_args): 

202 """The recompute dimension of a search: --recompute, else a hyper_v2 yaml's context.recompute. 

203 

204 A search config's ``parallelism.recompute`` is not a dimension: there the 

205 search chooses a mode for each strategy, and --recompute states it 

206 (:func:`_apply_cli_overrides`). 

207 

208 Args: 

209 cli_args: The parsed CLI namespace. 

210 

211 Returns: 

212 The modes, or None where nothing states the dimension. 

213 

214 Raises: 

215 ValueError: A value that is neither a mode nor auto. 

216 """ 

217 if cli_args.recompute is not None: 

218 return read_recompute_modes(cli_args.recompute, "--recompute") 

219 if cli_args.framework != "hyper_v2" or cli_args.search_config or not cli_args.yaml_config: 

220 return None 

221 if not os.path.isfile(cli_args.yaml_config): 

222 return None 

223 with open(cli_args.yaml_config, encoding="utf-8") as handle: 

224 raw = yaml.safe_load(handle) 

225 context = raw.get("context") if isinstance(raw, dict) else None 

226 stated = context.get("recompute") if isinstance(context, dict) else None 

227 return read_recompute_modes(stated, f"{cli_args.yaml_config}: context.recompute") 

228 

229 

230def _recompute_kwargs(auto_recompute, modes): 

231 """The search's recompute arguments for the modes --recompute or a yaml's context states. 

232 

233 Without -ar the modes are a dimension, every candidate ranked under each 

234 of them. With it they are the modes each layer of a candidate chooses 

235 among, the choice the trainer runs as activation_checkpoint.layers, as a 

236 search config's ``recompute: per_layer`` makes it. 

237 

238 Args: 

239 auto_recompute: Whether -ar chooses each layer's recompute. 

240 modes: The modes stated, or None. 

241 

242 Returns: 

243 The keyword arguments for Parallelize. 

244 """ 

245 if modes is None: 

246 return {} 

247 if auto_recompute: 

248 return {"recompute_modes": modes, "recompute_mode_per_layer": True} 

249 return {"recompute_dimension": modes} 

250 

251 

252def _check_recompute_potentials(cli_parser, cli_args, modes) -> None: 

253 """Refuse -rp/--recompute_potentials where there is no plan per layer to extend, or nowhere to write it. 

254 

255 Args: 

256 cli_parser: The argument parser, which reports the refusal. 

257 cli_args: The parsed CLI namespace. 

258 modes: The recompute modes stated, or None. 

259 """ 

260 if not cli_args.auto_recompute or modes is None: 

261 cli_parser.error("-rp/--recompute_potentials extends a recompute plan per layer: it needs " 

262 "-ar/--auto_recompute with the modes each layer chooses among, from --recompute or a " 

263 "hyper_v2 yaml's context.recompute") 

264 if cli_args.real_csv is not None or cli_args.search_config or cli_args.verify: 

265 cli_parser.error("-rp/--recompute_potentials applies to the standard ND search, not --real_csv, " 

266 "--search-config, or --verify") 

267 if cli_args.verbosity < 2: 

268 cli_parser.error("-rp/--recompute_potentials writes a CSV and two plots, which a verbosity of at least 2 " 

269 "turns on") 

270 

271 

272def _priced_train_yaml(search_config: str) -> str: 

273 """The train.yaml a search config names, which its search prices. 

274 

275 Args: 

276 search_config: The search config's path. 

277 

278 Returns: 

279 Its ``train_yaml``, or an empty string where it names none, as a 

280 standalone search config does. 

281 """ 

282 with open(search_config, encoding="utf-8") as handle: 

283 raw = yaml.safe_load(handle) 

284 named = raw.get("train_yaml") if isinstance(raw, dict) else None 

285 return named if isinstance(named, str) else "" 

286 

287 

288def _check_train_yaml(cli_parser: argparse.ArgumentParser, cli_args: argparse.Namespace) -> None: 

289 """Refuse a -y that is not the train.yaml the search config prices. 

290 

291 The search prices the train.yaml its config names, and -y is the one 

292 the resolved strategy is written over: two files would put one run's 

293 strategy over another's run. 

294 

295 Args: 

296 cli_parser: The parser, whose ``error()`` exits. 

297 cli_args: The parsed CLI namespace, both files found. 

298 """ 

299 priced = _priced_train_yaml(cli_args.search_config) 

300 if priced and os.path.isfile(priced) and not os.path.samefile(priced, cli_args.yaml_config): 

301 cli_parser.error( 

302 f"-y names {cli_args.yaml_config}, but the search config prices its train_yaml, {priced}: " 

303 "pass the same file" 

304 ) 

305 

306 

307def _compare_with_real_csv(runner, cli_args): 

308 """Print ND's estimate next to the configurations measured in a classified CSV. 

309 

310 Args: 

311 runner: The ND runner built from the CLI arguments. 

312 cli_args: The parsed CLI namespace. Requires ``real_csv``; ND's 

313 real-versus-estimate plot goes to ``output_dir`` when it is set. 

314 """ 

315 if cli_args.output_dir is not None: 

316 os.makedirs(cli_args.output_dir, exist_ok=True) 

317 # Keep debug.csv beside the plot rather than inside the installed package. 

318 Debug.set_output_dir(cli_args.output_dir) 

319 configs_estimated, metrics = runner.compare_with_csv( 

320 cli_args.real_csv, output_path=cli_args.output_dir, plot_idle=True 

321 ) 

322 logger.output("%s", Debug.format_classified_comparison(configs_estimated)) 

323 Debug.print_correlations_classified([metrics]) 

324 if cli_args.write_ratios is not None: 

325 ratios = Ratios.fit_ratios(configs_estimated) 

326 Ratios.write_ratios(cli_args.write_ratios, ratios, configs_estimated, cli_args.real_csv) 

327 for text in Ratios.report(configs_estimated, ratios): 

328 logger.output("%s", text) 

329 logger.output("Ratios written to %s; run_nd -c reads them", cli_args.write_ratios) 

330 

331 

332def _run_hyper_v2_search(cli_parser, cli_args): 

333 """Run the HyperParallel V2 strategy search via ``config_adapter``. 

334 

335 This branch is activated when ``-f hyper_v2`` is combined with 

336 ``-s/--search-config``. It reads the Search Config YAML, validates 

337 it, runs the ND search engine through :func:`search_strategies`, 

338 and writes the resolved strategy back into a copy of the original 

339 ``train.yaml``. 

340 

341 Args: 

342 cli_parser: The :class:`argparse.ArgumentParser` (used for ``error()``). 

343 cli_args: The parsed CLI namespace. Requires ``yaml_config``, 

344 ``search_config``, and optionally ``output_dir``. 

345 

346 Raises: 

347 SystemExit: If validation fails (via ``parser.error``). 

348 """ 

349 # pylint: disable=import-outside-toplevel 

350 from hyper_parallel.auto_parallel.config_adapter import ( 

351 read_search_config, 

352 validate, 

353 search_strategies, 

354 write_resolved_yaml, 

355 ) 

356 

357 if not os.path.isfile(cli_args.search_config): 

358 cli_parser.error(f"search-config not found: {cli_args.search_config}") 

359 if not os.path.isfile(cli_args.yaml_config): 

360 cli_parser.error(f"yaml-config not found: {cli_args.yaml_config}") 

361 _check_train_yaml(cli_parser, cli_args) 

362 

363 set_verbose_level(cli_args.verbosity) 

364 Debug.set_output_dir(cli_args.output_dir, fresh=True) 

365 

366 search_cfg = read_search_config(cli_args.search_config) 

367 try: 

368 _apply_cli_overrides(search_cfg, cli_args) 

369 except ValueError as exc: 

370 cli_parser.error(str(exc)) 

371 offload = None 

372 if getattr(cli_args, "auto_offload", False): 

373 if search_cfg.estimator.get("recompute_strategy") != "per_layer": 

374 cli_parser.error("-ao offloads in a choice of a recompute mode per layer: give the search config " 

375 "recompute: per_layer, and leave out -ar, which chooses one mode for every layer") 

376 offload = _link_figures(cli_parser, cli_args) 

377 

378 if getattr(search_cfg, "parallelism_summary", ""): 

379 logger.output("Parallelism: %s", search_cfg.parallelism_summary) 

380 

381 errors = validate(search_cfg) 

382 hard_errors = [e for e in errors if e.severity == "error"] 

383 warnings = [e for e in errors if e.severity == "warning"] 

384 for w in warnings: 

385 logger.warning("%s: %s", w.field_path, w.message) 

386 if hard_errors: 

387 for e in hard_errors: 

388 logger.error("%s: %s", e.field_path, e.message) 

389 cli_parser.error( 

390 f"Search config validation failed with {len(hard_errors)} error(s)." 

391 ) 

392 

393 result = search_strategies(search_cfg) if offload is None else search_strategies(search_cfg, offload=offload) 

394 search_cfg.resolved_strategy = result 

395 

396 output_dir = cli_args.output_dir or "." 

397 if not os.path.isdir(output_dir): 

398 os.makedirs(output_dir, exist_ok=True) 

399 resolve_path = os.path.join(output_dir, "resolved.yaml") 

400 write_resolved_yaml(search_cfg, cli_args.yaml_config, resolve_path) 

401 logger.output("Resolved strategy written to %s", resolve_path) 

402 if "activation_checkpoint" in result: 

403 _log_activation_checkpoint(result) 

404 logger.output( 

405 "Optimal strategy: dp=%(dp)s tp=%(tp)s pp=%(pp)s " 

406 "cp=%(cp)s ep=%(ep)s mb_num=%(micro_batch_num)s " 

407 "recompute=%(activation_checkpoint)s " 

408 "mem=%(memory_estimate_mb).0f MB score=%(score).2e", 

409 result, 

410 ) 

411 

412 

413def _log_activation_checkpoint(result): 

414 """Log the activation checkpointing a search chose: its mode, the layers that run another, and each layer's own. 

415 

416 A plan per layer the installed trainer cannot run, as it declares no 

417 ``activation_checkpoint.layers``, is logged as advice: the resolved yaml 

418 runs every layer the mode. 

419 

420 Args: 

421 result: The search's result, which states ``activation_checkpoint``. 

422 """ 

423 plan = result.get("activation_checkpoint_layers") 

424 if plan and trainer_runs_layer_plans(): 

425 logger.output("Activation checkpoint mode %s, and per layer: %s", result["activation_checkpoint"], plan) 

426 elif plan: 

427 logger.output("Activation checkpoint mode for every layer: %s, as the resolved yaml states it. The plan per " 

428 "layer %s is advice: the installed trainer's activation_checkpoint has no layers field", 

429 result["activation_checkpoint"], plan) 

430 else: 

431 logger.output("Activation checkpoint mode for every layer: %s", result["activation_checkpoint"]) 

432 if result.get("offloaded_layers"): 

433 moved = result.get("offloaded_gib", {}) 

434 shares = result.get("offloaded_share", {}) 

435 layers = ", ".join( 

436 f"{name} ({moved[name]:.2f} GiB a micro-batch, {shares[name]:.0%} of what it keeps)" 

437 if name in moved and name in shares else name 

438 for name in result["offloaded_layers"] 

439 ) 

440 logger.output("Offloaded to the host, priced as offload should run: layers %s. The trainer runs no offload " 

441 "yet, so the plan keeps them off", layers) 

442 per_layer = result.get("recompute_per_layer") 

443 if per_layer: 

444 logger.output( 

445 "With each layer run its own way, the score would be %.2e at %.0f MB: %s", 

446 per_layer["score"], per_layer["memory_estimate_mb"], per_layer["ranges"], 

447 ) 

448 

449 

450if __name__ == "__main__": 

451 parser = argparse.ArgumentParser( 

452 prog="python run_nd.py", 

453 description=("Provides a degree to *N* parallelism dimensions"), 

454 epilog="", 

455 ) 

456 

457 parser.add_argument( 

458 "-y", 

459 "--yaml_config", 

460 type=str, 

461 required=True, 

462 help="Path to yaml configuration file", 

463 ) 

464 parser.add_argument( 

465 "-f", 

466 "--framework", 

467 default="mindformers", 

468 type=str, 

469 required=False, 

470 help="Framework to evaluate in " 

471 "[mindformers, mindspeed, hyperparallel, hyper_v2, torchtitan]", 

472 ) 

473 parser.add_argument( 

474 "-d", 

475 "--devices", 

476 type=int, 

477 default=None, 

478 help="Number of devices. Takes yaml value if unspecified", 

479 ) 

480 parser.add_argument( 

481 "-b", 

482 "--global_batch_size", 

483 type=int, 

484 default=None, 

485 help="Global batch size. Takes yaml value if unspecified", 

486 ) 

487 parser.add_argument( 

488 "-m", 

489 "--model", 

490 type=str, 

491 default=None, 

492 help="Model Name to use. Takes yaml value if unspecified", 

493 ) 

494 # parser.add_argument( 

495 # "-g", 

496 # "--generate_yaml_in", 

497 # type=str, 

498 # default=None, 

499 # help="Generate all fitting yaml configurations in the given folder", 

500 # ) 

501 # parser.add_argument( 

502 # "-c", 

503 # "--csv", 

504 # type=str, 

505 # default=None, 

506 # help="Computes correlation coefficient from csv results file", 

507 # ) 

508 parser.add_argument( 

509 "-l", 

510 "--dimensions", 

511 nargs="*", 

512 type=str, 

513 default=None, 

514 help="list of varying (output) dimensions", 

515 ) 

516 parser.add_argument( 

517 "-j", 

518 "--jobs", 

519 type=int, 

520 default=None, 

521 help="Worker processes that price the search's candidates; 1 prices them in this one. " 

522 "Default: ND_JOBS, else half the logical cores and at most 4, and 1 above -v 2, whose " 

523 "line per estimate only this process prints. A search with -ar, --recompute, a recompute " 

524 "dimension in its yaml or a multimodal model, every mode but the plain search, and a " 

525 "search of fewer than 2000 candidates price them here", 

526 ) 

527 parser.add_argument( 

528 "-v", 

529 "--verbosity", 

530 type=int, 

531 default=2, 

532 help="Level of verbosity in range [0,6], " 

533 "0 being no output and 6 being debug level output. " 

534 "Plot and debug csv are generated from 2", 

535 ) 

536 parser.add_argument( 

537 "-k", 

538 "--ppb_k", 

539 type=int, 

540 default=None, 

541 help="Write the pipeline balancer's layer description, with per-layer " 

542 "times for every recompute option, of the k-th ranked configuration " 

543 "(0 is the best) to the output directory.", 

544 ) 

545 parser.add_argument( 

546 "-A", 

547 "--device_type", 

548 default=None, 

549 help="choose device type between A2 or A3: A2 unless given, or a search config states one", 

550 ) 

551 parser.add_argument( 

552 "-swap_os", 

553 "--swap_opt_state", 

554 action=argparse.BooleanOptionalAction, 

555 default=False, 

556 help="Activate swap optimiezr state", 

557 ) 

558 # parser.add_argument( 

559 # "-lm", 

560 # "--less_memory", 

561 # action=argparse.BooleanOptionalAction, 

562 # default=False, 

563 # help="Activate less memory schedule", 

564 # ) 

565 parser.add_argument( 

566 "-mppb", 

567 "-–manual_pipeline_balance", 

568 action=argparse.BooleanOptionalAction, 

569 default=False, 

570 help="Takes offset and recompute from yaml", 

571 ) 

572 parser.add_argument( 

573 "-ar", 

574 "--auto_recompute", 

575 action=argparse.BooleanOptionalAction, 

576 default=False, 

577 help="Give every layer of each configuration the fastest recompute " 

578 "option that fits, instead of scoring it fully recomputed. With " 

579 "--recompute, or a hyper_v2 yaml's context.recompute, each layer " 

580 "chooses among those modes, a plan the trainer runs as " 

581 "activation_checkpoint.layers; without, among sets of recomputed ops " 

582 "the trainer does not run", 

583 ) 

584 parser.add_argument( 

585 "-ao", 

586 "--auto_offload", 

587 action=argparse.BooleanOptionalAction, 

588 default=False, 

589 help="With -ar, let each pipeline stage's first layers offload their " 

590 "activations to the host instead of recomputing them, over the " 

591 "device's host link", 

592 ) 

593 parser.add_argument( 

594 "--host_link_gibps", 

595 type=float, 

596 default=None, 

597 help="The host link's sustained copy bandwidth in GiB/s, for -ao, as " 

598 "hyper_offload's profile_transfer_bandwidth measures it; the device's " 

599 "placeholder when omitted", 

600 ) 

601 parser.add_argument( 

602 "--sustained_tflops", 

603 type=float, 

604 default=None, 

605 help="The device's sustained TFLOP/s at the training precision, for " 

606 "-ao; the device's placeholder when omitted, and unused with -c, whose " 

607 "FORWARD or COMPUTE ratio converts a copy's seconds instead", 

608 ) 

609 parser.add_argument( 

610 "--host_link_overlap", 

611 type=float, 

612 default=None, 

613 help="The share of the forward the copies may take, for -ao: 1 where " 

614 "they meet nothing, less where the step's collectives share what they " 

615 "need; the device's 0.8 when omitted", 

616 ) 

617 parser.add_argument( 

618 "--host_link_cost", 

619 type=float, 

620 default=None, 

621 help="What the copies cost the step, in ms a GiB moved either way, for " 

622 "-ao; the device's when omitted (A3: 2.8, measured inside the Demo 2 " 

623 "step), 0 to price offload as free", 

624 ) 

625 parser.add_argument( 

626 "-t", 

627 "--top_config_number", 

628 type=int, 

629 default=None, 

630 help="Number of top configs to print & plot", 

631 ) 

632 parser.add_argument( 

633 "-e", 

634 "--exhaustive", 

635 type=_non_negative_int, 

636 default=0, 

637 metavar="N", 

638 help=( 

639 "Ensure each requested dimension is > 1 in at least N plotted " 

640 "results; default: 0" 

641 ), 

642 ) 

643 parser.add_argument( 

644 "-ee", 

645 "--force_exhaustive", 

646 type=_non_negative_int, 

647 default=0, 

648 metavar="N", 

649 help=( 

650 "When N >= --exhaustive, append N further results per requested " 

651 "dimension with degree > 1; otherwise append at least N and top up " 

652 "to --exhaustive total; default: 0" 

653 ), 

654 ) 

655 parser.add_argument( 

656 "-eee", 

657 "--fforce-exhaustive", 

658 type=_non_negative_int, 

659 default=0, 

660 metavar="N", 

661 help=( 

662 "Append N distinct results per requested dimension with degree > 1; " 

663 "each added result is reserved for one dimension and this option " 

664 "takes precedence over -e/-ee; default: 0" 

665 ), 

666 ) 

667 parser.add_argument( 

668 "-mem", 

669 "--mem_for_ppb", 

670 type=str, 

671 default="0GB", 

672 help="Memory to reserve for pipeline balancing, taken out of the " 

673 "memory budget ND allows (default 0GB).", 

674 ) 

675 parser.add_argument( 

676 "-c", 

677 "--cache_file", 

678 type=str, 

679 default=None, 

680 help="Cache file with ratios to recalibrate ND scores. With -ao its " 

681 "FORWARD ratio, where a round measured the forward on its own, else " 

682 "its COMPUTE ratio, also turns a copy's seconds into the estimate's " 

683 "units, and its COMPUTE ratio the copies' cost; with a search config " 

684 "(-s) that is all it does. " 

685 "Will be defaulted to 'None'.", 

686 ) 

687 

688 parser.add_argument( 

689 "-M", 

690 "--max_mem", 

691 type=str, 

692 default=None, 

693 help="Device memory budget the search must fit in, e.g. '58GB'. " 

694 "Overrides the yaml capacity and cluster.device_memory_gb. " 

695 "To reserve memory instead of capping it, use -mem/--mem_for_ppb.", 

696 ) 

697 parser.add_argument( 

698 "--train-yaml", 

699 type=str, 

700 default=None, 

701 help="Path to training configuration yaml file (for hyperparallel2)", 

702 ) 

703 parser.add_argument( 

704 "--accelerate-yaml", 

705 type=str, 

706 default=None, 

707 help="Path to accelerate configuration yaml file (for hyperparallel2)", 

708 ) 

709 parser.add_argument( 

710 "-s", 

711 "--search-config", 

712 type=str, 

713 default=None, 

714 help="Path to Search Config YAML for fine-grained search-space control " 

715 "(hyper_v2 only). Scalar=fixed, list=candidates, 'auto'=ND decides.", 

716 ) 

717 parser.add_argument( 

718 "-V", 

719 "--verify", 

720 action="store_true", 

721 help="Verify mode (hyper_v2 only): set the parameters and FLOPs ND prices " 

722 "of each part of the model -y trains beside those of the Transformers " 

723 "layers its checkpoint builds, and the model spec the resolver reads " 

724 "beside the one the census measures, with the run priced with each; " 

725 "with -o, write the census's spec there; and exit.", 

726 ) 

727 parser.add_argument( 

728 "-o", 

729 "--output-dir", 

730 type=str, 

731 default=None, 

732 help="Directory for output files when using --search-config " 

733 "(default: current directory), and for the real-versus-estimate " 

734 "plot of --real_csv (no plot when omitted).", 

735 ) 

736 parser.add_argument( 

737 "--real_csv", 

738 type=str, 

739 default=None, 

740 help="Instead of searching, compare ND's estimate with the configurations " 

741 "measured in a classified profiling CSV (see nd.trace_classify).", 

742 ) 

743 parser.add_argument( 

744 "--ranking_csv", 

745 type=str, 

746 default=None, 

747 help="Also write every configuration the search keeps, in ND's order, " 

748 "to this CSV: rank, degrees, memory in MB, score and its parts.", 

749 ) 

750 parser.add_argument( 

751 "--write_ratios", 

752 type=str, 

753 default=None, 

754 help="With --real_csv, fit a ratio per part, measured over ND's estimate, " 

755 "on the configurations the CSV measured, write them to this JSON file " 

756 "(the file -c reads) and print how well they predict each configuration " 

757 "when fitted on the others.", 

758 ) 

759 parser.add_argument( 

760 "--recompute", 

761 nargs="+", 

762 default=None, 

763 metavar="MODE", 

764 help="Search recompute as a dimension: the activation checkpoint modes " 

765 "(off, selective, full) a candidate may run, or auto for all three. " 

766 "Every candidate is priced under each and the pairs are ranked together, " 

767 "so a mode left out is never proposed. Overrides context.recompute of a " 

768 "hyper_v2 yaml and the search config's parallelism.recompute. Without " 

769 "either, every candidate keeps the recompute it derives, full unless " 

770 "-mppb. With -ar, the modes each layer chooses among instead. With " 

771 "--real_csv, the mode of rows that state none.", 

772 ) 

773 parser.add_argument( 

774 "-rp", 

775 "--recompute_potentials", 

776 action="store_true", 

777 help="With -ar and a mode chosen per layer, also report what the cost " 

778 "model estimates each further recomputed layer would cost: for every " 

779 "configuration that fits, the sequence of plans that fully recompute one " 

780 "more layer at the same degrees and memory budget, written to " 

781 "recompute_potentials.csv and to two comparison PDFs beside the results " 

782 "plot. Off by default, and the search does none of that work without it: " 

783 "the report prices far more plans than the search itself and takes " 

784 "accordingly longer. Every figure is the cost model's estimate of the " 

785 "step, which a measured run can disagree with.", 

786 ) 

787 parser.add_argument( 

788 "--strict", 

789 action="store_true", 

790 help="Refuse a config field no parser set where the estimate reads it, " 

791 "instead of pricing it as 0. Without it, a search names the fields it " 

792 "read that way under its ranking.", 

793 ) 

794 

795 args = parser.parse_args() 

796 set_strict(args.strict) 

797 if args.jobs is not None and args.jobs < 1: 

798 parser.error(f"-j/--jobs counts processes, at least 1, got {args.jobs}") 

799 if args.real_csv is not None and not os.path.isfile(args.real_csv): 

800 parser.error(f"real_csv not found: {args.real_csv}") 

801 if args.write_ratios is not None and args.real_csv is None: 

802 parser.error("--write_ratios fits the ratios on a comparison: it needs --real_csv") 

803 try: 

804 recompute_modes = _recompute_modes(args) 

805 except ValueError as exc: 

806 parser.error(str(exc)) 

807 

808 exhaustive_count = args.exhaustive 

809 force_exhaustive = args.force_exhaustive 

810 fforce_exhaustive = args.fforce_exhaustive 

811 if (exhaustive_count or force_exhaustive or fforce_exhaustive) and ( 

812 args.real_csv is not None or args.search_config or args.verify 

813 ): 

814 parser.error( 

815 "--exhaustive, --force_exhaustive, and --fforce-exhaustive apply to " 

816 "the standard ND search, not --real_csv, --search-config, or --verify" 

817 ) 

818 

819 max_mem = ( 

820 Memory.from_string(args.max_mem.strip()) 

821 if args.max_mem is not None 

822 else None 

823 ) 

824 

825 if args.cache_file is not None: 

826 if not os.path.exists(args.cache_file): 

827 logger.error( 

828 f"cache file not found:" 

829 f" {args.cache_file}" 

830 "\nProceeding without cache file..." 

831 ) 

832 args.cache_file = None 

833 

834 if args.auto_recompute and args.mppb: 

835 parser.error("-ar/--auto_recompute chooses the recompute, so it cannot take it from the yaml (-mppb)") 

836 if args.auto_offload and args.search_config and args.framework != "hyper_v2": 

837 parser.error("-ao/--auto_offload with a search config (-s) needs -f hyper_v2, whose search config chooses " 

838 "a recompute mode per layer") 

839 if args.auto_offload and not args.auto_recompute and not args.search_config: 

840 parser.error("-ao/--auto_offload offloads in the choice per layer -ar/--auto_recompute makes") 

841 if not args.auto_offload and (args.host_link_gibps is not None or args.sustained_tflops is not None 

842 or args.host_link_overlap is not None or args.host_link_cost is not None): 

843 parser.error("--host_link_gibps, --sustained_tflops, --host_link_overlap and --host_link_cost price " 

844 "offload, which needs -ao/--auto_offload") 

845 if args.recompute_potentials: 

846 _check_recompute_potentials(parser, args, recompute_modes) 

847 if args.verify: 

848 if args.framework != "hyper_v2": 

849 parser.error("-V/--verify needs -f hyper_v2: it builds the Transformers checkpoint -y names") 

850 set_verbose_level(args.verbosity) 

851 logger.output("Parameters") 

852 for line in report(verify_parameters(args.yaml_config)): 

853 logger.output(line) 

854 logger.output("Forward FLOPs of one sequence; the time model prices the backward at twice them") 

855 for line in report(verify_flops(args.yaml_config)): 

856 logger.output(line) 

857 logger.output("Activations a layer keeps for its backward, bytes a token by op: the records' and the census's") 

858 for line in report(verify_activations(args.yaml_config)): 

859 logger.output(line) 

860 logger.output("Model spec as ND prices it: read from the checkpoint's config, and measured by the census") 

861 for line in report_spec(verify_spec(args.yaml_config)): 

862 logger.output(line) 

863 logger.output("The run priced with each spec: the resolver's in the ND column, the census's beside it") 

864 for line in report(verify_estimate(args.yaml_config, args.device_type or "A2")): 

865 logger.output(line) 

866 if args.output_dir: 

867 logger.output(f"census spec written to {census_spec_yaml(args.yaml_config, args.output_dir)}") 

868 logger.output("Bytes a forward of one sequence moves, the whole layer, beside the FLOPs the time model " 

869 "prices of it") 

870 for line in traffic_report(verify_traffic(args.yaml_config)): 

871 logger.output(line) 

872 sys.exit(0) 

873 

874 if args.framework == "hyper_v2" and args.search_config: 

875 _run_hyper_v2_search(parser, args) 

876 sys.exit(0) 

877 

878 if args.framework == "hyper_v2" and args.devices is None: 

879 # An AutoModels train.yaml carries no world size: the runtime derives 

880 # the data-parallel replicate degree from it at launch. Without -d the 

881 # cluster would be inferred as d*t*cp*p, which understates HSDP runs 

882 # and silently invalidates every candidate in the search. 

883 parser.error( 

884 "-d/--devices is required for hyper_v2: the device count is not " 

885 "expressible in an AutoModels train.yaml. Alternatively set " 

886 "context.device_num in the config." 

887 ) 

888 

889 set_verbose_level(args.verbosity) 

890 Debug.set_output_dir(args.output_dir, fresh=True) 

891 dims = Dim.get_dims(args.dimensions) 

892 YAML_FOLDER = None # args.generate_yaml_in 

893 machine = Hard.Machine(args.devices, args.device_type or "A2") 

894 host_link = _host_link(parser, args, machine.device) 

895 

896 if args.framework == "hyperparallel2": 

897 if args.yaml_config is None or args.train_yaml is None or args.accelerate_yaml is None: 

898 parser.error("-y (model yaml), --train-yaml, and --accelerate-yaml are required for hyperparallel2") 

899 input_config = { 

900 "model": args.yaml_config, 

901 "train": args.train_yaml, 

902 "accelerate": args.accelerate_yaml, 

903 "machine": args.devices 

904 } 

905 elif args.framework == "torchtitan":