| """Standalone SHARED inference after a bounded visible-response Fable SFT stage.""" |
|
|
| import argparse |
| import hashlib |
| import importlib.util |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| FORMAT = "native-loom-fable-response-sft-v1" |
| PAIR_HELPER_SHA = "62f3432b3174d6825b39aca1ac9fbab38f924f8b5e6c04f7e45c17e1a8c69e6e" |
| PARENT_SHA = "809a3836e2ad6a62c0bc95ae95431b4a18ba50bc839d25998a03bfa62eadd915" |
| INITIAL_SHA = "18c3eac60d7b51d7211cd943eb646bf1cb50d0dacfb587bd182cd9fd9d88b935" |
| DATA_SHA = "8dc518ae52704a7b24d6f2d5b0e43fa9f6e4374b6a9284a4a5db9f377a07936c" |
| PROMPT_TEMPLATE = "### Context\n{prompt}\n\n### Response\n" |
|
|
|
|
| def require(condition, message): |
| if not condition: |
| raise ValueError(message) |
|
|
|
|
| def digest(path): |
| with Path(path).open("rb") as stream: |
| return hashlib.file_digest(stream, "sha256").hexdigest() |
|
|
|
|
| def pair_helper(root): |
| """Import only the exact reviewed, bundled pure helper before model construction.""" |
| path = root / "runtime/loom_melt_pair_portable.py" |
| require( |
| path.is_file() |
| and not path.is_symlink() |
| and not path.parent.is_symlink() |
| and path.resolve().is_relative_to(root), |
| "Reviewed portable helper is missing or changed", |
| ) |
| source = path.read_bytes() |
| require( |
| hashlib.sha256(source).hexdigest() == PAIR_HELPER_SHA, |
| "Reviewed portable helper is missing or changed", |
| ) |
| spec = importlib.util.spec_from_file_location("_loom_fable_pair_reference", path) |
| module = importlib.util.module_from_spec(spec) |
| previous = sys.dont_write_bytecode |
| try: |
| sys.dont_write_bytecode = True |
| |
| exec(compile(source, str(path), "exec"), module.__dict__) |
| finally: |
| sys.dont_write_bytecode = previous |
| return module |
|
|
|
|
| def verify_bundle(directory): |
| root = Path(directory).resolve() |
| helper = pair_helper(root) |
| manifest = helper.read(helper.local_file(root, "manifest.json")) |
| require( |
| manifest["format"] == FORMAT and manifest["variant"] == "shared", |
| "Expected the single SHARED Fable response-SFT export", |
| ) |
| actual = set() |
| for path in root.rglob("*"): |
| require( |
| not path.is_symlink() |
| and "__pycache__" not in path.relative_to(root).parts |
| and path.suffix not in {".pyc", ".pyo"}, |
| "Unbound executable cache or symlink", |
| ) |
| if path.is_file(): |
| actual.add(str(path.relative_to(root))) |
| require(actual == set(manifest["files"]) | {"manifest.json"}, "Bundle inventory differs") |
| for name, binding in manifest["files"].items(): |
| require( |
| helper.file_row(helper.local_file(root, name)) == binding, |
| "Bundled file changed: " + name, |
| ) |
| return root, manifest, helper |
|
|
|
|
| def sft_contract(manifest): |
| sft = manifest["sft"] |
| require( |
| manifest["parent_manifest_sha256"] == PARENT_SHA |
| and manifest["initial_shared_state_sha256"] == INITIAL_SHA |
| and manifest["data_receipt_sha256"] == DATA_SHA |
| and manifest["adapted_updates"] == 32 |
| and manifest["adapted_targets_per_arm"] == 4087238 |
| and manifest["prompt_template"] == PROMPT_TEMPLATE |
| and sft["target_kind"] == "visible_explanatory_response" |
| and sft["private_cot_claim"] is False |
| and type(sft["epochs"]) is int |
| and sft["epochs"] in {1, 2, 3} |
| and sft["updates"] == 7 * sft["epochs"] |
| and sft["input_tokens"] == 54651 * sft["epochs"] |
| and sft["supervised_targets"] == 49199 * sft["epochs"] |
| and sft["optimizer_restart"] is True, |
| "Saved parent, data, formatting or complete SFT epoch contract differs", |
| ) |
|
|
|
|
| def load_bundle(directory, device="cpu"): |
| require(device in {"cpu", "cuda"}, "Unknown inference device") |
| require(torch.__version__.split("+")[0] == "2.13.0", "Use pinned Torch 2.13.0") |
| require( |
| all( |
| os.environ.get(name) in {None, "0"} |
| for name in ("NVIDIA_TF32_OVERRIDE", "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE") |
| ), |
| "TF32 overrides conflict with FP32 routing", |
| ) |
| root, manifest, helper = verify_bundle(directory) |
| sft_contract(manifest) |
| config = helper.read(helper.local_file(root, "model_config.json")) |
| |
| helper.native_contract(manifest, config) |
| require( |
| digest(helper.local_file(root, "source-integrity.json")) == helper.INTEGRITY_SHA, |
| "Frozen native source receipt changed", |
| ) |
| reviewed = helper.read(root / "source-integrity.json")["child_hashes"] |
| require( |
| { |
| name[7:]: binding["sha256"] |
| for name, binding in manifest["files"].items() |
| if name.startswith("source/") |
| } |
| == reviewed, |
| "Frozen source inventory differs", |
| ) |
| adapter, exporter = helper.runtime_modules(root) |
| contract = sys.modules["research.probes.loom_melt_v2_contract"] |
| require( |
| contract.matches_compute_contract(manifest["compute_contract"]), |
| "Saved BF16 training compute contract changed", |
| ) |
| model = helper.construct_model(root, config, manifest["files"]) |
| require( |
| sum(p.numel() for p in model.parameters()) == 112237056 |
| and sum(".melt_gate." in name for name, _ in model.named_parameters()) == 18, |
| "Native parameter or gate inventory differs", |
| ) |
| helper.load_state(model, config, helper.local_file(root, "model.pt"), exporter) |
| packed = adapter.pack_experts(model) |
| require( |
| packed["roundtrip_bit_equal"] |
| and packed["parameters"] == 112237056 |
| and packed["parameter_tensors"] == 92 |
| and helper.state_digest(model) == manifest["packed_state_sha256"], |
| "Saved SFT state differs from its validated checkpoint", |
| ) |
| if device == "cuda": |
| require( |
| torch.version.cuda == "13.0" and torch.cuda.is_bf16_supported(), |
| "CUDA inference requires pinned CUDA 13.0 and BF16", |
| ) |
| model.to(device).eval() |
| from tokenizers import Tokenizer |
|
|
| tokenizer_path = helper.local_file(root, "tokenizer/tokenizer.json") |
| require( |
| digest(tokenizer_path) |
| == "9ca9acddb6525a194ec8ac7a87f24fbba7232a9a15ffa1af0c1224fcd888e47c", |
| "Unchanged native tokenizer required", |
| ) |
| tokenizer = Tokenizer.from_file(str(tokenizer_path)) |
| engine = adapter.CudaMeltEngine( |
| model, |
| "shared", |
| 50, |
| 50, |
| precision="fp32" if device == "cpu" else "bf16", |
| cpu_reference=device == "cpu", |
| ) |
| return engine, tokenizer, config, helper |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--bundle", type=Path, default=Path(__file__).resolve().parent) |
| parser.add_argument("--prompt", required=True) |
| parser.add_argument("--max-new-tokens", type=int, default=64) |
| parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu") |
| args = parser.parse_args() |
| torch.set_num_threads(1) |
| engine, tokenizer, config, helper = load_bundle(args.bundle, args.device) |
| ids = tokenizer.encode(PROMPT_TEMPLATE.format(prompt=args.prompt), add_special_tokens=False).ids |
| generated = helper.greedy( |
| engine, |
| ids, |
| args.max_new_tokens, |
| config["model"]["max_seq_len"], |
| config["tokenizer"]["eos_token_id"], |
| ) |
| print( |
| json.dumps( |
| { |
| "text": tokenizer.decode(generated, skip_special_tokens=True), |
| "new_tokens": len(generated), |
| "variant": "shared", |
| "target_kind": "visible_explanatory_response", |
| } |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
