Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / weight_verifier.py: 90%

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"""Check live model shards against local HF weights before training updates. 

16 

17Inverse weight conversions run symbolically; actual model values are read back 

18in bounded tiles. No full model, expert tensor or Host embedding is gathered. 

19All ranks must call ``verify_hf_weights``. Run it before the first forward to 

20avoid confusing router/buffer updates with weight-loading errors. 

21 

22This is a check of live values under the recorded conversion rules, not a 

23mathematical oracle for those rules. The reference decoder reuses the shared 

24Transformers-based quantized reader; it does not recover pre-quantization FP32 values. 

25""" 

26 

27from __future__ import annotations 

28 

29import csv 

30import json 

31import logging 

32import math 

33from collections import defaultdict 

34from pathlib import Path 

35from typing import Any 

36 

37import torch 

38import torch.distributed as dist 

39from safetensors import safe_open 

40from torch import nn 

41 

42from hyper_parallel import DTensor 

43from hyper_parallel.components.checkpoint.load_groups import ( 

44 build_load_targets, local_target_tensor, checkpoint_logical_shapes, checkpoint_padding_boxes, reject_partial_layout, 

45 resolve_checkpoint_index, target_layout, 

46) 

47from hyper_parallel.components.checkpoint.quantized_hf_storage import build_quantized_hf_reader 

48from hyper_parallel.components.checkpoint.region_tensor import RegionTensor 

49from hyper_parallel.components.checkpoint.weight_conversion import ( 

50 WeightConverter, WeightRenaming, get_model_conversion_mapping, revert_weight_conversion, 

51) 

52from hyper_parallel.core.distributed_checkpoint.utils import create_chunk_list_for_tensor 

53from hyper_parallel.models.external_state import CheckpointTensorShard, get_model_external_state 

54from hyper_parallel.components.checkpoint._weight_check_geometry import box_union_volume, local_block_box, tile_box 

55 

56logger = logging.getLogger(__name__) 

57 

58 

59class _Reference: 

60 """Read selected disk tensors without running model-side conversions.""" 

61 

62 def __init__(self, path: str) -> None: 

63 """Read indexed tensor headers and optional quantization metadata.""" 

64 location = Path(path).expanduser().resolve() 

65 if not location.exists(): 

66 raise ValueError(f"Weight checking requires a local HF checkpoint: {location}") 

67 self.index = resolve_checkpoint_index(str(location)) 

68 self.quantized = build_quantized_hf_reader(location) 

69 self.shapes = {} 

70 if self.quantized is not None: 

71 self.shapes = { 

72 name: tuple(entry.size) 

73 for name, entry in self.quantized.load_metadata().state_dict_metadata.items() 

74 if name in self.index.files_by_key 

75 } 

76 else: 

77 by_file = defaultdict(list) 

78 for name, filename in self.index.files_by_key.items(): 

79 by_file[filename].append(name) 

80 for filename, names in by_file.items(): 

81 with safe_open(str(filename), framework="pt", device="cpu") as reader: 

82 for name in names: 

83 self.shapes[name] = tuple(reader.get_slice(name).get_shape()) 

84 if not self.shapes: 

85 raise ValueError("No logical HF tensors were selected for checking") 

86 

87 def read(self, name: str, offsets: tuple[int, ...], lengths: tuple[int, ...]) -> torch.Tensor: 

88 """Decode a bounded source rectangle, keeping ordinary source dtypes.""" 

89 if self.quantized is not None and len(lengths) == 2: 

90 return self.quantized.read_region(name, offsets, lengths) 

91 with safe_open(str(self.index.files_by_key[name]), framework="pt", device="cpu") as reader: 

92 if not lengths: 

93 return reader.get_tensor(name) 

94 region = tuple(slice(start, start + count) for start, count in zip(offsets, lengths)) 

95 return reader.get_slice(name)[region] 

96 

97 

98class _InverseModel: 

99 """Expose export conversions without mutating the live model's metadata.""" 

100 

101 def __init__(self, model: nn.Module, conversions: Any) -> None: 

