Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / case_compare.py: 85%

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"""Comparators across validation-case evidence directories.""" 

16# pylint: disable=forbidden-backend-import 

17 

18from __future__ import annotations 

19 

20import json 

21import math 

22import statistics 

23from pathlib import Path 

24from typing import Any, Iterable, Mapping, Optional 

25 

26import torch 

27 

28from hyper_parallel.tools.model_integration.schemas import tolerance_value 

29 

30 

31def _optional_tolerance(tolerance: Mapping[str, Any], name: str) -> Optional[float]: 

32 """Return one declared tolerance without inventing an absent zero bound.""" 

33 value = tolerance.get(name) 

34 return None if value is None else float(value) 

35 

36 

37def _relative_error(reference: float, candidate: float) -> float: 

38 """Return reference-relative error with a finite zero-safe denominator.""" 

39 return abs(candidate - reference) / max(abs(reference), 1.0e-30) 

40 

41 

42def _combine_tolerance_checks( 

43 checks: Iterable[bool], 

44 combination: Any, 

45 *, 

46 owner: str, 

47) -> bool: 

48 """Combine explicitly declared numerical checks with validated semantics.""" 

49 if combination not in ("all", "any"): 

50 raise ValueError( 

51 f"{owner} tolerance combination must be 'all' or 'any', got " 

52 f"{combination!r}" 

53 ) 

54 resolved_checks = tuple(checks) 

55 if not resolved_checks: 

56 raise ValueError(f"{owner} tolerance must declare at least one numerical bound") 

57 return all(resolved_checks) if combination == "all" else any(resolved_checks) 

58 

59 

60def _scalar_metric_delta( 

61 reference: float, 

62 candidate: float, 

63 tolerance: Mapping[str, Any], 

64 prefix: str, 

65) -> dict[str, Any]: 

66 """Compare one scalar using explicitly declared absolute/relative bounds.""" 

67 max_abs = abs(candidate - reference) 

68 max_rel = _relative_error(reference, candidate) 

69 max_abs_limit = _optional_tolerance(tolerance, f"{prefix}_max_abs") 

70 max_rel_limit = _optional_tolerance(tolerance, f"{prefix}_max_rel") 

71 if max_abs_limit is None and max_rel_limit is None: 

72 # Preserve the schema-v1 contract: an omitted scalar policy means 

73 # exact absolute equality, rather than silently disabling the gate. 

74 max_abs_limit = 0.0 

75 checks = [] 

76 if max_abs_limit is not None: 

77 checks.append(max_abs <= max_abs_limit) 

78 if max_rel_limit is not None: 

79 checks.append(max_rel <= max_rel_limit) 

80 combination = tolerance.get( 

81 f"{prefix}_combination", 

82 tolerance.get("combination", "all"), 

83 ) 

84 return { 

85 "status": _combine_tolerance_checks( 

86 checks, 

87 combination, 

88 owner=f"{prefix} scalar", 

89 ), 

90 "max_abs": max_abs, 

91 "max_rel": max_rel, 

92 "combination": combination, 

93 } 

94 

95 

96def _align_scalar_rows( 

97 baseline: list[Mapping[str, Any]], 

98 candidate: list[Mapping[str, Any]], 

99 allow_candidate_subset: bool, 

100) -> tuple[list[Mapping[str, Any]], Optional[dict[str, Any]]]: 

101 """Align candidate rows with baseline steps or explain why not.""" 

102 if allow_candidate_subset: 

103 baseline_by_step = {row.get("step"): row for row in baseline} 

104 missing_steps = [row.get("step") for row in candidate if row.get("step") not in baseline_by_step] 

105 if missing_steps: 

106 return baseline, { 

107 "status": "FAIL", 

108 "reason": f"resume steps are absent from uninterrupted baseline: {missing_steps}", 

109 } 

110 return [baseline_by_step[row.get("step")] for row in candidate], None 

111 if len(baseline) != len(candidate): 

