| |
| """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("<i4", copy=False).tobytes(order="C") |
|
|
|
|
| def tensor_sha256(tensor: torch.Tensor) -> 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() |
|
|