102 """Bind one inverse-conversion stage to a read-only model view.""" 

103 self.model = model 

104 self.config = model.config 

105 self._weight_conversions = conversions 

106 

107 def __getattr__(self, name: str) -> Any: 

108 return getattr(self.model, name) 

109 

110 

111def _model_shards(model: nn.Module) -> dict[str, CheckpointTensorShard]: 

112 """Describe every materialized local model shard for comparison.""" 

113 logical_shapes = checkpoint_logical_shapes(model) 

114 external = get_model_external_state(model) 

115 external_shards = {} 

116 if external is not None and external.parameters: 

117 provider = getattr(external, "checkpoint_tensor_shards", None) 

118 if not callable(provider): 

119 raise ValueError("External state must expose checkpoint_tensor_shards() for complete weight checking") 

120 external_shards = provider() 

121 if set(external_shards) != set(external.parameters): 

122 raise ValueError("External checkpoint shard descriptors do not cover exactly the external parameters") 

123 shards = {} 

124 for name, target in build_load_targets(model).items(): 

125 if name in external_shards: 

126 shards[name] = external_shards[name] 

127 continue 

128 layout = target_layout(target) 

129 view = target 

130 if layout is not None: 

131 reject_partial_layout(name, layout) 

132 if not isinstance(view, DTensor): 

133 shape = getattr(layout, "tensor_shape", None) 

134 view = DTensor.from_local_with_layout(target.detach(), layout, shape=shape) 

135 chunks = create_chunk_list_for_tensor(view) 

136 if len(chunks) != 1: 

137 raise ValueError(f"Weight checking requires one rectangular local shard: {name}") 

138 local = local_target_tensor(target) 

139 if local.is_meta or tuple(local.shape) != tuple(chunks[0].sizes): 

140 raise ValueError(f"Live storage does not match its checkpoint shard metadata: {name}") 

141 logical_shape = logical_shapes.get(name) 

142 if logical_shape == tuple(view.shape): 

143 logical_shape = None 

144 shards[name] = CheckpointTensorShard( 

145 local, tuple(view.shape), tuple(chunks[0].offsets), logical_shape, 

146 ) 

147 if not shards: 

148 raise ValueError("The model has no checkpoint tensors to check") 

149 return shards 

150 

151 

152def _inverse_regions(model: nn.Module, shards: dict[str, CheckpointTensorShard]) -> dict[str, RegionTensor]: 

153 """Reverse recorded load conversions into symbolic HF tensor regions.""" 

154 regions = { 

155 name: RegionTensor.leaf(name, shard.logical_shape or shard.global_shape, shard.tensor.dtype) 

156 for name, shard in shards.items() 

157 } 

158 replacements = getattr(model, "_hp_used_replacement_weight_conversions", None) 

159 if replacements is not None: 

160 stages = [replacements, getattr(model, "_hp_used_base_weight_conversions", [])] 

161 else: 

162 conversions = getattr(model, "_weight_conversions", None) 

163 if conversions is None: 

164 conversions = get_model_conversion_mapping(model, add_legacy=False) 

165 stages = [conversions] 

166 for conversions in stages: 

167 converters = [item for item in conversions if isinstance(item, WeightConverter)] 

168 renamings = [item for item in conversions if isinstance(item, WeightRenaming)] 

169 if len(converters) + len(renamings) != len(conversions): 

170 raise ValueError("Unsupported recorded weight-conversion type") 

171 if converters: 

172 regions = _reverse_stage(model, regions, converters) 

173 for renaming in reversed(renamings): 

174 regions = _reverse_stage(model, regions, [renaming]) 

175 if any(not isinstance(value, RegionTensor) for value in regions.values()): 

176 raise ValueError("Inverse conversion produced non-symbolic weights; bounded verification is unsupported") 

177 return regions 

178 

179 

180def _reverse_stage(model: nn.Module, regions: dict[str, RegionTensor], 

181 conversions: list[Any]) -> dict[str, RegionTensor]: 

