Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _strategy_output.py: 90%

210 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"""Strategy output writer for auto parallel strategy search. 

16 

17Serializes :class:`NormalizedConfig` instances to JSON or YAML files, 

18generates human-readable summaries, and provides a PPB (PR524) 

19configuration stub. 

20""" 

21 

22import json 

23import logging 

24import os 

25from typing import Any, Dict, Optional 

26 

27import yaml # type: ignore[import-untyped] 

28 

29from hyper_parallel.auto_parallel._hf_model_spec import is_auto_models_schema 

30from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig 

31 

32 

33logger = logging.getLogger(__name__) 

34 

35 

36def _resolve_pp_degree(pp_cfg: Dict) -> int: 

37 """Resolve pp_degree to a scalar from pp_config.""" 

38 pp_degree_val = pp_cfg.get("pp_degree", 1) 

39 if isinstance(pp_degree_val, list): 

40 return pp_degree_val[0] if pp_degree_val else 1 

41 return pp_degree_val 

42 

43 

44def _write_json(data: Dict[str, Any], output_path: str) -> None: 

45 """Write a dict as a pretty-printed JSON file.""" 

46 with open(output_path, "w", encoding="utf-8") as fh: 

47 json.dump(data, fh, indent=2, default=str) 

48 fh.write("\n") 

49 

50 

51# Mapping from resolved strategy keys to HP YAML accelerator field names. 

52_YAML_KEY_MAP: Dict[str, str] = { 

53 "dp_shard": "dp_shard", 

54 "dp_replicate": "dp_replicate", 

55 "data_parallel_shard_degree": "dp_shard", 

56 "data_parallel_replicate_degree": "dp_replicate", 

57 "data_parallel_degree": "dp_replicate", 

58 "dp": "dp_replicate", 

59 "tensor_parallel_degree": "tp_degree", 

60 "tp_degree": "tp_degree", 

61 "tp": "tp_degree", 

62 "pipeline_parallel_degree": "pipeline_parallel_degree", 

63 "context_parallel_degree": "context_parallel_degree", 

64 "expert_parallel_degree": "expert_parallel_degree", 

65 "expert_tensor_parallel_degree": "expert_tensor_parallel_degree", 

66 "pp_degree": "pipeline_parallel_degree", 

67 "pp": "pipeline_parallel_degree", 

68 "cp_degree": "context_parallel_degree", 

69 "cp": "context_parallel_degree", 

70 "ep_degree": "expert_parallel_degree", 

71 "ep": "expert_parallel_degree", 

72 "etp_degree": "expert_tensor_parallel_degree", 

73 "etp": "expert_tensor_parallel_degree", 

74} 

75 

76_AUTO_MODELS_YAML_KEY_MAP: Dict[str, str] = { 

77 "tensor_parallel_degree": "tp_size", 

78 "tp_degree": "tp_size", 

79 "tp": "tp_size", 

80 "pipeline_parallel_degree": "pp_size", 

81 "pp_degree": "pp_size", 

82 "pp": "pp_size", 

83 "context_parallel_degree": "cp_size", 

84 "cp_degree": "cp_size", 

85 "cp": "cp_size", 

86 "expert_parallel_degree": "ep_size", 

87 "ep_degree": "ep_size", 

88 "ep": "ep_size", 

89} 

90 

91 

92def _validate_strategy_and_yaml( 

93 config: NormalizedConfig, 

94 original_yaml_path: str, 

95) -> None: 

96 """Validate that resolved_strategy is set and the original YAML exists.""" 

97 if config.resolved_strategy is None: 

98 raise ValueError( 

99 "config.resolved_strategy is None — no strategy to write. " 

100 "Set config.resolved_strategy first." 

101 ) 

102 if not os.path.isfile(original_yaml_path): 

103 raise FileNotFoundError( 

104 f"Original YAML not found: {original_yaml_path}" 

105 ) 

106 

107 

108def _load_yaml_to_inject(original_yaml_path: str) -> Dict[str, Any]: 

109 """Load the original YAML and initialize its strategy sections.""" 

110 with open(original_yaml_path, "r", encoding="utf-8") as fh: 

111 data = yaml.safe_load(fh) 

112 if data is None or not isinstance(data, dict): 

113 raise ValueError( 

114 f"Original YAML {original_yaml_path} must contain a top-level mapping." 

115 ) 

116 if is_auto_models_schema(data): 

117 if not isinstance(data.get("training"), dict): 

118 data["training"] = {} 

