diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 97e4945b..2319a90a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -50,6 +50,21 @@ jobs: python -m pip install pytest PYTHONPATH=src python -m pytest tests/clef -q + kev-export: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v5 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: CPU-only export dependencies + run: | + python -m pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install -r recipe/kev/requirements.txt + - name: Export contract tests (no weights or GPU) + run: python -m unittest discover -s tests/kev -p 'test_*.py' -v + cuda-contract: # The CUDA backend contract and its checker. No CUDA toolkit and no GPU: both # read files as text. diff --git a/Cargo.lock b/Cargo.lock index 7151d686..9f5d30e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1263,6 +1263,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-kev-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "safetensors 0.8.0", + "serde_json", +] + [[package]] name = "omni-laya" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 7910b3dd..1ebf8e22 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/omnijev/native", "src/models/laya", "src/backends/cuda", "src/models/decider/native", "src/models/jemm/native", "src/models/startlux/native", "src/models/clef/native"] +members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/omnijev/native", "src/models/laya", "src/backends/cuda", "src/models/decider/native", "src/models/jemm/native", "src/models/startlux/native", "src/models/clef/native", "src/models/kev/native"] resolver = "3" diff --git a/README.md b/README.md index 1502a8fb..1be04e5d 100644 --- a/README.md +++ b/README.md @@ -191,6 +191,7 @@ checkpoint. Public231 concurrency-1 latency for heads `1ba07d1`, `f968e576`, | CLM-v0.1-8B | [External worker with a CPU stub encoder](recipe/clm/README.md); contract checks only, real Qwen3-8B decisions unverified by this recipe | | LFM2.5-350M | [Transformers worker](recipe/lfm2/README.md); text state and independent choice questions | | OmniJev-4B v1.1 | [Native Rust/CUDA worker](recipe/omnijev/native.md), eager: one image, Choice, Noul and Score; [reference comparison on sm_89](recipe/omnijev/native.md#comparison-with-the-reference) | +| Kev-4B | [Pinned CPU export and Rust pointer head](recipe/kev/README.md); no worker or full-model parity yet ([#27](https://github.com/InferMatrixLab/system1-omni/issues/27)) | | D3 (Decision 3.0) | [Python reference worker](recipe/d3/reference.md) on the checkpoint's pinned runtime; `choice`/`noul`/`score` over text, image and video state; author-run on A100 for text and media serving (see the integration and media PRs) | | Vela 2.0 0.3B | [Python reference worker](recipe/vela/reference.md) on the checkpoint's pinned runtime; all five kinds (`choice`/`noul`/`score`/`span`/`set`) for text state; CPU FP32 validated (author-run, evidence in the integration PR) | | Vela 2.0 0.8B hybrid | [Python reference worker](recipe/vela/reference.md) on the checkpoint's pinned runtime; same surface plus `span_heads`; CPU FP32 parity recorded in the integration PR | diff --git a/docs/supported-models.md b/docs/supported-models.md index ca929b0e..36fa0907 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -34,6 +34,10 @@ Models that are being added are also tracked in issues labeled [new model](https | Clef (27B) | Same `omni-clef` server and text contract as Clef-flash, selected with `--variant clef` | Contract tests only; no inference | Not run | Not supported | Same stack; pinned `Cloudflare/clef` snapshot (~55.0 GB). Weights and joint-head config differ from flash | | pplx-decider-v1.1-27b | [Python reference worker](../recipe/pplx_decider/README.md) on the checkpoint's own runtime. Text and up to four images; `choice`/`noul`/`score`. Native CUDA executor is not this worker. Lab L20X-labeled GPUs are the H200 target. The causal-conv1d build keeps sm_90 and sm_89. Capability 8.9 is ready on FLA chunk. Capability 9.0 uses all-sequence FlashInfer. The kept primary rotates 48 `cu_seqlens` buffers per device and length, passes `use_cp` false, skips the q/k repeat, and packs q and k before the L2-norm view. `801ef7b` fused the 48 linear gated RMSNorms and missed the median (p50 89.172 ms, p95 274.790 ms), so that kernel stays off. `2bbbe4e` replaced `silu(gate) * up` with FLA `swiglu_fwd` and the owner closeout rejected it (p50 97.613 ms, p95 292.795 ms), so that kernel stays off. Gate and up GEMMs stay separate. Residual RMSNorm and the gated output norm stay eager. An unknown `PPLX_GDN_PREFILL` fails closed on every device, including capability 8.9. `3ecdf5a` crashed in warmup (0/231). Always-on fusion `3afb472` and the 512-token gate `705ae089` missed the median. The 48-slot FlashInfer state ring `17a7b120` missed the gate. `PPLX_GDN_PREFILL=fla` reproduces the prior `f968e576` chunk. `PPLX_GDN_PREFILL=short` is the rejected ≤256 dispatch. SDPA stays on the checkpoint order. The fused gate-up GEMM is refused. The forward stays eager. `PPLX_COMPILE=dynamic` and `PPLX_COMPILE=dynamic-tail` are rejected opt-ins | Contract checks only; no inference | Reported load on 1× lab L20X (sm_89, ~140 GiB) at head `1ba07d1`: BF16 resident ~50333 MiB, no OOM. Historical primary: 231/231, p50 98.777 ms, p95 341.202 ms. Prior primary is the `f968e576` remeasure on the lab L20X-labeled GPU, `compute_capability=9.0`, `fla-chunk`, `flashinfer` null: 231/231, p50 94.077 ms, p95 300.200 ms (p50 −4.76%, p95 −12.02%). Prior primary `ffae182`: 231/231, p50 92.844 ms, p95 296.458 ms (p50 −1.31%, p95 −1.25%). Prior primary `4b99b73`: 231/231, p50 88.942 ms, p95 295.797 ms (p50 −4.20%, p95 −0.22%), 48-slot `cu_seqlens`, quality `choice_acc=0.9065` (139), `noul_acc=0.8784` (74), `score_mae=0.2035` (18). Kept primary is `fcf5e2b`: 231/231, p50 87.782 ms, p95 291.183 ms (p50 −1.30%, p95 −1.56%), grouped-value FlashInfer with q and k packed, same quality digits. Uncached all-sequence FlashInfer `3c72f02`: p50 88.808 ms, p95 300.751 ms, gate failed. Short-sequence FlashInfer `a50b9aa`: p50 95.446 ms, p95 300.871 ms, gate failed. cuDNN-first SDPA `0dcaa0ac`: p50 138.501 ms, p95 332.859 ms, gate failed. Fused gate-up `8b4ee326`: p50 102.437 ms, p95 304.970 ms, gate failed. Always-on dynamic compile `e36b32ed`: p50 97.751 ms, p95 263.124 ms, gate failed (p95 improved). 512-token compile `7e3b1be`: p50 100.037 ms, p95 257.084 ms, gate failed. Isolated-median compile `34ba5de`: p50 97.084 ms, p95 256.679 ms, gate failed (median probe was eager). FlashInfer state-workspace ring `17a7b120`: p50 89.165 ms, p95 295.723 ms, gate failed. Always-on FLA residual RMSNorm `3afb472`: p50 93.935 ms, p95 274.902 ms, gate failed. Length-gated FLA residual RMSNorm `705ae089`: p50 94.604 ms, p95 274.835 ms, gate failed. `use_cp=auto` `a6f4d9a`: p50 112.294 ms, p95 299.419 ms, gate failed. Grouped-value `3ecdf5a` crashed in warmup (0/231, no p50 or p95). The fused linear gated-norm candidate has no sample. Does not enter the H800 board | Not supported | Python 3.12, pinned checkpoint revision `14c9d90137962d300f3de95633f7b6535e80c5ce`, checkpoint `requirements.txt`, `causal-conv1d` 1.7.0, `flash-linear-attention` 0.5.2, optional `flashinfer-python` 0.6.16.post3. Shard sizes in the release manifest sum with the readout to about 48.6 GiB before activations | +Kev-4B has [CPU checkpoint preparation and a Rust pointer head](../recipe/kev/README.md) +only. No Kev worker, backbone parity, GPU serving or full-model accuracy is +established by that preparation slice; native integration remains [#27](https://github.com/InferMatrixLab/system1-omni/issues/27). + - **Validated:** covered by the recipe on `main` or by the checks in the linked merged pull request. - **Author-run:** recorded by the recipe author on the stated hardware and not independently reproduced; the row names its owner. - **Unverified:** the worker accepts this device, but no recipe or merged pull request covers it. diff --git a/mkdocs.yml b/mkdocs.yml index 86abd35a..e011fc70 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -81,6 +81,7 @@ nav: - JEV-27B-VL experimental worker: recipe/jev_vl/README.md - JEV-27B-VL validation status: recipe/jev_vl/validation.md - OmniJev preparation: recipe/omnijev/README.md + - Kev-4B preparation: recipe/kev/README.md - OmniJev native worker: recipe/omnijev/native.md - CLM stub-encoder contract: recipe/clm/README.md - LFM2.5-350M choice worker: recipe/lfm2/README.md diff --git a/recipe/README.md b/recipe/README.md index 22179d21..6250a7cd 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -42,6 +42,8 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting - [OmniJev-4B v1.1 preparation](omnijev/README.md): export the checkpoint and regenerate the CPU fixtures; [native worker](omnijev/native.md): serve it on CUDA and compare it with the reference. +- [Kev-4B preparation](kev/README.md): export all LoRA projections and validate + the Rust FP32 pointer head on CPU; worker integration remains pending. - [CLM behind the frontend](clm/README.md): run CLM's own server behind the frontend on CPU with a stub encoder, and what the response comparison has to allow for. - [pplx-decider reference worker](pplx_decider/README.md): pin diff --git a/recipe/kev/README.md b/recipe/kev/README.md new file mode 100644 index 00000000..50993c42 --- /dev/null +++ b/recipe/kev/README.md @@ -0,0 +1,77 @@ +# Kev-4B: CPU export and pointer-head validation + +This is a preparation recipe for [#27](https://github.com/InferMatrixLab/system1-omni/issues/27), +not a runnable Kev server. The [model contract](../../src/models/kev/README.md) +records the pinned artifacts, numerical boundaries and missing worker pieces. +All commands below run from the repository root. No GPU is needed. + +## Environment and downloads + +Use Python 3.12 and the stable Rust toolchain used by this repository. The CPU +export reads about 9.5 GB of inputs and writes about 8.5 GB; reserve at least 25 GB +of free disk and 24 GB of host RAM. It processes one source shard at a time. + +```sh +python3.12 -m venv .venv-kev +.venv-kev/bin/pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu +.venv-kev/bin/pip install -r recipe/kev/requirements.txt huggingface-hub==1.33.0 + +.venv-kev/bin/hf download Qwen/Qwen3.5-4B-Base \ + --revision 1001bb4d826a52d1f399e183466143f4da7b741b \ + --include 'config.json' 'tokenizer.json' 'tokenizer_config.json' \ + 'model.safetensors.index.json' '*.safetensors' --local-dir weights/kev-base +.venv-kev/bin/hf download jaredpalmer/kev-4b \ + --revision 139fdd94f1b6a6ad80cc15e08fcb99cac885a101 \ + --include 'adapter_config.json' 'adapter_model.safetensors' 'head.pt' \ + --local-dir weights/kev-adapter + +.venv-kev/bin/python recipe/kev/export.py \ + --base weights/kev-base --checkpoint weights/kev-adapter --out weights/kev-native +``` + +Every required input must match its pinned SHA-256. The output directory must not +already exist. `kev_export.json` lists merged targets, output hashes, FP32 head +dimensions and calibration, plus the base tokenizer's five delimiter IDs. +Do not replace these artifacts with another Kev size or a bare Qwen checkpoint. + +## CPU checks + +```sh +.venv-kev/bin/python -m unittest discover -s tests/kev -p 'test_*.py' -v +cargo test --locked -p omni-kev-native -- --list +cargo test --locked -p omni-kev-native +KEV_HEAD="$PWD/weights/kev-native/head.safetensors" \ + cargo test --release --locked -p omni-kev-native --test head \ + -- --ignored --nocapture +``` + +The final test is opt-in because it needs the released head, but it is still +CPU-only. It compares the Rust implementation against the committed small oracle +for 1/2/7/31/255 options. Default CI executes the ordinary head tests without +downloading weights. Full backbone inference is not part of these checks. + +To regenerate the oracle using the official class rather than this implementation: + +```sh +git clone https://github.com/jaredpalmer/kev.git /tmp/kev-reference +git -C /tmp/kev-reference checkout --detach 2074cc44256200bb37540103fa4572901faa10f3 +.venv-kev/bin/python tests/kev/generate_head_fixture.py \ + --reference-model /tmp/kev-reference/kev/model.py \ + --head weights/kev-adapter/head.pt --out /tmp/kev-head-reference.json +diff -u tests/kev/data/head.json /tmp/kev-head-reference.json +``` + +The generator verifies the complete source and head hashes, extracts only the +upstream `PointerHead` class, and uses CPU PyTorch. It does not need Transformers, +PEFT or the backbone. The fixture's Torch version includes the wheel's local +suffix, so a different wheel can change provenance even when numbers agree. + +## Next integration slice + +There is no worker binary, frontend launch command, health endpoint or example +HTTP request yet. The next slice adds exact prompt/row preparation and typed +answers, verified export loading, shared Qwen execution and serial admission. +It must qualify the BF16 merge against the full reference, then test real warmup +and first inference directly and through the frontend. Record hardware, full +dependency pins and numerical gates before those runs. No latency or accuracy +claim follows from export success or isolated head agreement. diff --git a/recipe/kev/export.py b/recipe/kev/export.py new file mode 100644 index 00000000..d00fea7d --- /dev/null +++ b/recipe/kev/export.py @@ -0,0 +1,258 @@ +"""Export pinned Kev-4B language weights and pointer head on CPU (no serving).""" + +import argparse +import hashlib +import json +import math +import os +import shutil +import tempfile +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +MODEL_ID = "jaredpalmer/kev-4b" +REVISION = "139fdd94f1b6a6ad80cc15e08fcb99cac885a101" +BASE_ID = "Qwen/Qwen3.5-4B-Base" +BASE_REVISION = "1001bb4d826a52d1f399e183466143f4da7b741b" +REFERENCE_REVISION = "2074cc44256200bb37540103fa4572901faa10f3" +SPECIAL = { + "<|fim_prefix|>": 248060, + "<|fim_middle|>": 248061, + "<|box_start|>": 248049, + "<|box_end|>": 248050, + "<|fim_suffix|>": 248062, +} +BASE_FILES = { + "config.json": "ddc63e1c717afa86c865bb5e01313d89d72bb53b97ad4a8a03ba8510c0621670", + "tokenizer.json": "fe000e3ed39ed12b8d2481d527d44f93c65d37e87645d2dcc80d1bf9d50d2927", + "tokenizer_config.json": "3891e840d7dc5fca0af33d3a25083a735e36fe06214e3f707024820cb6b9f89c", + "model.safetensors.index.json": "eae340074abb0a5f31a6621f7ae8e8248a7c1790df04a722c4e4b70c2a6d1dbb", + "model.safetensors-00001-of-00002.safetensors": "df547074dce70532a0493e5433152bd17a65efb89088cfabc2e7e2371a93d712", + "model.safetensors-00002-of-00002.safetensors": "590fbaac095dd31db886c322d9d2f7df47777966391acf306ddddc3e4e3a15ef", +} +CHECKPOINT_FILES = { + "adapter_config.json": "8a05dfd6c5e7f61a62e6db5d8de8093a8dbbd105aa7ee7b5e7b73f6249f45617", + "adapter_model.safetensors": "90e817356246e7f18bfa7ca3d31794cd4fbeb3332a66a84cb51d9ceae925f2b2", + "head.pt": "dd633435998ecc751ac538717a3742e32149500fabf7d7276287dbf0693f347c", +} + + +def expected_targets(): + """The pinned 32-layer model: 96 MLP, 32 attention and 120 GDN targets.""" + targets = set() + for layer in range(32): + modules = [f"mlp.{p}_proj" for p in ("gate", "up", "down")] + if layer % 4 == 3: + modules += [f"self_attn.{p}_proj" for p in ("q", "k", "v", "o")] + else: + modules += [f"linear_attn.{p}" for p in ( + "in_proj_qkv", "in_proj_a", "in_proj_z", "in_proj_b", "out_proj" + )] + targets.update(f"model.language_model.layers.{layer}.{m}.weight" for m in modules) + return targets + + +def validate_targets(targets): + expected = expected_targets() + if set(targets) != expected: + raise ValueError( + f"LoRA target mismatch: missing={sorted(expected - set(targets))}, " + f"unexpected={sorted(set(targets) - expected)}" + ) + + +def sha256(path): + h = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1 << 24), b""): + h.update(block) + return h.hexdigest() + + +def verify_files(directory, expected): + for name, digest in expected.items(): + if sha256(directory / name) != digest: + raise ValueError(f"{name} does not match the pinned release") + + +def validate_tokenizer(path): + tokens = { + t["content"]: t["id"] for t in json.loads(path.read_text())["added_tokens"] + } + if any(tokens.get(token) != value for token, value in SPECIAL.items()): + raise ValueError("base tokenizer has unexpected delimiter IDs") + + +def adapter_pairs(tensors): + pairs = {} + for name, value in tensors.items(): + prefix, marker, side = name.partition(".lora_") + if ( + not prefix.startswith("base_model.model.layers.") + or not marker + or side not in ("A.weight", "B.weight") + ): + raise ValueError(f"unsupported adapter tensor: {name}") + if ( + value.dtype != torch.float32 + or value.ndim != 2 + or not torch.isfinite(value).all() + ): + raise ValueError(f"invalid FP32 adapter tensor: {name}") + target = ( + "model.language_model." + + prefix.removeprefix("base_model.model.") + + ".weight" + ) + pairs.setdefault(target, {})[side[0]] = value + if any(set(pair) != {"A", "B"} for pair in pairs.values()): + raise ValueError("incomplete LoRA pair") + return pairs + + +def merge_weight(weight, a, b): + """Kev's alpha/r = 2: FP32 delta and addition, one final BF16 rounding.""" + if ( + weight.ndim != 2 + or a.shape != (16, weight.shape[1]) + or b.shape != (weight.shape[0], 16) + ): + raise ValueError("LoRA shape mismatch") + merged = (weight.float() + (b @ a) * 2.0).to(torch.bfloat16) + if not torch.isfinite(merged).all(): + raise ValueError("non-finite merged weight") + return merged.contiguous() + + +def pointer_head(meta): + expected = { + "base": BASE_ID, + "base_revision": BASE_REVISION, + "lora": 16, + "head_dim": 256, + "option_isolation": False, + "special_embeddings": False, + "weights_dtype": "fp32", + } + if any(meta.get(key) != value for key, value in expected.items()): + raise ValueError("unsupported Kev head metadata") + temperature = meta.get("temperature") + if ( + not isinstance(temperature, (int, float)) + or not math.isfinite(temperature) + or temperature <= 0 + ): + raise ValueError("invalid calibration temperature") + shapes = { + "q.weight": (256, 2560), + "k.weight": (256, 2560), + "q.bias": (256,), + "k.bias": (256,), + } + head = meta["head"] + if set(head) != set(shapes): + raise ValueError("unexpected pointer-head keys") + for name, shape in shapes.items(): + value = head[name] + if ( + value.dtype != torch.float32 + or tuple(value.shape) != shape + or not torch.isfinite(value).all() + ): + raise ValueError(f"invalid pointer-head tensor: {name}") + return {name: value.contiguous() for name, value in head.items()}, temperature + + +def export(base, checkpoint, out): + if out.exists(): + raise ValueError("output already exists; choose a new directory") + verify_files(base, BASE_FILES) + verify_files(checkpoint, CHECKPOINT_FILES) + validate_tokenizer(base / "tokenizer.json") + # Hash the pickle before parsing and use the restricted weights-only loader. + head, temperature = pointer_head( + torch.load(checkpoint / "head.pt", map_location="cpu", weights_only=True) + ) + pairs = adapter_pairs( + load_file(checkpoint / "adapter_model.safetensors", device="cpu") + ) + validate_targets(pairs) + out.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{out.name}-", dir=out.parent)) + used, weight_map = set(), {} + try: + for name in BASE_FILES: + if not name.endswith(".safetensors"): + continue + weights = {} + with safe_open(base / name, framework="pt", device="cpu") as source: + for key in source.keys(): + if not key.startswith("model.language_model."): + continue # Kev is text-only; exclude vision and MTP. + weight = source.get_tensor(key) + if key in pairs: + pair = pairs[key] + weight = merge_weight(weight, pair["A"], pair["B"]) + used.add(key) + else: + weight = weight.to(torch.bfloat16).contiguous() + if not torch.isfinite(weight).all(): + raise ValueError(f"non-finite base tensor: {key}") + weights[key] = weight + if key in weight_map: + raise ValueError(f"duplicate base tensor: {key}") + weight_map[key] = name + if weights: + save_file(weights, staging / name) + del weights + if used != set(pairs): + raise ValueError(f"unmerged adapter targets: {sorted(set(pairs) - used)}") + (staging / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}, indent=2) + "\n" + ) + for name in ("config.json", "tokenizer.json", "tokenizer_config.json"): + shutil.copyfile(base / name, staging / name) + save_file(head, staging / "head.safetensors") + manifest = { + "format": "kev-4b-merged/1", + "model_id": MODEL_ID, + "revision": REVISION, + "base_id": BASE_ID, + "base_revision": BASE_REVISION, + "reference_revision": REFERENCE_REVISION, + "backbone_dtype": "bfloat16", + "head_dtype": "float32", + "hidden_size": 2560, + "pointer_size": 256, + "temperature": temperature, + "special_tokens": SPECIAL, + "merged_targets": sorted(used), + "inputs": {"base": BASE_FILES, "checkpoint": CHECKPOINT_FILES}, + "files": {p.name: sha256(p) for p in sorted(staging.iterdir())}, + } + (staging / "kev_export.json").write_text(json.dumps(manifest, indent=2) + "\n") + os.rename(staging, out) + finally: + if staging.exists(): + shutil.rmtree(staging) + return manifest + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True, type=Path) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args() + torch.set_num_threads(8) + manifest = export(args.base, args.checkpoint, args.out) + print( + f"Exported {len(manifest['merged_targets'])} merged projections to {args.out}" + ) + + +if __name__ == "__main__": + main() diff --git a/recipe/kev/requirements.txt b/recipe/kev/requirements.txt new file mode 100644 index 00000000..227972ef --- /dev/null +++ b/recipe/kev/requirements.txt @@ -0,0 +1,2 @@ +torch==2.8.0 +safetensors==0.8.0 diff --git a/src/models/kev/README.md b/src/models/kev/README.md new file mode 100644 index 00000000..c5a2f004 --- /dev/null +++ b/src/models/kev/README.md @@ -0,0 +1,84 @@ +# Kev-4B preparation and pointer head + +First independent slice of [#27](https://github.com/InferMatrixLab/system1-omni/issues/27): +a pinned CPU exporter and a Rust FP32 pointer head. **There is no Kev worker or +end-to-end inference in this slice.** The [recipe](../../../recipe/kev/README.md) +reproduces preparation and head validation. Other Kev sizes are not supported. + +## Pinned artifacts + +| Artifact | Immutable revision | +| --- | --- | +| [jaredpalmer/kev-4b](https://huggingface.co/jaredpalmer/kev-4b/tree/139fdd94f1b6a6ad80cc15e08fcb99cac885a101), adapter and head | `139fdd94f1b6a6ad80cc15e08fcb99cac885a101` | +| [Qwen/Qwen3.5-4B-Base](https://huggingface.co/Qwen/Qwen3.5-4B-Base/tree/1001bb4d826a52d1f399e183466143f4da7b741b) | `1001bb4d826a52d1f399e183466143f4da7b741b` | +| [Kev reference code](https://github.com/jaredpalmer/kev/tree/2074cc44256200bb37540103fa4572901faa10f3) | `2074cc44256200bb37540103fa4572901faa10f3` | + +All are Apache-2.0; the adapted head's license is retained in +[`native/LICENSE.kev`](native/LICENSE.kev). No model weights are committed. +The source hash in the fixture generator covers the complete pinned `kev/model.py`. + +## Export contract + +[`export.py`](../../../recipe/kev/export.py) verifies SHA-256 for every required +base/checkpoint file **before deserializing `head.pt`**, which uses PyTorch's +weights-only CPU loader. It merges all 496 FP32 LoRA tensors into 248 language +projections as `BF16(FP32(W) + 2 * (B @ A))`. This includes all five GDN projection +types: `in_proj_qkv`, `in_proj_a`, `in_proj_z`, `in_proj_b` and `out_proj`. +Missing pairs or unused targets fail the export. The output excludes vision/MTP +weights and preserves the original language tensor names for the shared Qwen loader. +Other language tensors are converted to BF16, including originally FP32 norms. + +The four pointer-head tensors remain FP32: `q.weight`/`k.weight` are `[256, 2560]` +and their biases are `[256]`. The manifest records temperature +`2.406050072164233`, dimensions, delimiter IDs, exact input hashes, every merged +target and all output hashes. It is written after the artifacts, and the private +staging directory is renamed only on success. Existing outputs are refused. + +The tokenizer comes from the **base**: `<|fim_prefix|>` = 248060 (state), +`<|fim_middle|>` = 248061 (question), `<|box_start|>` = 248049 (option start), +`<|box_end|>` = 248050 (option readout), `<|fim_suffix|>` = 248062 (decision +readout). The adapter's older added-token metadata is not copied. + +## Executor and postprocessing boundary + +`head::PointerHead` owns the trained CPU weights and accepts one final-normalized +decision hidden state plus 1–255 option hidden states, each with 2560 FP32 values +for this checkpoint. It computes `k(h_option) dot q(h_decide) / sqrt(256)` in FP32, +including both learned biases, and preserves candidate order. Inputs, shapes, +weights and results must be finite. Returned logits own their storage. +`postprocess::probabilities` divides these logits by the saved temperature and +normalizes across the **whole question**. This follows the eval-mode reference; +temperature must not be applied twice. Projection and calibration have separate +modules even though both currently run on CPU. + +The future processor must reproduce `encode`/`rows_of`: Qwen's hybrid layers require +one causal row containing state plus one question's complete options. Each option +is read at its closing marker, and the decision is read at the final marker. +Questions must not leak into one another's GDN state. The reference rewrites +caller-supplied `<|name|>` to `<¦name¦>` before tokenization. No HTTP mapping, +prompt compiler, token accounting, request budgets, scheduler or device ownership +is implemented here. Those remain explicit requirements for the worker slice. + +## Validation and remaining gates + +Tests live under [`tests/kev/`](../../../tests/kev/). Ordinary Rust tests cover +learned biases, scale, candidate ordering, one/255 candidates, malformed weights, +hidden states and temperature. CPU Python tests cover pinned-input rejection, +no overwrite, all five GDN target names, one-rounding merge arithmetic, incomplete +pairs, delimiter IDs and head metadata. + +The opt-in **CPU** test loads the real exported head and compares 1, 2, 7, 31 and +255 synthetic hidden-state candidates with the actual pinned upstream +`PointerHead` class (PyTorch 2.8.0 CPU). Logit tolerance is +`1e-4 + 1e-5 * abs(reference)`; probability absolute tolerance is `1e-5`, declared +before comparison. The small committed fixture contains expected outputs and +provenance, not weights; its deterministic binary-fraction inputs are regenerated +by the generator and Rust test. + +Passing export/head checks does **not** establish BF16 merge decision parity, +backbone parity, real prompt/token parity, native serving, GPU support or a speedup. +The worker follow-up must qualify BF16 merge versus the official unmerged/FP32 +reference, pin full runtime dependencies, reuse selected-position Qwen readouts +(already carried in [#129](https://github.com/InferMatrixLab/system1-omni/pull/129)), +implement model-specific processing, and validate readiness plus direct/frontend +inference for `choice`, `score` and `noul`. Prefix reuse is a later measured change. diff --git a/src/models/kev/native/Cargo.toml b/src/models/kev/native/Cargo.toml new file mode 100644 index 00000000..6020a043 --- /dev/null +++ b/src/models/kev/native/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "omni-kev-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "Kev pointer head; native worker integration pending" + +[dependencies] +anyhow = "1.0.100" +safetensors = "0.8.0" + +[dev-dependencies] +serde_json = "1.0.149" + +[[test]] +name = "head" +path = "../../../../tests/kev/head.rs" diff --git a/src/models/kev/native/LICENSE.kev b/src/models/kev/native/LICENSE.kev new file mode 100644 index 00000000..e1ece45b --- /dev/null +++ b/src/models/kev/native/LICENSE.kev @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 Jared Palmer + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/src/models/kev/native/src/head.rs b/src/models/kev/native/src/head.rs new file mode 100644 index 00000000..dc3eea2e --- /dev/null +++ b/src/models/kev/native/src/head.rs @@ -0,0 +1,102 @@ +//! Bilinear pointer readout from jaredpalmer/kev at 2074cc44256200bb37540103fa4572901faa10f3. +//! Learned projection stays separate from calibration and response interpretation. +use anyhow::{Context, Result, ensure}; +use safetensors::{Dtype, SafeTensors}; +use std::path::Path; + +pub struct PointerHead { + hidden: usize, + pointer: usize, + q: Vec, + k: Vec, + q_bias: Vec, + k_bias: Vec, +} + +fn tensor(tensors: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let view = tensors + .tensor(name) + .with_context(|| format!("missing {name}"))?; + ensure!( + view.dtype() == Dtype::F32 && view.shape() == shape, + "invalid {name} shape or dtype" + ); + let values: Vec = view + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|&b| f32::from_le_bytes(b)) + .collect(); + ensure!(values.iter().all(|v| v.is_finite()), "non-finite {name}"); + Ok(values) +} + +fn project(weights: &[f32], bias: &[f32], input: &[f32]) -> Vec { + weights + .chunks_exact(input.len()) + .zip(bias) + .map(|(row, &bias)| row.iter().zip(input).map(|(&w, &x)| w * x).sum::() + bias) + .collect() +} + +impl PointerHead { + /// Load exactly q/k matrices and biases, preserving their FP32 representation. + /// Dimensions come from the head; Kev-4B exports use 2560 hidden / 256 pointer. + pub fn load(path: &Path) -> Result { + let bytes = std::fs::read(path)?; + let tensors = SafeTensors::deserialize(&bytes)?; + ensure!( + tensors.names().len() == 4, + "expected exactly four pointer-head tensors" + ); + let q = tensors.tensor("q.weight")?; + ensure!( + q.shape().len() == 2 && q.shape().iter().all(|&d| d > 0), + "invalid q.weight shape" + ); + let (pointer, hidden) = (q.shape()[0], q.shape()[1]); + Ok(Self { + hidden, + pointer, + q: tensor(&tensors, "q.weight", &[pointer, hidden])?, + k: tensor(&tensors, "k.weight", &[pointer, hidden])?, + q_bias: tensor(&tensors, "q.bias", &[pointer])?, + k_bias: tensor(&tensors, "k.bias", &[pointer])?, + }) + } + + /// h_decide is the final-normalized state at <|fim_suffix|>; each h_opt is + /// at an option's closing <|box_end|> in the SAME complete question row. + /// Returns one raw FP32 logit per option, in caller order. + pub fn logits(&self, h_decide: &[f32], h_opts: &[&[f32]]) -> Result> { + ensure!( + (1..=255).contains(&h_opts.len()), + "expected 1..=255 options" + ); + ensure!( + std::iter::once(h_decide) + .chain(h_opts.iter().copied()) + .all(|row| row.len() == self.hidden && row.iter().all(|v| v.is_finite())), + "invalid hidden-state shape or values" + ); + let query = project(&self.q, &self.q_bias, h_decide); + let scale = (self.pointer as f64).sqrt().recip() as f32; + let logits: Vec = h_opts + .iter() + .map(|row| { + project(&self.k, &self.k_bias, row) + .iter() + .zip(&query) + .map(|(&k, &q)| k * q) + .sum::() + * scale + }) + .collect(); + ensure!( + logits.iter().all(|v| v.is_finite()), + "non-finite pointer logits" + ); + Ok(logits) + } +} diff --git a/src/models/kev/native/src/lib.rs b/src/models/kev/native/src/lib.rs new file mode 100644 index 00000000..16fed8da --- /dev/null +++ b/src/models/kev/native/src/lib.rs @@ -0,0 +1,3 @@ +//! Kev's model-specific CPU readout. This crate does not yet serve requests. +pub mod head; +pub mod postprocess; diff --git a/src/models/kev/native/src/postprocess.rs b/src/models/kev/native/src/postprocess.rs new file mode 100644 index 00000000..8346ecc4 --- /dev/null +++ b/src/models/kev/native/src/postprocess.rs @@ -0,0 +1,21 @@ +//! Kev checkpoint-temperature calibration; wire answers remain a later slice. +use anyhow::{Result, ensure}; + +pub fn probabilities(logits: &[f32], temperature: f64) -> Result> { + ensure!((1..=255).contains(&logits.len()), "expected 1..=255 logits"); + let temperature = temperature as f32; + ensure!( + temperature.is_finite() && temperature > 0.0, + "invalid temperature" + ); + let scaled: Vec = logits.iter().map(|v| v / temperature).collect(); + ensure!( + scaled.iter().all(|v| v.is_finite()), + "non-finite calibrated logits" + ); + let max = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut values: Vec = scaled.iter().map(|v| (v - max).exp()).collect(); + let sum: f32 = values.iter().sum(); + values.iter_mut().for_each(|v| *v /= sum); + Ok(values) +} diff --git a/tests/kev/data/head.json b/tests/kev/data/head.json new file mode 100644 index 00000000..68c8cff8 --- /dev/null +++ b/tests/kev/data/head.json @@ -0,0 +1,638 @@ +{ + "reference_revision": "2074cc44256200bb37540103fa4572901faa10f3", + "reference_model_sha256": "2ff3b2e67c048e3bc5a0d1a32aafe50804ed315a5564fce06bb4b95697c5271c", + "head_pt_sha256": "dd633435998ecc751ac538717a3742e32149500fabf7d7276287dbf0693f347c", + "torch": "2.8.0+cpu", + "device": "cpu", + "temperature": 2.406050072164233, + "hidden_rule": "((i * 37 + row * 19) % 101 - 50) / 64; row 0 = decide", + "cases": [ + { + "options": 1, + "logits": [ + 0.16330909729003906 + ], + "probabilities": [ + 1.0 + ] + }, + { + "options": 2, + "logits": [ + 0.16330908238887787, + -0.04296588525176048 + ], + "probabilities": [ + 0.5214198231697083, + 0.47858014702796936 + ] + }, + { + "options": 7, + "logits": [ + 0.16330906748771667, + -0.04296588897705078, + -0.1620718538761139, + 0.012366419658064842, + 0.047717608511447906, + 0.15346293151378632, + -0.011503779329359531 + ], + "probabilities": [ + 0.1512979418039322, + 0.1388673633337021, + 0.13216042518615723, + 0.14209792017936707, + 0.1442011296749115, + 0.1506800651550293, + 0.14069513976573944 + ] + }, + { + "options": 31, + "logits": [ + 0.16330909729003906, + -0.04296589270234108, + -0.1620718389749527, + 0.012366367504000664, + 0.04771760478615761, + 0.1534629762172699, + -0.011503748595714569, + -0.15352031588554382, + -0.12729023396968842, + -0.018939895555377007, + 0.07874026149511337, + 0.1336236596107483, + -0.07686153799295425, + -0.1321595162153244, + -0.038695480674505234, + 0.08788273483514786, + 0.18759208917617798, + -0.036322422325611115, + -0.18305128812789917, + 0.025450605899095535, + 0.07632403075695038, + 0.15722697973251343, + 0.00885787419974804, + -0.14430277049541473, + -0.09872071444988251, + 0.013050932437181473, + 0.10318844020366669, + 0.11492683738470078, + -0.07907889038324356, + -0.12160822004079819, + -0.05631116405129433 + ], + "probabilities": [ + 0.03454456478357315, + 0.031706396490335464, + 0.030175060033798218, + 0.032444000244140625, + 0.03292420879006386, + 0.0344034880399704, + 0.03212371841073036, + 0.03028249740600586, + 0.03061443567276001, + 0.03202458843588829, + 0.03335146605968475, + 0.03412097692489624, + 0.03126285597681999, + 0.030552541837096214, + 0.03176271915435791, + 0.03347843885421753, + 0.03489496931433678, + 0.03179406002163887, + 0.02991309203207493, + 0.03262091055512428, + 0.033317990601062775, + 0.03445735201239586, + 0.03239672631025314, + 0.03039873205125332, + 0.030980117619037628, + 0.032453231513500214, + 0.03369208425283432, + 0.03385685756802559, + 0.031234057620167732, + 0.030686818063259125, + 0.03153102099895477 + ] + }, + { + "options": 255, + "logits": [ + 0.16330909729003906, + -0.04296589270234108, + -0.1620718389749527, + 0.012366367504000664, + 0.04771760478615761, + 0.1534629762172699, + -0.011503748595714569, + -0.15352031588554382, + -0.12729023396968842, + -0.018939895555377007, + 0.07874026149511337, + 0.1336236596107483, + -0.07686153799295425, + -0.1321595162153244, + -0.038695480674505234, + 0.08788273483514786, + 0.18759208917617798, + -0.036322422325611115, + -0.18305128812789917, + 0.025450605899095535, + 0.07632403075695038, + 0.15722697973251343, + 0.00885787419974804, + -0.14430277049541473, + -0.09872071444988251, + 0.013050932437181473, + 0.10318844020366669, + 0.11492683738470078, + -0.07907890528440475, + -0.12160823494195938, + -0.05631116405129433, + 0.07068727165460587, + 0.167240172624588, + -0.055517807602882385, + -0.15244919061660767, + -0.0018523912876844406, + 0.05228900536894798, + 0.18918664753437042, + 0.003996687941253185, + -0.14683139324188232, + -0.09803099185228348, + 0.07003703713417053, + 0.15613742172718048, + 0.09519751369953156, + -0.06439848989248276, + -0.12586374580860138, + -0.05452625826001167, + 0.10998698323965073, + 0.1803295910358429, + -0.1095396876335144, + -0.17051587998867035, + -0.058045294135808945, + 0.06769002228975296, + 0.17876161634922028, + 0.02370007336139679, + -0.14835882186889648, + -0.04036284238100052, + 0.024112947285175323, + 0.14640827476978302, + 0.09839790314435959, + -0.08899127691984177, + -0.1127111166715622, + -0.05864831432700157, + 0.09408044815063477, + 0.15783268213272095, + -0.08195436745882034, + -0.14064112305641174, + -0.06312527507543564, + 0.0818595364689827, + 0.1974586397409439, + -0.008748650550842285, + -0.15902179479599, + -0.009172530844807625, + 0.020293492823839188, + 0.1351587176322937, + 0.04563114047050476, + -0.07360135018825531, + -0.11956068873405457, + -0.01191377080976963, + 0.07325820624828339, + 0.14599785208702087, + -0.08439307659864426, + -0.11306572705507278, + -0.07357587665319443, + 0.10550788789987564, + 0.19497254490852356, + -0.04602048918604851, + -0.17397882044315338, + -0.01272417139261961, + 0.05303587391972542, + 0.12375271320343018, + 0.0177279245108366, + -0.11302593350410461, + -0.10215660184621811, + 0.006327805109322071, + 0.05703119561076164, + 0.13509738445281982, + -0.07641758024692535, + -0.13544845581054688, + -0.04186152666807175, + 0.09860523045063019, + 0.16330909729003906, + -0.04296589270234108, + -0.1620718389749527, + 0.012366367504000664, + 0.04771760478615761, + 0.1534629762172699, + -0.011503748595714569, + -0.15352031588554382, + -0.12729023396968842, + -0.018939895555377007, + 0.07874026149511337, + 0.1336236596107483, + -0.07686153799295425, + -0.1321595162153244, + -0.038695480674505234, + 0.08788273483514786, + 0.18759208917617798, + -0.036322422325611115, + -0.18305128812789917, + 0.025450605899095535, + 0.07632403075695038, + 0.15722697973251343, + 0.00885787419974804, + -0.14430277049541473, + -0.09872071444988251, + 0.013050932437181473, + 0.10318844020366669, + 0.11492683738470078, + -0.07907890528440475, + -0.12160823494195938, + -0.05631116405129433, + 0.07068727165460587, + 0.167240172624588, + -0.055517807602882385, + -0.15244919061660767, + -0.0018523912876844406, + 0.05228900536894798, + 0.18918664753437042, + 0.003996687941253185, + -0.14683139324188232, + -0.09803099185228348, + 0.07003703713417053, + 0.15613742172718048, + 0.09519751369953156, + -0.06439848989248276, + -0.12586374580860138, + -0.05452625826001167, + 0.10998698323965073, + 0.1803295910358429, + -0.1095396876335144, + -0.17051587998867035, + -0.058045294135808945, + 0.06769002228975296, + 0.17876161634922028, + 0.02370007336139679, + -0.14835882186889648, + -0.04036284238100052, + 0.024112947285175323, + 0.14640827476978302, + 0.09839790314435959, + -0.08899127691984177, + -0.1127111166715622, + -0.05864831432700157, + 0.09408044815063477, + 0.15783268213272095, + -0.08195436745882034, + -0.14064112305641174, + -0.06312527507543564, + 0.0818595364689827, + 0.1974586397409439, + -0.008748650550842285, + -0.15902179479599, + -0.009172530844807625, + 0.020293492823839188, + 0.1351587176322937, + 0.04563114047050476, + -0.07360135018825531, + -0.11956068873405457, + -0.01191377080976963, + 0.07325820624828339, + 0.14599785208702087, + -0.08439307659864426, + -0.11306572705507278, + -0.07357587665319443, + 0.10550788789987564, + 0.19497254490852356, + -0.04602048918604851, + -0.17397882044315338, + -0.01272417139261961, + 0.05303587391972542, + 0.12375271320343018, + 0.0177279245108366, + -0.11302593350410461, + -0.10215660184621811, + 0.006327805109322071, + 0.05703119561076164, + 0.13509738445281982, + -0.07641758024692535, + -0.13544845581054688, + -0.04186152666807175, + 0.09860523045063019, + 0.16330909729003906, + -0.04296589270234108, + -0.1620718389749527, + 0.012366367504000664, + 0.04771760478615761, + 0.1534629762172699, + -0.011503748595714569, + -0.15352031588554382, + -0.12729023396968842, + -0.018939895555377007, + 0.07874026149511337, + 0.1336236596107483, + -0.07686153799295425, + -0.1321595162153244, + -0.038695480674505234, + 0.08788273483514786, + 0.18759208917617798, + -0.036322422325611115, + -0.18305128812789917, + 0.025450605899095535, + 0.07632403075695038, + 0.15722697973251343, + 0.00885787419974804, + -0.14430277049541473, + -0.09872071444988251, + 0.013050932437181473, + 0.10318844020366669, + 0.11492683738470078, + -0.07907890528440475, + -0.12160823494195938, + -0.05631116405129433, + 0.07068727165460587, + 0.167240172624588, + -0.055517807602882385, + -0.15244919061660767, + -0.0018523912876844406, + 0.05228900536894798, + 0.18918664753437042, + 0.003996687941253185, + -0.14683139324188232, + -0.09803099185228348, + 0.07003703713417053, + 0.15613742172718048, + 0.09519751369953156, + -0.06439848989248276, + -0.12586374580860138, + -0.05452625826001167, + 0.10998698323965073, + 0.1803295910358429, + -0.1095396876335144, + -0.17051589488983154, + -0.058045294135808945, + 0.06769000738859177 + ], + "probabilities": [ + 0.004190933890640736, + 0.0038466083351522684, + 0.0036608269438147545, + 0.003936094231903553, + 0.003994353115558624, + 0.00417381851002574, + 0.0038972378242760897, + 0.0036738612689077854, + 0.0037141318898648024, + 0.00388521165587008, + 0.004046187736093998, + 0.004139544442296028, + 0.0037927983794361353, + 0.003706623101606965, + 0.003853441448882222, + 0.004061591811478138, + 0.004233445040881634, + 0.003857244038954377, + 0.0036290453281253576, + 0.003957557491958141, + 0.004042126704007387, + 0.004180353134870529, + 0.003930358681827784, + 0.0036879628896713257, + 0.003758496604859829, + 0.003937214147299528, + 0.004087510984390974, + 0.004107501823455095, + 0.003789304755628109, + 0.0037229135632514954, + 0.003825332038104534, + 0.004032668191939592, + 0.004197786562144756, + 0.003826593514531851, + 0.003675497369840741, + 0.003912901971489191, + 0.004001948982477188, + 0.0042362515814602375, + 0.003922426141798496, + 0.0036840890534222126, + 0.003759574145078659, + 0.004031578078866005, + 0.0041784606873989105, + 0.00407395837828517, + 0.003812495619058609, + 0.0037163347005844116, + 0.003828170709311962, + 0.004099077079445124, + 0.004220685921609402, + 0.0037416343111544847, + 0.0036480019334703684, + 0.0038225757889449596, + 0.004027647897601128, + 0.0042179361917078495, + 0.003954679239541292, + 0.0036817509680986404, + 0.0038507720455527306, + 0.0039553577080369, + 0.00416159862652421, + 0.004079381003975868, + 0.003773725824430585, + 0.0037367057520896196, + 0.0038216179236769676, + 0.004072067327797413, + 0.004181405529379845, + 0.0037847787607461214, + 0.0036935796961188316, + 0.003814513562247157, + 0.004051436670124531, + 0.004250840749591589, + 0.00390170281752944, + 0.003665470751002431, + 0.0039010157343000174, + 0.003949083853513002, + 0.004142186604440212, + 0.003990890458226204, + 0.003797941142693162, + 0.003726082853972912, + 0.003896573791280389, + 0.004036979284137487, + 0.004160888958722353, + 0.0037809445057064295, + 0.0037361548747867346, + 0.003797981422394514, + 0.004091453738510609, + 0.004246450960636139, + 0.0038417279720306396, + 0.0036427550949156284, + 0.0038952615577727556, + 0.004003191366791725, + 0.004122596699744463, + 0.003944875206798315, + 0.003736216574907303, + 0.0037531331181526184, + 0.003926228266209364, + 0.004009844269603491, + 0.004142080899327993, + 0.003793498268350959, + 0.0037015597335994244, + 0.0038483741227537394, + 0.004079732578247786, + 0.004190933890640736, + 0.0038466083351522684, + 0.0036608269438147545, + 0.003936094231903553, + 0.003994353115558624, + 0.00417381851002574, + 0.0038972378242760897, + 0.0036738612689077854, + 0.0037141318898648024, + 0.00388521165587008, + 0.004046187736093998, + 0.004139544442296028, + 0.0037927983794361353, + 0.003706623101606965, + 0.003853441448882222, + 0.004061591811478138, + 0.004233445040881634, + 0.003857244038954377, + 0.0036290453281253576, + 0.003957557491958141, + 0.004042126704007387, + 0.004180353134870529, + 0.003930358681827784, + 0.0036879628896713257, + 0.003758496604859829, + 0.003937214147299528, + 0.004087510984390974, + 0.004107501823455095, + 0.003789304755628109, + 0.0037229135632514954, + 0.003825332038104534, + 0.004032668191939592, + 0.004197786562144756, + 0.003826593514531851, + 0.003675497369840741, + 0.003912901971489191, + 0.004001948982477188, + 0.0042362515814602375, + 0.003922426141798496, + 0.0036840890534222126, + 0.003759574145078659, + 0.004031578078866005, + 0.0041784606873989105, + 0.00407395837828517, + 0.003812495619058609, + 0.0037163347005844116, + 0.003828170709311962, + 0.004099077079445124, + 0.004220685921609402, + 0.0037416343111544847, + 0.0036480019334703684, + 0.0038225757889449596, + 0.004027647897601128, + 0.0042179361917078495, + 0.003954679239541292, + 0.0036817509680986404, + 0.0038507720455527306, + 0.0039553577080369, + 0.00416159862652421, + 0.004079381003975868, + 0.003773725824430585, + 0.0037367057520896196, + 0.0038216179236769676, + 0.004072067327797413, + 0.004181405529379845, + 0.0037847787607461214, + 0.0036935796961188316, + 0.003814513562247157, + 0.004051436670124531, + 0.004250840749591589, + 0.00390170281752944, + 0.003665470751002431, + 0.0039010157343000174, + 0.003949083853513002, + 0.004142186604440212, + 0.003990890458226204, + 0.003797941142693162, + 0.003726082853972912, + 0.003896573791280389, + 0.004036979284137487, + 0.004160888958722353, + 0.0037809445057064295, + 0.0037361548747867346, + 0.003797981422394514, + 0.004091453738510609, + 0.004246450960636139, + 0.0038417279720306396, + 0.0036427550949156284, + 0.0038952615577727556, + 0.004003191366791725, + 0.004122596699744463, + 0.003944875206798315, + 0.003736216574907303, + 0.0037531331181526184, + 0.003926228266209364, + 0.004009844269603491, + 0.004142080899327993, + 0.003793498268350959, + 0.0037015597335994244, + 0.0038483741227537394, + 0.004079732578247786, + 0.004190933890640736, + 0.0038466083351522684, + 0.0036608269438147545, + 0.003936094231903553, + 0.003994353115558624, + 0.00417381851002574, + 0.0038972378242760897, + 0.0036738612689077854, + 0.0037141318898648024, + 0.00388521165587008, + 0.004046187736093998, + 0.004139544442296028, + 0.0037927983794361353, + 0.003706623101606965, + 0.003853441448882222, + 0.004061591811478138, + 0.004233445040881634, + 0.003857244038954377, + 0.0036290453281253576, + 0.003957557491958141, + 0.004042126704007387, + 0.004180353134870529, + 0.003930358681827784, + 0.0036879628896713257, + 0.003758496604859829, + 0.003937214147299528, + 0.004087510984390974, + 0.004107501823455095, + 0.003789304755628109, + 0.0037229135632514954, + 0.003825332038104534, + 0.004032668191939592, + 0.004197786562144756, + 0.003826593514531851, + 0.003675497369840741, + 0.003912901971489191, + 0.004001948982477188, + 0.0042362515814602375, + 0.003922426141798496, + 0.0036840890534222126, + 0.003759574145078659, + 0.004031578078866005, + 0.0041784606873989105, + 0.00407395837828517, + 0.003812495619058609, + 0.0037163347005844116, + 0.003828170709311962, + 0.004099077079445124, + 0.004220685921609402, + 0.0037416343111544847, + 0.0036480019334703684, + 0.0038225757889449596, + 0.00402764743193984 + ] + } + ] +} diff --git a/tests/kev/generate_head_fixture.py b/tests/kev/generate_head_fixture.py new file mode 100644 index 00000000..4b3eff3e --- /dev/null +++ b/tests/kev/generate_head_fixture.py @@ -0,0 +1,91 @@ +"""Run the pinned upstream PointerHead on deterministic hidden states, on CPU.""" + +import argparse +import ast +import hashlib +import json +import math +from pathlib import Path + +import torch + +SOURCE_SHA256 = "2ff3b2e67c048e3bc5a0d1a32aafe50804ed315a5564fce06bb4b95697c5271c" +HEAD_SHA256 = "dd633435998ecc751ac538717a3742e32149500fabf7d7276287dbf0693f347c" + + +def hidden(row): + return torch.tensor( + [((i * 37 + row * 19) % 101 - 50) / 64 for i in range(2560)], + dtype=torch.float32, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference-model", type=Path, required=True) + parser.add_argument("--head", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + source = args.reference_model.read_bytes() + if ( + hashlib.sha256(source).hexdigest() != SOURCE_SHA256 + or hashlib.sha256(args.head.read_bytes()).hexdigest() != HEAD_SHA256 + ): + raise ValueError("reference source or head does not match the pinned release") + # Execute the actual pinned class without importing the unrelated backbone and + # training dependencies. The hash above covers the entire upstream model.py. + definition = next( + n + for n in ast.parse(source).body + if isinstance(n, ast.ClassDef) and n.name == "PointerHead" + ) + namespace = {"nn": torch.nn, "math": math} + exec( + compile( + ast.Module(body=[definition], type_ignores=[]), + str(args.reference_model), + "exec", + ), + namespace, + ) + meta = torch.load(args.head, map_location="cpu", weights_only=True) + head = namespace["PointerHead"](2560, 256) + head.load_state_dict(meta["head"], strict=True) + head.temperature = meta["temperature"] + torch.set_num_threads(1) + cases = [] + with torch.no_grad(): + for n in (1, 2, 7, 31, 255): + options = torch.stack([hidden(i + 1) for i in range(n)]) + head.train() + logits = head(hidden(0), options) + head.eval() + probabilities = head(hidden(0), options).softmax(-1) + cases.append( + { + "options": n, + "logits": logits.tolist(), + "probabilities": probabilities.tolist(), + } + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps( + { + "reference_revision": "2074cc44256200bb37540103fa4572901faa10f3", + "reference_model_sha256": SOURCE_SHA256, + "head_pt_sha256": HEAD_SHA256, + "torch": torch.__version__, + "device": "cpu", + "temperature": meta["temperature"], + "hidden_rule": "((i * 37 + row * 19) % 101 - 50) / 64; row 0 = decide", + "cases": cases, + }, + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/kev/head.rs b/tests/kev/head.rs new file mode 100644 index 00000000..b7eefb83 --- /dev/null +++ b/tests/kev/head.rs @@ -0,0 +1,127 @@ +use omni_kev_native::{head::PointerHead, postprocess::probabilities}; +use safetensors::{ + Dtype, + tensor::{TensorView, serialize}, +}; +use std::{ + path::PathBuf, + sync::atomic::{AtomicU64, Ordering}, +}; + +static SERIAL: AtomicU64 = AtomicU64::new(0); + +struct Fixture(PathBuf); +impl Fixture { + fn new(weights: &[f32]) -> Self { + let path = std::env::temp_dir().join(format!( + "kev-head-{}-{}.safetensors", + std::process::id(), + SERIAL.fetch_add(1, Ordering::Relaxed) + )); + let data: Vec = weights.iter().flat_map(|v| v.to_le_bytes()).collect(); + let bias: Vec = [0.5f32, -0.25] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let tensors = [ + ( + "q.weight", + TensorView::new(Dtype::F32, vec![2, 2], &data).unwrap(), + ), + ( + "k.weight", + TensorView::new(Dtype::F32, vec![2, 2], &data).unwrap(), + ), + ( + "q.bias", + TensorView::new(Dtype::F32, vec![2], &bias).unwrap(), + ), + ( + "k.bias", + TensorView::new(Dtype::F32, vec![2], &bias).unwrap(), + ), + ]; + std::fs::write(&path, serialize(tensors, None).unwrap()).unwrap(); + Self(path) + } +} +impl Drop for Fixture { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } +} + +#[test] +fn trained_biases_scaling_calibration_and_candidate_order() { + let fixture = Fixture::new(&[1., 0., 0., 1.]); + let head = PointerHead::load(&fixture.0).unwrap(); + let logits = head + .logits(&[1., 2.], &[&[3., 4.], &[1., 0.], &[3., 4.]]) + .unwrap(); + let expected = [11.8125f32, 1.8125, 11.8125]; // dot after both learned biases + for (actual, expected) in logits.iter().zip(expected) { + assert!((actual - expected / 2f32.sqrt()).abs() < 1e-6); + } + let probs = probabilities(&logits, 2.4).unwrap(); + assert_eq!(probs[0], probs[2]); + assert!(probs[0] > probs[1]); + assert!((probs.iter().sum::() - 1.).abs() < 1e-6); + assert_eq!(probabilities(&[123.], 2.4).unwrap(), vec![1.]); + assert!(head.logits(&[1., 2.], &vec![&[0., 1.][..]; 255]).is_ok()); +} + +#[test] +fn rejects_invalid_weights_inputs_and_calibration() { + let fixture = Fixture::new(&[1., 0., 0., 1.]); + let head = PointerHead::load(&fixture.0).unwrap(); + assert!(head.logits(&[1.], &[&[1., 2.]]).is_err()); + assert!(head.logits(&[1., 2.], &[&[f32::NAN, 2.]]).is_err()); + assert!(head.logits(&[1., 2.], &[]).is_err()); + assert!(head.logits(&[1., 2.], &vec![&[0., 1.][..]; 256]).is_err()); + assert!(PointerHead::load(&Fixture::new(&[f32::NAN; 4]).0).is_err()); + std::fs::write(&fixture.0, b"invalid safetensors").unwrap(); + assert!(PointerHead::load(&fixture.0).is_err()); + for temperature in [0., -1., f64::NAN, f64::INFINITY, 1e-100, 1e100] { + assert!(probabilities(&[1., 2.], temperature).is_err()); + } + assert!(probabilities(&[f32::INFINITY], 1.).is_err()); + assert!(probabilities(&[], 1.).is_err()); +} + +fn hidden(row: usize) -> Vec { + (0..2560) + .map(|i| (((i * 37 + row * 19) % 101) as f32 - 50.) / 64.) + .collect() +} + +#[test] +#[ignore = "needs KEV_HEAD pointing to the exported pinned head.safetensors; CPU only"] +fn released_head_matches_pinned_upstream_class() { + let path = PathBuf::from(std::env::var_os("KEV_HEAD").expect("set KEV_HEAD")); + let head = PointerHead::load(&path).unwrap(); + let fixture: serde_json::Value = serde_json::from_str(include_str!("data/head.json")).unwrap(); + let mut max_logit_error = 0f64; + let mut max_probability_error = 0f64; + for case in fixture["cases"].as_array().unwrap() { + let options: Vec<_> = (1..=case["options"].as_u64().unwrap() as usize) + .map(hidden) + .collect(); + let refs: Vec<_> = options.iter().map(Vec::as_slice).collect(); + let logits = head.logits(&hidden(0), &refs).unwrap(); + let probs = probabilities(&logits, fixture["temperature"].as_f64().unwrap()).unwrap(); + for (&actual, expected) in logits.iter().zip(case["logits"].as_array().unwrap()) { + let expected = expected.as_f64().unwrap(); + let error = (actual as f64 - expected).abs(); + max_logit_error = max_logit_error.max(error); + assert!(error <= 1e-4 + 1e-5 * expected.abs(), "logit error {error}"); + } + for (&actual, expected) in probs.iter().zip(case["probabilities"].as_array().unwrap()) { + let error = (actual as f64 - expected.as_f64().unwrap()).abs(); + max_probability_error = max_probability_error.max(error); + assert!(error <= 1e-5, "probability error {error}"); + } + } + eprintln!( + "max logit error={max_logit_error:.8e}; max probability error={max_probability_error:.8e}" + ); +} diff --git a/tests/kev/test_export.py b/tests/kev/test_export.py new file mode 100644 index 00000000..9659d45b --- /dev/null +++ b/tests/kev/test_export.py @@ -0,0 +1,146 @@ +"""CPU-only exporter gates and LoRA arithmetic; no checkpoint download.""" + +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import torch + +SPEC = importlib.util.spec_from_file_location( + "kev_export", Path(__file__).resolve().parents[2] / "recipe/kev/export.py" +) +exporter = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(exporter) + + +class ExportTests(unittest.TestCase): + def test_all_layers_and_projection_families_are_required(self): + targets = exporter.expected_targets() + self.assertEqual(len(targets), 248) + for name, count in (("gate_proj", 32), ("up_proj", 32), ("down_proj", 32), + ("q_proj", 8), ("k_proj", 8), ("v_proj", 8), ("o_proj", 8), + ("in_proj_qkv", 24), ("in_proj_a", 24), ("in_proj_z", 24), + ("in_proj_b", 24), ("out_proj", 24)): + self.assertEqual(sum(t.endswith(f".{name}.weight") for t in targets), count) + exporter.validate_targets(targets) + missing = "model.language_model.layers.0.linear_attn.in_proj_b.weight" + # The same count with a different target is also invalid. + with self.assertRaisesRegex(ValueError, "target mismatch"): + exporter.validate_targets((targets - {missing}) | {missing.replace(".0.", ".3.")}) + + def test_merge_includes_each_gdn_projection_with_one_rounding(self): + # Includes all five GDN projections omitted by a Cua-S1-only target list. + for module in ( + "in_proj_qkv", + "in_proj_a", + "in_proj_z", + "in_proj_b", + "out_proj", + ): + name = f"base_model.model.layers.0.linear_attn.{module}" + a = torch.arange(48, dtype=torch.float32).reshape(16, 3) / 31 + b = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 43 + pair = exporter.adapter_pairs( + {f"{name}.lora_A.weight": a, f"{name}.lora_B.weight": b} + ) + target = f"model.language_model.layers.0.linear_attn.{module}.weight" + self.assertEqual(set(pair), {target}) + weight = torch.tensor( + [[1.1, -2.2, 3.3], [4.4, 0.1, -0.7]], dtype=torch.bfloat16 + ) + actual = exporter.merge_weight(weight, a, b) + # Independent elementwise FP64 sums; final BF16 rounding is exact here. + expected = weight.double() + 2 * torch.stack( + [ + torch.stack( + [ + sum(b[i, k].double() * a[k, j].double() for k in range(16)) + for j in range(3) + ] + ) + for i in range(2) + ] + ) + self.assertTrue(torch.equal(actual, expected.bfloat16())) + + def test_missing_unknown_nonfinite_and_shape_mismatches_fail(self): + key = "base_model.model.layers.0.mlp.down_proj.lora_A.weight" + for tensors in ( + {key: torch.ones(16, 2)}, + {"unknown": torch.ones(2, 2)}, + {key: torch.full((16, 2), float("nan"))}, + {key: torch.ones(16, 2).half()}, + ): + with self.assertRaises(ValueError): + exporter.adapter_pairs(tensors) + with self.assertRaisesRegex(ValueError, "shape"): + exporter.merge_weight(torch.ones(2, 3), torch.ones(4, 3), torch.ones(2, 4)) + + def test_wrong_hash_precedes_pickle_loading_and_output_creation(self): + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) + (root / "config.json").write_text("wrong revision") + with patch.object( + exporter.torch, + "load", + side_effect=AssertionError("must not deserialize"), + ): + with self.assertRaisesRegex(ValueError, "pinned release"): + exporter.export(root, root, root / "out") + self.assertFalse((root / "out").exists()) + self.assertEqual(sorted(p.name for p in root.iterdir()), ["config.json"]) + + def test_never_overwrites_existing_output(self): + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) + with self.assertRaisesRegex(ValueError, "already exists"): + exporter.export(root, root, root) + + def test_delimiters_use_base_ids(self): + with tempfile.TemporaryDirectory() as temp: + path = Path(temp) / "tokenizer.json" + tokens = [{"content": k, "id": v} for k, v in exporter.SPECIAL.items()] + path.write_text(json.dumps({"added_tokens": tokens})) + exporter.validate_tokenizer(path) + tokens[0]["id"] = 151659 + path.write_text(json.dumps({"added_tokens": tokens})) + with self.assertRaisesRegex(ValueError, "delimiter"): + exporter.validate_tokenizer(path) + + def test_head_metadata_and_calibration_are_required(self): + meta = dict( + base=exporter.BASE_ID, + base_revision=exporter.BASE_REVISION, + lora=16, + head_dim=256, + option_isolation=False, + special_embeddings=False, + weights_dtype="fp32", + temperature=2.406050072164233, + head={ + "q.weight": torch.zeros(256, 2560), + "k.weight": torch.ones(256, 2560), + "q.bias": torch.zeros(256), + "k.bias": torch.ones(256), + }, + ) + _, temperature = exporter.pointer_head(meta) + self.assertEqual(temperature, meta["temperature"]) + for change in ( + dict(temperature=0), + dict(temperature=float("nan")), + dict(special_embeddings=True), + dict(base_revision="wrong"), + dict(head={}), + dict(weights_dtype="bf16"), + ): + with self.assertRaises(ValueError): + exporter.pointer_head({**meta, **change}) + + +if __name__ == "__main__": + torch.set_num_threads(2) + unittest.main()