182 """Undo one stage, respecting anchored names relative to the model root. 

183 

184 Loading resolves the base-model prefix after applying its renamings. Match 

185 full scoped names first, then relative names; symbolic leaves retain the 

186 original live FQNs so that changing lookup names cannot change the data read. 

187 """ 

188 prefix = getattr(model, "base_model_prefix", "") 

189 prefix = prefix + "." if prefix else "" 

190 probes = [conversion.reverse_transform() for conversion in conversions] 

191 normalized = {} 

192 for name, region in regions.items(): 

193 lookup = name 

194 if (prefix and name.startswith(prefix) 

195 and not any(probe.rename_source_key(name)[1] is not None for probe in probes)): 

196 relative = name[len(prefix):] 

197 if any(probe.rename_source_key(relative)[1] is not None for probe in probes): 

198 lookup = relative 

199 if lookup in normalized: 

200 raise ValueError(f"Ambiguous inverse conversion after resolving base-model prefix: {lookup}") 

201 normalized[lookup] = region 

202 return revert_weight_conversion(_InverseModel(model, conversions), normalized) 

203 

204 

205def _model_box(block: Any, offsets: tuple[int, ...], 

206 lengths: tuple[int, ...]) -> tuple[tuple[int, ...], tuple[int, ...]]: 

207 """Map a rectangular recovered HF region back to its live model coordinates.""" 

208 delta = tuple(start - origin for start, origin in zip(offsets, block.offsets)) 

209 base = tuple( 

210 origin + sum(step[dim] * shift for step, shift in zip(block.steps, delta)) 

211 for dim, origin in enumerate(block.base) 

212 ) 

213 model_lengths = tuple( 

214 1 + sum(step[dim] * (size - 1) for step, size in zip(block.steps, lengths)) 

215 for dim in range(len(base)) 

216 ) 

217 return base, model_lengths 

218 

219 

220def _live_tile(block: Any, shard: CheckpointTensorShard, offsets: tuple[int, ...], 

221 lengths: tuple[int, ...]) -> torch.Tensor: 

222 """Read one transformed tile from the live model shard on CPU.""" 

223 base, _ = _model_box(block, offsets, lengths) 

224 tensor = shard.tensor.detach() 

225 strides = tensor.stride() 

226 tile_strides = tuple(sum(step * stride for step, stride in zip(column, strides)) for column in block.steps) 

227 storage_offset = tensor.storage_offset() + sum( 

228 (origin - start) * stride for origin, start, stride in zip(base, shard.offsets, strides) 

229 ) 

230 values = tensor.as_strided(lengths, tile_strides, storage_offset).cpu() 

231 for operation in block.post: 

232 values = operation(values) 

233 return values 

234 

235 

236def _difference(actual: torch.Tensor, expected: torch.Tensor, storage_dtype: torch.dtype) -> dict[str, Any]: 

237 """Count raw and load-dtype mismatches for one pair of tensor tiles.""" 

238 if actual.is_complex() or expected.is_complex(): 

239 raise ValueError("Complex weight checking is not supported") 

240 if tuple(actual.shape) != tuple(expected.shape): 

241 raise ValueError(f"Compared weight tiles have different shapes: {actual.shape}, {expected.shape}") 

242 raw_bad = actual != expected 

243 if actual.is_floating_point() and expected.is_floating_point(): 

244 raw_bad |= (actual == 0) & (expected == 0) & (torch.signbit(actual) != torch.signbit(expected)) 

245 nonfinite = ~(torch.isfinite(actual) & torch.isfinite(expected)) 

246 raw_count = int(torch.count_nonzero(raw_bad)) 

247 nonfinite_count = int(torch.count_nonzero(nonfinite)) 

248 if not raw_count and not nonfinite_count: 

249 return { 

250 "elements": actual.numel(), "exact_reference_mismatches": 0, "load_dtype_mismatches": 0, 

251 "nonfinite_elements": 0, "max_abs_error": 0.0, "max_load_dtype_abs_error": 0.0, 

252 "first_difference": None, 

253 } 