119 if not isinstance(data.get("accelerator"), dict): 

120 data["accelerator"] = {} 

121 if not isinstance(data.get("fsdp_config"), dict): 

122 data["fsdp_config"] = {} 

123 else: 

124 if "train" not in data or not isinstance(data["train"], dict): 

125 data["train"] = {} 

126 if not isinstance(data["train"].get("accelerator"), dict): 

127 data["train"]["accelerator"] = {} 

128 return data 

129 

130 

131def _check_batch_derivation(data: Dict[str, Any], resolved: Dict[str, Any]) -> None: 

132 """Verify the resolved strategy matches the Trainer's batch derivation. 

133 

134 AutoModels has no micro-batch-number field: it derives the accumulation 

135 count as ``global_batch_size / (micro_batch_size * dp)``. ND applies the 

136 same identity, so the searched value needs no writing, but a mismatch 

137 would silently train a different schedule than the one that was scored. 

138 """ 

139 micro_batch_num = resolved.get("micro_batch_num") 

140 data_parallel = resolved.get("dp") 

141 training = data.get("training", {}) 

142 global_batch_size = training.get("global_batch_size") 

143 micro_batch_size = training.get("micro_batch_size", 1) 

144 if not (micro_batch_num and data_parallel and global_batch_size): 

145 return 

146 expected = micro_batch_num * micro_batch_size * data_parallel 

147 if expected != global_batch_size: 

148 raise ValueError( 

149 "resolved strategy is inconsistent with the trainer batch " 

150 f"derivation: global_batch_size={global_batch_size} != " 

151 f"micro_batch_num={micro_batch_num} * " 

152 f"micro_batch_size={micro_batch_size} * dp={data_parallel}" 

153 ) 

154 

155 

156def _inject_activation_checkpoint(data: Dict[str, Any], mode: str, 

157 layers: Optional[Dict[str, str]] = None) -> None: 

158 """Set the activation checkpoint mode the search chose for every layer. 

159 

160 The AutoModels schema states it as ``activation_checkpoint.mode`` 

161 (``off``, ``full`` or ``selective``), the older one as 

162 ``train.gradient_checkpointing.activation_checkpoint``, where ``off`` 

163 is spelled ``none``. The layers *layers* names run another mode, which 

164 the AutoModels schema states as ``activation_checkpoint.layers`` where 

165 the installed trainer declares that field; a plan the train yaml states 

166 gives way to the search's, which it was not priced with. A trainer that 

167 declares no such field refuses it, so there the yaml runs every layer 

168 *mode*, the mode the plan runs most layers of those that recompute, and 

169 the plan is logged as advice. 

170 

171 Raises: 

172 ValueError: For *layers* in the older schema, which has no field 

173 for them. 

174 """ 

175 if is_auto_models_schema(data): 

176 section = data.get("activation_checkpoint") 

177 if not isinstance(section, dict): 

178 section = data["activation_checkpoint"] = {} 

179 before = section.get("mode", "off") 

180 section["mode"] = mode 

181 stated = section.pop("layers", None) 

182 written = _runnable_plan(mode, layers) 

183 if written: 

184 section["layers"] = dict(written) 

185 if stated and stated != written: 

186 logger.info("activation_checkpoint.layers %s in the train yaml, %s in the resolved yaml", 

187 stated, written or "none") 

188 else: 

189 if layers: 

190 raise ValueError( 

191 "the search chose a mode per layer, which only the AutoModels schema's activation_checkpoint.layers " 

192 "states; the train yaml uses train.gradient_checkpointing" 

193 ) 

194 checkpointing = data["train"].get("gradient_checkpointing") 

195 if not isinstance(checkpointing, dict): 

196 checkpointing = data["train"]["gradient_checkpointing"] = {} 

197 before = checkpointing.get("activation_checkpoint", "none") 

198 checkpointing["activation_checkpoint"] = "none" if mode == "off" else mode 

199 if {"none": "off"}.get(before, before) != mode: 

200 logger.info( 

201 "activation checkpoint mode %s in the train yaml, %s in the search: writing the searched mode", 

202 before, mode, 

203 ) 

204 

205 

206def _trainer_runs_layer_plans() -> bool: 

207 """Whether the installed trainer runs a recompute mode per layer (:func:`trainer_runs_layer_plans`). 

208 

209 ND's recompute module loads the memory model, as long to import as the 

210 rest of this package, so it loads only where a plan per layer is written. 

211 """ 

212 # pylint: disable=import-outside-toplevel 

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