112 return baseline, { 

113 "status": "FAIL", 

114 "reason": f"step count differs: baseline={len(baseline)}, candidate={len(candidate)}", 

115 } 

116 return baseline, None 

117 

118 

119def _optional_pair_equal( 

120 baseline_row: Mapping[str, Any], 

121 candidate_row: Mapping[str, Any], 

122 field: str, 

123) -> bool: 

124 """Compare one optional field, treating absent values as unequal.""" 

125 baseline_value = baseline_row.get(field) 

126 candidate_value = candidate_row.get(field) 

127 return ( 

128 baseline_value is not None 

129 and candidate_value is not None 

130 and baseline_value == candidate_value 

131 ) 

132 

133 

134def _post_clip_norm_delta( 

135 baseline_row: Mapping[str, Any], 

136 candidate_row: Mapping[str, Any], 

137 tolerance: Mapping[str, Any], 

138) -> tuple[Optional[float], Optional[float], bool]: 

139 """Compare optional post-clip norms only when both sides recorded them.""" 

140 baseline_post_norm = baseline_row.get("grad_norm_post_clip") 

141 candidate_post_norm = candidate_row.get("grad_norm_post_clip") 

142 if baseline_post_norm is None or candidate_post_norm is None: 

143 return None, None, baseline_post_norm is None and candidate_post_norm is None 

144 comparison = _scalar_metric_delta( 

145 float(baseline_post_norm), 

146 float(candidate_post_norm), 

147 tolerance, 

148 "norm", 

149 ) 

150 return comparison["max_abs"], comparison["max_rel"], comparison["status"] 

151 

152 

153def _compare_scalar_row( 

154 baseline_row: Mapping[str, Any], 

155 candidate_row: Mapping[str, Any], 

156 tolerance: Mapping[str, Any], 

157) -> dict[str, Any]: 

158 """Compare one aligned step row against the declared scalar tolerances.""" 

159 baseline_loss = float(baseline_row["loss"]) 

160 candidate_loss = float(candidate_row["loss"]) 

161 baseline_norm = float(baseline_row["grad_norm_pre_clip"]) 

162 candidate_norm = float(candidate_row["grad_norm_pre_clip"]) 

163 loss_delta = _scalar_metric_delta(baseline_loss, candidate_loss, tolerance, "loss") 

164 norm_delta = _scalar_metric_delta(baseline_norm, candidate_norm, tolerance, "norm") 

165 input_equal = _optional_pair_equal(baseline_row, candidate_row, "global_input_sha256") 

166 learning_rate_equal = _optional_pair_equal(baseline_row, candidate_row, "lr") 

167 step_equal = _optional_pair_equal(baseline_row, candidate_row, "step") 

168 post_norm_delta, post_norm_relative_delta, post_norm_equal = _post_clip_norm_delta( 

169 baseline_row, 

170 candidate_row, 

171 tolerance, 

172 ) 

173 finite = all( 

174 math.isfinite(float(value)) 

175 for value in ( 

176 baseline_row["loss"], 

177 candidate_row["loss"], 

178 baseline_row["grad_norm_pre_clip"], 

179 candidate_row["grad_norm_pre_clip"], 

180 ) 

181 ) 

182 numerical_pass = ( 

183 loss_delta["status"] 

184 and norm_delta["status"] 

185 and post_norm_equal 

186 and learning_rate_equal 

187 and step_equal 

188 ) 

189 if not finite: 

190 row_status = "FAIL" 

191 elif not input_equal: 

192 # Different logical inputs make a numerical A/B verdict invalid; 

193 # the topology is not proven wrong, but this experiment is blocked. 

194 row_status = "BLOCKED" 

195 else: 

196 row_status = "PASS" if numerical_pass else "FAIL" 