254 adjusted = expected.to(storage_dtype).to(actual.dtype) 

255 bad = actual != adjusted 

256 if actual.is_floating_point() and expected.is_floating_point(): 

257 bad |= (actual == 0) & (adjusted == 0) & (torch.signbit(actual) != torch.signbit(adjusted)) 

258 error = (actual.double() - expected.double()).abs() 

259 adjusted_error = (actual.double() - adjusted.double()).abs() 

260 first = None 

261 mask = bad | nonfinite 

262 load_mismatch = bool(torch.any(mask)) 

263 if not load_mismatch: 

264 mask = raw_bad 

265 if bool(torch.any(mask)): 

266 flat_index = int(torch.argmax(mask.reshape(-1).to(torch.int8))) 

267 index = [] 

268 remainder = flat_index 

269 for size in reversed(actual.shape): 

270 index.append(remainder % size) 

271 remainder //= size 

272 first = { 

273 "kind": "load_dtype_mismatch" if load_mismatch else "dtype_cast", 

274 "tile_index": list(reversed(index)), 

275 "actual": _json_number(actual.reshape(-1)[flat_index]), 

276 "reference": _json_number(expected.reshape(-1)[flat_index]), 

277 "reference_in_load_dtype": _json_number(adjusted.reshape(-1)[flat_index]), 

278 } 

279 return { 

280 "elements": actual.numel(), 

281 "exact_reference_mismatches": raw_count, 

282 "load_dtype_mismatches": int(torch.count_nonzero(bad)), 

283 "nonfinite_elements": nonfinite_count, 

284 "max_abs_error": float(error.max()) if error.numel() and bool(torch.isfinite(error).all()) else None, 

285 "max_load_dtype_abs_error": ( 

286 float(adjusted_error.max()) 

287 if adjusted_error.numel() and bool(torch.isfinite(adjusted_error).all()) else None 

288 ), 

289 "first_difference": first, 

290 } 

291 

292 

293def _json_number(value: torch.Tensor) -> float | str: 

294 number = float(value) 

295 return number if math.isfinite(number) else str(number) 

296 

297 

298def _accumulate(record: dict[str, Any], tile: dict[str, Any], offsets: tuple[int, ...]) -> None: 

299 """Merge tile comparison counts and retain the first mismatch.""" 

300 for field in ("elements", "exact_reference_mismatches", "load_dtype_mismatches", "nonfinite_elements"): 

301 record[field] = record.get(field, 0) + tile[field] 

302 for field in ("max_abs_error", "max_load_dtype_abs_error"): 

303 error = tile[field] 

304 if error is not None: 

305 record[field] = max(record.get(field, 0.0), error) 

306 existing = record.get("first_difference") 

307 candidate = tile["first_difference"] 

308 if candidate is not None and (existing is None or ( 

309 existing["kind"] == "dtype_cast" and candidate["kind"] == "load_dtype_mismatch")): 

310 first = dict(tile["first_difference"]) 

311 first["hf_index"] = [start + index for start, index in zip(offsets, first.pop("tile_index"))] 

312 record["first_difference"] = first 

313 

314 

315def _padding_check(name: str, shard: CheckpointTensorShard, limit: int) -> dict[str, Any] | None: 

316 """Check that physical padding beyond the logical tensor is zero-filled.""" 

317 logical = shard.logical_shape 

318 if logical is None or logical == shard.global_shape: 

319 return None 

320 if (len(logical) != len(shard.global_shape) 

321 or any(left > right for left, right in zip(logical, shard.global_shape))): 

322 raise ValueError(f"Invalid logical checkpoint shape for {name}") 

323 record = {"model_key": name, "kind": "padding", "elements": 0} 

324 for box in checkpoint_padding_boxes(logical, shard.offsets, tuple(shard.tensor.shape)): 

325 for offsets, sizes in tile_box(*box, limit): 

326 slices = tuple(slice(left, left + size) for left, size in zip(offsets, sizes)) 

327 values = shard.tensor.detach()[slices].cpu() 

328 _accumulate(record, _difference(values, torch.zeros_like(values), values.dtype), offsets) 

329 if record.get("first_difference"): 

330 record["first_difference"]["local_index"] = record["first_difference"].pop("hf_index") 

331 return record 

332 

333 

334def _check_local(model: nn.Module, reference: _Reference, limit: int) -> dict[str, Any]: 

335 """Compare one rank's live storage and calculate its source coverage.""" 

336 shards = _model_shards(model) 

337 recovered = _inverse_regions(model, shards) 

338 source_coverage = defaultdict(list) 

339 model_coverage = defaultdict(list) 

340 records = [] 

341 unmatched = [] 

342 prefix = getattr(model, "base_model_prefix", "") + "." 

343 for number, (recovered_name, region) in enumerate(sorted(recovered.items()), start=1): 

344 name = recovered_name 

345 if name not in reference.shapes and name.startswith(prefix): 

346 name = name[len(prefix):] 

347 if name not in reference.shapes: 

348 unmatched.append(recovered_name) 

349 continue 

350 if tuple(region.shape) != reference.shapes[name]: 

351 raise ValueError(f"Recovered HF shape mismatch: {name}: {tuple(region.shape)} vs {reference.shapes[name]}") 

352 record = { 

353 "hf_key": name, "shape": list(region.shape), "elements": 0, 

354 "model_keys": sorted(region.sources()), 

355 } 

356 if math.prod(region.shape) >= 1 << 30 or number % 200 == 0: 

357 logger.info("Checking HF tensor %d/%d: %s, shape=%s", number, len(recovered), name, tuple(region.shape)) 

358 for block in region.blocks: 

359 shard = shards[block.source] 

360 box = local_block_box(block, shard.offsets, tuple(shard.tensor.shape)) 

361 if box is None: 

362 continue 

363 source_coverage[name].append(box) 

364 model_coverage[block.source].append(_model_box(block, *box)) 

365 for offsets, lengths in tile_box(*box, limit): 

366 actual = _live_tile(block, shard, offsets, lengths) 

367 expected = reference.read(name, offsets, lengths) 

368 difference = _difference(actual, expected, shard.tensor.dtype) 

369 if difference["first_difference"]: 

370 difference["first_difference"].update( 

371 model_key=block.source, storage_dtype=str(shard.tensor.dtype), 

372 reference_dtype=str(expected.dtype), 

373 ) 

374 _accumulate(record, difference, offsets) 

375 if record["elements"]: 

376 records.append(record) 

377 uncovered = [] 

378 for name, shard in shards.items(): 

379 padding = _padding_check(name, shard, limit) 

380 padding_elements = 0 

381 if padding is not None: 

382 records.append(padding) 

383 padding_elements = padding["elements"] 

384 covered = box_union_volume(model_coverage[name]) + padding_elements 

385 if covered != shard.tensor.numel(): 

386 uncovered.append({"model_key": name, "checked": covered, "local_elements": shard.tensor.numel()}) 

387 return { 

388 "tensors": records, "source_coverage": dict(source_coverage), "unmatched_recovered_keys": unmatched, 

389 "uncovered_model_shards": uncovered, "model_tensor_keys": list(shards), 

390 } 

391 

392 

393def _gather(value: Any) -> list[Any]: 

394 if not (dist.is_available() and dist.is_initialized()): 

395 return [value] 

396 values = [None] * dist.get_world_size() 

397 dist.all_gather_object(values, value) 

398 return values 

399 

400 

401def _summarize(reports: list[dict[str, Any]], shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]: 

402 """Combine rank reports into global mismatch and coverage results.""" 

403 coverage = defaultdict(list) 

404 errors = [] 

405 mismatches = 0 

406 raw_mismatches = 0 

407 nonfinite = 0 

408 for rank, report in enumerate(reports): 

409 if report.get("error"): 

410 errors.append({"rank": rank, "error": report["error"]}) 

411 for name, boxes in report.get("source_coverage", {}).items(): 

412 coverage[name].extend(boxes) 

413 if report.get("uncovered_model_shards") or report.get("unmatched_recovered_keys"): 

414 errors.append({"rank": rank, "error": "Model weights are not completely covered; see rank report"}) 

415 mismatches += report.get("load_dtype_mismatches", 0) 

416 nonfinite += report.get("nonfinite_elements", 0) 

417 raw_mismatches += report.get("exact_reference_mismatches", 0) 

418 uncovered = { 

419 name: {"checked": box_union_volume(coverage[name]), "expected": math.prod(shape)} 

420 for name, shape in shapes.items() if box_union_volume(coverage[name]) != math.prod(shape) 

421 } 

422 return { 

423 "status": "FAIL" if errors or uncovered or mismatches or nonfinite else "PASS", 

424 "exact_reference_match": raw_mismatches == 0 and not errors and not uncovered and not nonfinite, 

425 "load_dtype_mismatches_including_replicas": mismatches, 

426 "exact_reference_mismatches_including_replicas": raw_mismatches, 

427 "nonfinite_elements_including_replicas": nonfinite, 

428 "uncovered_hf_tensors": uncovered, "errors": errors, 

429 "ranks": len(reports), "hf_tensors": len(shapes), 

430 "hf_elements": sum(math.prod(shape) for shape in shapes.values()), 

431 "model_tensors": len({name for report in reports for name in report.get("model_tensor_keys", [])}), 

432 } 

433 

434 

435def _tensor_table(reports: list[dict[str, Any]], summary: dict[str, Any]) -> list[dict[str, Any]]: 

436 """Build per-tensor rows from rank reports and global coverage.""" 

437 rows = {} 

438 for report in reports: 

439 rank = report["rank"] 

440 for tensor in report.get("tensors", []): 

441 key = tensor.get("hf_key", "<padding>:" + tensor.get("model_key", "")) 

442 row = rows.setdefault(key, {"name": key, "kind": tensor.get("kind", "weight"), "ranks_with_mismatch": []}) 

443 for field in ("elements", "exact_reference_mismatches", "load_dtype_mismatches", "nonfinite_elements"): 

444 row[field] = row.get(field, 0) + tensor.get(field, 0) 

445 for field in ("max_abs_error", "max_load_dtype_abs_error"): 

446 row[field] = max(row.get(field, 0.0), tensor.get(field, 0.0)) 

447 if tensor.get("load_dtype_mismatches", 0) or tensor.get("nonfinite_elements", 0): 

448 row["ranks_with_mismatch"].append(rank) 

449 existing, candidate = row.get("first_difference"), tensor.get("first_difference") 

450 if candidate is not None and (existing is None or ( 

451 existing["kind"] == "dtype_cast" and candidate["kind"] == "load_dtype_mismatch")): 

452 row["first_difference"] = dict(candidate, rank=rank) 

453 for row in rows.values(): 

454 if row["load_dtype_mismatches"] or row["nonfinite_elements"]: 

455 row["result"] = "MISMATCH" 

456 elif row["exact_reference_mismatches"]: 

457 row["result"] = "DTYPE_CAST" 

458 else: 

459 row["result"] = "EXACT" 

460 for name, counts in summary["uncovered_hf_tensors"].items(): 

461 row = rows.setdefault(name, { 

462 "name": name, "kind": "weight", "elements": 0, "exact_reference_mismatches": 0, 

463 "load_dtype_mismatches": 0, "nonfinite_elements": 0, "max_abs_error": 0.0, 

464 "max_load_dtype_abs_error": 0.0, "ranks_with_mismatch": [], 

465 }) 

466 row["coverage"] = counts 

467 if row.get("result") != "MISMATCH": 

468 row["result"] = "INCOMPLETE" 

469 order = {"MISMATCH": 0, "INCOMPLETE": 1, "DTYPE_CAST": 2, "EXACT": 3} 

470 return sorted(rows.values(), key=lambda row: (order[row["result"]], row["name"])) 

471 

472 

473def _print_table(rows: list[dict[str, Any]], summary: dict[str, Any], root: Path) -> None: 

474 """Log the human-readable summary table on rank zero.""" 

475 lines = [ 

476 f"HF weight check: {summary['status']} | ranks={summary['ranks']} | " 

477 f"model tensors={summary['model_tensors']} | HF tensors={summary['hf_tensors']} | " 

478 f"weight elements={summary['hf_elements']:,}", 

479 f"EXACT={summary['exact_tensors']} DTYPE_CAST={summary['dtype_cast_tensors']} " 

480 f"MISMATCH={summary['mismatched_tensors']} INCOMPLETE={summary['incomplete_tensors']} " 

481 f"padding checks={summary['padding_tensors']}", 

482 "Result | HF tensor / padding table | Checked elems | Mismatches | Max abs error", 

483 "------------+---------------------------------------------------+---------------+------------+--------------", 

484 ] 

485 for row in rows[:50]: 

486 error = "nonfinite" if row["nonfinite_elements"] else f"{row['max_abs_error']:.6g}" 

487 lines.append(f"{row['result']:<11} | {row['name'][-49:]:<49} | {row['elements']:>13,} | " 

488 f"{row['load_dtype_mismatches']:>10,} | {error}") 

489 first = row.get("first_difference") 

490 if row["result"] == "MISMATCH" and first: 

491 index = first.get("hf_index", first.get("local_index")) 

492 lines.append(f" {row['name']}: rank={first['rank']} index={index} " 

493 f"actual={first['actual']} reference={first['reference']} " 

494 f"reference_in_load_dtype={first['reference_in_load_dtype']}") 

495 if len(rows) > 50: 

496 lines.append(f"Showing 50/{len(rows)} rows (mismatches first); complete table: {root / 'tensors.csv'}") 

497 lines.append("Checked elements and mismatch counts in table include replicated rank observations.") 

498 lines.append(f"Reports: {root}") 

499 logger.info("\n%s", "\n".join(lines)) 

500 

501 

502def _write_table(rows: list[dict[str, Any]], root: Path) -> None: 

503 """Write per-tensor comparison results as a CSV report.""" 

504 fields = [ 

505 "name", "kind", "result", "elements", "exact_reference_mismatches", "load_dtype_mismatches", 

506 "nonfinite_elements", "max_abs_error", "max_load_dtype_abs_error", "ranks_with_mismatch", "first_difference", 

507 "coverage", 

508 ] 

509 with (root / "tensors.csv").open("w", encoding="utf-8", newline="") as handle: 

510 writer = csv.DictWriter(handle, fieldnames=fields) 

511 writer.writeheader() 

512 for row in rows: 

513 writer.writerow({field: json.dumps(row[field]) if isinstance(row.get(field), (list, dict)) 

514 else row.get(field) for field in fields}) 

515 

516 

517@torch.no_grad() 

518def verify_hf_weights(model: nn.Module, pretrained_path: str, output_dir: str, *, 

519 chunk_elements: int = 1 << 20, phase: str = "caller_managed_before_optimizer") -> dict[str, Any]: 

520 """Reverse-convert live model weights and compare every element to local HF storage. 

521 

522 All ranks must call together, before parameter updates. The automatic builder 

523 hook runs before the first forward. The checker never writes model tensors. 

524 Quantized source weights are decoded to BF16 using Transformers; ordinary 

525 source dtypes are preserved. Both raw equality and exact equality after the 

526 declared load-dtype cast are reported. No tolerance is used. Padding is 

527 checked against zero separately. Unsupported inverse conversions, missing 

528 model/source coverage and nonfinite values fail rather than being skipped. 

529 

530 Args: 

531 model: Finalized, materialized and weight-loaded model. 

532 pretrained_path: Local checkpoint; its index defines the comparison scope. 

533 output_dir: Directory for per-rank reports and summary.json. 

534 chunk_elements: Maximum number of elements per comparison tile. 

535 phase: Description of the invocation point, recorded in the reports. 

536 

537 Returns: 

538 Global summary on every rank. PASS requires exact load-dtype agreement. 

539 

540 Raises: 

541 RuntimeError: Any rank finds mismatches, unsupported conversions or missing coverage. 

542 """ 

