| |
| """Convert the frozen VibeThinker J Lens checkpoint to Safetensors. |
| |
| The input path is intentionally required at runtime and is never copied into |
| the safetensors header or any generated public metadata. The conversion is |
| lossless: every FP16 bit pattern is compared after a safetensors round trip. |
| """ |
|
|
| 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 torch |
| from safetensors import safe_open |
|
|
| SOURCE_SHA256 = "f36a99447623e0d777c70951a9148a7a52e42e0df82942e22ce6326f63d8d664" |
| 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 bytes, including FP16 bit patterns.""" |
|
|
| array = tensor.detach().cpu().contiguous().view(torch.int16).numpy() |
| return array.astype("<i2", 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.float16, f"J[{layer}] is not FP16") |
| 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", |
| "d_model": str(D_MODEL), |
| "format": "pt", |
| "model_id": MODEL_ID, |
| "model_revision": MODEL_REVISION, |
| "n_prompts": str(N_PROMPTS), |
| "schema_version": "1", |
| "source_checkpoint_sha256": SOURCE_SHA256, |
| "source_layers": json.dumps(SOURCE_LAYERS, separators=(",", ":")), |
| "target_layer": str(TARGET_LAYER), |
| "tensor_dtype": "float16", |
| "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: |
| """Write the documented safetensors format with canonical key ordering. |
| |
| The upstream writer preserves tensor data exactly, but its Rust metadata |
| map can serialize keys in a process-random order. Canonical JSON ordering |
| makes the complete artifact reproducible byte for byte across runs. |
| """ |
|
|
| 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.float16, f"{key} is not FP16") |
| nbytes = tensor.numel() * tensor.element_size() |
| header[key] = { |
| "dtype": "F16", |
| "shape": list(tensor.shape), |
| "data_offsets": [offset, offset + nbytes], |
| } |
| offset += nbytes |
|
|
| encoded_header = json.dumps( |
| header, |
| ensure_ascii=False, |
| separators=(",", ":"), |
| ).encode("utf-8") |
| padding = (-len(encoded_header)) % 8 |
| encoded_header += b" " * padding |
|
|
| with path.open("wb") as handle: |
| handle.write(len(encoded_header).to_bytes(8, byteorder="little", signed=False)) |
| handle.write(encoded_header) |
| for key in sorted(tensors): |
| handle.write(tensor_storage_bytes(tensors[key])) |
|
|
|
|
| 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 / "model.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 / ".model.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 during serialization") |
| 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 == source_tensor.dtype, f"{key} dtype mismatch") |
| require(tuple(output_tensor.shape) == tuple(source_tensor.shape), 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.int16), source_tensor.view(torch.int16)), |
| f"{key} FP16 bit-pattern mismatch", |
| ) |
| source_tensor_digest = tensor_sha256(source_tensor) |
| output_tensor_digest = tensor_sha256(output_tensor) |
| require(source_tensor_digest == output_tensor_digest, f"{key} byte hash mismatch") |
| exact_matches += 1 |
| manifest_tensors[key] = { |
| "dtype": "float16", |
| "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 |
| tensor_manifest = { |
| "artifact": "model.safetensors", |
| "artifact_sha256": output_digest, |
| "artifact_size_bytes": output_size, |
| "schema_version": 1, |
| "tensor_count": len(manifest_tensors), |
| "tensor_storage_bytes": sum(item["nbytes"] for item in manifest_tensors.values()), |
| "tensors": manifest_tensors, |
| } |
| validation = { |
| "artifact": "model.safetensors", |
| "artifact_sha256": output_digest, |
| "artifact_size_bytes": output_size, |
| "checks": { |
| "all_source_tensors_contiguous": True, |
| "all_source_tensors_finite": True, |
| "all_source_tensors_fp16": True, |
| "all_source_tensors_shape_2048x2048": 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, |
| }, |
| "exact_tensor_matches": exact_matches, |
| "expected_tensor_matches": len(SOURCE_LAYERS), |
| "ok": exact_matches == len(SOURCE_LAYERS), |
| "schema_version": 1, |
| "source_checkpoint_sha256": source_digest, |
| "tensor_values_changed": 0, |
| } |
| write_json(output_dir / "tensor_manifest.json", tensor_manifest) |
| write_json(output_dir / "validation.json", validation) |
| (output_dir / "SHA256SUMS").write_text( |
| f"{output_digest} model.safetensors\n", |
| encoding="ascii", |
| ) |
| return validation |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--source", required=True, type=Path, help="Private source .pt checkpoint") |
| parser.add_argument("--output-dir", required=True, type=Path, help="Public artifact directory") |
| parser.add_argument("--overwrite", action="store_true") |
| args = parser.parse_args() |
|
|
| result = convert(args.source, args.output_dir, args.overwrite) |
| print(json.dumps(result, indent=2, sort_keys=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
