[CLM 2/2] Reach the encoder, answer requests, and check against CLM

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"

recipe/clm/native/compare_with_reference.py

@@ -0,0 +1,153 @@
+"""Compare omni-clm's decision with the CLM reference, over the same embeddings.
+
+Both sides call the same OpenAI-compatible `/v1/embeddings` endpoint, so the vectors are
+identical and the only difference left is the code between the vectors and the answer.
+
+    python recipe/clm/native/compare_with_reference.py \
+        --checkpoint /path/to/clm-export \
+        --bin /path/to/clm-run \
+        --emb-url http://127.0.0.1:8090/v1/embeddings \
+        --pt /path/to/CLM_v0.1-8B.pt
+
+## On the tolerance
+
+CLM scores with `exp(logit_scale) = 100`, so a difference of `d` in a cosine similarity
+becomes `100 d` in the logit. Two encoder forward passes that disagree by ~1e-4 in the
+cosine therefore disagree by ~1e-2 in the probabilities, and that is what two different
+implementations of Qwen3-8B produce -- vLLM's bf16 pooling against Transformers' forward.
+It is a property of how this model scores, not a defect on either side.
+
+The default tolerance reflects it. The engine's own arithmetic is checked separately and
+much more tightly: with the vectors held fixed the two implementations agree to 3e-06.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+CASES = [
+    ("choice_two", "choice", {"billing": "Charges and refunds", "technical": "Software problems"}),
+    ("choice_five", "choice", {"a": "alpha", "b": "beta", "c": "gamma", "d": "delta", "e": "epsilon"}),
+    ("score_three", "score", ["Not urgent", "Needs attention soon", "Needs attention immediately"]),
+    ("noul_stmt", "noul", None),
+    # A `noul` reads its two descriptions in `false`/`true` order, not in the order they
+    # were written, so a number here has to be found by the key it was written under.
+    # Counting positions gave `false: 2` / `true: 18446744073709551616` reversed, which
+    # is a different pair of candidate texts and so a different answer.
+    ("noul_numbers", "noul", {"true": 18446744073709551616, "false": 2}),
+]
+
+STATE = "I was charged twice for order 4411 and want the second charge refunded."
+
+# exp(logit_scale) with the published checkpoint, and what it does to a cosine difference.
+SCALE = 100.0
+TOLERANCE = 2.5e-2
+
+
+def request_for(name: str, kind: str, criteria) -> dict:
+    q = {"type": kind, "instructions": f"Question {name}"}
+    if criteria is not None:
+        q["criteria"] = criteria
+    return {"model": "clm-latest", "state": STATE, "questions": {name: q}}
+
+
+def reference(emb_url: str, emb_model: str, pt: Path, req: dict, temperature: float) -> dict:
+    """CLM's own heads and schema, applied to the same embeddings the Rust side gets."""
+    from clm.embedder import Embedder
+    from clm.engine import Engine
+
+    embedder = Embedder(url=emb_url, model=emb_model)
+    engine = Engine(embedder=embedder, checkpoint=str(pt))
+    return engine.answer(req["state"], req["questions"], "clm-latest", temperature)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--checkpoint", type=Path, required=True, help="converted export dir")
+    parser.add_argument("--bin", type=Path, required=True, help="clm-run")
+    parser.add_argument("--emb-url", required=True)
+    parser.add_argument("--emb-model", default="qwen3-8b")
+    parser.add_argument("--pt", type=Path, required=True, help="CLM_v0.1-8B.pt")
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument(
+        "--tolerance",
+        type=float,
+        default=TOLERANCE,
+        help=f"logits are scale*cos with scale={SCALE:g}, so two encoder implementations "
+        f"differ by about this much in the probabilities",
+    )
+    args = parser.parse_args()
+
+    requests_ = [(n, request_for(n, k, c)) for n, k, c in CASES]
+    payload = "".join(json.dumps(r) + "\n" for _, r in requests_)
+    proc = subprocess.run(
+        [
+            str(args.bin),
+            str(args.checkpoint),
+            "--emb-url",
+            args.emb_url,
+            "--model",
+            args.emb_model,
+            # Both sides must run the same temperature, or a non-default one compares two
+            # different configurations and reports a port regression that is not there.
+            "--temperature",
+            str(args.temperature),
+        ],
+        input=payload,
+        capture_output=True,
+        text=True,
+        timeout=600,
+    )
+    if proc.returncode != 0:
+        print(proc.stderr[-2000:], file=sys.stderr)
+        raise SystemExit(f"clm-run exited {proc.returncode}")
+    mine = [json.loads(line) for line in proc.stdout.splitlines() if line.strip()]
+    if len(mine) != len(requests_):
+        raise SystemExit(f"clm-run answered {len(mine)} of {len(requests_)}")
+
+    failures = 0
+    for (name, req), got in zip(requests_, mine):
+        want = reference(args.emb_url, args.emb_model, args.pt, req, args.temperature)
+        mine_p = got["answers"][name]
+        want_p = want["answers"][name]
+        kind = req["questions"][name]["type"]
+
+        if kind == "noul":
+            got_v, want_v = mine_p["noul"], want_p["noul"]
+            ok = abs(got_v - want_v) <= args.tolerance
+            label, detail = "noul", f"{got_v:.6f} vs {want_v:.6f}"
+        else:
+            gk, wk = mine_p["probabilities"], want_p["probabilities"]
+            if set(gk) != set(wk):
+                print(f"FAIL {name}: keys differ {sorted(gk)} vs {sorted(wk)}")
+                failures += 1
+                continue
+            worst = max(abs(gk[k] - wk[k]) for k in gk)
+            same_pick = (
+                mine_p.get("choice") == want_p.get("choice")
+                if kind == "choice"
+                else abs(mine_p["score"] - want_p["score"]) <= args.tolerance
+            )
+            ok = worst <= args.tolerance and same_pick
+            label = "choice" if kind == "choice" else "score"
+            detail = f"max|dp|={worst:.2e} " + (
+                f"choice={mine_p['choice']}/{want_p['choice']}"
+                if kind == "choice"
+                else f"score={mine_p['score']:.6f}/{want_p['score']:.6f}"
+            )
+        failures += not ok
+        print(f"{'PASS' if ok else 'FAIL'} {name:12} {label:7} {detail}")
+
+    print(
+        f"\n{'all cases agree' if not failures else str(failures) + ' FAILED'} "
+        f"(tolerance {args.tolerance:g}; the engine alone agrees to 3e-06 when the vectors are fixed)"
+    )
+    raise SystemExit(1 if failures else 0)
+
+
+if __name__ == "__main__":
+    main()

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()