From 4be35df43514b57636ab2926c63a52782f62b8c1 Mon Sep 17 00:00:00 2001 From: psv666 <2693925048@qq.com> Date: Sun, 25 Oct 2026 01:34:36 +0800 Subject: [PATCH 1/2] feat(kev): add pinned CPU export and FP32 pointer head Signed-off-by: psv666 <2693925048@qq.com> --- .github/workflows/ci.yml | 15 + Cargo.lock | 9 + Cargo.toml | 2 +- README.md | 1 + docs/supported-models.md | 4 + mkdocs.yml | 1 + recipe/README.md | 2 + recipe/kev/README.md | 77 +++ recipe/kev/export.py | 235 +++++++++ recipe/kev/requirements.txt | 2 + src/models/kev/README.md | 84 +++ src/models/kev/native/Cargo.toml | 17 + src/models/kev/native/LICENSE.kev | 202 +++++++ src/models/kev/native/src/head.rs | 102 ++++ src/models/kev/native/src/lib.rs | 3 + src/models/kev/native/src/postprocess.rs | 21 + tests/kev/data/head.json | 638 +++++++++++++++++++++++ tests/kev/generate_head_fixture.py | 91 ++++ tests/kev/head.rs | 127 +++++ tests/kev/test_export.py | 132 +++++ 20 files changed, 1764 insertions(+), 1 deletion(-) create mode 100644 recipe/kev/README.md create mode 100644 recipe/kev/export.py create mode 100644 recipe/kev/requirements.txt create mode 100644 src/models/kev/README.md create mode 100644 src/models/kev/native/Cargo.toml create mode 100644 src/models/kev/native/LICENSE.kev create mode 100644 src/models/kev/native/src/head.rs create mode 100644 src/models/kev/native/src/lib.rs create mode 100644 src/models/kev/native/src/postprocess.rs create mode 100644 tests/kev/data/head.json create mode 100644 tests/kev/generate_head_fixture.py create mode 100644 tests/kev/head.rs create mode 100644 tests/kev/test_export.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 306f3e39..6202d211 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -46,6 +46,21 @@ jobs: - name: Decider validation helpers run: python -m unittest discover -s tests/decider -p 'test_*.py' -v + 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 4affb469..6e6757f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1251,6 +1251,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 6c49c1f7..cd706ccd 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"] +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/kev/native", "src/models/laya", "src/backends/cuda", "src/models/decider/native", "src/models/jemm/native"] resolver = "3" diff --git a/README.md b/README.md index a1b55e87..80b83f68 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ pinned runtime for text, image and video state: | 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 | [Checkpoint export and CPU contract](recipe/omnijev/README.md) only; the native worker is in progress ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) | +| 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) | [Supported models and hardware](docs/supported-models.md) lists the devices diff --git a/docs/supported-models.md b/docs/supported-models.md index 4afd2514..d1778dcc 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -28,6 +28,10 @@ Models that are being added are also tracked in issues labeled [new model](https | OmniJev-4B v1.1 | None yet: [checkpoint export and CPU contract](../recipe/omnijev/README.md); the native worker is in progress ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) | Request contract, preparation, heads and finishing checked against the reference ([#120](https://github.com/ThinkFlowLab/system1-omni/pull/120)); no inference | Not yet | Not supported | Python 3.10 reference environment for the export and fixtures | | D3 (Decision 3.0) | [Python reference worker](../recipe/d3/reference.md) on the checkpoint's pinned runtime — text, image and video state over `choice`/`noul`/`score`, with the deliberate transport differences documented in the recipe; the native worker is in progress ([#136](https://github.com/ThinkFlowLab/system1-omni/issues/136)) | Unverified | Author-run on NVIDIA A100 (sm_80): serving smoke and reference parity for text and media (exact decisions, max delta 0.0); evidence in the integration and media PRs. Owner: @MrDongsls. | Unverified | Python 3.12, the [pinned runtime requirements](../recipe/d3/requirements.txt), the [pinned package](https://huggingface.co/vllm-sr/d3) (~52 GB); [media budgets and the 64 MiB total request cap](../recipe/d3/reference.md) | +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 ec30436b..e50a1ddd 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -78,6 +78,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 - CLM stub-encoder contract: recipe/clm/README.md - LFM2.5-350M choice worker: recipe/lfm2/README.md - D3 reference worker: recipe/d3/reference.md diff --git a/recipe/README.md b/recipe/README.md index 27778139..c62d1546 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -32,6 +32,8 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting image decisions; bounded H800 validation with explicit deployment limits. - [OmniJev-4B v1.1 preparation](omnijev/README.md): export the checkpoint and regenerate the CPU fixtures; the native worker is not available yet. +- [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. - [D3 reference worker](d3/reference.md): download the pinned vllm-sr/d3 package, 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..c9c92a5d --- /dev/null +++ b/recipe/kev/export.py @@ -0,0 +1,235 @@ +"""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 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") + ) + if len(pairs) != 248: + raise ValueError("expected all 248 Kev LoRA target matrices") + 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..93d0f32d --- /dev/null +++ b/tests/kev/test_export.py @@ -0,0 +1,132 @@ +"""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_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() From 8f765fdbf9eb782ab6f9d731df359699b70398ba Mon Sep 17 00:00:00 2001 From: psv666 <2693925048@qq.com> Date: Tue, 27 Oct 2026 16:54:45 +0800 Subject: [PATCH 2/2] fix(kev): validate every pinned LoRA projection target Signed-off-by: psv666 <2693925048@qq.com> --- recipe/kev/export.py | 27 +++++++++++++++++++++++++-- tests/kev/test_export.py | 14 ++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/recipe/kev/export.py b/recipe/kev/export.py index c9c92a5d..d00fea7d 100644 --- a/recipe/kev/export.py +++ b/recipe/kev/export.py @@ -40,6 +40,30 @@ } +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: @@ -155,8 +179,7 @@ def export(base, checkpoint, out): pairs = adapter_pairs( load_file(checkpoint / "adapter_model.safetensors", device="cpu") ) - if len(pairs) != 248: - raise ValueError("expected all 248 Kev LoRA target matrices") + 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(), {} diff --git a/tests/kev/test_export.py b/tests/kev/test_export.py index 93d0f32d..9659d45b 100644 --- a/tests/kev/test_export.py +++ b/tests/kev/test_export.py @@ -17,6 +17,20 @@ 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 (