543 rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 

544 root = Path(output_dir).expanduser().resolve() 

545 local, shapes = _check_rank_weights(model, pretrained_path, root, rank, chunk_elements, phase) 

546 reports = _gather(local) 

547 # All ranks use rank zero's reference scope, including when a local reader failed. 

548 scopes = _gather(shapes if rank == 0 else None) 

549 summary = _summarize(reports, scopes[0] or shapes) 

550 summary.update( 

551 phase=phase, reference="HF decoded weights; load-dtype equality, no tolerance", 

552 reference_scope="The selected HF index; unused modalities or cropped-out layers are not included", 

553 ) 

554 rows = _tensor_table(reports, summary) 

555 _add_result_counts(rows, summary) 

556 write_error = _write_summary_reports(rank, root, rows, summary) 

557 write_errors = _gather(write_error) 

558 _raise_check_failure(summary, write_errors, root) 

559 logger.info("HF weight check PASS on %d ranks; exact reference match=%s; reports=%s", 

560 summary["ranks"], summary["exact_reference_match"], root) 

561 return summary 

562 

563 

564def _check_rank_weights(model: nn.Module, pretrained_path: str, root: Path, rank: int, 

565 chunk_elements: int, phase: str) -> tuple[dict[str, Any], dict[str, tuple[int, ...]]]: 

566 """Run one rank's bounded check and capture errors for the later collective.""" 

567 shapes = {} 

568 local = {"rank": rank, "phase": phase, "checkpoint": str(pretrained_path)} 

569 try: 

570 if chunk_elements < 1: 

571 raise ValueError("chunk_elements must be positive") 

572 root.mkdir(parents=True, exist_ok=True) 

573 reference = _Reference(pretrained_path) 

574 shapes = reference.shapes 

575 if rank == 0: 

576 local["source_files"] = sorted({str(path.resolve()) for path in reference.index.files_by_key.values()}) 

577 logger.info("Starting HF weight check on rank %d: %d source tensors; phase=%s", rank, len(shapes), phase) 

578 local.update(_check_local(model, reference, chunk_elements)) 

579 for field in ("exact_reference_mismatches", "load_dtype_mismatches", "nonfinite_elements"): 

580 local[field] = sum(record.get(field, 0) for record in local["tensors"]) 

581 (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2, allow_nan=False), encoding="utf-8") 

582 except Exception as exc: # pylint: disable=broad-exception-caught 

583 # A local failure must reach the collective so other ranks cannot continue training. 

584 local["error"] = f"{type(exc).__name__}: {exc}" 

585 try: 

586 root.mkdir(parents=True, exist_ok=True) 

587 (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2), encoding="utf-8") 

588 except OSError: 

589 logger.exception("Cannot write HF weight-check rank report") 

590 return local, shapes 

591 

592 

593def _add_result_counts(rows: list[dict[str, Any]], summary: dict[str, Any]) -> None: 

594 """Add tensor-result totals to the global verification summary.""" 

595 weight_rows = [row for row in rows if row["kind"] == "weight"] 

596 summary.update( 

597 exact_tensors=sum(row["result"] == "EXACT" for row in weight_rows), 

598 dtype_cast_tensors=sum(row["result"] == "DTYPE_CAST" for row in weight_rows), 

599 mismatched_tensors=sum(row["result"] == "MISMATCH" for row in weight_rows), 

600 incomplete_tensors=len(summary["uncovered_hf_tensors"]), 

601 padding_tensors=len(rows) - len(weight_rows), 

602 ) 

603 

604 

605def _write_summary_reports(rank: int, root: Path, rows: list[dict[str, Any]], 

606 summary: dict[str, Any]) -> str | None: 

607 """Write rank-zero aggregate reports and return any filesystem error.""" 

608 if rank != 0: