Diff Coverage

def tile_box(offsets: tuple[int, ...], lengths: tuple[int, ...], limit: int) -> Iterator[Box]:
    """Split a rectangle into non-overlapping tiles of at most ``limit`` elements."""
    if limit < 1:
        raise ValueError("Tile limit must be positive")
    if any(size == 0 for size in lengths):
        return
    if math.prod(lengths) <= limit:
        yield offsets, lengths

    low = [0] * len(block.lengths)
    high = list(block.lengths)
    for size, column in zip(block.lengths, block.steps):
        if size > 1 and sum(step != 0 for step in column) != 1:
            raise ValueError("Inverse conversion has a non-rectangular shard mapping")
    for dim, origin in enumerate(block.base):
        steps = tuple(column[dim] for column in block.steps)
        moving = [(axis, step) for axis, step in enumerate(steps) if step]
        start, stop = shard_offsets[dim], shard_offsets[dim] + shard_shape[dim]

    def __init__(self, path: str) -> None:
        """Read indexed tensor headers and optional quantization metadata."""
        location = Path(path).expanduser().resolve()
        if not location.exists():
            raise ValueError(f"Weight checking requires a local HF checkpoint: {location}")
        self.index = resolve_checkpoint_index(str(location))
        self.quantized = build_quantized_hf_reader(location)
        self.shapes = {}
        if self.quantized is not None:

                with safe_open(str(filename), framework="pt", device="cpu") as reader:
                    for name in names:
                        self.shapes[name] = tuple(reader.get_slice(name).get_shape())
        if not self.shapes:
            raise ValueError("No logical HF tensors were selected for checking")

    def read(self, name: str, offsets: tuple[int, ...], lengths: tuple[int, ...]) -> torch.Tensor:
        """Decode a bounded source rectangle, keeping ordinary source dtypes."""
        if self.quantized is not None and len(lengths) == 2:

        if self.quantized is not None and len(lengths) == 2:
            return self.quantized.read_region(name, offsets, lengths)
        with safe_open(str(self.index.files_by_key[name]), framework="pt", device="cpu") as reader:
            if not lengths:
                return reader.get_tensor(name)
            region = tuple(slice(start, start + count) for start, count in zip(offsets, lengths))
            return reader.get_slice(name)[region]

        self.config = model.config
        self._weight_conversions = conversions

    def __getattr__(self, name: str) -> Any:
        return getattr(self.model, name)


def _model_shards(model: nn.Module) -> dict[str, CheckpointTensorShard]:
    """Describe every materialized local model shard for comparison."""

        if not callable(provider):
            raise ValueError("External state must expose checkpoint_tensor_shards() for complete weight checking")
        external_shards = provider()
        if set(external_shards) != set(external.parameters):
            raise ValueError("External checkpoint shard descriptors do not cover exactly the external parameters")
    shards = {}
    for name, target in build_load_targets(model).items():
        if name in external_shards:
            shards[name] = external_shards[name]

            continue
        layout = target_layout(target)
        view = target
        if layout is not None:
            reject_partial_layout(name, layout)
            if not isinstance(view, DTensor):
                shape = getattr(layout, "tensor_shape", None)
                view = DTensor.from_local_with_layout(target.detach(), layout, shape=shape)
        chunks = create_chunk_list_for_tensor(view)
        if len(chunks) != 1:
            raise ValueError(f"Weight checking requires one rectangular local shard: {name}")
        local = local_target_tensor(target)
        if local.is_meta or tuple(local.shape) != tuple(chunks[0].sizes):
            raise ValueError(f"Live storage does not match its checkpoint shard metadata: {name}")
        shards[name] = CheckpointTensorShard(local, tuple(view.shape), tuple(chunks[0].offsets))
    if not shards:
        raise ValueError("The model has no checkpoint tensors to check")
    return shards


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

        stages = [replacements, getattr(model, "_hp_used_base_weight_conversions", [])]
    else:
        conversions = getattr(model, "_weight_conversions", None)
        if conversions is None:
            conversions = get_model_conversion_mapping(model, add_legacy=False)
        stages = [conversions]
    for conversions in stages:
        converters = [item for item in conversions if isinstance(item, WeightConverter)]
        renamings = [item for item in conversions if isinstance(item, WeightRenaming)]

    for conversions in stages:
        converters = [item for item in conversions if isinstance(item, WeightConverter)]
        renamings = [item for item in conversions if isinstance(item, WeightRenaming)]
        if len(converters) + len(renamings) != len(conversions):
            raise ValueError("Unsupported recorded weight-conversion type")
        if converters:
            regions = _reverse_stage(model, regions, converters)
        for renaming in reversed(renamings):
            regions = _reverse_stage(model, regions, [renaming])

            regions = _reverse_stage(model, regions, converters)
        for renaming in reversed(renamings):
            regions = _reverse_stage(model, regions, [renaming])
    if any(not isinstance(value, RegionTensor) for value in regions.values()):
        raise ValueError("Inverse conversion produced non-symbolic weights; bounded verification is unsupported")
    return regions


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

            relative = name[len(prefix):]
            if any(probe.rename_source_key(relative)[1] is not None for probe in probes):
                lookup = relative
        if lookup in normalized:
            raise ValueError(f"Ambiguous inverse conversion after resolving base-model prefix: {lookup}")
        normalized[lookup] = region
    return revert_weight_conversion(_InverseModel(model, conversions), normalized)

        (origin - start) * stride for origin, start, stride in zip(base, shard.offsets, strides)
    )
    values = tensor.as_strided(lengths, tile_strides, storage_offset).cpu()
    for operation in block.post:
        values = operation(values)
    return values


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