197 return { 

198 "step": baseline_row.get("step"), 

199 "status": row_status, 

200 "loss_max_abs": loss_delta["max_abs"], 

201 "loss_max_rel": loss_delta["max_rel"], 

202 "loss_tolerance_combination": loss_delta["combination"], 

203 "norm_max_abs": norm_delta["max_abs"], 

204 "norm_max_rel": norm_delta["max_rel"], 

205 "norm_tolerance_combination": norm_delta["combination"], 

206 "post_clip_norm_max_abs": post_norm_delta, 

207 "post_clip_norm_max_rel": post_norm_relative_delta, 

208 "input_identity": input_equal, 

209 "learning_rate_equal": learning_rate_equal, 

210 "step_equal": step_equal, 

211 "finite": finite, 

212 } 

213 

214 

215def compare_scalar_curves( 

216 baseline_rows: Iterable[Mapping[str, Any]], 

217 candidate_rows: Iterable[Mapping[str, Any]], 

218 tolerance: Mapping[str, Any], 

219 *, 

220 allow_candidate_subset: bool = False, 

221) -> dict[str, Any]: 

222 """Compare aligned loss/norm/LR rows without parsing Trainer text logs.""" 

223 baseline = list(baseline_rows) 

224 candidate = list(candidate_rows) 

225 if not baseline or not candidate: 

226 return { 

227 "status": "BLOCKED", 

228 "reason": "structured scalar evidence is missing or empty", 

229 } 

230 baseline, alignment_error = _align_scalar_rows( 

231 baseline, 

232 candidate, 

233 allow_candidate_subset, 

234 ) 

235 if alignment_error is not None: 

236 return alignment_error 

237 comparisons = [] 

238 status = "PASS" 

239 for baseline_row, candidate_row in zip(baseline, candidate): 

240 comparison = _compare_scalar_row(baseline_row, candidate_row, tolerance) 

241 row_status = comparison["status"] 

242 if row_status == "FAIL": 

243 status = "FAIL" 

244 elif row_status == "BLOCKED" and status == "PASS": 

245 status = "BLOCKED" 

246 comparisons.append(comparison) 

247 return {"status": status, "steps": comparisons} 

248 

249 

250def load_jsonl(path: str | Path) -> list[dict[str, Any]]: 

251 """Load a structured JSON-lines evidence stream.""" 

252 input_path = Path(path) 

253 if not input_path.is_file(): 

254 return [] 

255 return [json.loads(line) for line in input_path.read_text(encoding="utf-8").splitlines() if line] 

256 

257 

258def _nested_value(value: Mapping[str, Any], path: str) -> Any: 

259 current: Any = value 

260 for name in path.split("."): 

261 if not isinstance(current, Mapping) or name not in current: 

262 return None 

263 current = current[name] 

264 return current 

265 

266 

267def compare_preflight_configs( 

268 baseline_dir: str | Path, 

269 candidate_dir: str | Path, 

270) -> dict[str, Any]: 

271 """Prove that topology cases retain one training/precision contract.""" 

272 baseline_path = Path(baseline_dir) / "trainer_config.resolved.json" 

273 candidate_path = Path(candidate_dir) / "trainer_config.resolved.json" 

274 if not baseline_path.is_file() or not candidate_path.is_file(): 

275 return {"status": "BLOCKED", "reason": "resolved Trainer config evidence is missing"} 

276 baseline = json.loads(baseline_path.read_text(encoding="utf-8")) 

277 candidate = json.loads(candidate_path.read_text(encoding="utf-8")) 

278 invariant_paths = ( 

279 "model_init_dtype", 

280 "training.global_batch_size", 

281 "training.max_grad_norm", 

282 "training.loss_aggregation", 

283 "training.seed", 

284 "fsdp_config.mix_precision.param_dtype", 

285 "fsdp_config.mix_precision.reduce_dtype", 

286 "fsdp_config.mix_precision.output_dtype", 

287 "fsdp_config.mix_precision.cast_forward_inputs", 

288 "optimizer.fp32_main_params", 

289 ) 

290 mismatches = [ 

291 { 

292 "path": path, 

293 "baseline": _nested_value(baseline, path), 

294 "candidate": _nested_value(candidate, path), 

295 } 

296 for path in invariant_paths 

297 if _nested_value(baseline, path) != _nested_value(candidate, path) 

298 ] 

