#!/usr/bin/env python3 """Convert the frozen FP32 evaluation lens to deterministic Safetensors.""" from __future__ import annotations import argparse import hashlib import json import os from collections.abc import Mapping from pathlib import Path from typing import Any import safetensors import torch from safetensors import safe_open SOURCE_SHA256 = "8f752032a26a5196c1cb447ef63f01e8a29820ff0c57178dd80c6e26d32b12a9" SOURCE_PROVENANCE_SHA256 = ( "b6e5764fa1580a142403a425ecd03cafe55d4cc36ca1e6b90da7fc19a35aad36" ) SOURCE_VALIDATION_SHA256 = ( "4aea71008a10ef2d043129f7767e5d387372fe1fd5022550b59017a44e0b965d" ) SOURCE_MATRIX_STATS_SHA256 = ( "a21fe7c1f661fa9b5455d89c0a415beae794012eab38f782b667b27fac942401" ) SOURCE_FIT_CHECKPOINT_SHA256 = ( "a1236cfe5d04601575b3de150ffe50e3a67e755ede1197ecf74c205a31bdc258" ) SOURCE_EXPORT_SCRIPT_SHA256 = ( "d1d9e0b7afd1d69771907a839f62b30936995fdd63b17ea2ac80f724b3cdce17" ) FP16_ARTIFACT_SHA256 = ( "089d776979408f23e5377539c15aa8025d171633718ccdce709bcd3372e7942c" ) MODEL_ID = "WeiboAI/VibeThinker-3B" MODEL_REVISION = "77bd2cced09193c8b9a59a32bd8577bbd1f3e01c" SOURCE_LAYERS = tuple(range(0, 36, 2)) TARGET_LAYER = 35 D_MODEL = 2048 N_PROMPTS = 1000 EXPECTED_TOP_LEVEL_KEYS = {"J", "n_prompts", "source_layers", "d_model"} FORBIDDEN_PUBLIC_FRAGMENTS = ( os.sep.join(("", "Users", "")), os.sep.join(("", "Volumes", "")), os.sep.join(("", "workspace")), ) def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def tensor_storage_bytes(tensor: torch.Tensor) -> bytes: """Return C-contiguous little-endian FP32 bytes without value conversion.""" array = tensor.detach().cpu().contiguous().view(torch.int32).numpy() return array.astype(" str: return hashlib.sha256(tensor_storage_bytes(tensor)).hexdigest() def write_json(path: Path, value: Any) -> None: path.write_text( json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True) + "\n", encoding="utf-8", ) def require(condition: bool, message: str) -> None: if not condition: raise ValueError(message) def validate_source(checkpoint: Any) -> dict[int, torch.Tensor]: require(isinstance(checkpoint, Mapping), "checkpoint must be a mapping") require( set(checkpoint) == EXPECTED_TOP_LEVEL_KEYS, f"unexpected checkpoint keys: {sorted(checkpoint)}", ) require(checkpoint["n_prompts"] == N_PROMPTS, "unexpected n_prompts") require(checkpoint["d_model"] == D_MODEL, "unexpected d_model") require( tuple(checkpoint["source_layers"]) == SOURCE_LAYERS, "unexpected source_layers", ) matrices = checkpoint["J"] require(isinstance(matrices, Mapping), "J must be a layer-to-tensor mapping") require(set(matrices) == set(SOURCE_LAYERS), "unexpected J layer keys") validated: dict[int, torch.Tensor] = {} for layer in SOURCE_LAYERS: tensor = matrices[layer] require(isinstance(tensor, torch.Tensor), f"J[{layer}] is not a tensor") require(tensor.device.type == "cpu", f"J[{layer}] is not on CPU") require(tensor.dtype == torch.float32, f"J[{layer}] is not FP32") require( tuple(tensor.shape) == (D_MODEL, D_MODEL), f"J[{layer}] shape mismatch", ) require(tensor.is_contiguous(), f"J[{layer}] is not contiguous") require( bool(torch.isfinite(tensor).all()), f"J[{layer}] contains non-finite values", ) validated[layer] = tensor return validated def public_header() -> dict[str, str]: return { "artifact_kind": "jacobian_lens_evaluation_fp32", "d_model": str(D_MODEL), "format": "pt", "model_id": MODEL_ID, "model_revision": MODEL_REVISION, "n_prompts": str(N_PROMPTS), "schema_version": "1", "source_fit_checkpoint_sha256": SOURCE_FIT_CHECKPOINT_SHA256, "source_fp32_checkpoint_sha256": SOURCE_SHA256, "source_layers": json.dumps(SOURCE_LAYERS, separators=(",", ":")), "target_layer": str(TARGET_LAYER), "tensor_dtype": "float32", "tensor_key_pattern": "J.{source_layer}", } def assert_public_header(metadata: Mapping[str, str]) -> None: encoded = json.dumps(dict(metadata), sort_keys=True) for fragment in FORBIDDEN_PUBLIC_FRAGMENTS: require( fragment not in encoded, f"private fragment found in Safetensors header: {fragment}", ) def save_deterministic_safetensors( tensors: Mapping[str, torch.Tensor], path: Path, metadata: Mapping[str, str], ) -> None: offset = 0 header: dict[str, Any] = { "__metadata__": {key: metadata[key] for key in sorted(metadata)} } for key in sorted(tensors): tensor = tensors[key] require(tensor.dtype == torch.float32, f"{key} is not FP32") nbytes = tensor.numel() * tensor.element_size() header[key] = { "dtype": "F32", "shape": list(tensor.shape), "data_offsets": [offset, offset + nbytes], } offset += nbytes encoded_header = json.dumps( header, ensure_ascii=False, separators=(",", ":"), ).encode("utf-8") encoded_header += b" " * ((-len(encoded_header)) % 8) with path.open("wb") as handle: handle.write(len(encoded_header).to_bytes(8, "little", signed=False)) handle.write(encoded_header) for key in sorted(tensors): handle.write(tensor_storage_bytes(tensors[key])) def compare_fp16( matrices: Mapping[int, torch.Tensor], fp16_path: Path, ) -> dict[str, Any]: require( sha256_file(fp16_path) == FP16_ARTIFACT_SHA256, "FP16 companion artifact SHA-256 mismatch", ) per_layer: dict[str, Any] = {} error_squared = 0.0 reference_squared = 0.0 with safe_open(fp16_path, framework="pt", device="cpu") as fp16_artifact: require( set(fp16_artifact.keys()) == {f"J.{layer}" for layer in SOURCE_LAYERS}, "FP16 companion key set mismatch", ) for layer in SOURCE_LAYERS: fp32 = matrices[layer] stored_fp16 = fp16_artifact.get_tensor(f"J.{layer}") cast_fp16 = fp32.to(torch.float16) require( torch.equal(cast_fp16.view(torch.int16), stored_fp16.view(torch.int16)), f"FP32-to-FP16 cast differs at J.{layer}", ) error = cast_fp16.float() - fp32 error_norm = float(torch.linalg.vector_norm(error)) reference_norm = float(torch.linalg.vector_norm(fp32)) layer_error_squared = error_norm**2 layer_reference_squared = reference_norm**2 error_squared += layer_error_squared reference_squared += layer_reference_squared per_layer[str(layer)] = { "cast_matches_model_safetensors_exactly": True, "max_absolute_error": float(error.abs().max().item()), "relative_frobenius_error": error_norm / reference_norm, } return { "schema_version": 1, "artifact_kind": "fp32_to_fp16_lens_compatibility", "fp32_source_checkpoint_sha256": SOURCE_SHA256, "fp16_artifact": "model.safetensors", "fp16_artifact_sha256": FP16_ARTIFACT_SHA256, "conversion": "IEEE_FP32_to_FP16_round_to_nearest_even", "all_layer_casts_match_exactly": True, "relative_frobenius_error": (error_squared / reference_squared) ** 0.5, "max_absolute_error": max( record["max_absolute_error"] for record in per_layer.values() ), "per_layer": per_layer, } def convert(source: Path, output_dir: Path, overwrite: bool) -> dict[str, Any]: source_digest = sha256_file(source) require(source_digest == SOURCE_SHA256, "source checkpoint SHA-256 mismatch") checkpoint = torch.load(source, map_location="cpu", weights_only=True) matrices = validate_source(checkpoint) tensors = {f"J.{layer}": matrices[layer] for layer in SOURCE_LAYERS} output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / "evaluation.safetensors" if output_path.exists() and not overwrite: raise FileExistsError( f"refusing to overwrite {output_path.name}; pass --overwrite" ) metadata = public_header() assert_public_header(metadata) temporary_path = output_dir / ".evaluation.safetensors.tmp" save_deterministic_safetensors(tensors, temporary_path, metadata) os.replace(temporary_path, output_path) manifest_tensors: dict[str, Any] = {} exact_matches = 0 with safe_open(output_path, framework="pt", device="cpu") as artifact: stored_metadata = artifact.metadata() or {} require(stored_metadata == metadata, "Safetensors metadata changed") assert_public_header(stored_metadata) require(set(artifact.keys()) == set(tensors), "Safetensors key set mismatch") for layer in SOURCE_LAYERS: key = f"J.{layer}" source_tensor = matrices[layer] output_tensor = artifact.get_tensor(key) require(output_tensor.dtype == torch.float32, f"{key} dtype mismatch") require( tuple(output_tensor.shape) == (D_MODEL, D_MODEL), f"{key} shape mismatch", ) require(output_tensor.is_contiguous(), f"{key} is not contiguous") require(torch.equal(output_tensor, source_tensor), f"{key} value mismatch") require( torch.equal( output_tensor.view(torch.int32), source_tensor.view(torch.int32), ), f"{key} FP32 bit-pattern mismatch", ) source_tensor_digest = tensor_sha256(source_tensor) require( tensor_sha256(output_tensor) == source_tensor_digest, f"{key} raw byte hash mismatch", ) exact_matches += 1 manifest_tensors[key] = { "dtype": "float32", "nbytes": source_tensor.numel() * source_tensor.element_size(), "numel": source_tensor.numel(), "sha256_c_contiguous_little_endian_bytes": source_tensor_digest, "shape": list(source_tensor.shape), "source_layer": layer, } output_digest = sha256_file(output_path) output_size = output_path.stat().st_size manifest = { "schema_version": 1, "artifact": "evaluation.safetensors", "artifact_sha256": output_digest, "artifact_size_bytes": output_size, "source_checkpoint_sha256": source_digest, "tensor_count": len(manifest_tensors), "tensor_storage_bytes": sum( record["nbytes"] for record in manifest_tensors.values() ), "tensors": manifest_tensors, } compatibility = compare_fp16(matrices, output_dir / "model.safetensors") compatibility["fp32_artifact"] = "evaluation.safetensors" compatibility["fp32_artifact_sha256"] = output_digest validation = { "schema_version": 1, "artifact": "evaluation.safetensors", "artifact_sha256": output_digest, "artifact_size_bytes": output_size, "source_checkpoint_sha256": source_digest, "exact_tensor_matches": exact_matches, "expected_tensor_matches": len(SOURCE_LAYERS), "tensor_values_changed": 0, "checks": { "all_source_tensors_contiguous": True, "all_source_tensors_finite": True, "all_source_tensors_fp32": True, "all_source_tensors_shape_2048x2048": True, "fp16_cast_matches_companion_artifact": True, "roundtrip_all_tensor_byte_hashes_equal": True, "roundtrip_all_tensor_dtypes_equal": True, "roundtrip_all_tensor_shapes_equal": True, "roundtrip_all_tensor_values_equal": True, "roundtrip_all_tensor_bit_patterns_equal": True, "roundtrip_key_set_exact": True, "safetensors_header_public_safe": True, "safetensors_header_roundtrip_exact": True, "source_checkpoint_sha256_exact": True, "source_metadata_exact": True, "source_top_level_key_set_exact": True, }, "ok": exact_matches == len(SOURCE_LAYERS), } provenance = { "schema_version": 1, "artifact_kind": "jacobian_lens_evaluation_fp32_provenance", "artifact": { "filename": "evaluation.safetensors", "format": "safetensors", "sha256": output_digest, "size_bytes": output_size, "tensor_conversion": "lossless_fp32_reserialization", }, "source_checkpoint": { "format": "pytorch", "sha256": source_digest, }, "derivation": { "formula": "jacobian_sum[layer] / n_done", "n_done": N_PROMPTS, "source_fit_checkpoint_sha256": SOURCE_FIT_CHECKPOINT_SHA256, "source_export_script_sha256": SOURCE_EXPORT_SCRIPT_SHA256, "source_matrix_stats_sha256": SOURCE_MATRIX_STATS_SHA256, "source_provenance_sha256": SOURCE_PROVENANCE_SHA256, "source_validation_sha256": SOURCE_VALIDATION_SHA256, }, "model": { "architecture": "Qwen2ForCausalLM", "d_model": D_MODEL, "id": MODEL_ID, "n_layers": 36, "revision": MODEL_REVISION, "revision_binding": "inferred_hub_head_unchanged_since_before_fit", "revision_last_modified": "2026-06-30T11:35:41+00:00", "tied_embeddings": True, }, "lens": { "d_model": D_MODEL, "dtype": "float32", "hook_convention": "post_transformer_block_output_residual", "n_prompts": N_PROMPTS, "source_layers": list(SOURCE_LAYERS), "target_layer": TARGET_LAYER, "estimator": { "dim_batch": 8, "exclude_final_position": True, "fit_dtype": "bfloat16", "max_seq_len": 128, "name": "causal_all_current_and_future_targets_mean_jacobian", "prompt_aggregation": "equal_weight_mean_over_prompts", "skip_first": 16, "source_position_aggregation": "mean_over_valid_source_positions", "target_position_aggregation": ( "sum_over_valid_targets_at_or_after_source" ), }, }, "compatibility": { "file": "evaluation_compatibility.json", "fp16_artifact": "model.safetensors", "fp16_artifact_sha256": FP16_ARTIFACT_SHA256, "all_layer_casts_match_exactly": True, }, "software": { "conversion_runtime": { "python_implementation": "CPython", "safetensors": safetensors.__version__, "torch": torch.__version__, }, "anthropic_jacobian_lens_commit": ( "581d398613e5602a5af361e1c34d3a92ea82ba8e" ), }, "limitations": [ "The original fit did not store the resolved Hugging Face commit; the revision binding was reconstructed from the Hub head and its last-modified timestamp.", "The recorded evaluation validates token readout and does not establish causal steering or a global workspace.", ], } write_json(output_dir / "evaluation_tensor_manifest.json", manifest) write_json(output_dir / "evaluation_compatibility.json", compatibility) write_json(output_dir / "evaluation_validation.json", validation) write_json(output_dir / "evaluation_provenance.json", provenance) return { "artifact": output_path.name, "artifact_sha256": output_digest, "artifact_size_bytes": output_size, "exact_tensor_matches": exact_matches, "fp16_cast_matches": compatibility["all_layer_casts_match_exactly"], "source_checkpoint_sha256": source_digest, } def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--source", required=True, type=Path) parser.add_argument( "--output-dir", type=Path, default=Path(__file__).resolve().parents[1], ) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() result = convert(args.source, args.output_dir.resolve(), args.overwrite) print(json.dumps(result, indent=2, sort_keys=True)) if __name__ == "__main__": main()