def _difference(actual: torch.Tensor, expected: torch.Tensor, storage_dtype: torch.dtype) -> dict[str, Any]:
    """Count raw and load-dtype mismatches for one pair of tensor tiles."""
    if actual.is_complex() or expected.is_complex():
        raise ValueError("Complex weight checking is not supported")
    if tuple(actual.shape) != tuple(expected.shape):
        raise ValueError(f"Compared weight tiles have different shapes: {actual.shape}, {expected.shape}")
    raw_bad = actual != expected
    if actual.is_floating_point() and expected.is_floating_point():
        raw_bad |= (actual == 0) & (expected == 0) & (torch.signbit(actual) != torch.signbit(expected))
    nonfinite = ~(torch.isfinite(actual) & torch.isfinite(expected))

    logical = shard.logical_shape
    if logical is None or logical == shard.global_shape:
        return None
    if not logical or logical[1:] != shard.global_shape[1:]:
        raise ValueError(f"External checkpoint padding must be on dimension zero: {name}")
    start = max(0, min(shard.tensor.shape[0], logical[0] - shard.offsets[0]))
    lengths = (shard.tensor.shape[0] - start, *shard.tensor.shape[1:])
    record = {"model_key": name, "kind": "padding", "elements": 0}
    for offsets, sizes in tile_box((start, *((0,) * (len(lengths) - 1))), lengths, limit):

    prefix = getattr(model, "base_model_prefix", "") + "."
    for number, (recovered_name, region) in enumerate(sorted(recovered.items()), start=1):
        name = recovered_name
        if name not in reference.shapes and name.startswith(prefix):
            name = name[len(prefix):]
        if name not in reference.shapes:
            unmatched.append(recovered_name)
            continue
        if tuple(region.shape) != reference.shapes[name]:
            raise ValueError(f"Recovered HF shape mismatch: {name}: {tuple(region.shape)} vs {reference.shapes[name]}")
        record = {
            "hf_key": name, "shape": list(region.shape), "elements": 0,
            "model_keys": sorted(region.sources()),
        }

            "hf_key": name, "shape": list(region.shape), "elements": 0,
            "model_keys": sorted(region.sources()),
        }
        if math.prod(region.shape) >= 1 << 30 or number % 200 == 0:
            logger.info("Checking HF tensor %d/%d: %s, shape=%s", number, len(recovered), name, tuple(region.shape))
        for block in region.blocks:
            shard = shards[block.source]
            box = local_block_box(block, shard.offsets, tuple(shard.tensor.shape))
            if box is None:

def _gather(value: Any) -> list[Any]:
    if not (dist.is_available() and dist.is_initialized()):
        return [value]
    values = [None] * dist.get_world_size()
    dist.all_gather_object(values, value)
    return values


def _summarize(reports: list[dict[str, Any]], shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]:
    """Combine rank reports into global mismatch and coverage results."""

    raw_mismatches = 0
    nonfinite = 0
    for rank, report in enumerate(reports):
        if report.get("error"):
            errors.append({"rank": rank, "error": report["error"]})
        for name, boxes in report.get("source_coverage", {}).items():
            coverage[name].extend(boxes)
        if report.get("uncovered_model_shards") or report.get("unmatched_recovered_keys"):
            errors.append({"rank": rank, "error": "Model weights are not completely covered; see rank report"})

            lines.append(f"  {row['name']}: rank={first['rank']} index={index} "
                         f"actual={first['actual']} reference={first['reference']} "
                         f"reference_in_load_dtype={first['reference_in_load_dtype']}")
    if len(rows) > 50:
        lines.append(f"Showing 50/{len(rows)} rows (mismatches first); complete table: {root / 'tensors.csv'}")
    lines.append("Checked elements and mismatch counts in table include replicated rank observations.")
    lines.append(f"Reports: {root}")
    logger.info("\n%s", "\n".join(lines))

    shapes = {}
    local = {"rank": rank, "phase": phase, "checkpoint": str(pretrained_path)}
    try:
        if chunk_elements < 1:
            raise ValueError("chunk_elements must be positive")
        root.mkdir(parents=True, exist_ok=True)
        reference = _Reference(pretrained_path)
        shapes = reference.shapes
        if rank == 0:

        local.update(_check_local(model, reference, chunk_elements))
        for field in ("exact_reference_mismatches", "load_dtype_mismatches", "nonfinite_elements"):
            local[field] = sum(record.get(field, 0) for record in local["tensors"])
        (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2, allow_nan=False), encoding="utf-8")
    except Exception as exc:  # pylint: disable=broad-exception-caught
        # A local failure must reach the collective so other ranks cannot continue training.
        local["error"] = f"{type(exc).__name__}: {exc}"
        try:
            root.mkdir(parents=True, exist_ok=True)
            (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2), encoding="utf-8")
        except OSError:
            logger.exception("Cannot write HF weight-check rank report")
    return local, shapes


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

def _write_summary_reports(rank: int, root: Path, rows: list[dict[str, Any]],
                           summary: dict[str, Any]) -> str | None:
    """Write rank-zero aggregate reports and return any filesystem error."""
    if rank != 0:
        return None
    try:
        root.mkdir(parents=True, exist_ok=True)
        _write_table(rows, root)
        (root / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

        root.mkdir(parents=True, exist_ok=True)
        _write_table(rows, root)
        (root / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
        _print_table(rows, summary, root)
    except OSError as exc:
        return str(exc)
    return None


def _raise_check_failure(summary: dict[str, Any], write_errors: list[Any], root: Path) -> None: