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..9a6cd308 --- /dev/null +++ b/src/models/clm/Cargo.toml @@ -0,0 +1,23 @@ +[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" + +# 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/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..ded0580b --- /dev/null +++ b/src/models/clm/src/config.rs @@ -0,0 +1,109 @@ +//! 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")?; + // `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}"))?; + 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")? + .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..935c38ac --- /dev/null +++ b/src/models/clm/src/scoring.rs @@ -0,0 +1,309 @@ +//! 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(), + // 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() + .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(), + } + } +} + +/// 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..59160167 --- /dev/null +++ b/src/models/clm/src/weights.rs @@ -0,0 +1,185 @@ +//! 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 or sets `layernorm: false`; + /// the 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 { + // `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])?, + norm_weight, + norm_bias, + ) + } 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/tests/clm/checkpoint.rs b/tests/clm/checkpoint.rs new file mode 100644 index 00000000..25f8ac5f --- /dev/null +++ b/tests/clm/checkpoint.rs @@ -0,0 +1,276 @@ +//! CPU checks for the CLM head loader. No GPU and no encoder, but the frozen export. +use omni_clm::{ + Config, Kind, Question, Weights, answer, confidence, distribution, head_tensors, 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); +} + +/// `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() +} + +/// 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"); +}