214 return trainer_runs_layer_plans() 

215 

216 

217def _runnable_plan(mode: str, layers: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]: 

218 """*layers*, the search's plan per layer, where the installed trainer runs one; else None, with a warning. 

219 

220 A trainer whose ``activation_checkpoint`` declares no ``layers`` refuses 

221 a yaml stating them, so it runs every layer *mode*, and the plan is 

222 advice it cannot run. 

223 """ 

224 if not layers or _trainer_runs_layer_plans(): 

225 return layers or None 

226 logger.warning("the search chose activation_checkpoint mode %s with layers %s, and the installed trainer's " 

227 "activation_checkpoint has no layers field: the resolved yaml runs every layer %s, and the " 

228 "plan per layer is advice the trainer cannot run", mode, layers, mode) 

229 return None 

230 

231 

232def _drop_cost_model_context(data: Dict[str, Any]) -> Dict[str, Any]: 

233 """Leave out the cost model's own ``context`` section, which the trainer's AutoModels schema refuses. 

234 

235 A train.yaml the search reads states its pricing options there, such as 

236 ``context.census``; the trainer has no such section and refuses the 

237 whole file over it. 

238 

239 Returns: 

240 The section left out, empty when there was none: some of its options 

241 decide what else the resolved yaml must state. 

242 """ 

243 if not is_auto_models_schema(data): 

244 return {} 

245 context = data.pop("context", None) 

246 if context is not None: 

247 logger.info("context states how the cost model prices the run, and the trainer refuses it: " 

248 "left out of the resolved yaml") 

249 return context if isinstance(context, dict) else {} 

250 

251 

252def _inject_expert_shard(data: Dict[str, Any], resolved: Dict[str, Any]) -> None: 

253 """Shard each routed expert over the whole expert data-parallel group, as the search priced it. 

254 

255 A train.yaml stating ``context.expert_shard: group`` is priced with each 

256 strategy's routed experts sharded, under expert parallelism, over its 

257 whole expert data-parallel group: the stage's DP x CP x TP ranks over EP, 

258 as the sweep launches it. The trainer shards an expert over 

259 ``fsdp_config.edp_shard_size`` ranks of that group and replicates it over 

260 the rest, so the value the train.yaml states for its own EP would launch 

261 another sharding than the one priced. Without expert parallelism the 

262 trainer builds no expert mesh, and the stated value is left as it is. 

263 """ 

264 accelerator = data["accelerator"] 

265 expert_parallel = int(accelerator.get("ep_size") or 1) 

266 if expert_parallel <= 1: 

267 return 

268 if "dp" not in resolved: 

269 logger.warning("context.expert_shard is group but the search states no DP: " 

270 "fsdp_config.edp_shard_size left as the train yaml states it") 

271 return 

272 ranks = int(resolved["dp"]) * int(accelerator.get("cp_size") or 1) * int(accelerator.get("tp_size") or 1) 

273 group = ranks // expert_parallel 

274 stated = data["fsdp_config"].get("edp_shard_size") 

275 if stated is not None and int(stated) != group: 

276 logger.info("edp_shard_size %s in the train yaml, the whole expert group of %d in the search: " 

277 "writing the group", stated, group) 

278 data["fsdp_config"]["edp_shard_size"] = group 

279 

280 

281def _inject_resolved_strategy(data: Dict[str, Any], resolved: Dict[str, Any]) -> None: 

282 """Inject resolved strategy values into the YAML data dict.""" 

283 if resolved.get("activation_checkpoint"): 

284 _inject_activation_checkpoint(data, resolved["activation_checkpoint"], 

285 resolved.get("activation_checkpoint_layers")) 

286 if is_auto_models_schema(data): 

287 accelerator = data["accelerator"] 

288 for src_key, dst_key in _AUTO_MODELS_YAML_KEY_MAP.items(): 

289 if src_key in resolved: 

290 accelerator[dst_key] = int(resolved[src_key]) 

291 

292 dp_shard_size = resolved.get("dp_shard") 

293 if dp_shard_size is None: 

294 dp_shard_size = resolved.get( 

295 "data_parallel_shard_degree", 

296 resolved.get("dp"), 

297 ) 

298 if dp_shard_size is not None: 

299 data["fsdp_config"]["dp_shard_size"] = int(dp_shard_size) 

300 if "global_batch_size" in resolved: 

301 data["training"]["global_batch_size"] = int( 

302 resolved["global_batch_size"] 

303 ) 