299 return { 

300 "status": "FAIL" if mismatches else "PASS", 

301 "mismatches": mismatches, 

302 } 

303 

304 

305def _probe_rows(case_dir: Path) -> list[dict[str, Any]]: 

306 candidates = tuple(case_dir.glob("parameter_probes/rank*.jsonl")) 

307 if not candidates: 

308 candidates = tuple(case_dir.glob("**/parameter_probes/rank*.jsonl")) 

309 return [row for path in sorted(candidates) for row in load_jsonl(path)] 

310 

311 

312def _numeric_summary_delta( 

313 baseline: Mapping[str, Any], 

314 candidate: Mapping[str, Any], 

315 tolerance: Mapping[str, Any], 

316) -> dict[str, Any]: 

317 max_abs_limit = tolerance_value(tolerance, "max_abs", "atol") 

318 relative_l2_limit = tolerance_value(tolerance, "relative_l2", "rtol") 

319 tolerance_combination = tolerance.get("combination", "all") 

320 baseline_values = baseline.get("values") 

321 candidate_values = candidate.get("values") 

322 exact_values = baseline_values is not None and candidate_values is not None 

323 summary_policy_source = None 

324 if exact_values: 

325 baseline_tensor = torch.tensor(baseline_values, dtype=torch.float64) 

326 candidate_tensor = torch.tensor(candidate_values, dtype=torch.float64) 

327 difference = candidate_tensor - baseline_tensor 

328 max_abs = float(difference.abs().max()) if difference.numel() else 0.0 

329 denominator = max( 

330 float(torch.linalg.vector_norm(baseline_tensor)), # pylint: disable=not-callable 

331 1.0e-30, 

332 ) 

333 relative_l2 = ( 

334 float(torch.linalg.vector_norm(difference)) # pylint: disable=not-callable 

335 / denominator 

336 ) 

337 within_tolerance = _combine_tolerance_checks( 

338 (max_abs <= max_abs_limit, relative_l2 <= relative_l2_limit), 

339 tolerance_combination, 

340 owner="parameter", 

341 ) 

342 comparison_mode = "exact_values" 

343 l2_norm_max_abs = None 

344 l2_norm_relative_error = None 

345 aggregate_deltas = None 

346 else: 

347 max_abs = None 

348 relative_l2 = None 

349 comparison_mode = "aggregate_summary" 

350 baseline_l2 = float(baseline.get("l2") or 0.0) 

351 candidate_l2 = float(candidate.get("l2") or 0.0) 

352 l2_norm_max_abs = abs(baseline_l2 - candidate_l2) 

353 l2_norm_relative_error = _relative_error(baseline_l2, candidate_l2) 

354 aggregate_deltas = { 

355 name: abs(float(baseline.get(name) or 0.0) - float(candidate.get(name) or 0.0)) 

356 for name in ("sum", "min", "max") 

357 } 

358 summary_tolerance = tolerance.get("summary", {}) 

359 if not isinstance(summary_tolerance, Mapping): 

360 raise ValueError("parameter summary tolerance must be a mapping") 

361 summary_relative_limit = _optional_tolerance( 

362 summary_tolerance, 

363 "l2_norm_relative", 

364 ) 

365 summary_policy_source = "summary.l2_norm_relative" 

366 if summary_relative_limit is None: 

367 summary_relative_limit = relative_l2_limit 

368 summary_policy_source = "parameters.relative_l2 (compatibility fallback)" 

369 summary_abs_limit = _optional_tolerance( 

370 summary_tolerance, 

371 "l2_norm_max_abs", 

372 ) 

373 summary_checks = [] 

374 if summary_abs_limit is not None: 

375 summary_checks.append(l2_norm_max_abs <= summary_abs_limit) 

376 if summary_relative_limit is not None: 

377 summary_checks.append(l2_norm_relative_error <= summary_relative_limit) 

378 summary_combination = summary_tolerance.get("combination", "all") 

