From 999dc1eb7564576fd0508837a967f51ec4f8c902 Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz Date: Mon, 28 Sep 2026 20:04:18 +0800 Subject: [PATCH 1/4] [CLM] Add the second model engine: checkpoint and typed decision scoring CLM is the second model tracked in #9, and it decides differently in a way LAYA does not cover: the engine does not compute embeddings. A frozen Qwen3-8B encoder runs as its own process, and the engine owns everything after it. This is the half that can be checked without a GPU -- the checkpoint, the head geometry, and the decision arithmetic. A CLM checkpoint is a torch.save dict, so recipe/clm/native/export_weights.py converts it first: tensors to safetensors with the head name as a prefix, and cfg, hidden_size, projection_dim and logit_scale into the metadata. The published checkpoint is 75 MB and holds the two heads, not the 8B encoder; the export is 16 tensors and 18.9 M parameters. - config reads the head geometry from the metadata and checks it. - weights lists the expected inventory from that geometry, checks shapes before reading, and loads FP32. - scoring projects a state and each candidate, L2-normalises both, scores them by cosine under exp(logit_scale), softmaxes across the question's candidates, and assembles choice, noul and score. The three types differ only after the distribution exists, so they share one path. confidence is the top probability minus the mean of the rest, clamped, and 1.0 for a single candidate, which is the definition src/clm/schema.py uses. exp(logit_scale) is capped at 100, as the reference caps it in heads.py. The published logit_scale is 4.6132, whose exponential is 100.82, so the cap binds: without it every probability is about 0.8 % off. Checks. head_oracle.py is an independent NumPy implementation of the same arithmetic, and tests/checkpoint.rs checks 16 tensor conversion hashes against the export oracle and five decisions against head_oracle.py. Default cargo test needs no checkpoint. fmt, clippy -D warnings. --- Cargo.lock | 15 ++ Cargo.toml | 2 +- recipe/clm/native/export_weights.py | 87 ++++++++ recipe/clm/native/head_oracle.py | 162 +++++++++++++++ src/models/clm/Cargo.toml | 17 ++ src/models/clm/README.md | 59 ++++++ src/models/clm/src/config.rs | 105 ++++++++++ src/models/clm/src/lib.rs | 20 ++ src/models/clm/src/scoring.rs | 303 ++++++++++++++++++++++++++++ src/models/clm/src/weights.rs | 174 ++++++++++++++++ src/models/clm/tests/checkpoint.rs | 149 ++++++++++++++ 11 files changed, 1092 insertions(+), 1 deletion(-) create mode 100644 recipe/clm/native/export_weights.py create mode 100644 recipe/clm/native/head_oracle.py create mode 100644 src/models/clm/Cargo.toml create mode 100644 src/models/clm/README.md create mode 100644 src/models/clm/src/config.rs create mode 100644 src/models/clm/src/lib.rs create mode 100644 src/models/clm/src/scoring.rs create mode 100644 src/models/clm/src/weights.rs create mode 100644 src/models/clm/tests/checkpoint.rs diff --git a/Cargo.lock b/Cargo.lock index 11b8ba25..dcba5bb4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -452,6 +452,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -974,6 +975,19 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "omni-clm" +version = "0.1.0" +dependencies = [ + "anyhow", + "memmap2", + "reqwest", + "safetensors 0.6.2", + "serde", + "serde_json", + "sha2", +] + [[package]] name = "omni-cua-s1-native" version = "0.1.0" @@ -1346,6 +1360,7 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", + "futures-channel", "futures-core", "futures-util", "http", diff --git a/Cargo.toml b/Cargo.toml index d5abbe60..242d51ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/runtime", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/laya", "src/backends/cuda"] +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/laya", "src/backends/cuda"] resolver = "3" diff --git a/recipe/clm/native/export_weights.py b/recipe/clm/native/export_weights.py new file mode 100644 index 00000000..ab1215fc --- /dev/null +++ b/recipe/clm/native/export_weights.py @@ -0,0 +1,87 @@ +"""Export a CLM head checkpoint to safetensors, and record the conversion oracle. + +A CLM checkpoint is a ``torch.save`` dict (``state_head``/``action_head`` state dicts, +``logit_scale``, ``cfg``), so it is a pickle and no non-Python reader can open it. This +writes the tensors to safetensors with the head name as a prefix, keeps the scalar and +config entries in the safetensors metadata, and emits the FP32/FP16/BF16 hash of every +tensor so a reader can be checked without comparing floats directly. + + python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR + python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR --oracle oracle.json + +Writes ``model.safetensors`` and, unless ``--no-oracle``, ``oracle.json``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from safetensors.torch import save_file + +HEADS = ("state_head", "action_head") + + +def tensors(ckpt: dict) -> dict[str, torch.Tensor]: + """Every parameter, prefixed by its head, in a stable order.""" + out: dict[str, torch.Tensor] = {} + for head in HEADS: + state = ckpt.get(head) + if not isinstance(state, dict): + raise SystemExit(f"checkpoint has no {head!r} state dict") + for name, value in state.items(): + if not isinstance(value, torch.Tensor): + raise SystemExit(f"{head}.{name} is {type(value).__name__}, not a tensor") + out[f"{head}.{name}"] = value.detach().to(torch.float32).contiguous() + return out + + +def oracle(weights: dict[str, torch.Tensor]) -> list[dict]: + rows = [] + for name, x in weights.items(): + row: dict = {"name": name, "shape": list(x.shape), "source_dtype": str(x.dtype)} + for key, dtype in (("f32", torch.float32), ("f16", torch.float16), ("bf16", torch.bfloat16)): + y = x.to(torch.float32).to(dtype).contiguous() + row[key] = hashlib.sha256(y.view(torch.uint8).numpy().tobytes()).hexdigest() + rows.append(row) + return rows + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("checkpoint", type=Path, help="CLM_v0.1-8B.pt") + parser.add_argument("output", type=Path, help="directory to write into") + parser.add_argument("--oracle", type=Path, help="where to write the conversion oracle") + parser.add_argument("--no-oracle", action="store_true", help="skip the oracle") + args = parser.parse_args() + + torch.set_num_threads(4) + ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + cfg = dict(ckpt.get("cfg") or {}) + weights = tensors(ckpt) + args.output.mkdir(parents=True, exist_ok=True) + + metadata = { + "format": "clm-heads", + "logit_scale": repr(float(ckpt["logit_scale"])), + "hidden_size": str(int(ckpt.get("hidden_size") or cfg.get("hidden_size"))), + "projection_dim": str(int(ckpt.get("projection_dim") or cfg.get("projection_dim"))), + "cfg": json.dumps(cfg, sort_keys=True), + } + path = args.output / "model.safetensors" + save_file(weights, str(path), metadata=metadata) + + params = sum(v.numel() for v in weights.values()) + print(f"WROTE {path} tensors={len(weights)} params={params}", flush=True) + + if not args.no_oracle: + where = args.oracle or (args.output / "oracle.json") + where.write_text(json.dumps(oracle(weights), indent=2) + "\n") + print(f"ORACLE {where} rows={len(weights)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/recipe/clm/native/head_oracle.py b/recipe/clm/native/head_oracle.py new file mode 100644 index 00000000..38c865d9 --- /dev/null +++ b/recipe/clm/native/head_oracle.py @@ -0,0 +1,162 @@ +"""Reference decisions for the CLM heads, for the Rust port to be checked against. + +Reads the exported safetensors (not the .pt) so both sides load byte-identical weights, +computes a decision for each fixed case with plain NumPy, and writes the result as JSON. + + python recipe/clm/native/head_oracle.py EXPORT_DIR OUT.json + +The embeddings here are synthesised from the case name, so the numbers are meaningless +as model output -- the point is that two implementations of the same arithmetic agree. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +from safetensors import safe_open + +HIDDEN = 4096 + + +def embedding(text: str, dim: int) -> np.ndarray: + """A stable unit vector per text, so the cases are reproducible without an encoder.""" + out: list[float] = [] + counter = 0 + while len(out) < dim: + digest = hashlib.sha256(f"{counter}:{text}".encode()).digest() + for i in range(0, len(digest) - 3, 4): + if len(out) == dim: + break + out.append(int.from_bytes(digest[i:i + 4], "big") / 2**31 - 1.0) + counter += 1 + x = np.asarray(out, dtype=np.float32) + return x / np.linalg.norm(x) + + +def load_head(f, prefix: str, width: int, proj: int, hidden: int, blocks: int) -> dict: + def get(name: str) -> np.ndarray: + return np.asarray(f.get_tensor(f"{prefix}.{name}"), dtype=np.float32) + + head = { + "inp_w": get("inp.weight"), "inp_b": get("inp.bias"), + "out_w": get("out.weight"), "out_b": get("out.bias"), + } + for i in range(blocks): + head[f"hidden{i}_w"] = get(f"hidden.{i}.weight") + head[f"hidden{i}_b"] = get(f"hidden.{i}.bias") + head[f"norm{i}_w"] = get(f"norms.{i}.weight") + head[f"norm{i}_b"] = get(f"norms.{i}.bias") + return head + + +def gelu(x: np.ndarray) -> np.ndarray: + # The erf form, which is what torch.nn.GELU() uses by default. + from math import sqrt + + return 0.5 * x * (1.0 + _erf(x / sqrt(2.0))) + + +def _erf(x: np.ndarray) -> np.ndarray: + """Vectorised erf via the same A&S 7.1.26 form the Rust side uses.""" + sign = np.sign(x) + x = np.abs(x) + t = 1.0 / (1.0 + 0.3275911 * x) + y = 1.0 - (((((1.0614054 * t - 1.45315203) * t) + 1.42141374) * t - 0.284496736) * t + + 0.2548296) * t * np.exp(-x * x) + return sign * y + + +def layernorm(x: np.ndarray, w: np.ndarray, b: np.ndarray, eps: float = 1e-5) -> np.ndarray: + mean = x.mean(axis=-1, keepdims=True) + var = x.var(axis=-1, keepdims=True) + return (x - mean) / np.sqrt(var + eps) * w + b + + +def project(head: dict, cfg: dict, x: np.ndarray) -> np.ndarray: + h = gelu(x @ head["inp_w"].T + head["inp_b"]) + blocks = cfg["depth"] - 2 + for i in range(blocks): + z = h @ head[f"hidden{i}_w"].T + head[f"hidden{i}_b"] + if cfg["layernorm"]: + z = layernorm(z, head[f"norm{i}_w"], head[f"norm{i}_b"]) + z = gelu(z) + h = h + z if cfg["residual"] else z + out = h @ head["out_w"].T + head["out_b"] + return out / np.linalg.norm(out, axis=-1, keepdims=True) + + +def softmax(v: np.ndarray) -> np.ndarray: + e = np.exp(v - v.max()) + return e / e.sum() + + +def confidence(probs: np.ndarray) -> float: + if probs.size < 2: + return 1.0 + j = int(np.argmax(probs)) + rest = np.delete(probs, j).mean() + return float(min(1.0, max(0.0, probs[j] - rest))) + + +CASES = [ + ("choice_binary", "choice", ["billing", "technical"]), + ("choice_five", "choice", ["a", "b", "c", "d", "e"]), + ("score_three", "score", ["0", "1", "2"]), + ("noul", "noul", ["false", "true"]), + ("choice_single", "choice", ["only"]), +] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("export", type=Path, help="directory holding model.safetensors") + parser.add_argument("output", type=Path) + args = parser.parse_args() + + with safe_open(args.export / "model.safetensors", framework="np") as f: + md = f.metadata() + cfg = json.loads(md["cfg"]) + hidden = int(md["hidden_size"]) + proj = int(md["projection_dim"]) + width = cfg["width"] + blocks = cfg["depth"] - 2 + # The reference caps this: heads.py does exp(logit_scale).clamp(max=100.0), + # and the published checkpoint's 4.6132 exponentiates to 100.82, so the cap + # binds. Without it every probability is about 0.8 % off. + scale = min(float(np.exp(np.float32(md["logit_scale"]))), 100.0) + state = load_head(f, "state_head", width, proj, hidden, blocks) + action = load_head(f, "action_head", width, proj, hidden, blocks) + + rows = [] + for name, kind, keys in CASES: + state_vec = embedding(f"state::{name}", hidden) + cand_vecs = np.stack([embedding(f"cand::{name}::{k}", hidden) for k in keys]) + zs = project(state, cfg, state_vec[None, :])[0] + zc = project(action, cfg, cand_vecs) + cos = zc @ zs + temperature = 1.0 if name != "choice_five" else 2.5 + probs = softmax((np.float32(scale) * cos / np.float32(temperature)).astype(np.float32)) + row = { + "name": name, "kind": kind, "keys": keys, "temperature": temperature, + "probabilities": [float(p) for p in probs], + } + if kind == "choice": + row["choice"] = keys[int(np.argmax(probs))] + row["confidence"] = confidence(probs) + elif kind == "score": + row["score"] = float(sum(i * float(p) for i, p in enumerate(probs))) + row["confidence"] = confidence(probs) + else: + row["noul"] = float(probs[keys.index("true")]) + rows.append(row) + + args.output.write_text(json.dumps({"cases": rows}, indent=2) + "\n") + print(f"ORACLE {args.output} cases={len(rows)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/models/clm/Cargo.toml b/src/models/clm/Cargo.toml new file mode 100644 index 00000000..6fa26413 --- /dev/null +++ b/src/models/clm/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "omni-clm" +version = "0.1.0" +edition = "2024" +publish = false +description = "CLM: projection heads, typed decision scoring and the embeddings client" + +[dependencies] +anyhow = "1" +memmap2 = "0.9" +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +safetensors = "0.6" +serde = { version = "1", features = ["derive"] } +# preserve_order keeps object key order, which to_text depends on: the heads are +# trained on the caller's field order, so re-sorting a state would change the text. +serde_json = { version = "1", features = ["preserve_order"] } +sha2 = "0.10" diff --git a/src/models/clm/README.md b/src/models/clm/README.md new file mode 100644 index 00000000..41d123cf --- /dev/null +++ b/src/models/clm/README.md @@ -0,0 +1,59 @@ +# CLM model engine + +CLM is the second model the project tracks ([#9](https://github.com/ThinkFlowLab/system1-omni/issues/9)), after LAYA. It decides differently in a way LAYA does not cover: **the engine does not compute embeddings.** A frozen `Qwen/Qwen3-8B` encoder runs as its own process behind an OpenAI-compatible `/v1/embeddings` endpoint, and the engine owns everything after it — two projection heads, the cosine score, and the typed answer. + +That split is the point of implementing it second. LAYA's engine owns one forward pass; this one owns a client to someone else's server, plus a candidate-vector cache that persists across requests. + +## What is here + +`omni-clm` reads a converted checkpoint and computes decisions. It does not call an embeddings endpoint and does not serve HTTP yet; those belong with the runtime that owns the request path. + +| module | responsibility | +| --- | --- | +| `config` | the head geometry, read from the safetensors metadata | +| `weights` | tensor inventory, shape checks, FP32 loading | +| `scoring` | projection, cosine score, softmax, and the three answer types | + +## The checkpoint is converted first + +A CLM checkpoint is a `torch.save` dict, so it is a pickle and no non-Python reader can open it. `recipe/clm/native/export_weights.py` writes the tensors to safetensors with the head name as a prefix and keeps `cfg`, `hidden_size`, `projection_dim` and `logit_scale` in the metadata. + +```sh +python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR +``` + +`CLM_v0.1-8B.pt` is the published checkpoint from `Contrastive-LM/CLM-v0.1-8B`; it is 75 MB and holds the two heads, not the 8B encoder. The export is 16 tensors, 18.9 M parameters. + +## The decision, in one pass + +A decision is `softmax(exp(logit_scale) * cos(state_head(s), action_head(c)) / temperature)` over a question's candidates. Both heads are `inp → [LayerNorm →] hidden → out` with GELU, and both projections are L2-normalised before the dot product; the published checkpoint sets `layernorm: true` and `residual: false` with one hidden block. + +The three question types differ only after the distribution exists, which is why they share one scoring path: + +| type | answer | +| --- | --- | +| `choice` | the argmax key, with `confidence` and the full distribution | +| `noul` | the `true` entry of a two-candidate distribution | +| `score` | the expected level index, `sum(i * p_i)`, with `confidence` | + +`confidence` is the top probability minus the mean of the rest, clamped to `[0, 1]`, and `1.0` for a single candidate — the TypeSafe-style definition `schema.py` uses. + +## CPU checks + +The default tests need no checkpoint: + +```sh +cargo test -p omni-clm +``` + +To check the full checkpoint and the decision arithmetic, export it first and point `CLM_EXPORT` at the directory: + +```sh +python recipe/clm/native/export_weights.py /path/to/CLM_v0.1-8B.pt /tmp/clm-export +python recipe/clm/native/head_oracle.py /tmp/clm-export /tmp/clm-export/head-oracle.json +CLM_EXPORT=/tmp/clm-export cargo test -p omni-clm -- --ignored +``` + +Two tests run there: every tensor's FP32 conversion hash against the export oracle, and **five decisions checked against an independent NumPy implementation of the same arithmetic** (`head_oracle.py`) on synthesised embeddings, so the two sides need no encoder to disagree. The embeddings are hash-derived and carry no meaning as model output — the check is that two implementations of the same maths agree. + +The default CI job skips these because it does not download the checkpoint. diff --git a/src/models/clm/src/config.rs b/src/models/clm/src/config.rs new file mode 100644 index 00000000..77e7eac7 --- /dev/null +++ b/src/models/clm/src/config.rs @@ -0,0 +1,105 @@ +//! CLM head configuration, read from the exported safetensors metadata. +//! +//! A CLM checkpoint is a `torch.save` dict, so `recipe/clm/native/export_weights.py` +//! converts it first; everything this crate reads is safetensors. The head geometry is +//! not a file of its own, so it travels in the safetensors metadata as `cfg`. +use anyhow::{Context, Result, ensure}; +use serde::Deserialize; + +/// The head shape the checkpoint was trained with. `depth` counts `inp`, the hidden +/// blocks and `out`, so `depth - 2` is the number of `hidden.N` / `norms.N` pairs. +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +pub struct HeadConfig { + pub hidden_size: usize, + pub projection_dim: usize, + pub width: usize, + pub depth: usize, + pub activation: String, + pub layernorm: bool, + pub residual: bool, + #[serde(default)] + pub model: Option, +} + +impl HeadConfig { + /// Number of hidden blocks; `depth` includes the input and output projections. + pub fn hidden_blocks(&self) -> Result { + ensure!( + self.depth >= 2, + "depth {} cannot cover an input and an output projection", + self.depth + ); + Ok(self.depth - 2) + } +} + +/// The reference's cap on `exp(logit_scale)`. +pub const MAX_SCALE: f32 = 100.0; + +#[derive(Debug, Clone, PartialEq)] +pub struct Config { + pub head: HeadConfig, + /// `exp(logit_scale)` is the inverse InfoNCE temperature the heads were trained with. + pub logit_scale: f32, +} + +impl Config { + /// Build from the safetensors metadata written by the export tool. + pub fn from_metadata( + metadata: Option<&std::collections::HashMap>, + ) -> Result { + let metadata = metadata.context("the checkpoint carries no metadata")?; + ensure!( + metadata.get("format").map(String::as_str) == Some("clm-heads"), + "not a converted CLM head checkpoint: format is {:?}", + metadata.get("format") + ); + let cfg = metadata.get("cfg").context("metadata has no cfg")?; + let mut head: HeadConfig = + serde_json::from_str(cfg).with_context(|| format!("parse cfg {cfg}"))?; + // hidden_size and projection_dim are repeated at the top level; prefer those + // when the cfg omits them so an older export still loads. + head.hidden_size = metadata + .get("hidden_size") + .and_then(|v| v.parse().ok()) + .unwrap_or(head.hidden_size); + head.projection_dim = metadata + .get("projection_dim") + .and_then(|v| v.parse().ok()) + .unwrap_or(head.projection_dim); + let logit_scale: f32 = metadata + .get("logit_scale") + .context("metadata has no logit_scale")? + .parse() + .context("logit_scale is not a number")?; + + ensure!(head.hidden_size > 0, "hidden_size must be positive"); + ensure!(head.width > 0, "width must be positive"); + ensure!(head.projection_dim > 0, "projection_dim must be positive"); + ensure!( + head.activation == "gelu" || head.activation == "relu" || head.activation == "silu", + "unsupported activation {:?}", + head.activation + ); + // The upstream head applies LayerNorm before every hidden block and adds the + // residual after it; the published 0.1 checkpoint sets both to true/false + // respectively, and the export keeps them so the two paths stay distinguishable. + head.hidden_blocks()?; + ensure!( + logit_scale.is_finite(), + "logit_scale {logit_scale} is not finite" + ); + Ok(Self { head, logit_scale }) + } + + /// `exp(logit_scale)`, capped as the reference caps it. + /// + /// `heads.py` computes `exp(logit_scale).clamp(max=100.0)`, and the published + /// checkpoint's `logit_scale` is 4.6132, whose exponential is 100.82 — so the cap + /// binds and the effective scale is 100, not 100.82. Without it every probability is + /// off by about 0.8 %, which is what the end-to-end comparison against the reference + /// caught; no CPU-side oracle can, because both sides of those share this constant. + pub fn scale(&self) -> f32 { + self.logit_scale.exp().min(MAX_SCALE) + } +} diff --git a/src/models/clm/src/lib.rs b/src/models/clm/src/lib.rs new file mode 100644 index 00000000..8530d0e3 --- /dev/null +++ b/src/models/clm/src/lib.rs @@ -0,0 +1,20 @@ +//! CLM: the second System1-Omni model engine. +//! +//! A CLM decision is made over embeddings the engine does not compute. A frozen +//! Qwen3-8B encoder sits behind an OpenAI-compatible `/v1/embeddings` endpoint and the +//! engine owns everything after it: load the two projection heads, normalise their +//! output, score candidates by cosine similarity under the trained temperature, and +//! assemble `choice`, `noul` and `score` answers. +//! +//! That split is why this model is implemented second — LAYA's engine owns one forward +//! pass, while this one owns a client to someone else's server. +//! +//! Nothing here needs a GPU: the tests drive the heads from hash-derived vectors that +//! match the Python oracle's, so the decision path is checked on a CPU-only machine. +pub mod config; +pub mod scoring; +pub mod weights; + +pub use config::{Config, HeadConfig}; +pub use scoring::{Answer, Kind, Question, answer, confidence, distribution}; +pub use weights::{Head, Heads, Weights, head_tensors}; diff --git a/src/models/clm/src/scoring.rs b/src/models/clm/src/scoring.rs new file mode 100644 index 00000000..02f7198e --- /dev/null +++ b/src/models/clm/src/scoring.rs @@ -0,0 +1,303 @@ +//! The typed decision computation: project embeddings, score candidates, assemble answers. +//! +//! This mirrors `src/clm/engine.py` and `src/clm/schema.py`: a state and each candidate +//! are embedded elsewhere (Qwen3-8B behind `/v1/embeddings`) and arrive here as vectors. +//! The state head sees the state, the action head sees every candidate, both projections +//! are L2-normalised, and the score of a pair is `exp(logit_scale) * cos(state, candidate)`, +//! divided by the request temperature and softmaxed across the question's candidates. +//! +//! The three question types differ only after the distribution exists. +use anyhow::{Result, ensure}; + +use crate::weights::{Head, Heads}; + +/// One question: its type and its candidate keys in the order the caller offered them. +#[derive(Debug, Clone, PartialEq)] +pub struct Question { + pub id: String, + pub kind: Kind, + /// Candidate keys, in order. For `Noul` these are `["false", "true"]`. + pub keys: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Kind { + Choice, + Noul, + Score, +} + +/// One answer, matching the shapes `client.py` parses. +#[derive(Debug, Clone, PartialEq)] +pub enum Answer { + Choice { + choice: String, + confidence: f32, + probabilities: Vec<(String, f32)>, + }, + Noul { + noul: f32, + }, + Score { + score: f32, + confidence: f32, + probabilities: Vec<(String, f32)>, + }, +} + +impl Answer { + /// The discrete label, as `schema.label_of` defines it. + pub fn label(&self) -> String { + match self { + Answer::Choice { choice, .. } => choice.clone(), + Answer::Noul { noul } => if *noul >= 0.5 { "true" } else { "false" }.to_string(), + Answer::Score { probabilities, .. } => probabilities + .iter() + .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(k, _)| k.clone()) + .unwrap_or_default(), + } + } +} + +/// L2-normalise one row in place, the way `torch.nn.functional.normalize` does. +/// +/// Shared with [`crate::embedding`], which normalises what the encoder hands back. +pub fn normalize(row: &mut [f32]) { + let norm = row.iter().map(|v| v * v).sum::().sqrt(); + if norm > 0.0 { + for v in row.iter_mut() { + *v /= norm; + } + } +} + +fn gelu(x: f32) -> f32 { + 0.5 * x * (1.0 + erf(x * std::f32::consts::FRAC_1_SQRT_2)) +} + +fn relu(x: f32) -> f32 { + x.max(0.0) +} + +fn silu(x: f32) -> f32 { + x / (1.0 + (-x).exp()) +} + +/// Abramowitz & Stegun 7.1.26. The coefficients are kept at full precision through named +/// constants: rounding one to a shorter f32 literal changes the value, and +/// `recipe/clm/native/head_oracle.py` carries the same digits so that both sides remain +/// the same arithmetic rather than merely similar. +#[allow(clippy::excessive_precision)] +const ERF_P: f32 = 0.3275911; +#[allow(clippy::excessive_precision)] +const ERF_A1: f32 = 0.254829592; +#[allow(clippy::excessive_precision)] +const ERF_A2: f32 = -0.284496736; +#[allow(clippy::excessive_precision)] +const ERF_A3: f32 = 1.421413741; +#[allow(clippy::excessive_precision)] +const ERF_A4: f32 = -1.453152027; +#[allow(clippy::excessive_precision)] +const ERF_A5: f32 = 1.061405429; + +fn erf(x: f32) -> f32 { + let sign = if x < 0.0 { -1.0 } else { 1.0 }; + let x = x.abs(); + let t = 1.0 / (1.0 + ERF_P * x); + let y = 1.0 + - (((((ERF_A5 * t + ERF_A4) * t + ERF_A3) * t + ERF_A2) * t + ERF_A1) * t * (-x * x).exp()); + sign * y +} + +fn activate(kind: &str, x: f32) -> f32 { + match kind { + "relu" => relu(x), + "silu" => silu(x), + _ => gelu(x), + } +} + +#[inline] +fn linear_into(x: &[f32], w: &[f32], b: &[f32], n: usize, k: usize, out: &mut [f32]) { + for i in 0..n { + let row = &w[i * k..(i + 1) * k]; + let mut acc = b[i]; + for (a, wv) in x.iter().zip(row) { + acc += a * wv; + } + out[i] = acc; + } +} + +fn layer_norm(x: &mut [f32], weight: &[f32], bias: &[f32]) { + let n = x.len() as f32; + let mean = x.iter().sum::() / n; + let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::() / n; + let inv = 1.0 / (var + 1e-5).sqrt(); + for i in 0..x.len() { + x[i] = (x[i] - mean) * inv * weight[i] + bias[i]; + } +} + +/// Project one embedding through a head, returning the L2-normalised vector. +pub fn project(head: &Head, cfg: &crate::config::HeadConfig, x: &[f32]) -> Result> { + ensure!( + x.len() == cfg.hidden_size, + "embedding has {} values, the head expects {}", + x.len(), + cfg.hidden_size + ); + let w = cfg.width; + let mut h = vec![0.0f32; w]; + linear_into( + x, + &head.inp_weight, + &head.inp_bias, + w, + cfg.hidden_size, + &mut h, + ); + for v in h.iter_mut() { + *v = activate(&cfg.activation, *v); + } + if !head.hidden_weight.is_empty() { + let mut hidden = vec![0.0f32; w]; + linear_into( + &h, + &head.hidden_weight, + &head.hidden_bias, + w, + w, + &mut hidden, + ); + if let (Some(nw), Some(nb)) = (&head.norm_weight, &head.norm_bias) { + layer_norm(&mut hidden, nw, nb); + } + for v in hidden.iter_mut() { + *v = activate(&cfg.activation, *v); + } + if cfg.residual { + for i in 0..w { + hidden[i] += h[i]; + } + } + h = hidden; + } + let p = cfg.projection_dim; + let mut out = vec![0.0f32; p]; + linear_into(&h, &head.out_weight, &head.out_bias, p, w, &mut out); + normalize(&mut out); + Ok(out) +} + +/// `softmax(scale * cos / temperature)`, the distribution every question type starts from. +pub fn distribution( + heads: &Heads, + state: &[f32], + candidates: &[Vec], + temperature: f32, +) -> Result> { + ensure!( + !candidates.is_empty(), + "a question needs at least one candidate" + ); + ensure!( + temperature > 0.0 && temperature <= 100.0, + "temperature must be in (0, 100], got {temperature}" + ); + let cfg = &heads.config.head; + let zs = project(&heads.state, cfg, state)?; + let mut logits = Vec::with_capacity(candidates.len()); + for candidate in candidates { + let zc = project(&heads.action, cfg, candidate)?; + let cos: f32 = zs.iter().zip(&zc).map(|(a, b)| a * b).sum(); + logits.push(heads.config.scale() * cos / temperature); + } + let max = logits.iter().copied().fold(f32::MIN, f32::max); + let exp: Vec = logits.iter().map(|v| (v - max).exp()).collect(); + let sum: f32 = exp.iter().sum(); + ensure!(sum > 0.0, "softmax denominator is zero"); + Ok(exp.iter().map(|v| v / sum).collect()) +} + +/// `schema.confidence`: the top probability minus the mean of the rest, clamped to [0, 1]. +/// A single candidate is fully decided by definition. +pub fn confidence(probs: &[f32]) -> f32 { + if probs.len() < 2 { + return 1.0; + } + let (j, top) = + probs.iter().enumerate().fold( + (0usize, f32::MIN), + |best, (i, p)| { + if *p > best.1 { (i, *p) } else { best } + }, + ); + let rest = (probs.len() - 1) as f32; + let mean_rest = probs + .iter() + .enumerate() + .filter(|(i, _)| *i != j) + .map(|(_, p)| *p) + .sum::() + / rest; + (top - mean_rest).clamp(0.0, 1.0) +} + +/// Assemble the answer for one question from its distribution. +pub fn answer(question: &Question, probs: &[f32]) -> Result { + ensure!( + question.keys.len() == probs.len(), + "{}: {} keys but {} probabilities", + question.id, + question.keys.len(), + probs.len() + ); + let pairs: Vec<(String, f32)> = question + .keys + .iter() + .cloned() + .zip(probs.iter().copied()) + .collect(); + Ok(match question.kind { + Kind::Noul => { + // `answer_from_probs` reads the `true` entry, which is the last key by + // convention; a caller that orders them differently still gets a stable + // answer because the key is looked up by name. + let p = pairs + .iter() + .find(|(k, _)| k == "true") + .map(|(_, p)| *p) + .unwrap_or_else(|| *probs.last().expect("checked non-empty")); + Answer::Noul { noul: p } + } + Kind::Choice => { + let (choice, _) = pairs + .iter() + .fold(None, |best: Option<&(String, f32)>, kv| match best { + Some(b) if b.1 >= kv.1 => Some(b), + _ => Some(kv), + }) + .expect("checked non-empty"); + Answer::Choice { + choice: choice.clone(), + confidence: confidence(probs), + probabilities: pairs, + } + } + Kind::Score => { + // The expected level index, `sum(i * p_i)`. + let score = probs + .iter() + .enumerate() + .map(|(i, p)| i as f32 * p) + .sum::(); + Answer::Score { + score, + confidence: confidence(probs), + probabilities: pairs, + } + } + }) +} diff --git a/src/models/clm/src/weights.rs b/src/models/clm/src/weights.rs new file mode 100644 index 00000000..9a42da3c --- /dev/null +++ b/src/models/clm/src/weights.rs @@ -0,0 +1,174 @@ +//! The CLM head tensors: inventory, loading and the expected shapes. +use anyhow::{Context, Result, ensure}; +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; +use std::{collections::BTreeSet, fs::File, path::Path}; + +use crate::config::{Config, HeadConfig}; + +/// The two projection heads. +pub const HEADS: [&str; 2] = ["state_head", "action_head"]; + +/// One head's parameters, all FP32 and row-major. +#[derive(Debug, Clone)] +pub struct Head { + pub inp_weight: Vec, + pub inp_bias: Vec, + pub hidden_weight: Vec, + pub hidden_bias: Vec, + /// Absent when the configuration has no hidden blocks; LayerNorm weight when present. + pub norm_weight: Option>, + pub norm_bias: Option>, + pub out_weight: Vec, + pub out_bias: Vec, +} + +pub struct Heads { + pub state: Head, + pub action: Head, + pub config: Config, +} + +pub struct Weights { + data: Mmap, +} + +impl Weights { + /// The checkpoint must remain immutable while the mapping exists. + pub fn open(path: &Path) -> Result { + let file = File::open(path).with_context(|| format!("open {}", path.display()))?; + // SAFETY: model files are read-only inputs; no mutable mapping is created. + let data = unsafe { Mmap::map(&file)? }; + SafeTensors::deserialize(&data)?; + Ok(Self { data }) + } + + fn tensors(&self) -> Result> { + Ok(SafeTensors::deserialize(&self.data)?) + } + + /// The safetensors metadata, which carries the head configuration. + /// + /// `SafeTensors` keeps its metadata private, so this reads the header directly from + /// the mapping. No tensor data is touched. + pub fn metadata(&self) -> Result>> { + let (_, header) = + SafeTensors::read_metadata(&self.data).context("read the safetensors header")?; + Ok(header.metadata().clone()) + } + + /// Reject omitted, extra or duplicate names in the expected inventory. + pub fn validate_names<'a>(&self, names: impl IntoIterator) -> Result<()> { + let tensors = self.tensors()?; + let mut expected = BTreeSet::new(); + for name in names { + ensure!(expected.insert(name), "duplicate expected tensor: {name}"); + } + let actual: BTreeSet<&str> = tensors.names().into_iter().collect(); + if expected != actual { + let joined = |set: BTreeSet<&str>| set.into_iter().collect::>().join(", "); + let missing = joined(expected.difference(&actual).copied().collect()); + let extra = joined(actual.difference(&expected).copied().collect()); + anyhow::bail!( + "tensor inventory does not match checkpoint: missing [{missing}], unexpected [{extra}]" + ); + } + Ok(()) + } + + /// One tensor as FP32, checking the shape first so a mismatch is an error rather + /// than a reinterpreted buffer. + pub fn f32(&self, name: &str, shape: &[usize]) -> Result> { + let tensors = self.tensors()?; + let t = tensors + .tensor(name) + .with_context(|| format!("read tensor {name}"))?; + ensure!( + t.shape() == shape, + "{name}: expected {shape:?}, got {:?}", + t.shape() + ); + let out = match t.dtype() { + Dtype::F32 => t + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect(), + dt => anyhow::bail!("{name}: unsupported dtype {dt:?}, the export writes FP32"), + }; + Ok(out) + } +} + +/// Names and shapes of every tensor the heads are built from. +pub fn head_tensors(cfg: &HeadConfig) -> Result)>> { + let blocks = cfg.hidden_blocks()?; + let (h, w, p) = (cfg.hidden_size, cfg.width, cfg.projection_dim); + let mut out = Vec::new(); + for head in HEADS { + out.push((format!("{head}.inp.weight"), vec![w, h])); + out.push((format!("{head}.inp.bias"), vec![w])); + for i in 0..blocks { + out.push((format!("{head}.hidden.{i}.weight"), vec![w, w])); + out.push((format!("{head}.hidden.{i}.bias"), vec![w])); + if cfg.layernorm { + out.push((format!("{head}.norms.{i}.weight"), vec![w])); + out.push((format!("{head}.norms.{i}.bias"), vec![w])); + } + } + out.push((format!("{head}.out.weight"), vec![p, w])); + out.push((format!("{head}.out.bias"), vec![p])); + } + Ok(out) +} + +fn load_head(weights: &Weights, name: &str, cfg: &HeadConfig) -> Result { + let (h, w, p) = (cfg.hidden_size, cfg.width, cfg.projection_dim); + let blocks = cfg.hidden_blocks()?; + ensure!( + blocks <= 1, + "the published checkpoint has {blocks} hidden blocks; the loader implements one" + ); + let (hidden_weight, hidden_bias, norm_weight, norm_bias) = if blocks == 1 { + ( + weights.f32(&format!("{name}.hidden.0.weight"), &[w, w])?, + weights.f32(&format!("{name}.hidden.0.bias"), &[w])?, + Some(weights.f32(&format!("{name}.norms.0.weight"), &[w])?), + Some(weights.f32(&format!("{name}.norms.0.bias"), &[w])?), + ) + } else { + (Vec::new(), Vec::new(), None, None) + }; + Ok(Head { + inp_weight: weights.f32(&format!("{name}.inp.weight"), &[w, h])?, + inp_bias: weights.f32(&format!("{name}.inp.bias"), &[w])?, + hidden_weight, + hidden_bias, + norm_weight, + norm_bias, + out_weight: weights.f32(&format!("{name}.out.weight"), &[p, w])?, + out_bias: weights.f32(&format!("{name}.out.bias"), &[p])?, + }) +} + +impl Heads { + /// Read the config from the metadata, check the inventory, and load every tensor. + pub fn load(weights: &Weights) -> Result { + let config = Config::from_metadata(weights.metadata()?.as_ref())?; + let expected = head_tensors(&config.head)?; + weights.validate_names(expected.iter().map(|(n, _)| n.as_str()))?; + let mut loaded = Vec::new(); + for head in HEADS { + loaded.push(load_head(weights, head, &config.head)?); + } + let action = loaded.pop().expect("two heads"); + let state = loaded.pop().expect("two heads"); + Ok(Self { + state, + action, + config, + }) + } +} diff --git a/src/models/clm/tests/checkpoint.rs b/src/models/clm/tests/checkpoint.rs new file mode 100644 index 00000000..94e67e23 --- /dev/null +++ b/src/models/clm/tests/checkpoint.rs @@ -0,0 +1,149 @@ +//! CPU checks for the CLM head loader. No GPU and no encoder, but the frozen export. +use omni_clm::{Kind, Question, Weights, answer, confidence, distribution, weights::Heads}; +use std::path::PathBuf; + +fn export_dir() -> PathBuf { + PathBuf::from(std::env::var_os("CLM_EXPORT").expect("set CLM_EXPORT to the export directory")) +} + +fn load() -> Heads { + let weights = Weights::open(&export_dir().join("model.safetensors")).unwrap(); + Heads::load(&weights).unwrap() +} + +#[test] +#[ignore = "requires CLM_EXPORT at a converted checkpoint; CPU only"] +fn every_tensor_conversion_matches_the_oracle() { + use sha2::{Digest, Sha256}; + + let dir = export_dir(); + let oracle: Vec = + serde_json::from_slice(&std::fs::read(dir.join("oracle.json")).unwrap()).unwrap(); + let weights = Weights::open(&dir.join("model.safetensors")).unwrap(); + + assert_eq!(oracle.len(), 16, "oracle must cover both heads"); + for row in oracle { + let name = row["name"].as_str().unwrap(); + let shape: Vec = serde_json::from_value(row["shape"].clone()).unwrap(); + let values = weights.f32(name, &shape).unwrap(); + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + format!("{:x}", Sha256::digest(&bytes)), + row["f32"].as_str().unwrap(), + "{name} f32" + ); + } +} + +#[test] +#[ignore = "requires CLM_EXPORT at a converted checkpoint; CPU only"] +fn decisions_match_the_reference_implementation() { + let heads = load(); + let oracle: serde_json::Value = + serde_json::from_slice(&std::fs::read(export_dir().join("head-oracle.json")).unwrap()) + .unwrap(); + + for case in oracle["cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let keys: Vec = serde_json::from_value(case["keys"].clone()).unwrap(); + let temperature = case["temperature"].as_f64().unwrap() as f32; + let expected: Vec = serde_json::from_value(case["probabilities"].clone()).unwrap(); + + let state = embedding(&format!("state::{name}"), heads.config.head.hidden_size); + let candidates: Vec> = keys + .iter() + .map(|k| embedding(&format!("cand::{name}::{k}"), heads.config.head.hidden_size)) + .collect(); + + let probs = distribution(&heads, &state, &candidates, temperature).unwrap(); + assert_eq!(probs.len(), expected.len(), "{name}: length"); + for (got, want) in probs.iter().zip(&expected) { + assert!( + (got - want).abs() < 1e-4, + "{name}: probability {got} vs {want}" + ); + } + + let kind = match case["kind"].as_str().unwrap() { + "choice" => Kind::Choice, + "noul" => Kind::Noul, + _ => Kind::Score, + }; + let question = Question { + id: name.to_string(), + kind, + keys, + }; + let answer = answer(&question, &probs).unwrap(); + match (&answer, case.get("choice")) { + ( + omni_clm::Answer::Choice { + choice, + confidence: c, + .. + }, + Some(want), + ) => { + assert_eq!(choice, want.as_str().unwrap(), "{name}: choice"); + let want_c = case["confidence"].as_f64().unwrap() as f32; + assert!( + (c - want_c).abs() < 1e-4, + "{name}: confidence {c} vs {want_c}" + ); + } + ( + omni_clm::Answer::Score { + score, + confidence: c, + .. + }, + None, + ) => { + let want_s = case["score"].as_f64().unwrap() as f32; + assert!( + (score - want_s).abs() < 1e-4, + "{name}: score {score} vs {want_s}" + ); + let want_c = case["confidence"].as_f64().unwrap() as f32; + assert!( + (c - want_c).abs() < 1e-4, + "{name}: confidence {c} vs {want_c}" + ); + } + (omni_clm::Answer::Noul { noul }, None) => { + let want = case["noul"].as_f64().unwrap() as f32; + assert!((noul - want).abs() < 1e-4, "{name}: noul {noul} vs {want}"); + } + (other, _) => panic!("{name}: unexpected answer {other:?}"), + } + } +} + +/// The same synthesised embedding the oracle uses, so both sides see identical vectors. +fn embedding(text: &str, dim: usize) -> Vec { + use sha2::{Digest, Sha256}; + + let mut out: Vec = Vec::with_capacity(dim); + let mut counter = 0u32; + while out.len() < dim { + let digest = Sha256::digest(format!("{counter}:{text}").as_bytes()); + for chunk in digest.as_chunks::<4>().0 { + if out.len() == dim { + break; + } + out.push(u32::from_be_bytes(*chunk) as f64 as f32 / 2f64.powi(31) as f32 - 1.0); + } + counter += 1; + } + let norm = out.iter().map(|v| v * v).sum::().sqrt(); + out.iter().map(|v| v / norm).collect() +} + +#[test] +fn confidence_matches_the_reference_definition() { + // Top minus the mean of the rest, clamped; a single candidate is fully decided. + assert_eq!(confidence(&[1.0]), 1.0); + assert!((confidence(&[0.75, 0.25]) - 0.5).abs() < 1e-6); + assert!((confidence(&[0.5, 0.3, 0.2]) - 0.25).abs() < 1e-6); + assert_eq!(confidence(&[0.4, 0.4, 0.4]), 0.0); +} From e1644f50b583d683493f36076e21b6834d5173e4 Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz Date: Wed, 30 Sep 2026 08:48:45 +0800 Subject: [PATCH 2/4] [CLM] Load a checkpoint that declares no LayerNorm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `head_tensors` lists `norms.N.*` only when `cfg.layernorm` is set, so a `layernorm: false` checkpoint is complete without them — but `load_head` read them unconditionally. A valid checkpoint passed the inventory check and then failed to load. The new test writes both shapes plus one that declares LayerNorm and omits the tensors, and checks the loader accepts the first two and rejects the third. Found in review. --- src/models/clm/src/weights.rs | 17 ++++-- src/models/clm/tests/checkpoint.rs | 86 +++++++++++++++++++++++++++++- 2 files changed, 99 insertions(+), 4 deletions(-) diff --git a/src/models/clm/src/weights.rs b/src/models/clm/src/weights.rs index 9a42da3c..59160167 100644 --- a/src/models/clm/src/weights.rs +++ b/src/models/clm/src/weights.rs @@ -16,7 +16,8 @@ pub struct Head { pub inp_bias: Vec, pub hidden_weight: Vec, pub hidden_bias: Vec, - /// Absent when the configuration has no hidden blocks; LayerNorm weight when present. + /// Absent when the configuration has no hidden blocks or sets `layernorm: false`; + /// the LayerNorm weight when present. pub norm_weight: Option>, pub norm_bias: Option>, pub out_weight: Vec, @@ -132,11 +133,21 @@ fn load_head(weights: &Weights, name: &str, cfg: &HeadConfig) -> Result { "the published checkpoint has {blocks} hidden blocks; the loader implements one" ); let (hidden_weight, hidden_bias, norm_weight, norm_bias) = if blocks == 1 { + // `head_tensors` only lists `norms.N.*` when `cfg.layernorm` is set, so a + // `layernorm: false` checkpoint is complete without them and loads without them. + let (norm_weight, norm_bias) = if cfg.layernorm { + ( + Some(weights.f32(&format!("{name}.norms.0.weight"), &[w])?), + Some(weights.f32(&format!("{name}.norms.0.bias"), &[w])?), + ) + } else { + (None, None) + }; ( weights.f32(&format!("{name}.hidden.0.weight"), &[w, w])?, weights.f32(&format!("{name}.hidden.0.bias"), &[w])?, - Some(weights.f32(&format!("{name}.norms.0.weight"), &[w])?), - Some(weights.f32(&format!("{name}.norms.0.bias"), &[w])?), + norm_weight, + norm_bias, ) } else { (Vec::new(), Vec::new(), None, None) diff --git a/src/models/clm/tests/checkpoint.rs b/src/models/clm/tests/checkpoint.rs index 94e67e23..370d2bb6 100644 --- a/src/models/clm/tests/checkpoint.rs +++ b/src/models/clm/tests/checkpoint.rs @@ -1,5 +1,7 @@ //! CPU checks for the CLM head loader. No GPU and no encoder, but the frozen export. -use omni_clm::{Kind, Question, Weights, answer, confidence, distribution, weights::Heads}; +use omni_clm::{ + Kind, Question, Weights, answer, confidence, distribution, head_tensors, weights::Heads, +}; use std::path::PathBuf; fn export_dir() -> PathBuf { @@ -147,3 +149,85 @@ fn confidence_matches_the_reference_definition() { assert!((confidence(&[0.5, 0.3, 0.2]) - 0.25).abs() < 1e-6); assert_eq!(confidence(&[0.4, 0.4, 0.4]), 0.0); } + +/// `layernorm: false` is a configuration `head_tensors` honours, so the loader must too: +/// a checkpoint that declares no LayerNorm is complete without `norms.*`, and one that +/// declares it is still rejected when the tensors are absent. +#[test] +fn the_loader_follows_the_layernorm_flag() { + let dir = std::env::temp_dir().join("omni-clm-synthetic-heads"); + std::fs::create_dir_all(&dir).unwrap(); + + let plain = synthetic(&dir.join("no-layernorm.safetensors"), false, false); + let heads = Heads::load(&Weights::open(&plain).unwrap()).unwrap(); + assert!(heads.state.norm_weight.is_none()); + assert!(heads.action.norm_weight.is_none()); + assert_eq!(head_tensors(&heads.config.head).unwrap().len(), 12); + + let normed = synthetic(&dir.join("layernorm.safetensors"), true, true); + let heads = Heads::load(&Weights::open(&normed).unwrap()).unwrap(); + assert!(heads.state.norm_weight.is_some()); + assert_eq!(head_tensors(&heads.config.head).unwrap().len(), 16); + + let claiming = synthetic(&dir.join("claims-layernorm.safetensors"), true, false); + let Err(err) = Heads::load(&Weights::open(&claiming).unwrap()) else { + panic!("a checkpoint that declares layernorm but omits norms.* was accepted"); + }; + assert!(err.to_string().contains("norms.0.weight"), "{err}"); +} + +/// A two-head checkpoint small enough to write here, with the LayerNorm tensors present +/// or absent independently of what the configuration declares. +fn synthetic(path: &std::path::Path, layernorm: bool, include_norms: bool) -> PathBuf { + use safetensors::Dtype; + use safetensors::tensor::{TensorView, serialize_to_file}; + use std::collections::HashMap; + + let metadata: HashMap = HashMap::from([ + ("format".to_string(), "clm-heads".to_string()), + ( + "cfg".to_string(), + format!( + r#"{{"hidden_size":4,"projection_dim":2,"width":3,"depth":3,"activation":"gelu","layernorm":{layernorm},"residual":false}}"# + ), + ), + ("hidden_size".to_string(), "4".to_string()), + ("projection_dim".to_string(), "2".to_string()), + ("logit_scale".to_string(), "1.0".to_string()), + ]); + + let (mut names, mut shapes, mut buffers) = (Vec::new(), Vec::new(), Vec::new()); + for head in ["state_head", "action_head"] { + let mut spec = vec![ + (format!("{head}.inp.weight"), vec![3, 4]), + (format!("{head}.inp.bias"), vec![3]), + (format!("{head}.hidden.0.weight"), vec![3, 3]), + (format!("{head}.hidden.0.bias"), vec![3]), + (format!("{head}.out.weight"), vec![2, 3]), + (format!("{head}.out.bias"), vec![2]), + ]; + if include_norms { + spec.push((format!("{head}.norms.0.weight"), vec![3])); + spec.push((format!("{head}.norms.0.bias"), vec![3])); + } + for (name, shape) in spec { + buffers.push(vec![0u8; shape.iter().product::() * 4]); + shapes.push(shape); + names.push(name); + } + } + let tensors: Vec<(String, TensorView)> = names + .iter() + .zip(&shapes) + .zip(&buffers) + .map(|((name, shape), buffer)| { + ( + name.clone(), + TensorView::new(Dtype::F32, shape.clone(), buffer).unwrap(), + ) + }) + .collect(); + + serialize_to_file(tensors, Some(metadata), path).unwrap(); + path.to_path_buf() +} From 2a2c489ad52ef2e47e762282d15494c85515d87e Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz Date: Sat, 3 Oct 2026 19:59:50 +0800 Subject: [PATCH 3/4] [CLM] Answer review: a reachable geometry fallback and first-maximum labels - `Config::from_metadata` applied the top-level `hidden_size`/`projection_dim` fallback *after* deserializing `cfg`, but both are required fields, so a `cfg` that omits them failed to parse and the fallback never ran. The geometry is now filled in before deserialization, and a `cfg` that omits it loads while one that omits it with no top-level copy is still an error. - `Answer::label` used `max_by`, which keeps the *last* of several equal values. `schema.label_of` is `max(p, key=p.__getitem__)`, which keeps the first, so a tie in a score distribution named a different level. Found in review. --- src/models/clm/src/config.rs | 26 +++++++++-------- src/models/clm/src/scoring.rs | 8 +++++- src/models/clm/tests/checkpoint.rs | 45 +++++++++++++++++++++++++++++- 3 files changed, 66 insertions(+), 13 deletions(-) diff --git a/src/models/clm/src/config.rs b/src/models/clm/src/config.rs index 77e7eac7..ded0580b 100644 --- a/src/models/clm/src/config.rs +++ b/src/models/clm/src/config.rs @@ -55,18 +55,22 @@ impl Config { metadata.get("format") ); let cfg = metadata.get("cfg").context("metadata has no cfg")?; - let mut head: HeadConfig = + // `hidden_size` and `projection_dim` are repeated at the top level, and `cfg` may + // omit them: the exporter takes them from the checkpoint's own entries. They are + // required fields of `HeadConfig`, so they have to be filled in *before* + // deserializing — a fallback applied afterwards never runs, because the parse + // fails on the missing field first. + let mut raw: serde_json::Map = serde_json::from_str(cfg).with_context(|| format!("parse cfg {cfg}"))?; - // hidden_size and projection_dim are repeated at the top level; prefer those - // when the cfg omits them so an older export still loads. - head.hidden_size = metadata - .get("hidden_size") - .and_then(|v| v.parse().ok()) - .unwrap_or(head.hidden_size); - head.projection_dim = metadata - .get("projection_dim") - .and_then(|v| v.parse().ok()) - .unwrap_or(head.projection_dim); + for key in ["hidden_size", "projection_dim"] { + if !raw.contains_key(key) + && let Some(value) = metadata.get(key).and_then(|v| v.parse::().ok()) + { + raw.insert(key.to_string(), value.into()); + } + } + let head: HeadConfig = serde_json::from_value(serde_json::Value::Object(raw)) + .with_context(|| format!("parse cfg {cfg}"))?; let logit_scale: f32 = metadata .get("logit_scale") .context("metadata has no logit_scale")? diff --git a/src/models/clm/src/scoring.rs b/src/models/clm/src/scoring.rs index 02f7198e..935c38ac 100644 --- a/src/models/clm/src/scoring.rs +++ b/src/models/clm/src/scoring.rs @@ -51,9 +51,15 @@ impl Answer { match self { Answer::Choice { choice, .. } => choice.clone(), Answer::Noul { noul } => if *noul >= 0.5 { "true" } else { "false" }.to_string(), + // The first maximum, as Python's `max` keeps. `max_by` would return the + // last of several equal values, which is a different label for the same + // distribution. Answer::Score { probabilities, .. } => probabilities .iter() - .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) + .fold(None, |best: Option<&(String, f32)>, kv| match best { + Some(b) if b.1 >= kv.1 => Some(b), + _ => Some(kv), + }) .map(|(k, _)| k.clone()) .unwrap_or_default(), } diff --git a/src/models/clm/tests/checkpoint.rs b/src/models/clm/tests/checkpoint.rs index 370d2bb6..25f8ac5f 100644 --- a/src/models/clm/tests/checkpoint.rs +++ b/src/models/clm/tests/checkpoint.rs @@ -1,6 +1,6 @@ //! CPU checks for the CLM head loader. No GPU and no encoder, but the frozen export. use omni_clm::{ - Kind, Question, Weights, answer, confidence, distribution, head_tensors, weights::Heads, + Config, Kind, Question, Weights, answer, confidence, distribution, head_tensors, weights::Heads, }; use std::path::PathBuf; @@ -231,3 +231,46 @@ fn synthetic(path: &std::path::Path, layernorm: bool, include_norms: bool) -> Pa serialize_to_file(tensors, Some(metadata), path).unwrap(); path.to_path_buf() } + +/// The exporter writes the geometry at the top level and may leave it out of `cfg`, so a +/// `cfg` without it has to be completed before deserialization rather than after. +#[test] +fn the_top_level_geometry_fills_in_a_cfg_that_omits_it() { + use std::collections::HashMap; + + let geom = r#""width":3,"depth":3,"activation":"gelu","layernorm":false,"residual":false"#; + let meta = |cfg: String| -> HashMap { + HashMap::from([ + ("format".to_string(), "clm-heads".to_string()), + ("cfg".to_string(), cfg), + ("hidden_size".to_string(), "4".to_string()), + ("projection_dim".to_string(), "2".to_string()), + ("logit_scale".to_string(), "1.0".to_string()), + ]) + }; + + let omitted = Config::from_metadata(Some(&meta(format!("{{{geom}}}")))).unwrap(); + assert_eq!(omitted.head.hidden_size, 4); + assert_eq!(omitted.head.projection_dim, 2); + + let both = meta(format!(r#"{{"hidden_size":4,"projection_dim":2,{geom}}}"#)); + assert_eq!(Config::from_metadata(Some(&both)).unwrap().head.width, 3); + + let mut neither = meta(format!("{{{geom}}}")); + neither.remove("hidden_size"); + neither.remove("projection_dim"); + assert!(Config::from_metadata(Some(&neither)).is_err()); +} + +/// `schema.label_of` is `max(p, key=p.__getitem__)`, which keeps the first of several +/// equal values. +#[test] +fn a_tied_score_label_keeps_the_first_level() { + let question = Question { + id: "tie".into(), + kind: Kind::Score, + keys: vec!["0".into(), "1".into(), "2".into()], + }; + let tied = answer(&question, &[0.5, 0.5, 0.0]).unwrap(); + assert_eq!(tied.label(), "0"); +} From 9650e2abc9a28fe3d0cb16aa67c63a58fa6ceef6 Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz Date: Mon, 5 Oct 2026 14:35:33 +0800 Subject: [PATCH 4/4] [CLM] Move the tests under the repository-level tests/ tree CONTRIBUTING puts every test body in the repository-level `tests/` tree and asks for explicit `[[test]]` registration, which is what `src/models/laya` does. `src/models/clm/tests/` was a crate-local directory under `src/`, which that rule prohibits; the tests ran either way, so this is placement rather than breakage. Found in review. --- src/models/clm/Cargo.toml | 6 ++++++ {src/models/clm/tests => tests/clm}/checkpoint.rs | 0 2 files changed, 6 insertions(+) rename {src/models/clm/tests => tests/clm}/checkpoint.rs (100%) diff --git a/src/models/clm/Cargo.toml b/src/models/clm/Cargo.toml index 6fa26413..9a6cd308 100644 --- a/src/models/clm/Cargo.toml +++ b/src/models/clm/Cargo.toml @@ -15,3 +15,9 @@ serde = { version = "1", features = ["derive"] } # trained on the caller's field order, so re-sorting a state would change the text. serde_json = { version = "1", features = ["preserve_order"] } sha2 = "0.10" + +# Test bodies live under the repository-level tests/ tree; this is the explicit +# registration CONTRIBUTING asks for, matching src/models/laya/Cargo.toml. +[[test]] +name = "checkpoint" +path = "../../../tests/clm/checkpoint.rs" diff --git a/src/models/clm/tests/checkpoint.rs b/tests/clm/checkpoint.rs similarity index 100% rename from src/models/clm/tests/checkpoint.rs rename to tests/clm/checkpoint.rs