diff --git a/CLAUDE.md b/CLAUDE.md index 5fe78e9..c7d44a6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -17,8 +17,16 @@ the canonical entry point for reproducing the published table. ## Setup and common commands ```bash -# One-time: fetch + build vendored llama.cpp (pinned commit 9e58d4d6) -git submodule update --init --recursive +# One-time: fetch + build vendored llama.cpp. +# Pinned commit c84e85af = head of llama.cpp PR #24423 (DiffusionGemma: +# `diffusion-gemma` arch + converter + llama-diffusion-cli canvas mode + +# llama-diffusion-gemma-server). NOT on master yet — if `submodule update` +# can't resolve the gitlink, fetch the PR ref inside the submodule and retry. +# Re-pin to master once the PR merges. +git submodule update --init --recursive || ( + git -C vendor/llama.cpp fetch origin pull/24423/head && + git submodule update --recursive +) cmake -S vendor/llama.cpp -B vendor/llama.cpp/build -DGGML_METAL=ON # Linux+CUDA: -DGGML_CUDA=ON cmake --build vendor/llama.cpp/build -j @@ -123,6 +131,62 @@ Pipeline behaviors worth knowing: ### Tensor naming `models/hf_gguf_map.py` maps HF parameter names to GGUF tensor names. Anything that crosses the HF↔GGUF boundary (imatrix variants, AWQ apply) goes through this mapping. +MoE expert stacks map to fused 3D tensors (`ffn_gate_up_exps`/`ffn_down_exps`, +`[n_expert, n_out, n_in]`); `is_moe_expert` gates per-expert handling the way +`is_ssm` gates SSM passthrough. DiffusionGemma module paths (`model.decoder.*`, +`model.encoder.language_model.*`) are prefix-normalized to `model.*` before +matching — same rewrite llama.cpp's converter applies. + +### DiffusionGemma / Gemma-4 MoE (`diffusion-gemma` arch) +`google/diffusiongemma-26B-A4B-it`: 128-expert (8 active) Gemma-4 MoE whose +causal **encoder** (prompt prefill) and bidirectional diffusion **decoder** +(256-token canvas denoising) share all backbone weights — only the attention +mask differs. Requires the PR-#24423 llama.cpp pin (see Setup). Key facts the +code depends on: +- **Checkpoint shape**: arch `DiffusionGemmaForBlockDiffusion`, backbone under + `model.decoder.layers.N.*` (encoder tied via `_tied_weights_keys`), vision + tower on the encoder (the converter drops it), decoder-only + `self_conditioning` gated MLP, root-level `canvas_length`. + `extract_text_lm` passes the checkpoint through untouched (`ForBlockDiffusion` + predicate) — never rebuild a text-only CausalLM from it. +- **Each layer has BOTH a dense MLP and an MoE block**, each behind its own + pre-norm (`pre_feedforward_layernorm` / `pre_feedforward_layernorm_2`). + Experts are ONE fused 3D param per projection (`experts.gate_up_proj`, + `experts.down_proj`); the router (`router.proj`) reads the **raw residual + through its own internal norm** — not the experts' pre-norm output. +- **imatrix works** (encoder mode is a standard causal forward over the shared + weights) with two caveats: (1) llama-imatrix/llama-perplexity run the + UNIFIED graph, which treats the last `canvas_length` positions of every + chunk as a bidirectional canvas — quant-vs-quant KLD stays valid, absolute + PPL is not deployment-faithful; raise `imatrix_ctx` (recipes use 2048) to + keep the canvas fraction small. (2) MoE imatrix entries carry **per-expert + counts**; `_load_base_imatrix` normalizes per expert (zero-count experts → + neutral ones, mirroring llama-quantize), the analytic variants rerank **per + expert** (L1 normalization must not bleed across experts), and + `imatrix.warn_moe_coverage` flags expert tensors below 95% coverage — the + fix is more calibration tokens. +- **AWQ works**: `discover_groups` emits `L{i}_moe` groups (the fused gate_up + stack under `pre_feedforward_layernorm_2`, anchored on the `experts` module + which sees every token pre-routing). Folding there leaves routing logits + exactly invariant because the router never consumes that norm's output — + do NOT add the router to the group. `moe_expert_sample` (default 8) caps the + α-search cost; `rmsnorm_plus_one: true` (Gemma `(1+γ)` convention). The HF + load goes through `_hf.load_calibration_model` (falls back to the concrete + arch class when `AutoModelForCausalLM` refuses) and sanity logits through + `_hf.causal_logits` (encoder + tied lm_head). +- **GPTQ on the experts is unsupported** (deferred): 128 per-expert Hessians, + each seeing ~1/16 of tokens, is days of badly-conditioned CPU Cholesky. +- **Task eval cannot use llama-server** (no diffusion decoding). Use + `eval/diffusion_server.py`: `running_diffusion_server(model_path, hf_dir=…)` + yields a `base_url` whose `/v1/chat/completions` runs the entropy-bound + block-diffusion sampler (numpy port of llama.cpp's reference, driven through + the persistent `llama-diffusion-gemma-server` — one model load per eval) and + parses Gemma-4 `<|tool_call>call:NAME{…}` markup into OpenAI + `tool_calls`. Reps scripts take `--backend diffusion --hf-dir `. +- **Speed**: llama-bench decode tok/s is meaningless for diffusion generation; + use `scripts/bench_diffusion_speed.py` (wraps llama-diffusion-cli). +- Recipes: `q2_k_diffusiongemma_imatrix` (low-resource, no HF load), + `iq2_m_diffusiongemma_awq` (AWQ + hybrid_custom stacked). ### Bench `bench/runner.py` defines `BenchRow` and `CSV_COLUMNS`. Sub-modules: @@ -233,7 +297,9 @@ creates `model_extracted/`, `corpus/`, `calibration/`, `gguf/`, `eval/` and rese `hybrid_custom` imatrix stacked, codebook proxies auto-selected), `q2_k_gptq`, `iq3_s_gptq` (asym grid + relaxed guardrails auto-derived). - Model-specific: `q4_k_m_qwen3_5_4b`, `{q4_k_m,q5_k_s}_qwen3_6_mtp{,_awq,_none}`, - `iq3_s_9b_mtp` (MTP heads kept via `extract.keep_mtp`). + `iq3_s_9b_mtp` (MTP heads kept via `extract.keep_mtp`), + `q2_k_diffusiongemma_imatrix` + `iq2_m_diffusiongemma_awq` (DiffusionGemma + 26B-A4B MoE — see the DiffusionGemma section for the caveats). A unit test (`test_all_packaged_recipes_parse`) requires every shipped recipe to validate, and IQ1/IQ2 recipes to carry a calibration method — keep it green when adding recipes. diff --git a/scripts/bench_diffusion_speed.py b/scripts/bench_diffusion_speed.py new file mode 100644 index 0000000..35edaf1 --- /dev/null +++ b/scripts/bench_diffusion_speed.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Denoising-throughput bench for diffusion GGUFs (DiffusionGemma et al.). + +llama-bench measures autoregressive prefill/decode — decode tok/s is NOT how +a block-diffusion model generates. This script times `llama-diffusion-cli` +end-to-end instead: wall time per generated token across N repetitions +(mean ± stdev), the diffusion analog of bench/speed.py. + +Usage: + uv run python scripts/bench_diffusion_speed.py \ + --model out/run/gguf/IQ2_M-awq-hybrid_custom.gguf \ + --prompt "Explain the Fourier transform." \ + --n-predict 256 --reps 5 --out results.csv +""" + +from __future__ import annotations + +import argparse +import csv +import re +import statistics +import subprocess +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from quant_tuner.paths import llama_bin # noqa: E402 + +# entropy-bound CLI timing line: "total time: 1234.56ms, time per step: ..." +_TOTAL_RE = re.compile(r"total time:\s*([0-9.]+)ms.*\((\d+) steps over (\d+) blocks", re.S) + + +def run_once(args, rep: int) -> dict: + cmd = [ + str(llama_bin("llama-diffusion-cli")), + "-m", str(args.model), + "-p", args.prompt, + "-n", str(args.n_predict), + "-ngl", str(args.ngl), + "--seed", str(args.seed + rep), + ] + if args.diffusion_steps: + cmd += ["--diffusion-steps", str(args.diffusion_steps)] + t0 = time.time() + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=args.timeout) + wall_s = time.time() - t0 + if proc.returncode != 0: + raise RuntimeError(f"llama-diffusion-cli failed (rc={proc.returncode}):\n{proc.stderr[-2000:]}") + combined = proc.stdout + proc.stderr + row = {"rep": rep, "wall_s": round(wall_s, 3)} + m = _TOTAL_RE.search(combined) + if m: + row["gen_ms"] = float(m.group(1)) + row["steps"] = int(m.group(2)) + row["blocks"] = int(m.group(3)) + row["ms_per_step"] = round(row["gen_ms"] / max(row["steps"], 1), 2) + return row + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", type=Path, required=True) + p.add_argument("--prompt", default="Write a Python function that merges two sorted lists.") + p.add_argument("--n-predict", type=int, default=256) + p.add_argument("--reps", type=int, default=5) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--ngl", type=int, default=99) + p.add_argument("--diffusion-steps", type=int, default=0, + help="override step count (0 = model/GGUF default)") + p.add_argument("--timeout", type=int, default=1800) + p.add_argument("--out", type=Path, default=None, help="CSV output path") + args = p.parse_args() + + rows = [] + for rep in range(args.reps): + row = run_once(args, rep) + rows.append(row) + print(f"rep {rep}: {row}", flush=True) + + walls = [r["wall_s"] for r in rows] + mean = statistics.mean(walls) + stdev = statistics.stdev(walls) if len(walls) > 1 else 0.0 + tok_s = args.n_predict / mean if mean > 0 else 0.0 + print(f"\n{args.model.name}: wall {mean:.2f}s ± {stdev:.2f} " + f"(~{tok_s:.1f} tok/s incl. model load, n={args.reps})") + + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + cols = sorted({k for r in rows for k in r}) + with args.out.open("w", newline="") as f: + w = csv.DictWriter(f, fieldnames=["model", *cols]) + w.writeheader() + for r in rows: + w.writerow({"model": args.model.name, **r}) + print(f"wrote {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_mmlu_pro_reps.py b/scripts/run_mmlu_pro_reps.py index 20e4931..00fc074 100644 --- a/scripts/run_mmlu_pro_reps.py +++ b/scripts/run_mmlu_pro_reps.py @@ -58,6 +58,12 @@ def _build_arg_parser() -> argparse.ArgumentParser: p.add_argument("--ctx", type=int, default=8192) p.add_argument("--ngl", type=int, default=99) p.add_argument("--max-tokens", type=int, default=2048) + p.add_argument("--backend", choices=("server", "diffusion"), default="server", + help="generation backend: llama-server (default) or the " + "DiffusionGemma block-diffusion shim (requires --hf-dir)") + p.add_argument("--hf-dir", type=Path, default=None, + help="HF snapshot dir with the tokenizer/chat template " + "(required for --backend diffusion)") p.add_argument("--temperature", type=float, default=0.6, help="Default 0.6 — reps only diverge at T>0. Set 0.0 for greedy.") p.add_argument("--top-p", type=float, default=0.95) @@ -124,6 +130,16 @@ def on_model(mr) -> None: print(f"holdout: {args.holdout}", flush=True) print(f"results: {args.results} / {args.aggregated}", flush=True) + server_factory = None + if args.backend == "diffusion": + if args.hf_dir is None: + raise SystemExit("--backend diffusion requires --hf-dir (tokenizer source)") + from functools import partial + + from quant_tuner.eval.diffusion_server import running_diffusion_server + + server_factory = partial(running_diffusion_server, hf_dir=args.hf_dir) + t0 = time.time() by_model = run_reps_for_models( models=list(args.models), @@ -137,6 +153,7 @@ def on_model(mr) -> None: chat_template_kwargs=(args.chat_template_kwargs or None), per_rep_callback=on_rep, per_model_callback=on_model, + server_factory=server_factory, ) write_csvs( diff --git a/scripts/run_toolcall_reps.py b/scripts/run_toolcall_reps.py index 9cb7f3d..47c4c98 100644 --- a/scripts/run_toolcall_reps.py +++ b/scripts/run_toolcall_reps.py @@ -91,6 +91,12 @@ def _build_arg_parser() -> argparse.ArgumentParser: help="Override sampling.top_p") p.add_argument("--top-k", type=int, default=None, help="Override sampling.top_k") + p.add_argument("--backend", choices=("server", "diffusion"), default="server", + help="generation backend: llama-server (default) or the " + "DiffusionGemma block-diffusion shim (requires --hf-dir)") + p.add_argument("--hf-dir", type=Path, default=None, + help="HF snapshot dir with the tokenizer/chat template " + "(required for --backend diffusion)") p.add_argument("--min-p", type=float, default=None, help="Override sampling.min_p") return p @@ -158,6 +164,16 @@ def on_model(mr) -> None: print(f"results: {args.results} (per-rep) / {args.aggregated} (aggregated)", flush=True) + server_factory = None + if args.backend == "diffusion": + if args.hf_dir is None: + raise SystemExit("--backend diffusion requires --hf-dir (tokenizer source)") + from functools import partial + + from quant_tuner.eval.diffusion_server import running_diffusion_server + + server_factory = partial(running_diffusion_server, hf_dir=args.hf_dir) + t0 = time.time() by_model = run_reps_for_models( models=list(args.models), @@ -170,6 +186,7 @@ def on_model(mr) -> None: chat_template_kwargs=(args.chat_template_kwargs or None), per_rep_callback=on_rep, per_model_callback=on_model, + server_factory=server_factory, ) write_csvs( diff --git a/src/quant_tuner/calibrate/_hf.py b/src/quant_tuner/calibrate/_hf.py index 8afbc8a..0a28874 100644 --- a/src/quant_tuner/calibrate/_hf.py +++ b/src/quant_tuner/calibrate/_hf.py @@ -2,6 +2,75 @@ from __future__ import annotations +import json +from pathlib import Path +from typing import Any + +# Candidate dotted paths (relative to the top-level HF model) for the decoder +# layer stack. DiffusionGemma's causal *encoder* stack is listed rather than +# its diffusion decoder: calibration forwards run in encoder (causal) mode, +# and the two stacks share every backbone weight (tied parameters), so weight +# mutations applied through encoder paths land in the decoder too. +_TRUNK_CANDIDATES = ( + "model", + "model.encoder.language_model", + "encoder.language_model", +) + + +def _get_attr_path(obj: Any, dotted: str) -> Any | None: + for part in dotted.split("."): + obj = getattr(obj, part, None) + if obj is None: + return None + return obj + + +def resolve_text_trunk(model) -> tuple[str, Any]: + """Locate the text decoder trunk (the module owning ``.layers``). + + Returns ``(dotted_path, module)`` where ``dotted_path`` is relative to the + top-level model (e.g. ``"model"`` for standard CausalLMs, + ``"model.encoder.language_model"`` for DiffusionGemma). Raises if no + candidate matches. + """ + for path in _TRUNK_CANDIDATES: + trunk = _get_attr_path(model, path) + if trunk is not None and hasattr(trunk, "layers"): + return path, trunk + raise RuntimeError( + "no text trunk with .layers found (tried: " + + ", ".join(f"model.{p}" for p in _TRUNK_CANDIDATES) + + "); unsupported architecture" + ) + + +def load_calibration_model(model_dir: Path, *, torch_dtype, device: str): + """Load an HF checkpoint for calibration, handling non-CausalLM classes. + + ``AutoModelForCausalLM`` covers the standard architectures; checkpoints + whose class isn't registered there (e.g. ``DiffusionGemmaForBlockDiffusion``) + are loaded via their concrete ``architectures[0]`` class from + ``transformers`` instead. + """ + import transformers + from transformers import AutoModelForCausalLM + + try: + model = AutoModelForCausalLM.from_pretrained( + model_dir, torch_dtype=torch_dtype, trust_remote_code=True + ) + except (ValueError, KeyError): + config = json.loads((Path(model_dir) / "config.json").read_text()) + archs = config.get("architectures") or [] + cls = getattr(transformers, archs[0], None) if archs else None + if cls is None: + raise + model = cls.from_pretrained( + model_dir, torch_dtype=torch_dtype, trust_remote_code=True + ) + return model.to(device) + def forward_no_logits(model, input_ids) -> None: """Run a forward pass for hook-based stat collection, skipping the LM head. @@ -10,11 +79,27 @@ def forward_no_logits(model, input_ids) -> None: logits are discarded, yet ``model(ids)`` materializes a ``[1, ctx, vocab]`` float tensor (several GB at ctx 4096 on a 150k-vocab model) and pays the lm_head matmul for every chunk. Calling the decoder - trunk (``model.model``) directly avoids both. Falls back to the full - model for architectures without the standard ``.model.layers`` shape. + trunk (``model.model``, or DiffusionGemma's causal encoder stack) directly + avoids both. Falls back to the full model for architectures without a + recognizable trunk. """ - base = getattr(model, "model", None) - if base is not None and hasattr(base, "layers"): - base(input_ids=input_ids) - else: + try: + _, trunk = resolve_text_trunk(model) + except RuntimeError: model(input_ids) + return + trunk(input_ids=input_ids) + + +def causal_logits(model, input_ids): + """Causal-LM logits for sanity checks, on architectures where ``model(ids)`` + isn't a causal forward (DiffusionGemma's top-level forward wants canvas + inputs; its *encoder* + tied ``lm_head`` is the causal path).""" + trunk_path, trunk = resolve_text_trunk(model) + if trunk_path == "model": + return model(input_ids).logits + hidden = trunk(input_ids=input_ids).last_hidden_state + head = getattr(model, "lm_head", None) + if head is None: + raise RuntimeError(f"no lm_head on model with trunk {trunk_path!r}") + return head(hidden) diff --git a/src/quant_tuner/calibrate/_iq2_grids.py b/src/quant_tuner/calibrate/_iq2_grids.py index 61e1a69..44401cb 100644 --- a/src/quant_tuner/calibrate/_iq2_grids.py +++ b/src/quant_tuner/calibrate/_iq2_grids.py @@ -1,6 +1,6 @@ """E8-lattice codebooks for the IQ2_* quants — GENERATED, do not edit. -Extracted from llama.cpp ggml-common.h at commit 9e58d4d692ed3d350591cc86d06c73c61c122509 +Extracted from llama.cpp ggml-common.h at commit c84e85af61011f9fbfcf41479381d5ed1661a564 by scripts/gen_iq2_grids.py. Each grid entry is 8 magnitude bytes (component order); signs live outside the codebook. For IQ2_XXS and IQ2_XS only sign patterns with an EVEN number of negatives per group diff --git a/src/quant_tuner/calibrate/awq.py b/src/quant_tuner/calibrate/awq.py index 78257a6..a2cde7a 100644 --- a/src/quant_tuner/calibrate/awq.py +++ b/src/quant_tuner/calibrate/awq.py @@ -39,7 +39,12 @@ import torch from quant_tuner.calibrate._device import resolve_device -from quant_tuner.calibrate._hf import forward_no_logits +from quant_tuner.calibrate._hf import ( + causal_logits, + forward_no_logits, + load_calibration_model, + resolve_text_trunk, +) # --- Group discovery ------------------------------------------------------ # @@ -54,27 +59,37 @@ class ScaleGroup: leaves untouched. """ - group_id: str # e.g. "L3_attn", "L3_mlp", "L3_attn_out", "L3_mlp_out" - anchor: str # dotted module path; first member, used for the hook - members: tuple[str, ...] # dotted module paths to all sharing linears + group_id: str # e.g. "L3_attn", "L3_mlp", "L3_moe", "L3_attn_out" + anchor: str # dotted module path; used for the activation hook + members: tuple[str, ...] # dotted paths to all sharing linears (or fused + # 3D expert parameters — input channels last axis) prev_norm: str | None # dotted module path to the RMSNorm whose γ we fold, or None def discover_groups(model, *, include_output_proj: bool = False) -> list[ScaleGroup]: - """Find attention + MLP scale groups in a HuggingFace decoder model. + """Find attention + MLP (+ MoE expert) scale groups in a HuggingFace decoder model. With ``include_output_proj=True`` (exp-011), also emit per-layer ``L{i}_attn_out`` (``o_proj``) and ``L{i}_mlp_out`` (``down_proj``) groups. These have ``prev_norm=None`` since their input is not a normalized activation — scales apply to the weight only, no RMSNorm fold cancellation. + + Gemma-4 / DiffusionGemma MoE blocks emit a ``L{i}_moe`` group: the fused + 3D ``experts.gate_up_proj`` parameter under its dedicated pre-norm + (``pre_feedforward_layernorm_2``). The router is deliberately NOT a member: + its input is the *raw residual* through the router's own internal norm, not + this norm's output, so folding the scale here leaves routing logits exactly + invariant. The anchor is the ``experts`` module itself — its forward + receives every token's normed hidden state (experts see only routed subsets + *inside* that forward, so any single expert projection would give biased + ``mean(|x|)``). """ - layers = getattr(model.model, "layers", None) - if layers is None: - raise RuntimeError("model.model.layers not found; unsupported architecture") + trunk_path, trunk = resolve_text_trunk(model) + layers = trunk.layers out: list[ScaleGroup] = [] for i, layer in enumerate(layers): - prefix = f"model.layers.{i}" + prefix = f"{trunk_path}.layers.{i}" attn = getattr(layer, "self_attn", None) if attn is not None and isinstance(getattr(attn, "q_proj", None), torch.nn.Linear): @@ -115,6 +130,18 @@ def discover_groups(model, *, include_output_proj: bool = False) -> list[ScaleGr prev_norm=f"{prefix}.{norm_name}", )) + experts = getattr(layer, "experts", None) + gate_up = getattr(experts, "gate_up_proj", None) if experts is not None else None + if experts is not None and torch.is_tensor(gate_up) and gate_up.ndim == 3: + moe_norm = getattr(layer, "pre_feedforward_layernorm_2", None) + if moe_norm is not None and hasattr(moe_norm, "weight"): + out.append(ScaleGroup( + group_id=f"L{i}_moe", + anchor=f"{prefix}.experts", + members=(f"{prefix}.experts.gate_up_proj",), + prev_norm=f"{prefix}.pre_feedforward_layernorm_2", + )) + if include_output_proj: if attn is not None and isinstance(getattr(attn, "o_proj", None), torch.nn.Linear): out.append(ScaleGroup( @@ -133,6 +160,46 @@ def discover_groups(model, *, include_output_proj: bool = False) -> list[ScaleGr return out +def _member_weight(model, dotted: str) -> torch.Tensor: + """Weight tensor of a group member. + + A member is either an ``nn.Linear`` module path (weight ``[n_out, n_in]``) + or a direct parameter path for fused 3D expert stacks + (``[n_expert, n_out, n_in]``). Input channels are the LAST axis in both + cases, so a per-input-channel scale broadcasts as + ``mul.view((1,) * (W.ndim - 1) + (-1,))``. + """ + obj = _get_module(model, dotted) + if isinstance(obj, torch.nn.Linear): + return obj.weight + if torch.is_tensor(obj): + return obj + raise TypeError(f"group member {dotted!r} is neither a Linear nor a weight tensor") + + +def _in_channel_view(mul: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """Reshape a per-input-channel vector to broadcast over ``w``'s last axis.""" + return mul.view((1,) * (w.ndim - 1) + (-1,)) + + +def _scoring_weight(model, member: str, moe_expert_sample: int) -> torch.Tensor: + """A member's weight as a 2D ``[rows, n_in]`` f32 tensor for proxy scoring. + + Fused 3D expert stacks are subsampled to ``moe_expert_sample`` evenly + spaced experts (deterministic; ``0`` keeps all) and flattened — the + block/group fake-quantizers operate per row, so stacking experts' output + rows reproduces exactly how llama-quantize sees each expert slice. + """ + w = _member_weight(model, member).detach() + if w.ndim == 3: + n_experts = w.shape[0] + if 0 < moe_expert_sample < n_experts: + idx = torch.linspace(0, n_experts - 1, moe_expert_sample).round().long() + w = w[idx.to(w.device)] + return w.reshape(-1, w.shape[-1]).float() + return w.float() + + def _get_module(model, dotted: str): obj = model for part in dotted.split("."): @@ -518,6 +585,7 @@ def calibrate( holdout_text: Path | None = None, cv_strategy: Literal["off", "gate", "mixed"] = "off", cv_weight: float = 1.0, + moe_expert_sample: int = 8, ) -> Path: """Calibrate per-group AWQ scales and write them to ``out_path``. @@ -544,6 +612,11 @@ def calibrate( - ``"mixed"`` (exp-018): score each α candidate as ``proxy_loss(X_cal, α) + cv_weight · proxy_loss(X_ho, α)``. ``cv_weight`` > 1 over-weights the held-out signal to push back against over-fit. + + ``moe_expert_sample`` caps how many experts of a fused 3D stack are scored + in the α search (evenly-spaced, deterministic). The scale vector is shared + by every expert anyway; fake-quantizing all 128 experts of a 26B MoE per α + candidate is the cost driver, not a quality driver. ``0`` scores all. """ # Validate cheap preconditions BEFORE the (expensive) model load + forward # pass so misconfigured runs fail in milliseconds, not hours. @@ -552,7 +625,7 @@ def calibrate( if proxy not in _PROXIES: raise ValueError(f"unknown proxy {proxy!r}; choose one of {sorted(_PROXIES)}") - from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers import AutoTokenizer device = resolve_device(device) torch_dtype = { @@ -566,9 +639,7 @@ def calibrate( print(f"[awq] load {model_dir} -> {device}/{dtype}", file=sys.stderr) tok = AutoTokenizer.from_pretrained(model_dir, fix_mistral_regex=True) - model = AutoModelForCausalLM.from_pretrained( - model_dir, torch_dtype=torch_dtype, trust_remote_code=True - ).to(device) + model = load_calibration_model(model_dir, torch_dtype=torch_dtype, device=device) model.eval() for p in model.parameters(): p.requires_grad_(False) @@ -576,7 +647,11 @@ def calibrate( groups = discover_groups(model, include_output_proj=include_output_proj) n_attn = sum(1 for g in groups if g.group_id.endswith("_attn")) n_mlp = sum(1 for g in groups if g.group_id.endswith("_mlp")) - print(f"[awq] groups: {len(groups)} total ({n_attn} attn, {n_mlp} mlp)", file=sys.stderr) + n_moe = sum(1 for g in groups if g.group_id.endswith("_moe")) + print( + f"[awq] groups: {len(groups)} total ({n_attn} attn, {n_mlp} mlp, {n_moe} moe)", + file=sys.stderr, + ) if not groups: raise RuntimeError("no AWQ scale groups discovered") @@ -670,7 +745,9 @@ def pre(_module, in_args): # on CPU: it is what gets persisted in the bundle. s_dev = s.to(device) X = x_cached.to(device) - Ws = [_get_module(model, m).weight.detach().float() for m in g.members] + Ws = [ + _scoring_weight(model, m, moe_expert_sample) for m in g.members + ] if force_alpha is not None: best_alpha = force_alpha @@ -879,7 +956,7 @@ def apply( the round-trip injects ~0.8% relative drift per layer (mantissa noise); ``sanity_max_rel`` rejects anything well above that floor. """ - from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers import AutoTokenizer device = resolve_device(device) torch_dtype = { @@ -891,9 +968,7 @@ def apply( print(f"[awq.apply] load {model_dir} -> {device}/{dtype}", file=sys.stderr) tok = AutoTokenizer.from_pretrained(model_dir, fix_mistral_regex=True) - model = AutoModelForCausalLM.from_pretrained( - model_dir, torch_dtype=torch_dtype, trust_remote_code=True - ).to(device) + model = load_calibration_model(model_dir, torch_dtype=torch_dtype, device=device) model.eval() ref_logits = None @@ -902,7 +977,8 @@ def apply( text = "def hello_world():\n print('hi')\n" * 8 ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids[0] sanity_ids = ids[:sanity_tokens].unsqueeze(0).to(device) - ref_logits = model(sanity_ids).logits.detach().float().cpu() + with torch.no_grad(): + ref_logits = causal_logits(model, sanity_ids).detach().float().cpu() n_applied = 0 for g in bundle.groups: @@ -922,21 +998,22 @@ def apply( scaled_members: list[tuple[str, torch.Tensor]] = [] skip_group = False for m in g.members: - W = _get_module(model, m).weight + W = _member_weight(model, m) mul = member_mul.get(m, scale) - if W.shape[1] != mul.shape[0]: + if W.shape[-1] != mul.shape[0]: print( - f"[awq.apply] WARN: {m} in={W.shape[1]} vs scale={mul.shape[0]}; " + f"[awq.apply] WARN: {m} in={W.shape[-1]} vs scale={mul.shape[0]}; " f"skipping group {g.group_id}", file=sys.stderr, ) skip_group = True break - W.data.mul_(mul.unsqueeze(0)) + W.data.mul_(_in_channel_view(mul, W)) scaled_members.append((m, mul)) if skip_group: for m, mul in scaled_members: # revert any partial mutations - _get_module(model, m).weight.data.mul_((1.0 / mul).unsqueeze(0)) + W = _member_weight(model, m) + W.data.mul_(_in_channel_view(1.0 / mul, W)) continue if g.prev_norm is None: @@ -957,7 +1034,8 @@ def apply( file=sys.stderr, ) for m, mul in scaled_members: - _get_module(model, m).weight.data.mul_((1.0 / mul).unsqueeze(0)) + W = _member_weight(model, m) + W.data.mul_(_in_channel_view(1.0 / mul, W)) continue norm_w = _get_module(model, g.prev_norm).weight @@ -967,7 +1045,8 @@ def apply( file=sys.stderr, ) for m, mul in scaled_members: - _get_module(model, m).weight.data.mul_((1.0 / mul).unsqueeze(0)) + W = _member_weight(model, m) + W.data.mul_(_in_channel_view(1.0 / mul, W)) continue new_gain = fold_rmsnorm_gain(norm_w.data, inv_f, plus_one=rmsnorm_plus_one) @@ -977,7 +1056,7 @@ def apply( print(f"[awq.apply] {n_applied}/{len(bundle.groups)} groups folded", file=sys.stderr) if ref_logits is not None and sanity_ids is not None: - new_logits = model(sanity_ids).logits.detach().float().cpu() + new_logits = causal_logits(model, sanity_ids).detach().float().cpu() max_abs = (new_logits - ref_logits).abs().max().item() rel = max_abs / max(ref_logits.abs().max().item(), 1e-8) print(f"[awq.apply] sanity max|Δlogit|={max_abs:.3e} rel={rel:.3e}", file=sys.stderr) diff --git a/src/quant_tuner/calibrate/imatrix.py b/src/quant_tuner/calibrate/imatrix.py index e25d3f9..792b21e 100644 --- a/src/quant_tuner/calibrate/imatrix.py +++ b/src/quant_tuner/calibrate/imatrix.py @@ -28,7 +28,7 @@ from quant_tuner.calibrate._device import resolve_device from quant_tuner.calibrate._hf import forward_no_logits -from quant_tuner.models.hf_gguf_map import is_ssm, map_hf_to_gguf +from quant_tuner.models.hf_gguf_map import is_moe_expert, is_ssm, map_hf_to_gguf from quant_tuner.paths import LLAMA_CPP_DIR Variant = Literal["analytic", "mix_50", "hybrid_custom", "outlier_l4", "outlier_max"] @@ -55,7 +55,15 @@ def _col_l2_sq(weight: np.ndarray) -> np.ndarray: def _load_base_imatrix(path: Path) -> dict[str, np.ndarray]: - """Read an existing imatrix GGUF as ``{tensor: E[a^2]}`` (sum_sq divided by count).""" + """Read an existing imatrix GGUF as ``{tensor: E[a^2]}`` (sum_sq divided by count). + + Dense tensors yield a 1-D ``[n_in]`` vector. Stacked MoE expert tensors + (e.g. ``ffn_gate_up_exps``) carry one count *per expert*: llama-imatrix + stores ``in_sum2`` as ``[n_expert, n_in]`` (expert-major) and ``counts`` as + ``[n_expert, 1]``; these yield a 2-D ``[n_expert, n_in]`` array normalized + per expert. Experts the calibration corpus never routed to (count 0) get a + neutral all-ones row, mirroring llama-quantize's own fallback. + """ _ensure_gguf_py() from gguf import GGUFReader @@ -69,8 +77,22 @@ def _load_base_imatrix(path: Path) -> dict[str, np.ndarray]: cnt = raw.get(f"{base}.counts") if cnt is None: continue - c = max(int(cnt.flat[0]), 1) - out[base] = arr.astype(np.float32) / c + counts = cnt.astype(np.float64).reshape(-1) + if counts.size <= 1: + c = max(float(counts[0]) if counts.size else 1.0, 1.0) + out[base] = (arr.astype(np.float32) / c).reshape(-1) + continue + n_mat = counts.size + if arr.size % n_mat != 0: + raise ValueError( + f"imatrix entry {base}: in_sum2 size {arr.size} not divisible by " + f"counts size {n_mat} (unexpected llama-imatrix layout)" + ) + mat = arr.astype(np.float32).reshape(n_mat, arr.size // n_mat) + ea2 = np.ones_like(mat) + seen = counts > 0 + ea2[seen] = mat[seen] / counts[seen, None].astype(np.float32) + out[base] = ea2 return out @@ -94,6 +116,58 @@ def _l1_normalize(v: np.ndarray) -> np.ndarray: return (v.astype(np.float32) * (v.size / s)).astype(np.float32) +def check_moe_coverage(base_imatrix: Path) -> dict[str, float]: + """Fraction of experts with nonzero activation mass, per stacked-expert tensor. + + With 128 experts and a top-k router, a too-small calibration corpus leaves + some experts unobserved — llama-quantize then falls back to neutral + importance for those experts, leaving them unprotected at low bit-widths. + Returns ``{tensor: covered_fraction}`` for every multi-count imatrix entry. + """ + _ensure_gguf_py() + from gguf import GGUFReader + + reader = GGUFReader(str(base_imatrix)) + out: dict[str, float] = {} + for t in reader.tensors: + if not t.name.endswith(".counts"): + continue + counts = np.array(t.data).reshape(-1) + if counts.size <= 1: + continue + out[t.name[: -len(".counts")]] = float((counts > 0).mean()) + return out + + +def warn_moe_coverage(base_imatrix: Path, threshold: float = 0.95) -> dict[str, float]: + """Log a warning for every expert tensor below ``threshold`` coverage. + + Diagnostic only — an unreadable imatrix is reported, never fatal (the + quantize step will surface real corruption with a hard error). + """ + try: + coverage = check_moe_coverage(base_imatrix) + except Exception as e: # noqa: BLE001 + print(f"[imatrix] coverage check skipped ({e!r})", file=sys.stderr) + return {} + low = {t: c for t, c in coverage.items() if c < threshold} + if low: + worst = min(low.values()) + print( + f"[imatrix] WARN: {len(low)}/{len(coverage)} expert tensors below " + f"{threshold:.0%} expert coverage (worst {worst:.0%}). Unobserved " + f"experts quantize with neutral importance — increase the " + f"calibration corpus (data.train_max_tokens) for low-bit targets.", + file=sys.stderr, + ) + elif coverage: + print( + f"[imatrix] MoE expert coverage OK ({len(coverage)} tensors >= {threshold:.0%})", + file=sys.stderr, + ) + return coverage + + # --- Forward-pass stats (E[a^4], max|a|) --------------------------------- # @dataclass @@ -143,7 +217,9 @@ def collect_forward_stats( coverage matches what llama-quantize will actually consume. """ import torch - from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers import AutoTokenizer + + from quant_tuner.calibrate._hf import load_calibration_model _ensure_gguf_py() from gguf import GGUFReader @@ -153,9 +229,7 @@ def collect_forward_stats( print(f"[forward_stats] load {model_dir} -> {device}/{dtype}", file=sys.stderr) tok = AutoTokenizer.from_pretrained(model_dir, fix_mistral_regex=True) - model = AutoModelForCausalLM.from_pretrained( - model_dir, torch_dtype=torch_dtype, trust_remote_code=True - ).to(device) + model = load_calibration_model(model_dir, torch_dtype=torch_dtype, device=device) model.eval() for p in model.parameters(): p.requires_grad_(False) @@ -229,6 +303,29 @@ def pre(_module, in_args): # pass through the raw E[a^2] signal (memory note: don't rerank ``blk.*.ssm_*``). +def _combine_moe( + weight3d: np.ndarray, ea2: np.ndarray, combiner, tname: str +) -> np.ndarray: + """Apply ``combiner`` per expert and flatten back expert-major. + + ``weight3d`` is ``[n_expert, n_out, n_in]`` (GGUF order); ``ea2`` is + ``[n_expert, n_in]`` from :func:`_load_base_imatrix`. Each expert's row is + combined against that expert's own ``||W_e[:, c]||^2`` — the L1 + normalization inside mix_50/hybrid_custom must not bleed across experts, + since llama-quantize consumes each expert's imatrix slice independently. + """ + if weight3d.ndim != 3 or ea2.ndim != 2 or weight3d.shape[0] != ea2.shape[0] or weight3d.shape[2] != ea2.shape[1]: + raise ValueError( + f"{tname}: expert weight {weight3d.shape} does not line up with " + f"imatrix rows {ea2.shape} as [n_expert, n_out, n_in] vs [n_expert, n_in] " + f"(llama.cpp imatrix layout change?)" + ) + w = weight3d.astype(np.float32, copy=False) + wcol = (w * w).sum(axis=1) # [n_expert, n_in] + rows = [combiner(wcol[e], ea2[e]).astype(np.float32) for e in range(ea2.shape[0])] + return np.stack(rows).reshape(-1) + + def _build_simple( f16: Path, base_imatrix: Path, @@ -238,22 +335,33 @@ def _build_simple( base = _load_base_imatrix(base_imatrix) weights = _load_weights(f16) out: dict[str, np.ndarray] = {} - ssm = skipped = 0 + ssm = moe = skipped = 0 for tname, ea2 in base.items(): if tname not in weights: skipped += 1 continue if is_ssm(tname): - out[tname] = ea2.astype(np.float32) + out[tname] = ea2.astype(np.float32).reshape(-1) ssm += 1 continue + if ea2.ndim == 2: + if not is_moe_expert(tname): + # multi-count entry on a non-expert tensor: no per-mat weight + # slicing rule we trust — keep the raw signal rather than guess + out[tname] = ea2.astype(np.float32).reshape(-1) + skipped += 1 + continue + out[tname] = _combine_moe(weights[tname], ea2, combiner, tname) + moe += 1 + continue wcol = _col_l2_sq(weights[tname]) if wcol.shape != ea2.shape: skipped += 1 continue out[tname] = combiner(wcol, ea2).astype(np.float32) print( - f"[{label}] {len(out)} tensors ({ssm} SSM passthrough, {skipped} skipped)", + f"[{label}] {len(out)} tensors ({ssm} SSM passthrough, {moe} MoE per-expert, " + f"{skipped} skipped)", file=sys.stderr, ) return out @@ -293,12 +401,18 @@ def build_outlier_l4( """ base = _load_base_imatrix(base_imatrix) out: dict[str, np.ndarray] = {} - ssm = missing = mismatch = 0 + ssm = moe = missing = mismatch = 0 for tname, ea2 in base.items(): if is_ssm(tname): - out[tname] = ea2.astype(np.float32) + out[tname] = ea2.astype(np.float32).reshape(-1) ssm += 1 continue + if ea2.ndim == 2: + # MoE experts: hooks see only routed token subsets, so per-expert + # outlier stats aren't collected — pass through E[a^2] per expert. + out[tname] = ea2.astype(np.float32).reshape(-1) + moe += 1 + continue e_a4 = forward_stats.e_a4.get(tname) if e_a4 is None: out[tname] = ea2.astype(np.float32) @@ -310,8 +424,8 @@ def build_outlier_l4( continue out[tname] = np.sqrt(np.maximum(e_a4, 0.0)).astype(np.float32) print( - f"[outlier_l4] {len(out)} tensors ({ssm} SSM, {missing} no-stats, " - f"{mismatch} shape-mismatch fallback)", + f"[outlier_l4] {len(out)} tensors ({ssm} SSM, {moe} MoE passthrough, " + f"{missing} no-stats, {mismatch} shape-mismatch fallback)", file=sys.stderr, ) return out @@ -324,12 +438,17 @@ def build_outlier_max( """``s_c = E[a_c^2] * max|a_c|^2`` for non-SSM tensors; passthrough ``E[a^2]`` for SSM.""" base = _load_base_imatrix(base_imatrix) out: dict[str, np.ndarray] = {} - ssm = missing = mismatch = 0 + ssm = moe = missing = mismatch = 0 for tname, ea2 in base.items(): if is_ssm(tname): - out[tname] = ea2.astype(np.float32) + out[tname] = ea2.astype(np.float32).reshape(-1) ssm += 1 continue + if ea2.ndim == 2: + # MoE experts: see build_outlier_l4 — per-expert E[a^2] passthrough. + out[tname] = ea2.astype(np.float32).reshape(-1) + moe += 1 + continue mx = forward_stats.max_abs.get(tname) if mx is None: out[tname] = ea2.astype(np.float32) @@ -341,8 +460,8 @@ def build_outlier_max( continue out[tname] = (ea2 * (mx**2)).astype(np.float32) print( - f"[outlier_max] {len(out)} tensors ({ssm} SSM, {missing} no-stats, " - f"{mismatch} shape-mismatch fallback)", + f"[outlier_max] {len(out)} tensors ({ssm} SSM, {moe} MoE passthrough, " + f"{missing} no-stats, {mismatch} shape-mismatch fallback)", file=sys.stderr, ) return out @@ -375,6 +494,7 @@ def write_imatrix( writer.add_array(key, [dataset_label]) for tname, vec in importance.items(): + vec = vec.reshape(-1) # llama-quantize consumes flat [n_mat * n_in] if not np.all(np.isfinite(vec)): print(f" WARN: non-finite values in {tname}; clamping", file=sys.stderr) vec = np.nan_to_num(vec, nan=0.0, posinf=0.0, neginf=0.0) @@ -460,4 +580,11 @@ def calibrate( ) -__all__ = ["ForwardStats", "Variant", "calibrate", "collect_forward_stats"] +__all__ = [ + "ForwardStats", + "Variant", + "calibrate", + "check_moe_coverage", + "collect_forward_stats", + "warn_moe_coverage", +] diff --git a/src/quant_tuner/eval/diffusion_server.py b/src/quant_tuner/eval/diffusion_server.py new file mode 100644 index 0000000..6d2cf49 --- /dev/null +++ b/src/quant_tuner/eval/diffusion_server.py @@ -0,0 +1,483 @@ +"""DiffusionGemma generation backend: an OpenAI-compatible shim over llama.cpp's +persistent ``llama-diffusion-gemma-server``. + +``llama-server`` cannot run block-diffusion decoding, so the tool-call and +MMLU-Pro evals can't point at it for DiffusionGemma quants. This module keeps +their ``base_url`` contract intact instead of changing the evals: + + * :class:`LogitsServer` drives llama.cpp's ``llama-diffusion-gemma-server`` + (loads the GGUF **once**, then serves raw canvas logits for + ``[prompt | canvas]`` requests over a stdin/file protocol — the PR-provided + building block for Python diffusion drivers). + * :func:`entropy_bound_generate` is a faithful numpy port of llama.cpp's + ``diffusion_generate_entropy_bound`` (the DiffusionGemma reference + sampler): per-step argmax/entropy/multinomial over the canvas, + lowest-entropy acceptance within the entropy bound, renoise of the rest, + self-conditioning between steps (cached server-side), block-autoregressive + commit with EOG/repetition trimming. Deterministic per seed (not + bit-identical to the C++ RNG stream). + * :func:`running_diffusion_server` spawns a local HTTP shim implementing + ``/health`` + ``/v1/chat/completions``: chat-template rendering (with + tools) via the HF tokenizer, generation via the sampler, and Gemma-4 + tool-call markup (``<|tool_call>call:NAME{...}``) parsed back + into OpenAI ``tool_calls``. + +Sampling: the entropy-bound schedule (t_max → t_min from GGUF metadata) +governs temperature; per-request ``seed`` is honored, other OpenAI sampling +knobs are accepted and ignored. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import tempfile +import threading +import time +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import numpy as np + +from quant_tuner.eval.server import free_port +from quant_tuner.paths import LLAMA_CPP_DIR, llama_bin + +# --- entropy-bound parameters (GGUF metadata, CLI reference defaults) ------ # + +@dataclass +class EbParams: + canvas_length: int = 256 + max_steps: int = 48 + t_min: float = 0.4 + t_max: float = 0.8 + entropy_bound: float = 0.1 + stability_threshold: int = 1 + confidence_threshold: float = 0.005 + + +def read_eb_params(gguf_path: Path) -> EbParams: + """Entropy-bound sampler config from GGUF ``diffusion.*`` metadata.""" + p = str(LLAMA_CPP_DIR / "gguf-py") + if p not in sys.path: + sys.path.insert(0, p) + from gguf import GGUFReader + + reader = GGUFReader(str(gguf_path)) + + def field(key: str, default): + f = reader.fields.get(key) + if f is None: + return default + return type(default)(f.parts[f.data[0]][0]) + + return EbParams( + canvas_length=field("diffusion.canvas_length", 256), + max_steps=field("diffusion.eb_max_steps", 48), + t_min=field("diffusion.eb_t_min", 0.4), + t_max=field("diffusion.eb_t_max", 0.8), + entropy_bound=field("diffusion.eb_entropy_bound", 0.1), + stability_threshold=field("diffusion.eb_stability_threshold", 1), + confidence_threshold=field("diffusion.eb_confidence_threshold", 0.005), + ) + + +# --- the persistent logits server ------------------------------------------ # + +class LogitsServer: + """Drive ``llama-diffusion-gemma-server``: one model load, many forwards. + + Protocol (see the example's header): a request-file path per stdin line; + the file holds ``int32 P, C, use_sc, float32 temp`` then ``P+C`` int32 + token ids (canvas last). The server replies ``C × n_vocab`` float32 canvas + logits to ``.resp`` and ``OK `` on stdout. Self-conditioning + state (the previous step's logits) is cached server-side. + """ + + def __init__(self, model_path: Path, *, ngl: int = 99, max_tokens: int = 4096, + kv_cache: bool = True, log_path: Path | None = None, + startup_timeout: float = 300.0): + binary = llama_bin("llama-diffusion-gemma-server") # raises if not built + self._tmp = tempfile.TemporaryDirectory(prefix="dg-shim-") + self._req_idx = 0 + self._lock = threading.Lock() + env = { + "NGL": str(ngl), + "MAXTOK": str(max_tokens), + "DG_KVCACHE": "1" if kv_cache else "0", + } + stderr: int | object = subprocess.DEVNULL + if log_path is not None: + log_path.parent.mkdir(parents=True, exist_ok=True) + stderr = log_path.open("w") + self._proc = subprocess.Popen( + [str(binary), str(model_path)], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=stderr, + env={**os.environ, **env}, text=True, bufsize=1, + ) + assert self._proc.stdin is not None and self._proc.stdout is not None + self._stdin = self._proc.stdin + self._stdout = self._proc.stdout + deadline = time.time() + startup_timeout + line = "" + while time.time() < deadline: + if self._proc.poll() is not None: + raise RuntimeError( + f"llama-diffusion-gemma-server exited with code {self._proc.returncode} " + f"before READY (bad GGUF / unsupported arch?)" + ) + line = self._stdout.readline() + if line.startswith("READY"): + break + if not line.startswith("READY"): + self.close() + raise TimeoutError(f"logits server not READY within {startup_timeout}s") + self.n_vocab = int(line.split()[1]) + + def canvas_logits( + self, prompt_ids: np.ndarray, canvas_ids: np.ndarray, + *, use_sc: bool, temp: float, + ) -> np.ndarray: + """One forward; returns raw canvas logits ``[C, n_vocab]`` (float32).""" + with self._lock: + self._req_idx += 1 + req = Path(self._tmp.name) / f"req-{self._req_idx:06d}" + header = np.empty(4, dtype=np.int32) + header[0] = len(prompt_ids) + header[1] = len(canvas_ids) + header[2] = 1 if use_sc else 0 + header[3:4].view(np.float32)[0] = temp + payload = np.concatenate([ + header, + np.asarray(prompt_ids, dtype=np.int32), + np.asarray(canvas_ids, dtype=np.int32), + ]) + payload.tofile(req) + self._stdin.write(f"{req}\n") + self._stdin.flush() + line = self._stdout.readline().strip() + if not line.startswith("OK"): + raise RuntimeError(f"logits server error: {line!r}") + resp = Path(f"{req}.resp") + logits = np.fromfile(resp, dtype=np.float32).reshape( + len(canvas_ids), self.n_vocab + ) + req.unlink(missing_ok=True) + resp.unlink(missing_ok=True) + return logits + + def close(self) -> None: + try: + if self._proc.poll() is None: + self._stdin.write("QUIT\n") + self._stdin.flush() + self._proc.wait(timeout=10) + except Exception: # noqa: BLE001 + self._proc.kill() + finally: + self._tmp.cleanup() + + +# --- the entropy-bound sampler (numpy port) -------------------------------- # + +def _detect_repetition_loop(tokens: list[int], start: int) -> bool: + """llama.cpp trim_canvas loop rule: >= 6 consecutive stride-1/2 repeats.""" + n = len(tokens) + for stride in (1, 2): + reps = 0 + j = start + while j + stride < n and tokens[j] == tokens[j + stride]: + reps += 1 + j += stride + if reps >= 6: + return True + return False + + +def trim_canvas(tokens: list[int], eog_ids: frozenset[int]) -> int: + """Cut index for a denoised canvas: first EOG token or repetition loop.""" + cut = len(tokens) + for i, t in enumerate(tokens): + if t in eog_ids: + cut = i + break + for i in range(max(cut - 1, 0)): + if _detect_repetition_loop(tokens, i): + cut = i + break + return cut + + +def entropy_bound_generate( + server: LogitsServer, + prompt_ids: list[int], + *, + eb: EbParams, + eog_ids: frozenset[int], + seed: int = 0, + max_blocks: int = 4, + max_tokens: int = 4096, +) -> list[int]: + """Block-autoregressive entropy-bound denoising; returns generated ids. + + Faithful port of ``diffusion_generate_entropy_bound`` + + ``run_turn``'s block loop from llama.cpp's diffusion example: each block + denoises a fixed ``canvas_length`` canvas conditioned on the prefix, the + trimmed result is committed and the next block starts, until an EOG + token, a repetition loop, the block budget, or the token budget. + """ + rng = np.random.default_rng(seed) + C = eb.canvas_length + S = max(1, eb.max_steps) + prefix = list(prompt_ids) + response: list[int] = [] + + for _block in range(max_blocks): + if len(prefix) + C > max_tokens: + break + canvas = rng.integers(0, server.n_vocab, size=C).astype(np.int64) + argmax = np.full(C, -1, dtype=np.int64) + prev_argmax = np.full(C, -2, dtype=np.int64) + held = 0 + + for step_idx in range(S): + cur_step = S - step_idx + t = eb.t_min + (eb.t_max - eb.t_min) * (cur_step / S) + logits = server.canvas_logits( + np.asarray(prefix), canvas, use_sc=step_idx > 0, temp=t, + ) + z = logits.astype(np.float64) / t + z -= z.max(axis=1, keepdims=True) + p = np.exp(z) + p /= p.sum(axis=1, keepdims=True) + + argmax = logits.argmax(axis=1) + with np.errstate(divide="ignore", invalid="ignore"): + logp = np.where(p > 0, np.log(p), 0.0) + entropy = -(p * logp).sum(axis=1) + + # one multinomial draw per position (inverse-CDF, like the C++) + cum = p.cumsum(axis=1) + u = rng.random(C) + sampled = (cum >= u[:, None]).argmax(axis=1) + + # accept the lowest-entropy positions within the bound + order = np.argsort(entropy, kind="stable") + cum_e = np.concatenate([[0.0], entropy[order].cumsum()[:-1]]) + accepted = np.zeros(C, dtype=bool) + accepted[order[cum_e <= eb.entropy_bound]] = True + + renoise = rng.integers(0, server.n_vocab, size=C) + canvas = np.where(accepted, sampled, renoise).astype(np.int64) + + held = held + 1 if np.array_equal(argmax, prev_argmax) else 0 + prev_argmax = argmax + confident = entropy.mean() < eb.confidence_threshold + if held >= eb.stability_threshold and confident: + break + + out = [int(x) for x in argmax] + cut = trim_canvas(out, eog_ids) + response.extend(out[:cut]) + if cut < C: + break # end token or repetition loop: answer complete + prefix.extend(out[:cut]) + + return response + + +# --- Gemma-4 response markup ------------------------------------------------ # + +_TOOL_CALL_RE = re.compile( + r"<\|tool_call>call:\s*([\w.\-]+)\s*(\{.*?\})\s*", re.S +) +_CHANNEL_RE = re.compile(r"<\|channel>.*?(?:|$)", re.S) +_BARE_KEY_RE = re.compile(r"([{,]\s*)([A-Za-z_][\w\-]*)\s*:") + + +def _parse_gemma4_args(raw: str) -> dict: + """Parse a gemma4-dict argument blob: JSON first, then with bare keys quoted.""" + for candidate in (raw, _BARE_KEY_RE.sub(r'\1"\2":', raw)): + try: + parsed = json.loads(candidate) + if isinstance(parsed, dict): + return parsed + except json.JSONDecodeError: + continue + return {"_raw": raw} + + +def parse_gemma4_response(text: str) -> tuple[str, list[dict]]: + """Split a raw Gemma-4 completion into (content, OpenAI-style tool_calls). + + The model emits ``<|tool_call>call:NAME{args}`` segments and + optional ``<|channel>…`` reasoning blocks; both are stripped + from the content. + """ + tool_calls = [] + for i, m in enumerate(_TOOL_CALL_RE.finditer(text)): + tool_calls.append({ + "id": f"call_{i}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(_parse_gemma4_args(m.group(2))), + }, + }) + content = _TOOL_CALL_RE.sub("", text) + content = _CHANNEL_RE.sub("", content) + content = re.sub(r"<(?:end_of_turn|eos|\|im_end\|)>", "", content) + return content.strip(), tool_calls + + +# --- the OpenAI-compatible HTTP shim ---------------------------------------- # + +class _ShimState: + def __init__(self, server: LogitsServer, tokenizer, eb: EbParams, + eog_ids: frozenset[int], max_blocks: int, max_tokens: int): + self.server = server + self.tokenizer = tokenizer + self.eb = eb + self.eog_ids = eog_ids + self.max_blocks = max_blocks + self.max_tokens = max_tokens + self.model_name = "diffusion-gemma" + + +def _make_handler(state: _ShimState): + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): # silence per-request stderr noise + pass + + def _json(self, code: int, payload: dict) -> None: + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): # noqa: N802 — BaseHTTPRequestHandler API + if self.path == "/health": + self._json(200, {"status": "ok"}) + else: + self._json(404, {"error": "not found"}) + + def do_POST(self): # noqa: N802 + if not self.path.endswith("/chat/completions"): + self._json(404, {"error": "not found"}) + return + try: + length = int(self.headers.get("Content-Length", "0")) + req = json.loads(self.rfile.read(length)) + messages = req["messages"] + tools = req.get("tools") or None + seed = int(req.get("seed", 0) or 0) + + prompt = state.tokenizer.apply_chat_template( + messages, tools=tools, add_generation_prompt=True, + tokenize=False, + ) + ids = state.tokenizer( + prompt, add_special_tokens=False + ).input_ids + out_ids = entropy_bound_generate( + state.server, ids, eb=state.eb, eog_ids=state.eog_ids, + seed=seed, max_blocks=state.max_blocks, + max_tokens=state.max_tokens, + ) + raw = state.tokenizer.decode(out_ids, skip_special_tokens=False) + content, tool_calls = parse_gemma4_response(raw) + message: dict = {"role": "assistant", "content": content or None} + if tool_calls: + message["tool_calls"] = tool_calls + self._json(200, { + "id": "dgshim", + "object": "chat.completion", + "model": state.model_name, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": "tool_calls" if tool_calls else "stop", + }], + "usage": { + "prompt_tokens": len(ids), + "completion_tokens": len(out_ids), + "total_tokens": len(ids) + len(out_ids), + }, + }) + except Exception as e: # noqa: BLE001 — surfaced to the client + self._json(500, {"error": {"message": repr(e)}}) + + return Handler + + +def _gemma_eog_ids(tokenizer) -> frozenset[int]: + ids = set() + if tokenizer.eos_token_id is not None: + ids.add(int(tokenizer.eos_token_id)) + for tok in ("",): + tid = tokenizer.convert_tokens_to_ids(tok) + if tid is not None and tid >= 0: + ids.add(int(tid)) + return frozenset(ids) + + +@contextmanager +def running_diffusion_server( + model_path: Path, + *, + hf_dir: Path, + ctx: int = 4096, + ngl: int = 99, + log_path: Path | None = None, + startup_timeout: float = 300.0, + max_blocks: int = 4, + chat_template_kwargs: str | None = None, # accepted for interface parity; unused +): + """Drop-in sibling of :func:`quant_tuner.eval.server.running_server` for + DiffusionGemma GGUFs: yields a ``base_url`` whose ``/v1/chat/completions`` + runs block-diffusion generation. ``hf_dir`` must hold the HF tokenizer + (chat template + tools rendering) — the extracted snapshot + (``ws.model_extracted``) is the canonical source. + """ + del chat_template_kwargs + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(hf_dir) + eb = read_eb_params(model_path) + logits_server = LogitsServer( + model_path, ngl=ngl, max_tokens=ctx, + log_path=log_path, startup_timeout=startup_timeout, + ) + try: + state = _ShimState( + logits_server, tokenizer, eb, + _gemma_eog_ids(tokenizer), max_blocks, ctx, + ) + port = free_port() + httpd = ThreadingHTTPServer(("127.0.0.1", port), _make_handler(state)) + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{port}/v1" + finally: + httpd.shutdown() + thread.join(timeout=5) + finally: + logits_server.close() + + +__all__ = [ + "EbParams", + "LogitsServer", + "entropy_bound_generate", + "parse_gemma4_response", + "read_eb_params", + "running_diffusion_server", + "trim_canvas", +] diff --git a/src/quant_tuner/eval/reps.py b/src/quant_tuner/eval/reps.py index bc6cc10..b3c00ac 100644 --- a/src/quant_tuner/eval/reps.py +++ b/src/quant_tuner/eval/reps.py @@ -132,14 +132,21 @@ def run_reps_for_model( server_startup_timeout: float = 120.0, chat_template_kwargs: str | None = None, per_rep_callback: Callable[[RepResult], None] | None = None, + server_factory: Callable | None = None, ) -> ModelReps: """Spawn ``llama-server`` for ``model``, run ``eval_fn`` ``reps`` times. Each rep receives a fresh ``Sampling`` with ``seed = base_seed + rep_idx``. The server is torn down on exit (or on exception). + + ``server_factory`` swaps the backend: any context manager with + ``running_server``'s signature that yields a ``base_url`` works — e.g. + ``functools.partial(running_diffusion_server, hf_dir=…)`` from + :mod:`quant_tuner.eval.diffusion_server` for DiffusionGemma GGUFs. """ + factory = server_factory or running_server rep_results: list[RepResult] = [] - with running_server( + with factory( model, ctx=ctx, ngl=ngl, log_path=server_log_path, startup_timeout=server_startup_timeout, @@ -186,6 +193,7 @@ def run_reps_for_models( chat_template_kwargs: str | None = None, per_rep_callback: Callable[[str, RepResult], None] | None = None, per_model_callback: Callable[[ModelReps], None] | None = None, + server_factory: Callable | None = None, ) -> list[ModelReps]: """Run ``run_reps_for_model`` over a list of GGUFs, one server at a time.""" results: list[ModelReps] = [] @@ -209,6 +217,7 @@ def _cb(rr: RepResult, _name=model.name) -> None: server_startup_timeout=server_startup_timeout, chat_template_kwargs=chat_template_kwargs, per_rep_callback=_cb, + server_factory=server_factory, ) results.append(mr) if per_model_callback is not None: diff --git a/src/quant_tuner/models/extract.py b/src/quant_tuner/models/extract.py index 74ad6cc..05dbd01 100644 --- a/src/quant_tuner/models/extract.py +++ b/src/quant_tuner/models/extract.py @@ -71,7 +71,12 @@ def extract_text_lm( with open(model_path / "config.json") as f: config = json.load(f) arch = (config.get("architectures") or [""])[0] - if "ForCausalLM" in arch and not config.get("text_config"): + # DiffusionGemma is pass-through despite its text_config nesting: the + # llama.cpp converter consumes the full checkpoint (root canvas_length, + # model.decoder.* backbone, tied encoder) and drops the vision tower + # itself — rebuilding a text-only CausalLM here would break it. + is_block_diffusion = "ForBlockDiffusion" in arch + if is_block_diffusion or ("ForCausalLM" in arch and not config.get("text_config")): # Pass-through: still materialize into output_dir (via symlinks, no # copy) so callers that reference output_dir — the pipeline points # conversion/calibration at ws.model_extracted — find a real model diff --git a/src/quant_tuner/models/hf_gguf_map.py b/src/quant_tuner/models/hf_gguf_map.py index ba8ef01..431fdb8 100644 --- a/src/quant_tuner/models/hf_gguf_map.py +++ b/src/quant_tuner/models/hf_gguf_map.py @@ -1,16 +1,29 @@ """HuggingFace → GGUF tensor-name mapping for Linear layers. Covers attention projections (q/k/v/o + attn_output_gate + fused qkv), MLP -projections (gate/up/down), and the LM head. SSM tensors are intentionally -absent — they don't share the `y = W a` linear-projection structure, so -calibrators that derive activation-aware importance pass them through using -the standard E[a^2] signal instead. +projections (gate/up/down), MoE expert/router tensors (Gemma-4 style fused 3D +experts), and the LM head. SSM tensors are intentionally absent — they don't +share the `y = W a` linear-projection structure, so calibrators that derive +activation-aware importance pass them through using the standard E[a^2] +signal instead. + +DiffusionGemma note: the checkpoint's text backbone lives under +``model.decoder.*`` (encoder weights are tied; llama.cpp's converter rewrites +``model.decoder.`` → ``model.``). :func:`map_hf_to_gguf` applies the same +prefix normalization so HF module names from either the decoder or the tied +encoder stack resolve to the canonical GGUF tensor. """ from __future__ import annotations import re +# Trunk prefixes that alias `model.` (DiffusionGemma encoder/decoder stacks). +_TRUNK_PREFIXES: tuple[tuple[str, str], ...] = ( + ("model.decoder.", "model."), + ("model.encoder.language_model.", "model."), +) + _HF_TO_GGUF_RULES: list[tuple[re.Pattern[str], str]] = [ (re.compile(r"^model\.layers\.(\d+)\.self_attn\.q_proj$"), "blk.{bid}.attn_q.weight"), (re.compile(r"^model\.layers\.(\d+)\.self_attn\.k_proj$"), "blk.{bid}.attn_k.weight"), @@ -33,11 +46,27 @@ (re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_a$"), "blk.{bid}.ssm_alpha.weight"), (re.compile(r"^model\.layers\.(\d+)\.linear_attn\.in_proj_b$"), "blk.{bid}.ssm_beta.weight"), (re.compile(r"^model\.layers\.(\d+)\.linear_attn\.out_proj$"), "blk.{bid}.ssm_out.weight"), + # Gemma-4 / DiffusionGemma MoE. Experts are a single fused 3D parameter per + # projection ([n_expert, n_out, n_in]); the router is a plain Linear whose + # input is NOT the experts' pre-norm output (it has its own internal norm). + (re.compile(r"^model\.layers\.(\d+)\.experts\.gate_up_proj$"), "blk.{bid}.ffn_gate_up_exps.weight"), + (re.compile(r"^model\.layers\.(\d+)\.experts\.down_proj$"), "blk.{bid}.ffn_down_exps.weight"), + (re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj$"), "blk.{bid}.ffn_gate_up_exps.weight"), + (re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj$"), "blk.{bid}.ffn_down_exps.weight"), + (re.compile(r"^model\.layers\.(\d+)\.router\.proj$"), "blk.{bid}.ffn_gate_inp.weight"), + # DiffusionGemma decoder-only self-conditioning gated MLP (2D Linears). + (re.compile(r"^model\.self_conditioning\.gate_proj$"), "self_cond_gate.weight"), + (re.compile(r"^model\.self_conditioning\.up_proj$"), "self_cond_up.weight"), + (re.compile(r"^model\.self_conditioning\.down_proj$"), "self_cond_down.weight"), ] def map_hf_to_gguf(hf_name: str) -> str | None: """Map an HF module name (e.g. ``model.layers.3.mlp.up_proj``) to a GGUF tensor name.""" + for prefix, repl in _TRUNK_PREFIXES: + if hf_name.startswith(prefix): + hf_name = repl + hf_name[len(prefix):] + break for pat, tmpl in _HF_TO_GGUF_RULES: m = pat.match(hf_name) if m: @@ -51,3 +80,15 @@ def is_ssm(gguf_name: str) -> bool: """SSM tensors lack the y = W a structure; calibrators must not rerank them with an output-aware prior (E[a^2] passthrough is correct).""" return ".ssm_" in gguf_name + + +_MOE_EXPERT_MARKERS = (".ffn_gate_up_exps.", ".ffn_gate_exps.", ".ffn_up_exps.", ".ffn_down_exps.") + + +def is_moe_expert(gguf_name: str) -> bool: + """Stacked per-expert tensors (3D ``[n_expert, n_out, n_in]`` in GGUF). + + Their imatrix entries carry ``n_expert`` independent E[a^2] rows (one per + expert, expert-major); calibrators must rerank each expert's row against + that expert's weight slice, never across experts.""" + return any(m in gguf_name for m in _MOE_EXPERT_MARKERS) diff --git a/src/quant_tuner/pipeline.py b/src/quant_tuner/pipeline.py index a0c0311..dfa1275 100644 --- a/src/quant_tuner/pipeline.py +++ b/src/quant_tuner/pipeline.py @@ -198,6 +198,7 @@ def _calibrate_imatrix( step("llama-imatrix (base)", base_imatrix, lambda: llama_cpp.imatrix(f16, train_corpus, base_imatrix, ctx=imatrix_ctx, log=logs / "imatrix-base.log")) + imatrix.warn_moe_coverage(base_imatrix) variant = cfg.calibration.variant if variant == "default": @@ -259,6 +260,7 @@ def _calibrate_awq( step("llama-imatrix (on AWQ-folded F16)", awq_imatrix, lambda: llama_cpp.imatrix(f16_awq, train_corpus, awq_imatrix, ctx=imatrix_ctx, log=logs / "imatrix-awq.log")) + imatrix.warn_moe_coverage(awq_imatrix) if imatrix_variant == "default": return {"imatrix": awq_imatrix, "f16": f16_awq} diff --git a/src/quant_tuner/recipes/iq2_m_diffusiongemma_awq.yaml b/src/quant_tuner/recipes/iq2_m_diffusiongemma_awq.yaml new file mode 100644 index 0000000..c8c04c0 --- /dev/null +++ b/src/quant_tuner/recipes/iq2_m_diffusiongemma_awq.yaml @@ -0,0 +1,69 @@ +# DiffusionGemma 26B-A4B (Gemma-4 MoE block-diffusion) — AWQ → IQ2_M (~2.7 bpw). +# +# Targets google/diffusiongemma-26B-A4B-it: a 128-expert (8 active) Gemma-4 +# MoE whose causal *encoder* and bidirectional diffusion *decoder* share all +# backbone weights. Everything below runs in encoder (causal) mode on the HF +# side; the AWQ fold is attention-mask-agnostic, so it is valid for both +# modes. Requires the llama.cpp pin with `diffusion-gemma` support (see +# CLAUDE.md setup — PR #24423). +# +# MoE specifics handled by the pipeline: +# * `L{i}_moe` groups fold the scale into pre_feedforward_layernorm_2; the +# router reads the raw residual through its own norm, so routing logits +# are exactly invariant. +# * The α search scores a deterministic 8-expert subsample of each fused +# gate_up stack (moe_expert_sample) — the scale is shared group-wide. +# * The folded-F16 imatrix is re-weighted per expert with `hybrid_custom`; +# expert coverage is checked after llama-imatrix (raise train_max_tokens +# if the WARN fires: unobserved experts quantize unprotected). +# +# Bench/eval caveats (encoder-mode): +# * llama-imatrix/llama-perplexity run the UNIFIED forward — the last +# `canvas_length` (256) positions of each chunk are processed in +# bidirectional canvas mode. Quant-vs-quant KLD stays valid (identical +# computation on both sides); absolute PPL is not deployment-faithful. +# * llama-bench decode tok/s is NOT how this model generates; use +# scripts/bench_diffusion_speed.py for denoising throughput. +# * Tool-call/MMLU-Pro eval needs the diffusion backend +# (quant_tuner.eval.diffusion_server), not llama-server. +# +# Resources: bf16 HF load ≈ 52 GB (AWQ calibrate/apply) — 80GB-class GPU or +# ≥96 GB RAM with device: cpu. Disk ≥ 220 GB (HF + F16 + folded HF/F16 + +# quant + KLD baseline). +name: iq2_m_diffusiongemma_awq + +model: google/diffusiongemma-26B-A4B-it + +workspace: ./out/iq2_m_diffusiongemma_awq + +data: + logs: PLACEHOLDER + train_frac: 0.8 + test_frac: 0.1 + holdout_frac: 0.1 + # 128 experts x 8 active: expert coverage needs a larger corpus than dense + # models — the post-imatrix coverage check warns below 95%. + train_max_tokens: 500000 + eval_max_tokens: 30000 + supplement: ./calibration_supplement.txt + context_len: 16384 + +calibration: + method: awq + variant: best + params: + proxy_tokens: 256 + # proxy: iq2_s # auto-selected from quantize.type; uncomment to pin + imatrix_variant: hybrid_custom + rmsnorm_plus_one: true # Gemma (1+γ) RMSNorm convention + moe_expert_sample: 8 + sanity_max_rel: 0.03 + +quantize: + type: IQ2_M + +bench: + # KLD vs the F16 reference on external raw text; cap eval tokens — the + # 262k-token Gemma vocab makes the KLD baseline file large. + suite: full + eval_ctx: 8192 diff --git a/src/quant_tuner/recipes/q2_k_diffusiongemma_imatrix.yaml b/src/quant_tuner/recipes/q2_k_diffusiongemma_imatrix.yaml new file mode 100644 index 0000000..4cc832e --- /dev/null +++ b/src/quant_tuner/recipes/q2_k_diffusiongemma_imatrix.yaml @@ -0,0 +1,52 @@ +# DiffusionGemma 26B-A4B (Gemma-4 MoE block-diffusion) — hybrid_custom imatrix → Q2_K. +# +# The LOW-RESOURCE path for google/diffusiongemma-26B-A4B-it: no HF model +# load at all. llama-imatrix streams the F16 GGUF (CPU is fine), and the +# hybrid_custom re-weighting is closed-form — per-expert for the fused 3D +# `ffn_gate_up_exps` / `ffn_down_exps` stacks, standard for the dense tensors. +# Requires the llama.cpp pin with `diffusion-gemma` support (PR #24423). +# +# Why imatrix works here at all: llama-imatrix runs a standard forward, and +# DiffusionGemma's encoder mode IS a causal forward over the exact weights +# the diffusion decoder uses. Caveat: the UNIFIED graph treats the last +# `canvas_length` (256) positions of each chunk as a bidirectional canvas — +# with imatrix_ctx 2048 below, 87.5% of positions collect encoder-mode +# (causal) statistics and the rest cover the decoder's operating mode. +# Expert coverage is checked after llama-imatrix; raise train_max_tokens if +# the WARN fires (unobserved experts quantize with neutral importance). +# +# Same bench/eval caveats as iq2_m_diffusiongemma_awq: KLD is quant-vs-quant +# valid, absolute PPL is not deployment-faithful; decode tok/s from +# llama-bench is meaningless for diffusion generation; task eval needs the +# diffusion backend. +name: q2_k_diffusiongemma_imatrix + +model: google/diffusiongemma-26B-A4B-it + +workspace: ./out/q2_k_diffusiongemma_imatrix + +data: + logs: PLACEHOLDER + train_frac: 0.8 + test_frac: 0.1 + holdout_frac: 0.1 + # 128 experts x 8 active: expert coverage needs a larger corpus than dense + # models — the post-imatrix coverage check warns below 95%. + train_max_tokens: 500000 + eval_max_tokens: 30000 + supplement: ./calibration_supplement.txt + context_len: 16384 + +calibration: + method: imatrix + variant: hybrid_custom + params: + # 2048 keeps the canvas region (fixed 256) a small fraction of each chunk. + imatrix_ctx: 2048 + +quantize: + type: Q2_K + +bench: + suite: full + eval_ctx: 8192 diff --git a/tests/unit/test_awq.py b/tests/unit/test_awq.py index abbc96c..dfbea7c 100644 --- a/tests/unit/test_awq.py +++ b/tests/unit/test_awq.py @@ -468,3 +468,149 @@ def test_scale_group_prev_norm_optional(): prev_norm=None, ) assert g.prev_norm is None + + +# ----- DiffusionGemma / Gemma-4 MoE --------------------------------------- # + + +def _build_moe_model(*, n_experts: int = 4, hidden: int = 8, inter: int = 6, + diffusion_trunk: bool = False): + """Tiny stand-in with a Gemma-4 style MoE block: dense MLP + fused 3D + experts + router with its own internal norm (NOT a consumer of the + experts' pre-norm).""" + import torch.nn as nn + + class RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) # Gemma (1+γ) convention + + def forward(self, x): + v = x / x.norm(dim=-1, keepdim=True).clamp_min(1e-6) * (x.shape[-1] ** 0.5) + return v * (1.0 + self.weight) + + class Router(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(hidden, n_experts, bias=False) + + def forward(self, x): + # internal norm omitted: the test only needs "router input is the + # raw residual", which the Block forward provides + return self.proj(x) + + class Experts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.randn(n_experts, 2 * inter, hidden)) + self.down_proj = nn.Parameter(torch.randn(n_experts, hidden, inter)) + + def forward(self, x): + gu = torch.einsum("eoh,th->teo", self.gate_up_proj, x) + gate, up = gu.chunk(2, dim=-1) + h = torch.nn.functional.silu(gate) * up + return torch.einsum("ehi,tei->teh", self.down_proj, h).mean(dim=1) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(hidden, hidden, bias=False) + self.self_attn.k_proj = nn.Linear(hidden, hidden, bias=False) + self.self_attn.v_proj = nn.Linear(hidden, hidden, bias=False) + self.input_layernorm = RMSNorm(hidden) + self.mlp = nn.Module() + self.mlp.gate_proj = nn.Linear(hidden, inter, bias=False) + self.mlp.up_proj = nn.Linear(hidden, inter, bias=False) + self.pre_feedforward_layernorm = RMSNorm(hidden) + self.post_attention_layernorm = RMSNorm(hidden) + self.router = Router() + self.experts = Experts() + self.pre_feedforward_layernorm_2 = RMSNorm(hidden) + + class M(nn.Module): + def __init__(self): + super().__init__() + if diffusion_trunk: + # DiffusionGemma shape: model.encoder.language_model.layers + self.model = nn.Module() + self.model.encoder = nn.Module() + self.model.encoder.language_model = nn.Module() + self.model.encoder.language_model.layers = nn.ModuleList( + [Block() for _ in range(2)] + ) + else: + self.model = nn.Module() + self.model.layers = nn.ModuleList([Block() for _ in range(2)]) + + return M() + + +def test_discover_groups_emits_moe_group(): + m = _build_moe_model() + groups = {g.group_id: g for g in discover_groups(m)} + assert "L0_moe" in groups and "L1_moe" in groups + g = groups["L0_moe"] + # anchor = the experts module (sees every token's normed hidden state) + assert g.anchor == "model.layers.0.experts" + # the fused 3D gate_up stack is the only member: the router reads the raw + # residual through its own norm, so it must NOT be scaled + assert g.members == ("model.layers.0.experts.gate_up_proj",) + assert g.prev_norm == "model.layers.0.pre_feedforward_layernorm_2" + # the dense MLP group coexists in the same layer + assert "L0_mlp" in groups + assert groups["L0_mlp"].prev_norm == "model.layers.0.pre_feedforward_layernorm" + + +def test_discover_groups_diffusiongemma_trunk(): + m = _build_moe_model(diffusion_trunk=True) + groups = {g.group_id: g for g in discover_groups(m)} + assert groups["L0_moe"].members == ( + "model.encoder.language_model.layers.0.experts.gate_up_proj", + ) + assert groups["L0_attn"].anchor == ( + "model.encoder.language_model.layers.0.self_attn.q_proj" + ) + + +def test_moe_fold_preserves_outputs_and_routing_in_f32(): + """Folding 1/s into pre_feedforward_layernorm_2 ((1+γ) convention) and s + into every expert's input channels leaves (a) each expert's output and + (b) the routing logits exactly invariant in f32.""" + torch.manual_seed(0) + m = _build_moe_model().float() + layer = m.model.layers[0] + x = torch.randn(5, 8) + + norm_out = layer.pre_feedforward_layernorm_2(x) + ref_expert_out = layer.experts(norm_out) + ref_routing = layer.router(x) # router input: raw residual, NOT norm_out + + s = torch.rand(8) + 0.5 + from quant_tuner.calibrate.awq import _in_channel_view, _member_weight + + W = _member_weight(m, "model.layers.0.experts.gate_up_proj") + W.data.mul_(_in_channel_view(s, W)) + gamma = layer.pre_feedforward_layernorm_2.weight + gamma.data.copy_(fold_rmsnorm_gain(gamma.data, 1.0 / s, plus_one=True)) + + new_expert_out = layer.experts(layer.pre_feedforward_layernorm_2(x)) + new_routing = layer.router(x) + + torch.testing.assert_close(new_expert_out, ref_expert_out, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(new_routing, ref_routing) # untouched, exact + + +def test_scoring_weight_subsamples_experts_deterministically(): + from quant_tuner.calibrate.awq import _scoring_weight + + m = _build_moe_model(n_experts=4, hidden=8, inter=6) + member = "model.layers.0.experts.gate_up_proj" + full = _scoring_weight(m, member, 0) + assert full.shape == (4 * 12, 8) # all experts flattened to rows + sub = _scoring_weight(m, member, 2) + assert sub.shape == (2 * 12, 8) + torch.testing.assert_close(sub, _scoring_weight(m, member, 2)) # deterministic + # 2D linears pass through unchanged + lin = _scoring_weight(m, "model.layers.0.mlp.gate_proj", 2) + assert lin.shape == (6, 8) diff --git a/tests/unit/test_diffusion_server.py b/tests/unit/test_diffusion_server.py new file mode 100644 index 0000000..795ebf4 --- /dev/null +++ b/tests/unit/test_diffusion_server.py @@ -0,0 +1,95 @@ +"""Unit tests for the DiffusionGemma eval backend's pure pieces (no model, +no subprocess): Gemma-4 tool-call markup parsing, canvas trimming, and the +entropy-bound acceptance rule.""" + +from __future__ import annotations + +import json + +import numpy as np + +from quant_tuner.eval.diffusion_server import ( + EbParams, + _parse_gemma4_args, + parse_gemma4_response, + trim_canvas, +) + + +def test_parse_gemma4_single_tool_call(): + text = '<|tool_call>call:read_file{"path": "/etc/hosts"}' + content, calls = parse_gemma4_response(text) + assert content == "" + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "read_file" + assert json.loads(calls[0]["function"]["arguments"]) == {"path": "/etc/hosts"} + + +def test_parse_gemma4_bare_keys_and_parallel_calls(): + text = ( + "Let me check.\n" + "<|tool_call>call:search{query: \"llama\", limit: 3}" + "<|tool_call>call:fetch{url: \"http://x\"}" + ) + content, calls = parse_gemma4_response(text) + assert content == "Let me check." + assert [c["function"]["name"] for c in calls] == ["search", "fetch"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "llama", "limit": 3} + + +def test_parse_gemma4_strips_reasoning_channel(): + text = ( + "<|channel>thinking hard about itThe answer is B." + "" + ) + content, calls = parse_gemma4_response(text) + assert content == "The answer is B." + assert calls == [] + + +def test_parse_gemma4_unparseable_args_kept_raw(): + text = "<|tool_call>call:run{not json at all" + content, calls = parse_gemma4_response(text) + # no closing brace before the end marker: regex requires {...}, so this is + # treated as content (the eval scores it as a malformed/absent call) + assert calls == [] or "_raw" in json.loads(calls[0]["function"]["arguments"]) + + +def test_trim_canvas_cuts_at_eog(): + eog = frozenset({7}) + assert trim_canvas([1, 2, 3, 7, 4, 5], eog) == 3 + assert trim_canvas([1, 2, 3], eog) == 3 # no EOG -> full length + + +def test_trim_canvas_cuts_repetition_loop(): + # >= 6 consecutive stride-1 repeats + toks = [1, 2] + [9] * 9 + [3] + cut = trim_canvas(toks, frozenset()) + assert cut == 2 + + +def test_entropy_bound_acceptance_rule(): + """Mirror of the C++ rule: accept positions in ascending-entropy order + while the sum of strictly-earlier entropies stays within the bound.""" + entropy = np.array([0.5, 0.01, 0.02, 9.0]) + bound = 0.04 + order = np.argsort(entropy, kind="stable") + cum_e = np.concatenate([[0.0], entropy[order].cumsum()[:-1]]) + accepted = np.zeros(4, dtype=bool) + accepted[order[cum_e <= bound]] = True + # 0.01 (cum 0), 0.02 (cum 0.01), 0.5 (cum 0.03 <= 0.04!), 9.0 (cum 0.53 > bound) + assert accepted.tolist() == [True, True, True, False] + + +def test_eb_params_defaults_match_llama_cpp_reference(): + eb = EbParams() + assert (eb.max_steps, eb.t_min, eb.t_max) == (48, 0.4, 0.8) + assert (eb.entropy_bound, eb.stability_threshold, eb.confidence_threshold) == ( + 0.1, 1, 0.005, + ) + + +def test_parse_args_fallback_chain(): + assert _parse_gemma4_args('{"a": 1}') == {"a": 1} + assert _parse_gemma4_args('{a: 1, b-c: "x"}') == {"a": 1, "b-c": "x"} + assert "_raw" in _parse_gemma4_args("{a: }") diff --git a/tests/unit/test_imatrix.py b/tests/unit/test_imatrix.py index 130b33c..70e94d4 100644 --- a/tests/unit/test_imatrix.py +++ b/tests/unit/test_imatrix.py @@ -3,17 +3,20 @@ from pathlib import Path import numpy as np +import pytest from quant_tuner.calibrate import imatrix as imx from quant_tuner.calibrate.imatrix import ( ForwardStats, _col_l2_sq, + _combine_moe, _l1_normalize, build_analytic, + build_hybrid_custom, build_outlier_l4, ) from quant_tuner.models import llama_cpp -from quant_tuner.models.hf_gguf_map import is_ssm, map_hf_to_gguf +from quant_tuner.models.hf_gguf_map import is_moe_expert, is_ssm, map_hf_to_gguf def test_col_l2_sq_sums_over_rows(): @@ -60,6 +63,33 @@ def test_is_ssm(): assert not is_ssm("blk.3.ffn_down.weight") +def test_map_hf_to_gguf_moe_and_router(): + # Gemma-4 style fused 3D experts + router (a plain Linear). + assert map_hf_to_gguf("model.layers.5.experts.gate_up_proj") == "blk.5.ffn_gate_up_exps.weight" + assert map_hf_to_gguf("model.layers.5.experts.down_proj") == "blk.5.ffn_down_exps.weight" + assert map_hf_to_gguf("model.layers.5.mlp.experts.gate_up_proj") == "blk.5.ffn_gate_up_exps.weight" + assert map_hf_to_gguf("model.layers.5.router.proj") == "blk.5.ffn_gate_inp.weight" + + +def test_map_hf_to_gguf_diffusiongemma_trunk_prefixes(): + # DiffusionGemma: backbone under model.decoder.*; encoder stack is tied. + assert map_hf_to_gguf("model.decoder.layers.2.self_attn.q_proj") == "blk.2.attn_q.weight" + assert map_hf_to_gguf("model.decoder.layers.2.experts.gate_up_proj") == "blk.2.ffn_gate_up_exps.weight" + assert ( + map_hf_to_gguf("model.encoder.language_model.layers.2.mlp.up_proj") + == "blk.2.ffn_up.weight" + ) + assert map_hf_to_gguf("model.decoder.self_conditioning.gate_proj") == "self_cond_gate.weight" + + +def test_is_moe_expert(): + assert is_moe_expert("blk.0.ffn_gate_up_exps.weight") + assert is_moe_expert("blk.0.ffn_down_exps.weight") + assert is_moe_expert("blk.0.ffn_gate_exps.weight") + assert not is_moe_expert("blk.0.ffn_gate.weight") + assert not is_moe_expert("blk.0.ffn_gate_inp.weight") + + def test_forward_stats_roundtrip(tmp_path): stats = ForwardStats( e_a4={ @@ -129,6 +159,109 @@ def test_build_outlier_l4_falls_back_when_stats_missing(monkeypatch): np.testing.assert_allclose(out["blk.0.ssm_dt.weight"], [7.0, 7.0]) # SSM passthrough +def test_build_analytic_moe_per_expert(monkeypatch): + """3D expert tensors are reranked per expert and flattened expert-major.""" + tname = "blk.0.ffn_gate_up_exps.weight" + # 2 experts, n_out=2, n_in=3 + w = np.array( + [ + [[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]], # expert 0 col ssq: [1, 1, 2] + [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0]], # expert 1 col ssq: [4, 9, 0] + ], + dtype=np.float32, + ) + ea2 = np.array([[1.0, 4.0, 9.0], [2.0, 2.0, 2.0]], dtype=np.float32) + monkeypatch.setattr(imx, "_load_base_imatrix", lambda _p: {tname: ea2}) + monkeypatch.setattr(imx, "_load_weights", lambda _p: {tname: w}) + + out = build_analytic("dummy_f16", "dummy_base") + + # expert 0: [1*1, 1*4, 2*9]; expert 1: [4*2, 9*2, 0*2], expert-major flat + np.testing.assert_allclose(out[tname], [1.0, 4.0, 18.0, 8.0, 18.0, 0.0]) + + +def test_build_hybrid_custom_moe_normalizes_per_expert(monkeypatch): + """L1 normalization must not bleed across experts with different scales.""" + tname = "blk.0.ffn_down_exps.weight" + # Identical structure, expert 1's activations scaled 1000x: per-expert + # normalization makes both experts' outputs identical. + w = np.stack([np.eye(3, dtype=np.float32)] * 2) + ea2 = np.array([[1.0, 2.0, 3.0], [1000.0, 2000.0, 3000.0]], dtype=np.float32) + monkeypatch.setattr(imx, "_load_base_imatrix", lambda _p: {tname: ea2}) + monkeypatch.setattr(imx, "_load_weights", lambda _p: {tname: w}) + + out = build_hybrid_custom("dummy_f16", "dummy_base") + flat = out[tname] + np.testing.assert_allclose(flat[:3], flat[3:], rtol=1e-6) + + +def test_combine_moe_rejects_misaligned_shapes(): + w = np.zeros((2, 4, 3), dtype=np.float32) + ea2 = np.zeros((3, 3), dtype=np.float32) # wrong expert count + with pytest.raises(ValueError, match="does not line up"): + _combine_moe(w, ea2, lambda a, b: a * b, "blk.0.ffn_gate_up_exps.weight") + + +def test_build_outlier_l4_moe_passthrough(monkeypatch): + """Expert tensors pass through per-expert E[a^2] in outlier variants.""" + tname = "blk.0.ffn_gate_up_exps.weight" + ea2 = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + monkeypatch.setattr(imx, "_load_base_imatrix", lambda _p: {tname: ea2}) + + out = build_outlier_l4("dummy_base", ForwardStats(e_a4={}, max_abs={})) + np.testing.assert_allclose(out[tname], [1.0, 2.0, 3.0, 4.0]) + + +def _write_moe_imatrix_gguf(path): + """Write a minimal llama-imatrix-layout GGUF: one dense entry (scalar count) + and one stacked-expert entry (per-expert counts, one expert unobserved).""" + imx._ensure_gguf_py() + gguf = pytest.importorskip("gguf") + + w = gguf.GGUFWriter(str(path), "imatrix") + w.add_uint32("imatrix.chunk_count", 2) + w.add_uint32("imatrix.chunk_size", 512) + w.add_array("imatrix.datasets", ["test"]) + # dense: 1-D in_sum2, scalar count + w.add_tensor("blk.0.attn_q.weight.in_sum2", np.array([2.0, 4.0, 8.0], dtype=np.float32)) + w.add_tensor("blk.0.attn_q.weight.counts", np.array([2.0], dtype=np.float32)) + # MoE: [n_expert=3, n_in=2] sums, per-expert counts with expert 2 unseen + sums = np.array([[2.0, 4.0], [30.0, 60.0], [0.0, 0.0]], dtype=np.float32) + w.add_tensor("blk.0.ffn_gate_up_exps.weight.in_sum2", sums) + w.add_tensor( + "blk.0.ffn_gate_up_exps.weight.counts", + np.array([[2.0], [10.0], [0.0]], dtype=np.float32), + ) + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + + +def test_load_base_imatrix_moe_per_expert_counts(tmp_path): + path = tmp_path / "imatrix.gguf" + _write_moe_imatrix_gguf(path) + + base = imx._load_base_imatrix(path) + + np.testing.assert_allclose(base["blk.0.attn_q.weight"], [1.0, 2.0, 4.0]) + moe = base["blk.0.ffn_gate_up_exps.weight"] + assert moe.shape == (3, 2) + np.testing.assert_allclose(moe[0], [1.0, 2.0]) # / count 2 + np.testing.assert_allclose(moe[1], [3.0, 6.0]) # / count 10 + np.testing.assert_allclose(moe[2], [1.0, 1.0]) # unseen -> neutral ones + + +def test_check_moe_coverage(tmp_path): + path = tmp_path / "imatrix.gguf" + _write_moe_imatrix_gguf(path) + + cov = imx.check_moe_coverage(path) + # dense entries (single count) are not reported + assert list(cov) == ["blk.0.ffn_gate_up_exps.weight"] + np.testing.assert_allclose(cov["blk.0.ffn_gate_up_exps.weight"], 2 / 3) + + def _capture_imatrix_cmd(monkeypatch, **kwargs) -> list[str]: """Run llama_cpp.imatrix with run()/llama_bin stubbed; return the built cmd as strings.""" captured: dict[str, list] = {} diff --git a/vendor/llama.cpp b/vendor/llama.cpp index 9e58d4d..c84e85a 160000 --- a/vendor/llama.cpp +++ b/vendor/llama.cpp @@ -1 +1 @@ -Subproject commit 9e58d4d692ed3d350591cc86d06c73c61c122509 +Subproject commit c84e85af61011f9fbfcf41479381d5ed1661a564