379 within_tolerance = _combine_tolerance_checks( 

380 summary_checks, 

381 summary_combination, 

382 owner="parameter summary", 

383 ) 

384 tolerance_combination = summary_combination 

385 layout_equal = baseline.get("layout", {}).get("global_shape") == candidate.get( 

386 "layout", {} 

387 ).get("global_shape") 

388 finite = bool(baseline.get("finite")) and bool(candidate.get("finite")) 

389 status = "PASS" if finite and layout_equal and within_tolerance else "FAIL" 

390 result = { 

391 "status": status, 

392 "comparison_mode": comparison_mode, 

393 "max_abs": max_abs, 

394 "relative_l2": relative_l2, 

395 "l2_norm_max_abs": l2_norm_max_abs, 

396 "l2_norm_relative_error": l2_norm_relative_error, 

397 "aggregate_deltas": aggregate_deltas, 

398 "global_shape_equal": layout_equal, 

399 "finite": finite, 

400 "full_value_comparison": exact_values, 

401 "tolerance_combination": tolerance_combination, 

402 } 

403 if summary_policy_source is not None: 

404 result["summary_policy_source"] = summary_policy_source 

405 return result 

406 

407 

408def _first_rank_rows( 

409 rows: Iterable[Mapping[str, Any]], 

410) -> dict[tuple[Any, ...], Mapping[str, Any]]: 

411 """Index one canonical row per step, stage, and parameter.""" 

412 indexed = {} 

413 for row in rows: 

414 key = (row.get("step"), row.get("stage"), row.get("name")) 

415 if key not in indexed or int(row.get("rank", 0)) < int(indexed[key].get("rank", 0)): 

416 indexed[key] = row 

417 return indexed 

418 

419 

420def _zero_optimizer_state(row: Optional[Mapping[str, Any]]) -> bool: 

421 """Treat a load-primed zero state as equivalent to an absent lazy state.""" 

422 states = None if row is None else row.get("optimizer_state") 

423 return bool(states) and all( 

424 bool(summary.get("finite")) and float(summary.get("l2") or 0.0) == 0.0 

425 for summary in states.values() 

426 ) 

427 

428 

429def _unmatched_probe_keys( 

430 baseline: Mapping[tuple[Any, ...], Mapping[str, Any]], 

431 candidate: Mapping[tuple[Any, ...], Mapping[str, Any]], 

432 allow_candidate_subset: bool, 

433) -> set[tuple[Any, ...]]: 

434 """Diff the canonical key spaces of both evidence directories.""" 

435 if allow_candidate_subset: 

436 candidate_steps = {key[0] for key in candidate} 

437 expected_baseline = {key for key in baseline if key[0] in candidate_steps} 

438 return expected_baseline ^ set(candidate) 

439 return set(baseline) ^ set(candidate) 

440 

441 

442def _compare_optimizer_state_probes( 

443 key: tuple[Any, ...], 

444 baseline_row: Mapping[str, Any], 

445 candidate_row: Mapping[str, Any], 

446 tolerance: Mapping[str, Any], 

447) -> list[dict[str, Any]]: 

448 """Compare every declared optimizer-state summary for one parameter.""" 

449 baseline_states = baseline_row.get("optimizer_state", {}) 

450 candidate_states = candidate_row.get("optimizer_state", {}) 

451 comparisons = [] 

452 for state_name in sorted(set(baseline_states) | set(candidate_states)): 

453 if state_name not in baseline_states or state_name not in candidate_states: 

454 comparisons.append( 

455 {"key": (*key, state_name), "status": "FAIL", "reason": "state missing"} 

456 ) 

457 continue 

458 comparison = _numeric_summary_delta( 

459 baseline_states[state_name], 

460 candidate_states[state_name], 

461 tolerance, 

462 ) 

463 comparison["key"] = (*key, state_name) 

464 comparisons.append(comparison) 

465 return comparisons 

466 

467 

