#!/usr/bin/env python3 """Stream the official GLM-5.2 BF16 checkpoint into an all-SQG model. The materializer never reads a quantized carrier. It keeps exactly the 837 official tensors outside the SQG replacement domain, replaces every routed and non-routed matrix from sealed SQG artifacts, and publishes one atomic receipt per output shard. An interrupted run resumes only when the immutable plan and every already-published shard still match their receipts. """ from __future__ import annotations import argparse from copy import deepcopy import json import os from pathlib import Path import re import shutil import struct import sys import tempfile from typing import Any, Iterable, Mapping PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from bmmlaw_r7_encoder.safetensors_io import ( # noqa: E402 SafeTensorReader, TensorEntry, write_safetensors_atomic, ) from scripts.full_w4a8_production_lineage_v2 import ( # noqa: E402 LAYER_SCHEMA, _pretty_id, ) from scripts.validate_full_sqg_checkpoint import ( # noqa: E402 DIRECT_BF16_MATRICES, NONROUTED_K6_MATRICES, OFFICIAL_SOURCE_TENSORS, ROUTED_MATRICES, TOTAL_SQG_MATRICES, replacement_weights, validate as validate_full_checkpoint, ) from src.fresh_pipeline_artifacts import LAYER_ARTIFACT_SCHEMA # noqa: E402 from src.fresh_pipeline_common import ( # noqa: E402 SQG_MARKER, atomic_json, canonical_sha256, sha256_file, ) from src.glm52_fresh_sqg.manifest import validate_tensor_manifest # noqa: E402 from src.progressive_w4a8 import _derived_down_target # noqa: E402 SCHEMA = "glm52-bf16-only-all-sqg-materialization-v1" PLAN_SCHEMA = "glm52-bf16-only-all-sqg-materialization-plan-v1" SHARD_SCHEMA = "glm52-bf16-only-all-sqg-shard-receipt-v1" K6_SCHEMA = "glm52-dense-k6-layer-v1" CODEBOOK = "sqg_xor_cheb_t12" SOURCE_REPO = "zai-org/GLM-5.2" SOURCE_REVISION = "b4734de4facf877f85769a911abafc5283eab3d9" SOURCE_INDEX_SHA256 = "5fd47a926aefce0f2c917f42523e5e0f3c87e23e389e767c3681536a62f5cf5e" SOURCE_CONFIG_SHA256 = "185f93ee6d12548e16a847e279dc0c3c90b1524c970b0866b42fb545747d859a" FIRST_LAYER = 3 LAST_LAYER = 78 ROUTED_TENSORS_PER_LAYER = 768 ROUTED_SHARD_TENSORS = 2_306 K6_ROLES = ("shared_gate", "shared_up", "shared_down", "q_b", "o_proj") K6_SUFFIX = { "shared_gate": "mlp.shared_experts.gate_proj", "shared_up": "mlp.shared_experts.up_proj", "shared_down": "mlp.shared_experts.down_proj", "q_b": "self_attn.q_b_proj", "o_proj": "self_attn.o_proj", } ALLOWED_SQG_SUFFIXES = (".trellis", ".suh", ".svh", ".sqg") def _load_json(path: Path) -> dict[str, Any]: if not path.is_file() or path.is_symlink(): raise FileNotFoundError(f"JSON input is absent or unsafe: {path}") value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): raise TypeError(f"JSON root must be an object: {path}") return value def _regular(path: Path, *, label: str) -> Path: path = path.resolve() if not path.is_file() or path.is_symlink(): raise FileNotFoundError(f"{label} is absent or unsafe: {path}") return path def _safe_sibling(parent: Path, name: object, *, label: str) -> Path: if not isinstance(name, str) or Path(name).name != name: raise ValueError(f"{label} filename is unsafe: {name!r}") path = parent / name if not path.is_file() or path.is_symlink() or path.resolve().parent != parent.resolve(): raise FileNotFoundError(f"{label} is absent, symlinked, or escapes: {path}") return path.resolve() def _paths_overlap(first: Path, second: Path) -> bool: a = first.absolute() b = second.resolve() return a == b or a in b.parents or b in a.parents def _marker_is_sqg(info) -> bool: if info.dtype != "I32" or info.shape != () or info.nbytes != 4: return False raw = b"".join(info.payload.chunks()) return struct.unpack(" None: """Reject actual MCG authority while allowing explicit zero censuses.""" if isinstance(value, Mapping): for key, item in value.items(): lowered = str(key).lower() if ".mcg" in lowered: raise ValueError(f"{where}: forbidden .mcg key {key!r}") if lowered in {"codebook", "encoding", "source_kind"} and str(item).lower() == "mcg": raise ValueError(f"{where}: MCG authority is forbidden") if lowered in {"mcg_source", "legacy_mcg_input"} and item is not False: raise ValueError(f"{where}: non-false MCG lineage is forbidden") if lowered.startswith("mcg_") and isinstance(item, (int, float)) and item != 0: raise ValueError(f"{where}: nonzero MCG lineage is forbidden") _assert_no_mcg(item, where=where) elif isinstance(value, list): for item in value: _assert_no_mcg(item, where=where) def _expected_routed_names(layer: int) -> set[str]: if not FIRST_LAYER <= layer <= LAST_LAYER: raise ValueError("routed layer must lie in [3,78]") # Spell these names locally: the historical four-layer helper validates # against an environment-selected pilot and cannot represent MTP78. result = { f"model.layers.{layer}.mlp.experts.r7_shared.gate_up_suh", f"model.layers.{layer}.mlp.experts.r7_shared.down_svh", } for expert in range(256): for projection in ("gate_proj", "up_proj", "down_proj"): prefix = f"model.layers.{layer}.mlp.experts.{expert}.{projection}" result.update( { f"{prefix}.trellis", f"{prefix}.sqg", f"{prefix}.suh" if projection == "down_proj" else f"{prefix}.svh", } ) return result def _routed_receipt_path(root: Path, layer: int) -> Path: candidates = ( root / f"layer_{layer:03d}" / "final" / "full_w4a8_production_layer_v2.json", root / f"layer_{layer:03d}" / "full_w4a8_production_layer_v2.json", ) found = [path for path in candidates if path.is_file() and not path.is_symlink()] if len(found) != 1: raise FileNotFoundError( f"layer {layer}: expected exactly one production receipt; found={found}" ) return found[0].resolve() def load_routed_layer(root: Path, layer: int) -> dict[str, Any]: receipt_path = _routed_receipt_path(root, layer) receipt = _load_json(receipt_path) # Direct fixed-map production receipts are emitted by # ``encode_direct_assigned_w4a8_layer.py`` with that encoder's historical # pretty-JSON canonicalizer. Validate them with the matching lineage # helper; ``_compact_id`` is reserved for artifacts emitted by the newer # compact-JSON lineage path. _pretty_id(receipt, "production_layer_id") artifact = receipt.get("selected_layer_artifact") if ( receipt.get("schema") != LAYER_SCHEMA or receipt.get("complete") is not True or receipt.get("production_eligible") is not True or int(receipt.get("layer", -1)) != layer or receipt.get("numerical_device") != "cuda:0" or receipt.get("cpu_numerical_work") is not False or not isinstance(artifact, Mapping) ): raise ValueError(f"layer {layer}: production receipt contract differs") _assert_no_mcg(receipt, where=f"routed receipt L{layer}") manifest_path = _safe_sibling( receipt_path.parent, artifact.get("filename"), label=f"layer {layer} manifest" ) if sha256_file(manifest_path) != artifact.get("sha256"): raise ValueError(f"layer {layer}: production manifest hash differs") manifest = _load_json(manifest_path) expected_names = _expected_routed_names(layer) bit_map = manifest.get("bit_map") payload_hashes = manifest.get("payload_sha256") lineage = manifest.get("lineage") values = tuple(bit_map.values()) if isinstance(bit_map, Mapping) else () if ( manifest.get("schema") != LAYER_ARTIFACT_SCHEMA or manifest.get("complete") is not True or int(manifest.get("layer", -1)) != layer or not isinstance(payload_hashes, Mapping) or set(payload_hashes) != expected_names or len(expected_names) != ROUTED_SHARD_TENSORS or (len(values), values.count(3), values.count(4), sum(values)) != (768, 384, 384, 2688) or not isinstance(lineage, Mapping) or lineage.get("mcg_tensor_count") != 0 or lineage.get("sqg_tensor_count") != ROUTED_TENSORS_PER_LAYER ): raise ValueError(f"layer {layer}: routed SQG manifest differs") _assert_no_mcg(manifest, where=f"routed manifest L{layer}") shard = _safe_sibling( manifest_path.parent, manifest.get("shard"), label=f"layer {layer} SQG shard" ) if artifact.get("shard") != shard.name or artifact.get("shard_sha256") != manifest.get( "shard_sha256" ): raise ValueError(f"layer {layer}: receipt/shard binding differs") reader = SafeTensorReader(shard) if set(reader.tensors) != expected_names: raise ValueError(f"layer {layer}: routed shard tensor domain differs") if any(".mcg" in name.lower() for name in reader.tensors): raise ValueError(f"layer {layer}: routed shard contains MCG tensor") markers = [info for name, info in reader.tensors.items() if name.endswith(".sqg")] if len(markers) != ROUTED_TENSORS_PER_LAYER or any( not _marker_is_sqg(info) for info in markers ): raise ValueError(f"layer {layer}: routed SQG markers differ") return { "layer": layer, "receipt_path": receipt_path, "receipt_sha256": sha256_file(receipt_path), "receipt": receipt, "manifest_path": manifest_path, "manifest_sha256": sha256_file(manifest_path), "manifest": manifest, "shard": shard, "reader": reader, } def _k6_receipt_path(root: Path, layer: int) -> Path: candidates = ( root / f"layer_{layer:03d}" / "SQG_K6_LAYER.json", root / "k6" / f"layer_{layer:03d}" / "SQG_K6_LAYER.json", ) found = [path for path in candidates if path.is_file() and not path.is_symlink()] if len(found) != 1: raise FileNotFoundError( f"layer {layer}: expected exactly one K6 receipt; found={found}" ) return found[0].resolve() def load_k6_layer(root: Path, layer: int) -> dict[str, Any]: path = _k6_receipt_path(root, layer) value = _load_json(path) matrices = value.get("matrices") expected = { f"model.layers.{layer}.{suffix}": role for role, suffix in K6_SUFFIX.items() } if ( value.get("schema") != K6_SCHEMA or value.get("complete") is not True or int(value.get("layer", -1)) != layer or value.get("work_unit") != "finalize" or int(value.get("bits", -1)) != 6 or value.get("codebook") != CODEBOOK or value.get("w4a8") is not True or value.get("numerical_device") != "cuda:0" or value.get("cpu_numerical_work") is not False or value.get("fp32_accumulate") is not True or not isinstance(matrices, Mapping) or set(matrices) != set(expected) ): raise ValueError(f"layer {layer}: K6 final receipt differs") _assert_no_mcg(value, where=f"K6 receipt L{layer}") shard_records: dict[Path, dict[str, Any]] = {} for prefix, role in expected.items(): record = matrices[prefix] names = {f"{prefix}.{suffix}" for suffix in ("trellis", "suh", "svh", "sqg")} if ( not isinstance(record, Mapping) or record.get("tensor_id") != prefix or record.get("official_tensor_name") != f"{prefix}.weight" or record.get("role") != role or int(record.get("bits", -1)) != 6 or record.get("codebook") != CODEBOOK or record.get("weight_endpoint") not in { "native_e4m3", "native_e4m3_sqg", "native_finite_e4m3_sqg_labels", } or record.get("topology_neutral") is not True or record.get("physical_permutation") is not False or record.get("hadamard_side") != "activation" or set(record.get("payload_sha256", {})) != names or not isinstance(record.get("tensor_manifest_id"), str) ): raise ValueError(f"layer {layer}: K6 matrix {prefix} differs") validate_tensor_manifest(record["tensor_manifest"]) if canonical_sha256(record["tensor_manifest"]) != record["tensor_manifest_id"]: raise ValueError(f"layer {layer}: K6 tensor manifest ID differs: {prefix}") mcg_reads = record.get("mcg_reads") if not isinstance(mcg_reads, Mapping) or any(int(item) != 0 for item in mcg_reads.values()): raise ValueError(f"layer {layer}: K6 MCG reads are nonzero: {prefix}") shard = _safe_sibling(path.parent, record.get("shard"), label=f"K6 {prefix} shard") group = shard_records.setdefault( shard, { "expected_sha256": record.get("shard_sha256"), "expected_names": set(), "expected_payloads": {}, "prefixes": [], }, ) if group["expected_sha256"] != record.get("shard_sha256"): raise ValueError(f"layer {layer}: shared K6 shard hash binding differs") group["expected_names"].update(names) group["expected_payloads"].update(record["payload_sha256"]) group["prefixes"].append(prefix) for shard, group in shard_records.items(): reader = SafeTensorReader(shard) if set(reader.tensors) != group["expected_names"]: raise ValueError(f"layer {layer}: K6 shard tensor domain differs: {shard.name}") markers = [info for name, info in reader.tensors.items() if name.endswith(".sqg")] if len(markers) != len(group["prefixes"]) or any( not _marker_is_sqg(info) for info in markers ): raise ValueError(f"layer {layer}: K6 SQG marker differs: {shard.name}") group["reader"] = reader return { "layer": layer, "receipt_path": path, "receipt_sha256": sha256_file(path), "receipt": value, "matrices": dict(matrices), "shards": shard_records, } def load_source(source_model: Path) -> dict[str, Any]: source = source_model.resolve() if not source.is_dir() or source.is_symlink(): raise ValueError("official BF16 source must be a real directory") index_path = _regular(source / "model.safetensors.index.json", label="official index") config_path = _regular(source / "config.json", label="official config") if sha256_file(index_path) != SOURCE_INDEX_SHA256: raise ValueError("official GLM-5.2 index SHA256 differs") if sha256_file(config_path) != SOURCE_CONFIG_SHA256: raise ValueError("official GLM-5.2 config SHA256 differs") index = _load_json(index_path) weight_map = index.get("weight_map") if not isinstance(weight_map, Mapping) or len(weight_map) != OFFICIAL_SOURCE_TENSORS: raise ValueError("official BF16 tensor census differs") routed, nonrouted = replacement_weights() replaced = routed | nonrouted if not replaced <= set(weight_map): raise ValueError("official BF16 index lacks an SQG replacement matrix") ignored = set(weight_map) - replaced if len(ignored) != DIRECT_BF16_MATRICES: raise ValueError("official direct-BF16 ignored census differs") by_shard: dict[str, set[str]] = {} for name, shard in weight_map.items(): if not isinstance(name, str) or not isinstance(shard, str) or Path(shard).name != shard: raise ValueError("official index contains an unsafe entry") by_shard.setdefault(shard, set()).add(name) readers: dict[str, SafeTensorReader] = {} for shard, indexed_names in sorted(by_shard.items()): path = _safe_sibling(source, shard, label="official BF16 shard") reader = SafeTensorReader(path) if set(reader.tensors) != indexed_names: raise ValueError(f"official index/header binding differs: {shard}") readers[shard] = reader return { "root": source, "index_path": index_path, "index": index, "weight_map": dict(weight_map), "by_shard": by_shard, "readers": readers, "ignored": ignored, "replaced": replaced, } def _entry(info) -> TensorEntry: return TensorEntry(info.name, info.dtype, info.shape, info.payload) def _receipt_path(state_root: Path, output_name: str) -> Path: return state_root / "shards" / f"{output_name}.json" def _publish_shard( *, destination: Path, entries: Iterable[TensorEntry], state_root: Path, input_binding: Mapping[str, Any], role: str, ) -> dict[str, Any]: entries = tuple(entries) names = [item.name for item in entries] if not entries or len(names) != len(set(names)): raise ValueError(f"{destination.name}: empty or duplicate shard domain") if any(".mcg" in name.lower() for name in names): raise ValueError(f"{destination.name}: MCG tensor is forbidden") binding = { "output": destination.name, "role": role, "input": dict(input_binding), "tensor_count": len(names), "tensor_domain_sha256": canonical_sha256(sorted(names)), } binding_id = canonical_sha256(binding) receipt_path = _receipt_path(state_root, destination.name) if receipt_path.is_file(): receipt = _load_json(receipt_path) if ( receipt.get("schema") != SHARD_SCHEMA or receipt.get("complete") is not True or receipt.get("binding_id") != binding_id or receipt.get("binding") != binding or not destination.is_file() or destination.is_symlink() or sha256_file(destination) != receipt.get("output_sha256") ): raise ValueError(f"resumable shard receipt differs: {destination.name}") reader = SafeTensorReader(destination) if set(reader.tensors) != set(names): raise ValueError(f"resumed shard tensor domain differs: {destination.name}") return receipt if destination.exists(): raise FileExistsError(f"unreceipted output shard exists: {destination}") tensor_hashes, output_hash = write_safetensors_atomic( destination, entries, metadata={"format": "pt", "sqg_materialization": SCHEMA}, ) receipt = { "schema": SHARD_SCHEMA, "complete": True, "binding": binding, "binding_id": binding_id, "output_bytes": destination.stat().st_size, "output_sha256": output_hash, "payload_sha256": tensor_hashes, } atomic_json(receipt_path, receipt) return receipt def _publish_json(path: Path, value: Mapping[str, Any]) -> None: if path.is_file(): existing = _load_json(path) if existing != value: raise ValueError(f"existing sealed JSON differs: {path}") return if path.exists(): raise ValueError(f"sealed JSON path is not a regular file: {path}") atomic_json(path, value) def _copy_support_file(source: Path, destination: Path) -> dict[str, Any]: expected = sha256_file(source) if destination.is_file(): if destination.is_symlink() or sha256_file(destination) != expected: raise ValueError(f"existing support file differs: {destination}") else: destination.parent.mkdir(parents=True, exist_ok=True) fd, temporary = tempfile.mkstemp( prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent ) temp = Path(temporary) try: with source.open("rb") as inp, os.fdopen(fd, "wb") as out: shutil.copyfileobj(inp, out, length=32 << 20) out.flush() os.fsync(out.fileno()) os.replace(temp, destination) finally: temp.unlink(missing_ok=True) return {"bytes": destination.stat().st_size, "sha256": expected} def _quantization_config( *, tensor_storage: Mapping[str, Any] | None = None, down_targets: Mapping[str, Any] | None = None, ) -> dict[str, Any]: """Build the loader contract while keeping SQG as the codec identity. ``exl3`` is the vLLM loader/dispatch name, not the reconstruction law. The stored marker and codebook fields remain SQG-only. A materialized model supplies the complete per-logical-tensor storage map and the sealed derived down targets; the empty defaults keep this helper useful for static tests. """ return { "quant_method": "exl3", "quant_algo": "SQG", "format": "glm52_sqg_atoms_v2", "codebook": CODEBOOK, "tensor_storage": dict(tensor_storage or {}), "weight_endpoint": "native_e4m3_labels", "activation_endpoint": "full-w4a8", "activation_endpoint_scope": "routed_experts", "allow_a16_fallback": False, "r7_routed_experts": { "schema": "glm52-sqg-atoms-v2-routed-v1", "first_layer": FIRST_LAYER, "last_layer": LAST_LAYER, "moe_layers": [FIRST_LAYER, LAST_LAYER], "bits": [3, 4], "k_values": [3, 4], "independent_per_tensor_rates": True, "per_layer_k3": 384, "per_layer_k4": 384, "codebook": CODEBOOK, "rotation_layout": "shared_h_v1", "gate_up_suh_template": ( "model.layers.{layer}.mlp.experts.r7_shared.gate_up_suh" ), "down_svh_template": ( "model.layers.{layer}.mlp.experts.r7_shared.down_svh" ), "sidecar_template": "r7-experts-layer-{layer:03d}.json", }, "sqg_k6_nonrouted": { "bits": 6, "matrix_count": NONROUTED_K6_MATRICES, "manifest": "SQG_K6_NONROUTED.json", "codebook": CODEBOOK, "execution": "native_sqg_k6_w6a16", "activation_endpoint": "a16", }, "glm_sqg_w4a8": { "schema": "glm52_sqg_atoms_v2_w4a8_v1", "execution": "full_w4a8", "codebook": CODEBOOK, "rates": "independent_per_tensor_k3_k4", "direct_e4m3_weights": True, "allow_a16_fallback": False, "activation": "silu_gate_times_up", "topology": "topology_neutral", "per_layer_bit_census": {"k3": 384, "k4": 384, "total": 768}, "parallelism": { "tensor_parallel_size": 1, "pipeline_parallel_size": 8, "decode_context_parallel_size": 1, "mtp": True, "long_context_max_tokens": 1_048_576, }, "down_targets": dict(down_targets or {}), "h_a8": True, "act_a8": True, "activation_equation": "silu(gate)*up", "hadamard_side": "activation", "topology_neutral": True, "physical_permutation": False, }, } _TORCH_DTYPE_NAMES = { "BF16": "torch.bfloat16", "F16": "torch.float16", "F32": "torch.float32", "I16": "torch.int16", "I32": "torch.int32", "I64": "torch.int64", "U8": "torch.uint8", } def _tensor_storage_metadata( *, weight_map: Mapping[str, str], readers: Mapping[str, SafeTensorReader], bits_by_prefix: Mapping[str, int], ) -> dict[str, Any]: """Describe every SQG logical tensor using vLLM's EXL3 loader ABI. Routed atoms-v2 tensors deliberately keep one topology-neutral h-side rotation per layer: gate/up store only their private ``svh`` and down stores only its private ``suh``. The complementary vectors are the ``r7_shared`` tensors declared by ``r7_routed_experts``; duplicating them into 256 expert records would both misdescribe the bytes and destroy the topology-neutral storage contract. Dense K6 matrices remain ordinary four-component EXL3 records. """ result: dict[str, Any] = {} for prefix, bits in sorted(bits_by_prefix.items()): stored: dict[str, Any] = {} for suffix in ("trellis", "suh", "svh", "sqg"): name = f"{prefix}.{suffix}" shard = weight_map.get(name) if shard is None: continue info = readers[shard].tensors[name] dtype = _TORCH_DTYPE_NAMES.get(info.dtype) if dtype is None: raise ValueError(f"unsupported SQG storage dtype {info.dtype}: {name}") stored[name] = { "dtype": dtype, "shape": list(info.shape), "n_bytes": info.nbytes, } suffixes = {name.rsplit(".", 1)[-1] for name in stored} routed = ".mlp.experts." in prefix and ".r7_shared" not in prefix projection = prefix.rsplit(".", 1)[-1] if routed and projection in {"gate_proj", "up_proj"}: expected = {"trellis", "svh", "sqg"} shared_h = "gate_up_suh" elif routed and projection == "down_proj": expected = {"trellis", "suh", "sqg"} shared_h = "down_svh" else: expected = {"trellis", "suh", "svh", "sqg"} shared_h = None if suffixes != expected: raise ValueError(f"incomplete SQG tensor storage metadata: {prefix}") entry = { "bits_per_weight": int(bits), "quant_format": "exl3", "codec": "sqg", "codebook": CODEBOOK, "stored_tensors": stored, } if shared_h is not None: layer_match = re.search(r"\.layers\.(\d+)\.", prefix) if layer_match is None: raise ValueError(f"cannot resolve routed layer: {prefix}") layer = int(layer_match.group(1)) entry["shared_h_tensor"] = ( f"model.layers.{layer}.mlp.experts.r7_shared.{shared_h}" ) if entry["shared_h_tensor"] not in weight_map: raise ValueError( "missing topology-neutral shared-H tensor: " f"{entry['shared_h_tensor']}" ) result[prefix] = entry if len(result) != TOTAL_SQG_MATRICES: raise ValueError("SQG tensor_storage logical-matrix census differs") return result def _final_tensor_domain( weight_map: Mapping[str, str], ignored: set[str], sqg_prefixes: set[str] ) -> None: names = set(weight_map) shared_h = { f"model.layers.{layer}.mlp.experts.r7_shared.{field}" for layer in range(FIRST_LAYER, LAST_LAYER + 1) for field in ("gate_up_suh", "down_svh") } if names & {f"{prefix}.weight" for prefix in sqg_prefixes}: raise ValueError("final index retained a replaced BF16 weight") if any(name.endswith(".mcg") or ".mcg." in name for name in names): raise ValueError("final index contains an MCG tensor") sqg_names = names - ignored if not shared_h <= names: raise ValueError("final index is missing topology-neutral shared-H tensors") if any( not name.endswith(ALLOWED_SQG_SUFFIXES) and name not in shared_h for name in sqg_names ): raise ValueError("final index contains a non-SQG, non-ignored tensor") markers = {name.removesuffix(".sqg") for name in names if name.endswith(".sqg")} trellis = { name.removesuffix(".trellis") for name in names if name.endswith(".trellis") } if markers != sqg_prefixes or not sqg_prefixes <= trellis: raise ValueError("final SQG marker/trellis domain differs") def materialize(args: argparse.Namespace) -> dict[str, Any]: source = load_source(args.source_model) output = args.output.absolute() production_root = args.production_root.resolve() k6_root = args.k6_root.resolve() hessian_path = _regular(args.hessian_dataset_receipt, label="Hessian upload receipt") hessian = _load_json(hessian_path) commit = hessian.get("commit") revision = hessian.get("revision") if ( hessian.get("complete") is not True or hessian.get("public") is not True or hessian.get("repo") != "brandonmusic/GLM-5.2-BMM-Law-SQG-Hessians" or hessian.get("phase") != "final_publication_verified" or hessian.get("main_verified") is not True or hessian.get("accepted_tag_verified") is not True or revision != "main" or not isinstance(commit, str) or len(commit) != 40 ): raise ValueError("final public Hessian dataset binding differs") for protected in (source["root"], production_root, k6_root, hessian_path.parent): if _paths_overlap(output, protected): raise ValueError(f"output overlaps protected input: {protected}") routed = [load_routed_layer(production_root, layer) for layer in range(3, 79)] k6 = [load_k6_layer(k6_root, layer) for layer in range(3, 79)] plan: dict[str, Any] = { "schema": PLAN_SCHEMA, "complete": True, "official_source": { "repo": SOURCE_REPO, "revision": args.bf16_revision, "index_sha256": SOURCE_INDEX_SHA256, "config_sha256": SOURCE_CONFIG_SHA256, }, "routed_receipts": { str(item["layer"]): { "sha256": item["receipt_sha256"], "production_layer_id": item["receipt"]["production_layer_id"], "manifest_sha256": item["manifest_sha256"], "shard_sha256": item["manifest"]["shard_sha256"], } for item in routed }, "k6_receipts": { str(item["layer"]): { "sha256": item["receipt_sha256"], "selected_artifact_domain_sha256": item["receipt"].get( "selected_artifact_domain_sha256" ), } for item in k6 }, "hessian_dataset": { "receipt_sha256": sha256_file(hessian_path), "repo": hessian["repo"], "revision": revision, "commit": commit, }, "census": { "routed_k3_k4_matrices": ROUTED_MATRICES, "nonrouted_k6_matrices": NONROUTED_K6_MATRICES, "total_sqg_matrices": TOTAL_SQG_MATRICES, "direct_bf16_ignored_tensors": DIRECT_BF16_MATRICES, }, } plan["plan_id"] = canonical_sha256(plan) output.mkdir(parents=True, exist_ok=True) state_root = output / ".materialize_state" state_root.mkdir(parents=True, exist_ok=True) plan_path = state_root / "PLAN.json" if plan_path.is_file() and _load_json(plan_path) != plan: raise ValueError("existing materialization plan differs") atomic_json(plan_path, plan) weight_map: dict[str, str] = {} shard_receipts: dict[str, Any] = {} # Only the official tensors outside the 58,748-matrix replacement domain # are streamed from BF16. Empty official shards disappear from the model. for ordinal, (source_name, reader) in enumerate(sorted(source["readers"].items())): names = sorted(set(reader.tensors) & source["ignored"]) if not names: continue output_name = f"direct-bf16-{ordinal + 1:05d}-of-00282.safetensors" receipt = _publish_shard( destination=output / output_name, entries=(_entry(reader.tensors[name]) for name in names), state_root=state_root, input_binding={ "kind": "official_bf16_ignored", "source_shard": source_name, "source_index_sha256": SOURCE_INDEX_SHA256, }, role="direct_bf16_ignored", ) shard_receipts[output_name] = receipt weight_map.update({name: output_name for name in names}) routed_sidecars: dict[str, Any] = {} sqg_prefixes: set[str] = set() bits_by_prefix: dict[str, int] = {} down_targets: dict[str, Any] = {} for item in routed: layer = item["layer"] reader = item["reader"] output_name = f"r7-experts-layer-{layer:03d}.safetensors" receipt = _publish_shard( destination=output / output_name, entries=(_entry(info) for info in reader.tensors.values()), state_root=state_root, input_binding={ "kind": "production_routed_sqg", "layer": layer, "receipt_sha256": item["receipt_sha256"], "manifest_sha256": item["manifest_sha256"], "declared_shard_sha256": item["manifest"]["shard_sha256"], }, role="native_full_w4a8_routed_sqg", ) if receipt["payload_sha256"] != item["manifest"]["payload_sha256"]: raise ValueError(f"layer {layer}: streamed routed payload hashes differ") shard_receipts[output_name] = receipt weight_map.update({name: output_name for name in reader.tensors}) prefixes = set(item["manifest"]["bit_map"]) sqg_prefixes.update(prefixes) bits_by_prefix.update( {str(prefix): int(bits) for prefix, bits in item["manifest"]["bit_map"].items()} ) down_target = _derived_down_target(layer, item["manifest"]) down_targets[str(layer)] = { "derived_down_target_id": down_target["derived_down_target_id"], "down_target_beta": down_target["down_target_beta"], } sidecar: dict[str, Any] = { "schema": "glm52-native-sqg-w4a8-layer-sidecar-v1", "complete": True, "layer": layer, "codebook": CODEBOOK, "execution": "full_w4a8", "allow_a16_fallback": False, "shard": output_name, "shard_sha256": receipt["output_sha256"], "bit_map": item["manifest"]["bit_map"], "bit_histogram": {"3": 384, "4": 384}, "shared_vectors": item["manifest"]["shared_vectors"], "vector_refs": item["manifest"]["vector_refs"], "payload_sha256": receipt["payload_sha256"], "production": { "receipt_sha256": item["receipt_sha256"], "production_layer_id": item["receipt"]["production_layer_id"], "manifest_sha256": item["manifest_sha256"], }, "lineage": { "official_bf16_source": True, "sqg_tensors": 768, "topology_neutral": True, "native_e4m3": True, }, } sidecar["sidecar_id"] = canonical_sha256(sidecar) sidecar_name = f"r7-experts-layer-{layer:03d}.json" _publish_json(output / sidecar_name, sidecar) routed_sidecars[str(layer)] = { "filename": sidecar_name, "sha256": sha256_file(output / sidecar_name), "sidecar_id": sidecar["sidecar_id"], } k6_matrices: dict[str, Any] = {} for item in k6: layer = item["layer"] output_by_input: dict[Path, str] = {} for index, (input_shard, group) in enumerate(sorted(item["shards"].items())): output_name = f"sqg-k6-layer-{layer:03d}-{index:02d}.safetensors" reader = group["reader"] receipt = _publish_shard( destination=output / output_name, entries=(_entry(info) for info in reader.tensors.values()), state_root=state_root, input_binding={ "kind": "production_dense_k6_sqg", "layer": layer, "receipt_sha256": item["receipt_sha256"], "source_shard": input_shard.name, "declared_shard_sha256": group["expected_sha256"], }, role="native_full_w4a8_dense_k6_sqg", ) if receipt["payload_sha256"] != group["expected_payloads"]: raise ValueError(f"layer {layer}: streamed K6 payload hashes differ") shard_receipts[output_name] = receipt output_by_input[input_shard] = output_name weight_map.update({name: output_name for name in reader.tensors}) for prefix, record in item["matrices"].items(): input_shard = _safe_sibling( item["receipt_path"].parent, record["shard"], label=f"K6 matrix {prefix} shard", ) output_name = output_by_input[input_shard] sqg_prefixes.add(prefix) bits_by_prefix[prefix] = 6 k6_matrices[prefix] = { "bits": 6, "codebook": CODEBOOK, "role": record["role"], "shard": output_name, "tensor_manifest_id": record["tensor_manifest_id"], "calibration": record["calibration"], "topology_neutral": True, "physical_permutation": False, "hadamard_side": "activation", "source_receipt_sha256": item["receipt_sha256"], } if len(sqg_prefixes) != TOTAL_SQG_MATRICES or len(k6_matrices) != NONROUTED_K6_MATRICES: raise AssertionError("assembled SQG matrix census differs") if set(weight_map) & source["replaced"]: raise AssertionError("replaced BF16 tensors survived materialization") if set(weight_map) & set(source["ignored"]) != source["ignored"]: raise AssertionError("direct-BF16 ignored tensors were lost") _final_tensor_domain(weight_map, source["ignored"], sqg_prefixes) # Preserve official non-weight loader/tokenizer files. The index, config, # and quantization identity are generated below and can never be inherited. support_files: dict[str, Any] = {} excluded = { "model.safetensors.index.json", "config.json", "quantization_config.json", # Publication owns these files. Never inherit an upstream model card or # imply that the derived checkpoint relicenses GLM-5.2. "README.md", "LICENSE", "LICENSE.txt", "compose.yaml", "docker-compose.yml", "serve.sh", "server.sh", "RELEASE_PROVENANCE.json", } for path in sorted(source["root"].iterdir()): if path.name in excluded or path.name.endswith(".safetensors"): continue if path.is_symlink() or not path.is_file(): continue support_files[path.name] = _copy_support_file(path, output / path.name) output_readers = { shard: SafeTensorReader(output / shard) for shard in sorted(set(weight_map.values())) } total_size = sum( output_readers[shard].tensors[name].nbytes for name, shard in weight_map.items() ) index = { "metadata": {"total_size": total_size}, "weight_map": dict(sorted(weight_map.items())), } _publish_json(output / "model.safetensors.index.json", index) tensor_storage = _tensor_storage_metadata( weight_map=weight_map, readers=output_readers, bits_by_prefix=bits_by_prefix, ) quant = _quantization_config( tensor_storage=tensor_storage, down_targets=down_targets, ) config = deepcopy(_load_json(source["root"] / "config.json")) config["quantization_config"] = quant _publish_json(output / "quantization_config.json", quant) _publish_json(output / "config.json", config) k6_manifest: dict[str, Any] = { "schema": "glm52-sqg-k6-nonrouted-v1", "complete": True, "codebook": CODEBOOK, "bits": 6, "matrix_count": len(k6_matrices), "matrices": dict(sorted(k6_matrices.items())), } k6_manifest["manifest_id"] = canonical_sha256(k6_manifest) _publish_json(output / "SQG_K6_NONROUTED.json", k6_manifest) census = { "routed_k3_k4_matrices": ROUTED_MATRICES, "nonrouted_k6_matrices": NONROUTED_K6_MATRICES, "total_sqg_matrices": TOTAL_SQG_MATRICES, "mcg_marker_tensors": 0, "retained_replaced_bf16_weights": 0, "direct_bf16_ignored_tensors": len(source["ignored"]), "official_source_tensor_total": OFFICIAL_SOURCE_TENSORS, } manifest: dict[str, Any] = { "schema": SCHEMA, "complete": True, "official_bf16_repo": SOURCE_REPO, "official_bf16_revision": args.bf16_revision, "official_bf16_index_sha256": SOURCE_INDEX_SHA256, "activation_endpoint": "full-w4a8", "activation_endpoint_scope": "routed_experts", "execution_endpoints": { "routed_experts": "native_sqg_k3_k4_full_w4a8", "nonrouted_k6": "native_sqg_k6_w6a16", }, "weight_endpoint": "native_e4m3_labels", "codebook": CODEBOOK, "hessian_dataset": { "repo": hessian["repo"], "revision": revision, "commit": commit, "receipt_sha256": sha256_file(hessian_path), }, "census": census, "plan": {"sha256": sha256_file(plan_path), "plan_id": plan["plan_id"]}, "routed_sidecars": routed_sidecars, "k6_manifest": { "filename": "SQG_K6_NONROUTED.json", "sha256": sha256_file(output / "SQG_K6_NONROUTED.json"), "manifest_id": k6_manifest["manifest_id"], }, "shards": { name: { "bytes": (output / name).stat().st_size, "sha256": receipt["output_sha256"], "receipt_sha256": sha256_file(_receipt_path(state_root, name)), } for name, receipt in sorted(shard_receipts.items()) }, "support_files": support_files, "construction": { "bf16_only_source": True, "quantized_carrier_used": False, "topology_neutral": True, "physical_permutation": False, "hadamard_side": "activation", "resumable_shard_receipts": True, "checkpoint_tensor_layout": "whole_logical_tensor", "tensor_parallel_slicing": False, "acceptance_runtime": "tp1_pp8", }, } manifest["manifest_id"] = canonical_sha256(manifest) _publish_json(output / "FULL_SQG_NATIVE_MANIFEST.json", manifest) acceptance_args = argparse.Namespace( model=output, bf16_index=source["index_path"], bf16_revision=args.bf16_revision, hessian_dataset_receipt=hessian_path, output=output / "FULL_SQG_NATIVE_ACCEPTANCE.json", ) acceptance = validate_full_checkpoint(acceptance_args) _publish_json(output / "FULL_SQG_NATIVE_ACCEPTANCE.json", acceptance) return manifest def parser() -> argparse.ArgumentParser: value = argparse.ArgumentParser(description=__doc__) value.add_argument("--source-model", type=Path, required=True) value.add_argument("--production-root", type=Path, required=True) value.add_argument("--k6-root", type=Path, required=True) value.add_argument("--hessian-dataset-receipt", type=Path, required=True) value.add_argument("--bf16-revision", default=SOURCE_REVISION) value.add_argument("--output", type=Path, required=True) return value def main() -> int: args = parser().parse_args() if args.bf16_revision != SOURCE_REVISION: raise ValueError("official BF16 revision differs from the frozen campaign") result = materialize(args) print(json.dumps(result, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())