304 _check_batch_derivation(data, resolved) 

305 return 

306 

307 train = data["train"] 

308 accel = train["accelerator"] 

309 

310 for src_key, dst_key in _YAML_KEY_MAP.items(): 

311 if src_key in resolved: 

312 accel[dst_key] = int(resolved[src_key]) 

313 

314 if "global_batch_size" in resolved: 

315 train["global_batch_size"] = int(resolved["global_batch_size"]) 

316 

317 if "micro_batch_num" in resolved: 

318 train["micro_batch_num"] = int(resolved["micro_batch_num"]) 

319 

320 

321def _write_output_yaml( 

322 data: Dict[str, Any], 

323 output_path: str, 

324 overwrite: bool, 

325 original_yaml_path: str, 

326) -> None: 

327 """Write the data dict to the output YAML file.""" 

328 write_path = original_yaml_path if overwrite else output_path 

329 

330 parent_dir = os.path.dirname(os.path.abspath(write_path)) 

331 if parent_dir and not os.path.isdir(parent_dir): 

332 os.makedirs(parent_dir, exist_ok=True) 

333 

334 with open(write_path, "w", encoding="utf-8") as fh: 

335 yaml.dump(data, fh, default_flow_style=False, sort_keys=False) 

336 

337 logger.info( 

338 "Resolved YAML written to %s (overwrite=%s)", 

339 write_path, overwrite, 

340 ) 

341 

342 

343def write_strategy_config( 

344 config: NormalizedConfig, 

345 output_path: str, 

346 fmt: str = "json", 

347) -> None: 

348 """Write a normalized configuration to a file. 

349 

350 Args: 

351 config: The normalized configuration to serialize. 

352 output_path: Destination file path. 

353 fmt: Output format, ``"json"`` (default) only. 

354 

355 Raises: 

356 ValueError: If the output format is unsupported. 

357 OSError: If the file cannot be written. 

358 """ 

359 data = config.to_dict() 

360 

361 parent_dir = os.path.dirname(os.path.abspath(output_path)) 

362 if parent_dir and not os.path.isdir(parent_dir): 

363 os.makedirs(parent_dir, exist_ok=True) 

364 

365 try: 

366 if fmt == "json": 

367 _write_json(data, output_path) 

368 else: 

369 raise ValueError( 

370 f"Unsupported output format: {fmt!r}. " 

371 "Supported format: 'json'" 

372 ) 

373 except OSError as exc: 

374 raise OSError(f"Failed to write config to {output_path}: {exc}") from exc 

375 

376 logger.info("Strategy config written to %s", output_path) 

377 

378 

379def write_resolved_strategy( 

380 config: NormalizedConfig, 

381 output_path: str, 

382 fmt: str = "json", 

383) -> None: 

384 """Write only the resolved strategy to a file. 

385 

386 This produces a compact config snippet suitable for consumption 

387 by training scripts. 

388 

389 Args: 

390 config: The normalized configuration (must have ``resolved_strategy`` set). 

391 output_path: Destination file path. 

392 fmt: Output format, ``"json"`` (default) only. 

393 

394 Raises: 

395 ValueError: If ``resolved_strategy`` is ``None``. 

396 """ 

397 if config.resolved_strategy is None: 

398 raise ValueError("config.resolved_strategy is None — no strategy to write") 

399 

400 parent_dir = os.path.dirname(os.path.abspath(output_path)) 

401 if parent_dir and not os.path.isdir(parent_dir): 

402 os.makedirs(parent_dir, exist_ok=True) 

403 

404 data = {"resolved_strategy": dict(config.resolved_strategy)} 

405 

406 try: 

407 if fmt == "json": 

408 _write_json(data, output_path) 

409 else: 

410 raise ValueError( 

411 f"Unsupported output format: {fmt!r}. " 

412 "Supported format: 'json'" 

413 ) 

414 except OSError as exc: 

415 raise OSError(f"Failed to write resolved strategy to {output_path}: {exc}") from exc 

416 

417 logger.info("Resolved strategy written to %s", output_path) 

418 

419 

420def write_ppb_config( 

421 config: NormalizedConfig, 

422 output_path: str, 

423) -> None: 

424 """Write a PPB (PR524 pipeline balancer) input JSON file. 

425 

426 Writes the PP-relevant fields from ``NormalizedConfig`` 

427 in a structure compatible with ``args_for_pipeline_parallel.json``. 

428 

429 Args: 

430 config: The normalized configuration. 

431 output_path: Destination JSON file path. 

432 

433 Raises: 

434 OSError: If the file cannot be written. 

435 """ 

436 model = config.model_spec 

437 pp_cfg = config.pp_config 

438 constraint = config.constraint 

439 

440 ppb_data = { 

441 "llm_class": "0", 

442 "train_yaml": "", 

443 "mindformers_dir": "", 

444 "layer_ratio": 0.33, 

445 "backward_ratio": 2.0, 

446 "head_loss": 1.5, 

447 "recompute_ratio": 1, 

448 "time_limit": 2147483647, 

449 "dryrun": True, 

450 "check": True, 

451 "fit_level": 0, 

452 "extract": False, 

453 "env_json": "./config/env_config.json", 

454 "dryrun_lim": 16, 

455 "_hyper_model": { 

456 "num_hidden_layers": model.get("num_hidden_layers", 0), 

457 "hidden_size": model.get("hidden_size", 0), 

458 "num_attention_heads": model.get("num_attention_heads", 0), 

459 "vocab_size": model.get("vocab_size", 0), 

460 "max_position_embeddings": model.get("max_position_embeddings", 4096), 

461 "moe_enabled": model.get("moe_enabled", False), 

462 "num_experts": model.get("num_experts", 1), 

463 }, 

464 "_hyper_pp": { 

465 "pp_degree": _resolve_pp_degree(pp_cfg), 

466 "stage_partition_mode": pp_cfg.get("stage_partition_mode", "uniform"), 

467 "micro_batch_num": pp_cfg.get("micro_batch_num", 1), 

468 "layer_offset_range": list(pp_cfg.get("layer_offset_range", (0, 0))), 

469 "layer_recompute_layers": pp_cfg.get("layer_recompute_layers", []), 

470 }, 

471 "_hyper_constraint": { 

472 "global_batch_size": constraint.get("global_batch_size", 0), 

473 "memory_limit_gb": constraint.get("memory_limit_gb", 0.0), 

474 }, 

475 } 

476 

477 parent_dir = os.path.dirname(os.path.abspath(output_path)) 

478 if parent_dir and not os.path.isdir(parent_dir): 

479 os.makedirs(parent_dir, exist_ok=True) 

480 

481 try: 

482 with open(output_path, "w", encoding="utf-8") as fh: 

483 json.dump(ppb_data, fh, indent=4) 

484 fh.write("\n") 

485 except OSError as exc: 

486 raise OSError(f"Failed to write PPB config to {output_path}: {exc}") from exc 

487 

488 logger.info("PPB config stub written to %s", output_path) 

489 

490 

491def _apply_searched_batch_size(data: Dict[str, Any], config: NormalizedConfig) -> None: 

492 """Write the global batch size the search actually optimised for. 

493 

494 ``-b`` and ``constraint.global_batch_size`` override the train.yaml for 

495 the search, so leaving the original value in place would either emit a 

496 config the trainer derives a different schedule from, or fail the batch 

497 check below with two numbers and no hint of where either came from. 

498 """ 

499 searched = int(config.constraint.get("global_batch_size", 0) or 0) 

500 if not searched: 

501 return 

502 section = data.get("training") if is_auto_models_schema(data) else data.get("train") 

503 if not isinstance(section, dict): 

504 return 

505 # Only when the strategy really was scored at that batch size; otherwise 

506 # the check below is the right place for the disagreement to surface. 

507 resolved = config.resolved_strategy or {} 

508 implied = ( 

509 int(resolved.get("micro_batch_num", 0) or 0) 

510 * int(section.get("micro_batch_size", 1) or 1) 

511 * int(resolved.get("dp", 0) or 0) 

512 ) 

513 if implied != searched: 

514 return 

515 current = section.get("global_batch_size") 

516 if current is not None and int(current) != searched: 

517 logger.info( 

518 "global_batch_size %s in the train yaml, %s in the search: " 

519 "writing the searched value", current, searched, 

520 ) 

521 section["global_batch_size"] = searched 

522 

523 

524def write_resolved_yaml( 

525 config: NormalizedConfig, 

526 original_yaml_path: str, 

527 output_path: str, 

528 overwrite: bool = False, 

529) -> None: 