468def _compare_parameter_probe_row( 

469 key: tuple[Any, ...], 

470 baseline_row: Mapping[str, Any], 

471 candidate_row: Mapping[str, Any], 

472 tolerance: Mapping[str, Any], 

473) -> dict[str, Any]: 

474 """Compare one parameter lifecycle row including hash and replica gates.""" 

475 comparison = _numeric_summary_delta(baseline_row, candidate_row, tolerance) 

476 comparison["key"] = key 

477 if baseline_row.get("comparison") == "exact_hash": 

478 comparison["hash_equal"] = baseline_row.get("sha256") == candidate_row.get("sha256") 

479 if not comparison["hash_equal"]: 

480 comparison["status"] = "FAIL" 

481 baseline_replica = baseline_row.get("replica", {}).get("replicas_equal", True) 

482 candidate_replica = candidate_row.get("replica", {}).get("replicas_equal", True) 

483 comparison["replicas_equal"] = baseline_replica and candidate_replica 

484 if not comparison["replicas_equal"]: 

485 comparison["status"] = "FAIL" 

486 return comparison 

487 

488 

489def compare_parameter_probes( 

490 baseline_dir: str | Path, 

491 candidate_dir: str | Path, 

492 tolerance: Mapping[str, Any], 

493 *, 

494 allow_candidate_subset: bool = False, 

495 compare_optimizer_states: bool = True, 

496) -> dict[str, Any]: 

497 """Compare canonical parameter lifecycle summaries from two topologies. 

498 

499 Optimizer internals such as Muon momentum may be topology-local across HSDP 

500 replica groups even when synchronized gradients and model updates are 

501 identical. Cross-topology cases therefore validate their live layouts and 

502 checkpoint recovery, but reserve numerical state comparison for matching 

503 topologies. 

504 """ 

505 baseline_rows = _probe_rows(Path(baseline_dir)) 

506 candidate_rows = _probe_rows(Path(candidate_dir)) 

507 if not baseline_rows or not candidate_rows: 

508 return { 

509 "status": "BLOCKED", 

510 "reason": "strict parameter probe evidence is missing", 

511 } 

512 baseline = _first_rank_rows(baseline_rows) 

513 candidate = _first_rank_rows(candidate_rows) 

514 unmatched = _unmatched_probe_keys(baseline, candidate, allow_candidate_subset) 

515 zero_initialized = sorted( 

516 ( 

517 key 

518 for key in unmatched 

519 if _zero_optimizer_state(baseline.get(key)) 

520 or _zero_optimizer_state(candidate.get(key)) 

521 ), 

522 key=repr, 

523 ) 

524 missing = sorted(set(unmatched) - set(zero_initialized), key=repr) 

525 comparisons = [] 

526 skipped_optimizer_states = 0 

527 for key in sorted(set(baseline).intersection(candidate), key=repr): 

528 baseline_row = baseline[key] 

529 candidate_row = candidate[key] 

530 if "optimizer_state" in baseline_row or "optimizer_state" in candidate_row: 

531 if not compare_optimizer_states: 

532 skipped_optimizer_states += len( 

533 set(baseline_row.get("optimizer_state", {})) 

534 | set(candidate_row.get("optimizer_state", {})) 

535 ) 

536 continue 

537 comparisons.extend( 

538 _compare_optimizer_state_probes( 

539 key, 

540 baseline_row, 

541 candidate_row, 

542 tolerance, 

543 ) 

544 ) 

545 continue 

546 comparisons.append( 

547 _compare_parameter_probe_row(key, baseline_row, candidate_row, tolerance) 

548 ) 

549 status = "FAIL" if missing or any(row["status"] == "FAIL" for row in comparisons) else "PASS" 

550 return { 

551 "status": status, 

552 "missing": [repr(key) for key in missing], 

553 "zero_initialized_optimizer_states": [repr(key) for key in zero_initialized], 

554 "optimizer_state_numeric_comparison": ( 

555 "enabled" if compare_optimizer_states else "skipped_cross_topology" 

556 ), 

557 "skipped_optimizer_states": skipped_optimizer_states, 

558 "values": comparisons, 

559 } 