530 """Write a resolved strategy into a HyperParallel YAML file. 

531 

532 Copies the original ``train.yaml`` and replaces the parallel 

533 dimension fields with the resolved strategy values. This produces 

534 a complete, immediately launchable training configuration: an 

535 AutoModels one leaves out the cost model's ``context`` section, which 

536 the trainer refuses. 

537 

538 The resolved strategy is read from ``config.resolved_strategy``. 

539 Supported keys: ``dp_shard``, ``dp_replicate``, ``tp_degree``, 

540 ``pipeline_parallel_degree``, ``context_parallel_degree``, 

541 ``expert_parallel_degree``, ``global_batch_size``, 

542 ``activation_checkpoint``, the mode every layer runs but those 

543 ``activation_checkpoint_layers`` gives another, where the installed 

544 trainer declares ``activation_checkpoint.layers``. An AutoModels 

545 train.yaml stating ``context.expert_shard: group`` also gets the 

546 ``fsdp_config.edp_shard_size`` its experts were priced with. 

547 

548 Args: 

549 config: NormalizedConfig with ``resolved_strategy`` set. 

550 original_yaml_path: Path to the original ``train.yaml``. 

551 output_path: Destination file path for the resolved YAML. 

552 overwrite: If ``True``, write to ``original_yaml_path`` instead 

553 of ``output_path`` (default ``False``). 

554 

555 Raises: 

556 ValueError: If ``resolved_strategy`` is ``None``. 

557 FileNotFoundError: If ``original_yaml_path`` does not exist. 

558 """ 

559 _validate_strategy_and_yaml(config, original_yaml_path) 

560 data = _load_yaml_to_inject(original_yaml_path) 

561 context = _drop_cost_model_context(data) 

562 _apply_searched_batch_size(data, config) 

563 _inject_resolved_strategy(data, config.resolved_strategy) 

564 if context.get("expert_shard") == "group": 

565 _inject_expert_shard(data, config.resolved_strategy) 

566 _write_output_yaml(data, output_path, overwrite, original_yaml_path) 

567 

568 

569def normalized_to_summary(config: NormalizedConfig) -> Dict[str, Any]: 

570 """Generate a human-readable summary of the configuration. 

571 

572 Args: 

573 config: The normalized configuration to summarize. 

574 

575 Returns: 

576 A dictionary with summary fields suitable for logging or display. 

577 """ 

578 model = config.model_spec 

579 cluster = config.cluster_spec 

580 search = config.search_space 

581 constraint = config.constraint 

582 estimator = config.estimator 

583 pp_cfg = config.pp_config 

584 

585 return { 

586 "model": { 

587 "name": model.get("name", "unknown"), 

588 "num_hidden_layers": model.get("num_hidden_layers", 0), 

589 "hidden_size": model.get("hidden_size", 0), 

590 "intermediate_size": model.get("intermediate_size", 0), 

591 "num_attention_heads": model.get("num_attention_heads", 0), 

592 "num_key_value_heads": model.get("num_key_value_heads", 0), 

593 "vocab_size": model.get("vocab_size", 0), 

594 "max_position_embeddings": model.get("max_position_embeddings", 0), 

595 "moe_enabled": model.get("moe_enabled", False), 

596 "num_experts": model.get("num_experts", 0), 

597 }, 

598 "cluster": { 

599 "num_nodes": cluster.get("num_nodes", 0), 

600 "cards_per_node": cluster.get("cards_per_node", 0), 

601 "total_cards": ( 

602 cluster.get("num_nodes", 0) * cluster.get("cards_per_node", 0) 

603 ), 

604 "device_memory_gb": cluster.get("device_memory_gb", 0), 

605 "device_type": cluster.get("device_type", "unknown"), 

606 }, 

607 "search_space": dict(sorted(search.items())), 

608 "constraints": { 

609 "global_batch_size": constraint.get("global_batch_size", 0), 

610 "memory_limit_gb": constraint.get("memory_limit_gb", 0.0), 

611 "fixed_dimensions": { 

612 k: constraint.get(f"fixed_{k}_degree") 

613 for k in ("dp", "tp", "pp", "cp", "ep") 

614 if constraint.get(f"fixed_{k}_degree") is not None 

615 }, 

616 }, 

617 "estimator": { 

618 "type": estimator.get("type", "symbolic"), 

619 "recompute_strategy": estimator.get("recompute_strategy", "none"), 

620 }, 

621 "pipeline": { 

622 "pp_degree": _resolve_pp_degree(pp_cfg), 

623 "stage_partition_mode": pp_cfg.get("stage_partition_mode", "uniform"), 

624 "micro_batch_num": pp_cfg.get("micro_batch_num", 1), 

625 }, 

626 "resolved_strategy": config.resolved_strategy, 

627 }