560 

561 

562def summarize_performance( 

563 rows: Iterable[Mapping[str, Any]], 

564 warmup_steps: int = 1, 

565 memory_rows: Optional[Iterable[Mapping[str, Any]]] = None, 

566) -> dict[str, Any]: 

567 """Summarize steady-state wall time and throughput independently of precision.""" 

568 values = list(rows)[warmup_steps:] 

569 step_times = sorted(float(row["step_time_seconds"]) for row in values) 

570 if not step_times: 

571 return {"status": "BLOCKED", "reason": "no steady-state performance steps"} 

572 p90_index = min(math.ceil(0.9 * len(step_times)) - 1, len(step_times) - 1) 

573 token_rates = [ 

574 float(row["tokens_per_second"]) 

575 for row in values 

576 if row.get("tokens_per_second") is not None 

577 ] 

578 sample_rates = [ 

579 float(row["samples_per_second"]) 

580 for row in values 

581 if row.get("samples_per_second") is not None 

582 ] 

583 allocator_rows = list(memory_rows) if memory_rows is not None else values 

584 peak_allocated = [ 

585 int(row["peak_memory_allocated_bytes"]) 

586 for row in allocator_rows 

587 if row.get("peak_memory_allocated_bytes") is not None 

588 ] 

589 peak_reserved = [ 

590 int(row["peak_memory_reserved_bytes"]) 

591 for row in allocator_rows 

592 if row.get("peak_memory_reserved_bytes") is not None 

593 ] 

594 return { 

595 "status": "MEASURED", 

596 "steady_steps": len(step_times), 

597 "step_time_p50_seconds": statistics.median(step_times), 

598 "step_time_p90_seconds": step_times[p90_index], 

599 "tokens_per_second_p50": statistics.median(token_rates) if token_rates else None, 

600 "samples_per_second_p50": statistics.median(sample_rates) if sample_rates else None, 

601 "peak_memory_allocated_bytes": max(peak_allocated) if peak_allocated else None, 

602 "peak_memory_reserved_bytes": max(peak_reserved) if peak_reserved else None, 

603 } 

604 

605 

606def _group_checkpoint_events(files: list[Path]) -> dict[str, list[Path]]: 

607 """Group every rank file by its versioned checkpoint event name.""" 

608 grouped: dict[str, list[Path]] = {} 

609 for path in files: 

610 event_name = path.name.rsplit("_rank", 1)[0] 

611 grouped.setdefault(event_name, []).append(path) 

612 return grouped 

613 

614 

615def _checkpoint_event_order(event_name: str) -> tuple[int, str]: 

616 """Sort numbered checkpoint events after legacy unnumbered events.""" 

617 suffix = event_name.rsplit("_", 1)[-1] 

618 return (int(suffix), event_name) if suffix.isdigit() else (-1, event_name) 

619 

620 

621def _latest_checkpoint_event(events: Mapping[str, list[Path]], preferred: str) -> str: 

622 """Choose the newest numbered event, retaining legacy-name support.""" 

623 matching = [ 

624 event_name 

625 for event_name in events 

626 if event_name == preferred or event_name.startswith(f"{preferred}_") 

627 ] 

628 if not matching: 

629 return sorted(events)[-1] 

630 return max(matching, key=_checkpoint_event_order) 

631 

632 

633def _load_rank_layouts(files: list[Path]) -> list[dict[str, Any]]: 

634 """Load checkpoint tensor-layout maps for all selected ranks.""" 

635 return [ 

636 json.loads(path.read_text(encoding="utf-8")).get("tensor_layouts", {}) 

637 for path in files 

638 ] 

639 

640 

641def _missing_layout_names( 

642 baseline: Mapping[str, Any], 

643 candidate: Mapping[str, Any], 

644) -> tuple[list[str], list[str]]: 

645 """List one-sided keys, excluding lazily initialized optimizer state.""" 

646 baseline_only = set(baseline) - set(candidate) 

647 candidate_only = set(candidate) - set(baseline) 

648 initialized_optimizer_entries = sorted( 

649 name for name in candidate_only if name.startswith("optimizer.") 

650 ) 

651 missing = sorted( 

652 baseline_only 

653 | { 

654 name 

655 for name in candidate_only 

656 if name not in initialized_optimizer_entries 

657 } 

658 ) 

659 return missing, initialized_optimizer_entries 

660 

661 

662def _layout_mismatches( 

663 baseline: Mapping[str, Any], 

664 candidate: Mapping[str, Any], 

665 same_topology: bool, 

666) -> list[dict[str, Any]]: 

667 """Diff shared tensor layouts over the topology-appropriate field set.""" 

668 fields = ["global_shape"] 

669 if same_topology: 

670 fields.extend(("mesh_dim_names", "mesh_shape", "placements")) 

671 mismatches = [] 

672 for name in sorted(set(baseline).intersection(candidate)): 

673 baseline_layout = baseline[name] 

674 candidate_layout = candidate[name] 

675 differences = { 

676 field: { 

677 "baseline": baseline_layout.get(field), 

678 "candidate": candidate_layout.get(field), 

679 } 

680 for field in fields 

681 if baseline_layout.get(field) != candidate_layout.get(field) 

682 } 

683 if differences: 

684 mismatches.append({"name": name, "differences": differences}) 

685 return mismatches 

686 

687 

688def _rank_consistency_inconsistencies( 

689 side_name: str, 

690 rank_layouts: list[dict[str, Any]], 

691) -> list[dict[str, Any]]: 

692 """Flag keys whose rank-local payloads disagree on presence or shape.""" 

693 inconsistencies = [] 

694 all_names = set().union(*(set(layouts) for layouts in rank_layouts)) 

695 for name in sorted(all_names): 

696 values = [layouts.get(name) for layouts in rank_layouts] 

697 global_shapes = { 

698 tuple(value.get("global_shape", ())) 

699 for value in values 

700 if value is not None 

701 } 

702 if any(value is None for value in values) or len(global_shapes) != 1: 

703 inconsistencies.append( 

704 { 

705 "side": side_name, 

706 "name": name, 

707 "reason": "rank-local payloads disagree on key presence or global shape", 

708 } 

709 ) 

710 return inconsistencies 

711 

712 

713def _same_topology_rank_inconsistencies( 

714 baseline_ranks: list[dict[str, Any]], 

715 candidate_ranks: list[dict[str, Any]], 

716) -> list[dict[str, Any]]: 

717 """Require equal rank counts and stable local shard shapes.""" 

718 inconsistencies = [] 

719 if len(baseline_ranks) != len(candidate_ranks): 

720 inconsistencies.append( 

721 { 

722 "side": "comparison", 

723 "reason": "same-topology checkpoint evidence has different rank counts", 

724 "baseline_ranks": len(baseline_ranks), 

725 "candidate_ranks": len(candidate_ranks), 

726 } 

727 ) 

728 return inconsistencies 

729 for rank, (baseline_layouts, candidate_layouts) in enumerate( 

730 zip(baseline_ranks, candidate_ranks) 

731 ): 

732 for name in sorted(set(baseline_layouts).intersection(candidate_layouts)): 

733 if baseline_layouts[name].get("local_shape") != candidate_layouts[name].get( 

734 "local_shape" 

735 ): 

736 inconsistencies.append( 

737 { 

738 "side": "comparison", 

739 "rank": rank, 

740 "name": name, 

741 "reason": "same-topology local shard shape changed", 

742 } 

743 ) 

744 return inconsistencies 

745 

746 

747def compare_checkpoint_layouts( 

748 baseline_dir: str | Path, 

749 candidate_dir: str | Path, 

750 *, 

751 same_topology: bool, 

752) -> dict[str, Any]: 

753 """Compare DCP-bound logical layouts while allowing legal resharding.""" 

754 baseline_files = sorted(Path(baseline_dir).glob("**/checkpoint/*_rank*.json"))