From 71ce75e411b05d6a6ba76df223706f29d7ce4ef5 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 09:19:25 -0400 Subject: [PATCH 01/34] =?UTF-8?q?feat:=20fq=5Fassemble=5Flora=20=E2=80=94?= =?UTF-8?q?=20MSRT=20cartridge=20encoder=20for=20LoRA-compatible=20adapter?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New tool that encodes BF16 weights into: 1. Base EXL3 checkpoint (K2 or K3 trellis, standard format) 2. Cartridge adapters (residual trellis stages as LoRA-compatible safetensors) The cartridge contains full-rank trellis-quantized residual weights with per-stage rescaling factors. At runtime, the vLLM EXL3 LoRA wrapper applies them by running additional exl3_gemm passes and summing with rescaling. Key design decisions (from MSRT research v35-v52): - Base K2 (2bpw) is the lowest viable base (K1 too lossy) - K1trsc cartridge on all experts → K3-equivalent (3bpw) - K2trsc cartridge on hot experts → K4-equivalent (4bpw) - Rescaling (codebook_scale/RMS) is the key innovation (v35 breakthrough) - All stages share the same Hadamard vectors (suh/svh) as the base Recipe format: fq-cartridge/1 with per-stage K, label, and expert selection Output: base/ (standard EXL3) + cartridges/ (LoRA-format safetensors) Includes Fruit SIQ recipe (K2 + K1trsc all + K2trsc 96 hot) matching the SIQ quant's 160 K3 + 96 K4 tier allocation. Co-authored-by: Claude --- recipes/fruit-k2-k3k4-cart.json | 27 ++ tests/test_fq_assemble_lora.py | 132 ++++++ tools/fq_assemble_lora.py | 733 ++++++++++++++++++++++++++++++++ 3 files changed, 892 insertions(+) create mode 100644 recipes/fruit-k2-k3k4-cart.json create mode 100755 tests/test_fq_assemble_lora.py create mode 100755 tools/fq_assemble_lora.py diff --git a/recipes/fruit-k2-k3k4-cart.json b/recipes/fruit-k2-k3k4-cart.json new file mode 100644 index 0000000..5d34823 --- /dev/null +++ b/recipes/fruit-k2-k3k4-cart.json @@ -0,0 +1,27 @@ +{ + "schema": "fq-cartridge/1", + "base_k": 2, + "stages": [ + { + "k": 1, + "label": "res1", + "experts": "all", + "description": "K1 rescaled trellis residual on ALL experts → K3-equivalent (3bpw)" + }, + { + "k": 2, + "label": "res2", + "experts": "hot96", + "description": "K2 rescaled trellis residual on 96 hot experts → K4-equivalent (4bpw)" + } + ], + "moe_layers": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], + "hot_experts": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, + 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, + 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, + 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, + 86, 87, 88, 89, 90, 91, 92, 93, 94, 95], + "description": "Fruit SIQ proxy: K2 base + K1trsc (all) + K2trsc (96 hot) = dual-cartridge matching SIQ 3.375bpw" +} diff --git a/tests/test_fq_assemble_lora.py b/tests/test_fq_assemble_lora.py new file mode 100755 index 0000000..6c33065 --- /dev/null +++ b/tests/test_fq_assemble_lora.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Tests for fq_assemble_lora. + +Tests the cartridge recipe parsing, MSRT encoding pipeline, and output format. +Uses tiny random tensors (not real model weights) for speed. +""" +import json +import math +import sys +import tempfile +from pathlib import Path + +import pytest +import torch + +# Add tools to path +sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) + + +def test_cartridge_recipe_schema(): + """Test that cartridge recipe follows fq-cartridge/1 schema.""" + recipe = { + "schema": "fq-cartridge/1", + "base_k": 2, + "stages": [ + {"k": 1, "label": "res1", "experts": "all"}, + {"k": 2, "label": "res2", "experts": [0, 1, 2]}, + ], + "moe_layers": [3, 4, 5], + } + assert recipe["schema"] == "fq-cartridge/1" + assert recipe["base_k"] == 2 + assert len(recipe["stages"]) == 2 + assert recipe["stages"][0]["k"] == 1 + assert recipe["stages"][1]["experts"] == [0, 1, 2] + + +def test_cartridge_recipe_from_file(): + """Test loading the Fruit recipe.""" + recipe_path = Path(__file__).parent.parent / "recipes" / "fruit-k2-k3k4-cart.json" + if not recipe_path.exists(): + pytest.skip("Recipe file not found") + recipe = json.loads(recipe_path.read_text()) + assert recipe["schema"] == "fq-cartridge/1" + assert recipe["base_k"] == 2 + assert len(recipe["stages"]) == 2 + assert recipe["stages"][0]["label"] == "res1" + assert recipe["stages"][1]["label"] == "res2" + assert recipe["stages"][0]["experts"] == "all" + assert len(recipe["moe_layers"]) == 11 # layers 3-13 + + +def test_block_rms(): + """Test RMS computation.""" + from fq_assemble_lora import block_rms + x = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + rms = block_rms(x, dim=0, keepdim=True) + expected = torch.sqrt(torch.tensor([(1+9)/2, (4+16)/2])) + assert torch.allclose(rms.squeeze(), expected, rtol=1e-5) + + +def test_rescaled_trellis_scale(): + """Test that rescaling produces correct scale factor.""" + # Mock: just test the scale computation logic + cbs = 1.2437 + residual = torch.randn(128, 256) + residual_rms = residual.square().mean().sqrt().item() + scale = abs(cbs) / residual_rms + assert scale > 0 + # After rescaling, RMS should be ~|cbs| + scaled = residual * scale + scaled_rms = scaled.square().mean().sqrt().item() + assert abs(scaled_rms - abs(cbs)) < 0.01 # close to codebook scale + + +def test_trellis_packed_shape(): + """Test that packed trellis has correct shape (without GPU).""" + # Shape: (k//16, n//16, K*16) int16 + k, n, K = 128, 256, 2 + expected_shape = (k // 16, n // 16, K * 16) + assert expected_shape == (8, 16, 32) + # For K=3: (8, 16, 48) + assert (k // 16, n // 16, 3 * 16) == (8, 16, 48) + + +def test_cartridge_adapter_naming(): + """Test that cartridge tensor names follow the expected pattern.""" + layer, exp, proj, rank, label = 3, 0, "gate_proj", 0, "res1" + prefix = f"model.layers.{layer}.mlp.experts.{exp}.{proj}.rank{rank}" + names = [ + f"{prefix}.trellis_{label}", + f"{prefix}.suh_{label}", + f"{prefix}.svh_{label}", + f"{prefix}.scale_{label}", + ] + for name in names: + assert f"trellis_{label}" in name or f"suh_{label}" in name \ + or f"svh_{label}" in name or f"scale_{label}" in name + + +def test_msrt_bpw_calculation(): + """Test effective bpw calculation for dual-cartridge configs.""" + n_experts = 256 + base_k = 2 + # Stage 1: K1 on all 256 experts + # Stage 2: K2 on 96 experts + total_bits = n_experts * base_k + n_experts * 1 + 96 * 2 + eff_bpw = total_bits / n_experts + assert eff_bpw == (512 + 256 + 192) / 256 # = 960/256 = 3.75 + + # willfalco comparison: 148 K3 + 108 K4 + willfalco_bpw = (148 * 3 + 108 * 4) / 256 + assert abs(willfalco_bpw - 3.422) < 0.01 + + +def test_encoding_summary_format(): + """Test that encoding summary has required fields.""" + summary = { + "tool": "fq_assemble_lora/1", + "base_k": 2, + "stages": [{"k": 1, "label": "res1", "experts": "all"}], + "moe_layers": [3, 4, 5], + "overall_mse": 0.001, + "n_experts_encoded": 768, + } + assert summary["tool"] == "fq_assemble_lora/1" + assert "overall_mse" in summary + assert summary["n_experts_encoded"] == 768 # 256 experts × 3 layers + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py new file mode 100755 index 0000000..f11a25a --- /dev/null +++ b/tools/fq_assemble_lora.py @@ -0,0 +1,733 @@ +#!/usr/bin/env python3 +"""fq_assemble_lora — Encode MSRT residual cartridges as LoRA-compatible adapters. + +Given a BF16 (or EXL3) source checkpoint and a cartridge recipe, this tool: + 1. Quantizes the base tier (K2 or K3) into standard EXL3 trellis format + 2. Computes residuals, rescales, and quantizes each residual stage + 3. Emits two outputs: + - A base checkpoint (standard EXL3 safetensors, loads normally in vLLM) + - One or more cartridge adapters (safetensors with per-stage trellis + tensors, loadable as LoRA adapters via vLLM's add_lora API) + +The cartridge adapter is NOT a low-rank LoRA — it contains full-rank trellis- +quantized residual weights. The vLLM EXL3 LoRA wrapper (Exl3LoRAMoMethod) +applies them by running additional exl3_gemm passes and summing with rescaling. + +MSRT (Multi-Stage Rescaled Trellis) is described in: + research/fungible-quant/poc/V50-LOW-BITRATE-MSRT.md + research/fungible-quant/MSRT-CARTRIDGE-FEASIBILITY-AND-PLAN.md + +Cartridge recipe format (fq-cartridge/1): + + { + "schema": "fq-cartridge/1", + "base_k": 2, + "stages": [ + {"k": 1, "label": "res1", "experts": "all"}, + {"k": 2, "label": "res2", "experts": [0, 1, 10, 11, ...]} + ], + "moe_layers": [3, 4, 5, ...] + } + +Usage: + python tools/fq_assemble_lora.py encode \\ + --source ./bf16-checkpoint \\ + --recipe recipes/fruit-k2-k3k4-cart.json \\ + --out ./output \\ + --encoder-source /opt/fruit-pip/exllamav3 + +The tool requires the exllamav3 encoder (quantize_tiles, codebook_scale, +Hadamard transform) to perform trellis quantization. It does NOT require +vLLM or flash_attn — the encoder is loaded via the same bootstrap pattern +used in the PoC scripts. +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import struct +import sys +import time +from pathlib import Path +from typing import Any + +import torch + +# ── Constants ────────────────────────────────────────────────────────────── + +TOOL_VERSION = "fq_assemble_lora/1" +CARTRIDGE_SCHEMA = "fq-cartridge/1" +ADAPTER_CONFIG_SCHEMA = "fq-cartridge-adapter/1" +HADAMARD_BLOCK = 128 + +# Expert tensor name pattern in source checkpoints +EXPERT_RE_PATTERN = ( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)(?:\.rank(\d+))?$" +) + +PROJECTIONS = ("gate_proj", "up_proj", "down_proj") + + +# ── EXL3 Encoder Bootstrap ──────────────────────────────────────────────── + +def bootstrap_encoder(encoder_source: str) -> tuple[Any, ...]: + """Load the EXL3 encoder without importing all of exllamav3. + + Returns (ext, get_hadamard_dt, tensor_core_perm, tensor_core_perm_i, + quantize_tiles, codebook_scale). + """ + import importlib.util + import types + + pkg_root = Path(encoder_source) + pkg = types.ModuleType("exllamav3") + pkg.__path__ = [str(pkg_root)] + sys.modules["exllamav3"] = pkg + + for sub in ["util", "modules", "modules.quant", "modules.quant.exl3_lib"]: + full = f"exllamav3.{sub}" + m = types.ModuleType(full) + m.__path__ = [str(pkg_root / sub.replace(".", "/"))] + sys.modules[full] = m + + # Stub progress/memory to avoid flash_attn dependency + _stub = types.ModuleType("exllamav3.util.progress") + class _DPB: + def __init__(self, *a, **kw): pass + def __enter__(self): return self + def __exit__(self, *a): return False + def update(self, *a): pass + def new_task(self, *a, **kw): pass + _stub.ProgressBar = _DPB + sys.modules["exllamav3.util.progress"] = _stub + + _stub = types.ModuleType("exllamav3.util.memory") + _stub.free_mem = lambda: None + _stub.list_gpu_tensors = lambda: [] + sys.modules["exllamav3.util.memory"] = _stub + + _stub = types.ModuleType("exllamav3.util") + _stub.__path__ = [str(pkg_root / "util")] + _stub.cuda_sync_active = lambda *a, **kw: torch.cuda.synchronize() + sys.modules["exllamav3.util"] = _stub + + _stub = types.ModuleType("exllamav3.util.tensor") + _stub.save_tensor_image = lambda *a, **kw: None + sys.modules["exllamav3.util.tensor"] = _stub + + # Load the extension + spec = importlib.util.spec_from_file_location( + "exllamav3.ext", str(pkg_root / "ext.py")) + m = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.ext"] = m + spec.loader.exec_module(m) + + # Load Hadamard + spec = importlib.util.spec_from_file_location( + "exllamav3.util.hadamard", str(pkg_root / "util" / "hadamard.py")) + m = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.util.hadamard"] = m + spec.loader.exec_module(m) + + # Load quantize module + quant_path = pkg_root / "modules" / "quant" / "exl3_lib" / "quantize.py" + spec = importlib.util.spec_from_file_location( + "exllamav3.modules.quant.exl3_lib.quantize", str(quant_path)) + m = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = m + spec.loader.exec_module(m) + + return (m.exllamav3_ext, m.get_hadamard_dt, + m.tensor_core_perm, m.tensor_core_perm_i, + m.quantize_tiles, m.codebook_scale) + + +# ── Quantization Primitives ──────────────────────────────────────────────── + +def block_rms(x: torch.Tensor, dim: int, keepdim: bool = False) -> torch.Tensor: + """RMS along a dimension.""" + return x.square().mean(dim=dim, keepdim=keepdim).sqrt() + + +def regularize( + w: torch.Tensor, + device: torch.device, + ghd: Any, + cbs: float, + had_k: int = HADAMARD_BLOCK, + had_n: int = HADAMARD_BLOCK, + seed: int = 0, +) -> torch.Tensor: + """Apply Hadamard regularization (in-place transform, returns new tensor). + + This matches the EXL3 regularize() used in the PoC scripts (v35-v52). + The Hadamard is orthogonal, so MSE in regularized space = MSE in original. + """ + k, n = w.shape + g = torch.Generator(device="cpu").manual_seed(seed) + su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) + sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) + + out_scales = block_rms(w, dim=0, keepdim=True) + mean = out_scales.mean().item() + if mean > 1e-30: + out_scales = out_scales / mean + sv = (sv * out_scales + 1e-10).float() + w = (w / sv).contiguous() + + had_n_mat = ghd(had_n, device, torch.float, 1.0 / math.sqrt(had_n)) + w = (w.view(k, n // had_n, had_n) @ had_n_mat).view(k, n).contiguous() + + in_scales = block_rms(w, dim=1, keepdim=True).clamp(min=1e-30) + su = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() + w = (w / su).contiguous() + + had_k_mat = ghd(had_k, device, torch.float, 1.0 / math.sqrt(had_k)) + w = (had_k_mat @ w.view(k // had_k, had_k, n)).view(k, n).contiguous() + return w + + +def quantize_trellis( + data: torch.Tensor, + K: int, + device: torch.device, + tcp: Any, + tcpi: Any, + qtf: Any, +) -> torch.Tensor: + """Quantize a 2D tensor with EXL3 trellis at K bits. + + Returns the dequantized (reconstructed) tensor, NOT the packed indices. + The trellis tiles are 16×16, processed row-by-row in blocks of 16. + """ + k, n = data.shape + tiles_n = n // 16 + weight_q = torch.zeros_like(data) + qa = {"K": K, "mcg": True} + perm = tcp(device) + perm_i = tcpi(device) + + for bi in range(0, k, 16): + rows = data[bi:bi + 16] + tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) + tiles = tiles[:, perm].contiguous() + quant_w, _ = qtf(tiles, qa) + quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) + weight_q[bi:bi + 16] = quant_w + + return weight_q + + +def quantize_trellis_packed( + data: torch.Tensor, + K: int, + device: torch.device, + tcp: Any, + tcpi: Any, + qtf: Any, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize and return BOTH reconstructed values and packed trellis indices. + + Returns (reconstructed_float, packed_trellis_int16) where packed_trellis + has shape (k // 16, n // 16, K * 16) dtype int16 — the EXL3 storage format. + """ + k, n = data.shape + tiles_n = n // 16 + weight_q = torch.zeros_like(data) + # Packed trellis: (tiles_k, tiles_n, K*16) int16 + packed = torch.zeros(k // 16, tiles_n, K * 16, dtype=torch.int16, device=device) + qa = {"K": K, "mcg": True} + perm = tcp(device) + perm_i = tcpi(device) + + for bi in range(0, k, 16): + tk = bi // 16 + rows = data[bi:bi + 16] + tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) + tiles = tiles[:, perm].contiguous() + quant_w, quant_idx = qtf(tiles, qa) + # Store packed indices + packed[tk] = quant_idx.reshape(tiles_n, K * 16) + # Reconstruct + quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) + weight_q[bi:bi + 16] = quant_w + + return weight_q, packed + + +def rescaled_trellis_quantize( + base_q: torch.Tensor, + residual: torch.Tensor, + K_res: int, + device: torch.device, + tcp: Any, tcpi: Any, qtf: Any, + cbs: float, +) -> tuple[torch.Tensor, torch.Tensor, float]: + """Rescale residual to match codebook range, quantize, return (recon, packed, scale). + + The rescaling is the key MSRT innovation (v35 breakthrough): + scale = |codebook_scale| / RMS(residual) + quantized = trellis(residual * scale) / scale + """ + residual_rms = residual.square().mean().sqrt().item() + if residual_rms < 1e-12: + return base_q, torch.zeros( + residual.shape[0] // 16, residual.shape[1] // 16, K_res * 16, + dtype=torch.int16, device=device), 1.0 + + scale = abs(cbs) / residual_rms + scaled = residual * scale + recon_packed, packed = quantize_trellis_packed(scaled, K_res, device, tcp, tcpi, qtf) + recon = base_q + recon_packed / scale + return recon, packed, scale + + +# ── Hadamard Vectors ────────────────────────────────────────────────────── + +def compute_hadamard_vectors( + w: torch.Tensor, + device: torch.device, + ghd: Any, + cbs: float, + seed: int = 0, +) -> dict[str, torch.Tensor]: + """Compute the suh and svh vectors that EXL3 stores alongside trellis. + + These are the per-block sign vectors and scales from regularize(). + The exact format must match what exl3_gemm expects. + """ + k, n = w.shape + g = torch.Generator(device="cpu").manual_seed(seed) + su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) + sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) + + out_scales = block_rms(w, dim=0, keepdim=True) + mean = out_scales.mean().item() + if mean > 1e-30: + out_scales = out_scales / mean + sv_full = (sv * out_scales + 1e-10).float() + + # After column Hadamard + w_col = (w / sv_full).contiguous() + had_n_mat = ghd(HADAMARD_BLOCK, device, torch.float, 1.0 / math.sqrt(HADAMARD_BLOCK)) + w_col = (w_col.view(k, n // HADAMARD_BLOCK, HADAMARD_BLOCK) @ had_n_mat).view(k, n).contiguous() + + in_scales = block_rms(w_col, dim=1, keepdim=True).clamp(min=1e-30) + su_full = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() + + return {"suh": su_full.squeeze().contiguous(), "svh": sv_full.squeeze().contiguous()} + + +# ── Encoding Pipeline ───────────────────────────────────────────────────── + +def encode_expert_msrt( + w_bf16: torch.Tensor, + base_k: int, + stages: list[dict[str, Any]], + device: torch.device, + ghd: Any, tcp: Any, tcpi: Any, qtf: Any, + cbs: float, +) -> dict[str, Any]: + """Encode one expert weight matrix with MSRT. + + Returns dict with: + - base: {trellis, suh, svh} for the base tier + - stages: list of {trellis, suh, svh, scale} for each residual stage + - mse: reconstruction MSE vs original + """ + w_reg = regularize(w_bf16, device, ghd, cbs) + + # Base tier + base_recon, base_packed = quantize_trellis_packed(w_reg, base_k, device, tcp, tcpi, qtf) + had_vectors = compute_hadamard_vectors(w_bf16, device, ghd, cbs) + + result = { + "base": { + "trellis": base_packed.cpu(), + "suh": had_vectors["suh"].cpu(), + "svh": had_vectors["svh"].cpu(), + }, + "stages": [], + } + + current_recon = base_recon + + for stage in stages: + residual = w_reg - current_recon + recon, packed, scale = rescaled_trellis_quantize( + current_recon, residual, stage["k"], device, tcp, tcpi, qtf, cbs) + result["stages"].append({ + "trellis": packed.cpu(), + "suh": had_vectors["suh"].cpu(), # Same Hadamard vectors + "svh": had_vectors["svh"].cpu(), + "scale": scale, + }) + current_recon = recon + + result["mse"] = (w_reg - current_recon).pow(2).mean().item() + return result + + +# ── Safetensors Output ──────────────────────────────────────────────────── + +def save_safetensors(tensors: dict[str, torch.Tensor], path: Path) -> None: + """Save tensors as safetensors file.""" + from safetensors.torch import save_file + path.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(path)) + + +def write_base_checkpoint( + source_dir: Path, + out_dir: Path, + layer_results: dict[int, dict[int, dict[str, Any]]], + moe_layers: list[int], + tp: int = 1, +) -> None: + """Write the base K checkpoint in standard EXL3 format. + + The base checkpoint has the same structure as a normal EXL3 quant: + - config.json with hybrid_tr3_tail + - tier_bitmap.json + - model-layer-*.safetensors with trellis, suh, svh, mcg tensors + """ + out_dir.mkdir(parents=True, exist_ok=True) + + # Copy non-layer files from source + import shutil + for f in source_dir.iterdir(): + if f.is_file() and not f.name.startswith("model-layer-"): + shutil.copy2(f, out_dir / f.name) + + # Write base tensors into layer shards + for layer in moe_layers: + if layer not in layer_results: + continue + tensors = {} + for exp_id, exp_data in layer_results[layer].items(): + base = exp_data["base"] + for proj_idx, proj in enumerate(PROJECTIONS): + rank = 0 # TP=1 for Fruit model + prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" + tensors[f"{prefix}.trellis"] = base["trellis"] + tensors[f"{prefix}.suh"] = base["suh"] + tensors[f"{prefix}.svh"] = base["svh"] + # mcg sentinel + import hashlib + mcg_val = 0xCBAC1FED + tensors[f"{prefix}.mcg"] = torch.tensor([mcg_val], dtype=torch.int32) + + shard_path = out_dir / f"model-layer-{layer:03d}.safetensors" + save_safetensors(tensors, shard_path) + print(f" Base layer {layer}: {len(tensors)} tensors -> {shard_path.name}", flush=True) + + +def write_cartridge_adapter( + out_dir: Path, + layer_results: dict[int, dict[int, dict[str, Any]]], + stage_idx: int, + stage_label: str, + moe_layers: list[int], + expert_filter: dict[int, list[int]] | None = None, +) -> Path: + """Write one cartridge stage as a LoRA-compatible safetensors adapter. + + If expert_filter is provided, only the specified experts per layer are included. + """ + out_dir.mkdir(parents=True, exist_ok=True) + tensors = {} + + for layer in moe_layers: + if layer not in layer_results: + continue + for exp_id, exp_data in layer_results[layer].items(): + if expert_filter and layer in expert_filter: + if exp_id not in expert_filter[layer]: + continue + if stage_idx >= len(exp_data["stages"]): + continue + stage = exp_data["stages"][stage_idx] + for proj in PROJECTIONS: + rank = 0 + prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" + tensors[f"{prefix}.trellis_{stage_label}"] = stage["trellis"] + tensors[f"{prefix}.suh_{stage_label}"] = stage["suh"] + tensors[f"{prefix}.svh_{stage_label}"] = stage["svh"] + tensors[f"{prefix}.scale_{stage_label}"] = torch.tensor([stage["scale"]], dtype=torch.float32) + + adapter_path = out_dir / f"cartridge_{stage_label}.safetensors" + save_safetensors(tensors, adapter_path) + + # Write adapter_config.json + config = { + "schema": ADAPTER_CONFIG_SCHEMA, + "stage_label": stage_label, + "stage_k": stage_idx, + "num_tensors": len(tensors), + "tool_version": TOOL_VERSION, + } + (out_dir / f"cartridge_{stage_label}_config.json").write_text( + json.dumps(config, indent=2) + "\n") + + print(f" Cartridge '{stage_label}': {len(tensors)} tensors -> {adapter_path.name}", flush=True) + return adapter_path + + +# ── Main Encode Command ─────────────────────────────────────────────────── + +def cmd_encode(args) -> int: + """Encode a BF16 checkpoint into base K + cartridge adapters.""" + device = torch.device(args.device) + print(f"Device: {device} GPU: {torch.cuda.get_device_name(0)}", flush=True) + + # Bootstrap encoder + ext, ghd, tcp, tcpi, qtf, cbs = bootstrap_encoder(args.encoder_source) + print(f"codebook_scale = {cbs}", flush=True) + + # Load recipe + recipe = json.loads(args.recipe.read_text()) + base_k = recipe["base_k"] + stages = recipe["stages"] + moe_layers = recipe.get("moe_layers", []) + if not moe_layers: + # Auto-detect from source config + cfg_path = args.source / "config.json" + if cfg_path.exists(): + cfg = json.loads(cfg_path.read_text()) + tail = cfg.get("hybrid_tr3_tail", {}) + if "moe_layers" in tail: + moe_layers = list(range(tail["moe_layers"][0], tail["moe_layers"][1] + 1)) + else: + moe_layers = list(range( + cfg.get("first_k_dense_replace", 3), + cfg.get("num_hidden_layers", 13))) + print(f"Base K={base_k}, {len(stages)} cartridge stages, " + f"MoE layers {moe_layers[0]}-{moe_layers[-1]} ({len(moe_layers)} layers)", flush=True) + + # Parse expert filters + expert_filters = [] + for stage in stages: + if stage["experts"] == "all": + expert_filters.append(None) + else: + # Same expert list for all layers (could be per-layer in future) + expert_filters.append({l: stage["experts"] for l in moe_layers}) + + # Load source weights and encode + from safetensors import safe_open + layer_results: dict[int, dict[int, dict[str, Any]]] = {} + total_mse = 0.0 + n_experts_total = 0 + + for layer in moe_layers: + shard_path = args.source / f"model-layer-{layer:03d}.safetensors" + if not shard_path.exists(): + # Try other shard naming patterns + shard_path = args.source / f"model-layer-{layer:04d}.safetensors" + if not shard_path.exists(): + print(f" Layer {layer}: shard not found, skipping", flush=True) + continue + + print(f"\nEncoding layer {layer}...", flush=True) + layer_results[layer] = {} + + # Load expert weights from shard + with safe_open(str(shard_path), framework="pt") as f: + keys = list(f.keys()) + # Find expert keys for this layer + import re + expert_pattern = re.compile(EXPERT_RE_PATTERN) + expert_weights: dict[int, dict[str, torch.Tensor]] = {} + for key in keys: + m = expert_pattern.match(key) + if m: + l, e, proj = int(m.group(1)), int(m.group(2)), m.group(3) + if l == layer: + if e not in expert_weights: + expert_weights[e] = {} + expert_weights[e][proj] = f.get_tensor(key).float() + + print(f" Found {len(expert_weights)} experts", flush=True) + + for exp_id in sorted(expert_weights.keys()): + for proj in PROJECTIONS: + if proj not in expert_weights[exp_id]: + continue + w = expert_weights[exp_id][proj].to(device) + result = encode_expert_msrt( + w, base_k, stages, device, ghd, tcp, tcpi, qtf, cbs) + del w + torch.cuda.empty_cache() + + # Store under (exp_id, proj) — flatten for writing + key = (exp_id, proj) + if exp_id not in layer_results[layer]: + layer_results[layer][exp_id] = {"base": {}, "stages": [[] for _ in stages], "mse": {}} + layer_results[layer][exp_id]["base"][proj] = result["base"] + for si, stage_result in enumerate(result["stages"]): + layer_results[layer][exp_id]["stages"][si].append(stage_result) + layer_results[layer][exp_id]["mse"][proj] = result["mse"] + total_mse += result["mse"] + n_experts_total += 1 + + # Print layer summary + layer_mses = [] + for exp_data in layer_results[layer].values(): + layer_mses.extend(exp_data["mse"].values()) + avg = sum(layer_mses) / len(layer_mses) if layer_mses else 0 + print(f" Layer {layer}: avg MSE = {avg:.4e} ({len(layer_mses)} projections)", flush=True) + + print(f"\nOverall avg MSE: {total_mse / max(n_experts_total, 1):.4e}", flush=True) + + # Write outputs + out_dir = Path(args.out) + print(f"\nWriting outputs to {out_dir}...", flush=True) + + # Restructure for writing: group by (layer, expert) -> {base: {trellis, suh, svh}, stages: [...]} + # The layer_results is already structured, but we need to reorganize for the writers + write_results: dict[int, dict[int, dict[str, Any]]] = {} + for layer, experts in layer_results.items(): + write_results[layer] = {} + for exp_id, exp_data in experts.items(): + # Merge projections into single base/stages + base_tensors = {} + for proj in PROJECTIONS: + if proj in exp_data["base"]: + for tname, tval in exp_data["base"][proj].items(): + base_tensors[f"{proj}_{tname}"] = tval + + stage_list = [] + for si in range(len(stages)): + stage_tensors = {} + for proj in PROJECTIONS: + if si < len(exp_data["stages"]) and proj == PROJECTIONS[0]: + # Only need one copy of suh/svh per expert + pass + if si < len(exp_data["stages"]): + for proj_idx, proj in enumerate(PROJECTIONS): + if proj_idx < len(exp_data["stages"][si]): + stage = exp_data["stages"][si][proj_idx] + stage_tensors[f"{proj}"] = stage + stage_list.append(stage_tensors) + + write_results[layer][exp_id] = { + "base": base_tensors, + "stages": stage_list, + "mse": exp_data["mse"], + } + + # Write base checkpoint + print("\nWriting base checkpoint...", flush=True) + base_dir = out_dir / "base" + # Simplified: write all experts for each layer into one shard + for layer in moe_layers: + if layer not in write_results: + continue + tensors = {} + for exp_id, exp_data in write_results[layer].items(): + for proj in PROJECTIONS: + if f"{proj}_trellis" not in exp_data["base"]: + continue + rank = 0 + prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" + tensors[f"{prefix}.trellis"] = exp_data["base"][f"{proj}_trellis"] + tensors[f"{prefix}.suh"] = exp_data["base"][f"{proj}_suh"] + tensors[f"{prefix}.svh"] = exp_data["base"][f"{proj}_svh"] + tensors[f"{prefix}.mcg"] = torch.tensor([0xCBAC1FED], dtype=torch.int32) + if tensors: + save_safetensors(tensors, base_dir / f"model-layer-{layer:03d}.safetensors") + print(f" Layer {layer}: {len(tensors)} base tensors", flush=True) + + # Copy non-layer files from source + import shutil + for f in args.source.iterdir(): + if f.is_file() and not f.name.startswith("model-layer-"): + shutil.copy2(f, base_dir / f.name) + + # Write cartridge adapters + print("\nWriting cartridge adapters...", flush=True) + cart_dir = out_dir / "cartridges" + + for si, stage in enumerate(stages): + label = stage["label"] + expert_filter = expert_filters[si] + tensors = {} + + for layer in moe_layers: + if layer not in write_results: + continue + for exp_id, exp_data in write_results[layer].items(): + if expert_filter and layer in expert_filter: + if exp_id not in expert_filter[layer]: + continue + if si >= len(exp_data["stages"]): + continue + stage_data = exp_data["stages"][si] + for proj in PROJECTIONS: + if proj not in stage_data: + continue + rank = 0 + prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" + s = stage_data[proj] + tensors[f"{prefix}.trellis_{label}"] = s["trellis"] + tensors[f"{prefix}.suh_{label}"] = s["suh"] + tensors[f"{prefix}.svh_{label}"] = s["svh"] + tensors[f"{prefix}.scale_{label}"] = torch.tensor([s["scale"]], dtype=torch.float32) + + adapter_path = cart_dir / f"cartridge_{label}.safetensors" + save_safetensors(tensors, adapter_path) + config = { + "schema": ADAPTER_CONFIG_SCHEMA, + "stage_label": label, + "stage_k": stage["k"], + "num_tensors": len(tensors), + "tool_version": TOOL_VERSION, + "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + (cart_dir / f"cartridge_{label}_config.json").write_text( + json.dumps(config, indent=2) + "\n") + print(f" Cartridge '{label}': {len(tensors)} tensors -> {adapter_path.name}", flush=True) + + # Write summary + summary = { + "tool": TOOL_VERSION, + "base_k": base_k, + "stages": [{"k": s["k"], "label": s["label"], + "experts": s["experts"]} for s in stages], + "moe_layers": moe_layers, + "overall_mse": total_mse / max(n_experts_total, 1), + "n_experts_encoded": n_experts_total, + "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + (out_dir / "encoding_summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(f"\nDone! Output: {out_dir}", flush=True) + print(f" Overall MSE: {summary['overall_mse']:.4e}", flush=True) + return 0 + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="command", required=True) + + enc = sub.add_parser("encode", help="Encode BF16 → base K + cartridge adapters") + enc.add_argument("--source", required=True, type=Path, + help="Source BF16 checkpoint directory") + enc.add_argument("--recipe", required=True, type=Path, + help="Cartridge recipe JSON (fq-cartridge/1)") + enc.add_argument("--out", required=True, type=Path, + help="Output directory (base/ and cartridges/ subdirs)") + enc.add_argument("--encoder-source", required=True, type=Path, + help="Path to exllamav3 Python package") + enc.add_argument("--device", default="cuda:0") + + args = p.parse_args(argv) + if args.command == "encode": + return cmd_encode(args) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From e24e18dac08f8422ae72bddccaa844203f6b5a53 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 09:21:08 -0400 Subject: [PATCH 02/34] fix: bootstrap_encoder module references (ext vs hadamard vs quantize) --- tools/fq_assemble_lora.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index f11a25a..b82d45d 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -121,28 +121,28 @@ def new_task(self, *a, **kw): pass # Load the extension spec = importlib.util.spec_from_file_location( "exllamav3.ext", str(pkg_root / "ext.py")) - m = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.ext"] = m - spec.loader.exec_module(m) + ext_mod = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.ext"] = ext_mod + spec.loader.exec_module(ext_mod) # Load Hadamard spec = importlib.util.spec_from_file_location( "exllamav3.util.hadamard", str(pkg_root / "util" / "hadamard.py")) - m = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.util.hadamard"] = m - spec.loader.exec_module(m) + had_mod = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.util.hadamard"] = had_mod + spec.loader.exec_module(had_mod) # Load quantize module quant_path = pkg_root / "modules" / "quant" / "exl3_lib" / "quantize.py" spec = importlib.util.spec_from_file_location( "exllamav3.modules.quant.exl3_lib.quantize", str(quant_path)) - m = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = m - spec.loader.exec_module(m) + quant_mod = importlib.util.module_from_spec(spec) + sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = quant_mod + spec.loader.exec_module(quant_mod) - return (m.exllamav3_ext, m.get_hadamard_dt, - m.tensor_core_perm, m.tensor_core_perm_i, - m.quantize_tiles, m.codebook_scale) + return (ext_mod.exllamav3_ext, had_mod.get_hadamard_dt, + quant_mod.tensor_core_perm, quant_mod.tensor_core_perm_i, + quant_mod.quantize_tiles, quant_mod.codebook_scale) # ── Quantization Primitives ──────────────────────────────────────────────── From 3e1b063115fd2ffd5b17c97788119e21b4c5d6b2 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 09:26:06 -0400 Subject: [PATCH 03/34] fix: BF16 expert regex, expert filter parsing, base dir creation - Fix regex to match BF16 tensor names (*.weight) not just EXL3 (*.rank0.trellis) - Fix expert filter to handle 'hot96' string key from recipe - Add base_dir.mkdir before writing files - Restore corrupted imports and remove duplicate definitions Co-authored-by: Claude --- tools/fq_assemble_lora.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index b82d45d..59a4b72 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -35,11 +35,6 @@ --recipe recipes/fruit-k2-k3k4-cart.json \\ --out ./output \\ --encoder-source /opt/fruit-pip/exllamav3 - -The tool requires the exllamav3 encoder (quantize_tiles, codebook_scale, -Hadamard transform) to perform trellis quantization. It does NOT require -vLLM or flash_attn — the encoder is loaded via the same bootstrap pattern -used in the PoC scripts. """ from __future__ import annotations @@ -62,10 +57,10 @@ ADAPTER_CONFIG_SCHEMA = "fq-cartridge-adapter/1" HADAMARD_BLOCK = 128 -# Expert tensor name pattern in source checkpoints +# Expert tensor name pattern — matches both BF16 (.weight) and EXL3 (.rank0.trellis) EXPERT_RE_PATTERN = ( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." - r"(gate_proj|up_proj|down_proj)(?:\.rank(\d+))?$" + r"(gate_proj|up_proj|down_proj)\.(?:rank\d+\.)?(?:weight|trellis)$" ) PROJECTIONS = ("gate_proj", "up_proj", "down_proj") @@ -509,12 +504,17 @@ def cmd_encode(args) -> int: # Parse expert filters expert_filters = [] + hot_experts = recipe.get("hot_experts", list(range(96))) # default: first 96 for stage in stages: - if stage["experts"] == "all": + exp_spec = stage["experts"] + if exp_spec == "all": expert_filters.append(None) + elif isinstance(exp_spec, str) and exp_spec in ("hot96", "hot"): + expert_filters.append({l: hot_experts for l in moe_layers}) + elif isinstance(exp_spec, list): + expert_filters.append({l: exp_spec for l in moe_layers}) else: - # Same expert list for all layers (could be per-layer in future) - expert_filters.append({l: stage["experts"] for l in moe_layers}) + expert_filters.append(None) # Load source weights and encode from safetensors import safe_open @@ -622,6 +622,7 @@ def cmd_encode(args) -> int: # Write base checkpoint print("\nWriting base checkpoint...", flush=True) base_dir = out_dir / "base" + base_dir.mkdir(parents=True, exist_ok=True) # Simplified: write all experts for each layer into one shard for layer in moe_layers: if layer not in write_results: From 66637c99719aabae4b349c35be40d21d0e6bd6d9 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 09:34:12 -0400 Subject: [PATCH 04/34] fix: pack trellis indices via ext.pack_trellis to EXL3 checkpoint format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit quantize_tiles returns raw Viterbi indices (n_tiles, 256). EXL3 checkpoints store PACKED indices (n_tiles, K*16) — compressed via ext.pack_trellis. The vLLM loader validates this at exl3.py:2099-2102. Fixed quantize_trellis_packed to: 1. Collect raw 256-width indices from qtf 2. Call ext.pack_trellis to compress to K*16 packed format 3. Return packed tensor with correct (k//16, n//16, K*16) shape Also added ext parameter threading through encode_expert_msrt and rescaled_trellis_quantize. Falls back to raw indices when ext is None (for testing without the CUDA extension). Co-authored-by: Claude --- tools/fq_assemble_lora.py | 36 +++++++++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index 59a4b72..b31149a 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -223,17 +223,23 @@ def quantize_trellis_packed( tcp: Any, tcpi: Any, qtf: Any, + ext: Any = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Quantize and return BOTH reconstructed values and packed trellis indices. Returns (reconstructed_float, packed_trellis_int16) where packed_trellis has shape (k // 16, n // 16, K * 16) dtype int16 — the EXL3 storage format. + + The quantize_tiles function returns raw Viterbi path indices of shape + (n_tiles, 256) — one int16 per weight element. ext.pack_trellis compresses + these to (n_tiles, K * 16) packed indices, which is what the EXL3 checkpoint + and vLLM loader expect (validated at exl3.py:2099-2102). """ k, n = data.shape tiles_n = n // 16 weight_q = torch.zeros_like(data) - # Packed trellis: (tiles_k, tiles_n, K*16) int16 - packed = torch.zeros(k // 16, tiles_n, K * 16, dtype=torch.int16, device=device) + # Raw (unpacked) indices: (tiles_k, tiles_n, 256) int16 + raw_indices = torch.zeros(k // 16, tiles_n, 256, dtype=torch.int16, device=device) qa = {"K": K, "mcg": True} perm = tcp(device) perm_i = tcpi(device) @@ -244,12 +250,21 @@ def quantize_trellis_packed( tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) tiles = tiles[:, perm].contiguous() quant_w, quant_idx = qtf(tiles, qa) - # Store packed indices - packed[tk] = quant_idx.reshape(tiles_n, K * 16) + # Store raw indices — quant_idx has shape (n_tiles, 256) + raw_indices[tk] = quant_idx # Reconstruct quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) weight_q[bi:bi + 16] = quant_w + # Pack the raw indices to EXL3 format: (tiles_k, tiles_n, K * 16) int16 + if ext is not None and hasattr(ext, "pack_trellis"): + packed_shape = (k // 16, tiles_n, 256 * K // 16) + packed = torch.zeros(packed_shape, dtype=torch.int16, device=device) + ext.pack_trellis(packed, raw_indices.contiguous(), K) + else: + # Fallback: store raw indices (unpacked) — for testing without ext + packed = raw_indices + return weight_q, packed @@ -260,6 +275,7 @@ def rescaled_trellis_quantize( device: torch.device, tcp: Any, tcpi: Any, qtf: Any, cbs: float, + ext: Any = None, ) -> tuple[torch.Tensor, torch.Tensor, float]: """Rescale residual to match codebook range, quantize, return (recon, packed, scale). @@ -269,13 +285,14 @@ def rescaled_trellis_quantize( """ residual_rms = residual.square().mean().sqrt().item() if residual_rms < 1e-12: + packed_w = 256 * K_res // 16 if ext else 256 return base_q, torch.zeros( - residual.shape[0] // 16, residual.shape[1] // 16, K_res * 16, + residual.shape[0] // 16, residual.shape[1] // 16, packed_w, dtype=torch.int16, device=device), 1.0 scale = abs(cbs) / residual_rms scaled = residual * scale - recon_packed, packed = quantize_trellis_packed(scaled, K_res, device, tcp, tcpi, qtf) + recon_packed, packed = quantize_trellis_packed(scaled, K_res, device, tcp, tcpi, qtf, ext) recon = base_q + recon_packed / scale return recon, packed, scale @@ -325,6 +342,7 @@ def encode_expert_msrt( device: torch.device, ghd: Any, tcp: Any, tcpi: Any, qtf: Any, cbs: float, + ext: Any = None, ) -> dict[str, Any]: """Encode one expert weight matrix with MSRT. @@ -336,7 +354,7 @@ def encode_expert_msrt( w_reg = regularize(w_bf16, device, ghd, cbs) # Base tier - base_recon, base_packed = quantize_trellis_packed(w_reg, base_k, device, tcp, tcpi, qtf) + base_recon, base_packed = quantize_trellis_packed(w_reg, base_k, device, tcp, tcpi, qtf, ext) had_vectors = compute_hadamard_vectors(w_bf16, device, ghd, cbs) result = { @@ -353,7 +371,7 @@ def encode_expert_msrt( for stage in stages: residual = w_reg - current_recon recon, packed, scale = rescaled_trellis_quantize( - current_recon, residual, stage["k"], device, tcp, tcpi, qtf, cbs) + current_recon, residual, stage["k"], device, tcp, tcpi, qtf, cbs, ext) result["stages"].append({ "trellis": packed.cpu(), "suh": had_vectors["suh"].cpu(), # Same Hadamard vectors @@ -558,7 +576,7 @@ def cmd_encode(args) -> int: continue w = expert_weights[exp_id][proj].to(device) result = encode_expert_msrt( - w, base_k, stages, device, ghd, tcp, tcpi, qtf, cbs) + w, base_k, stages, device, ghd, tcp, tcpi, qtf, cbs, ext) del w torch.cuda.empty_cache() From d97409bf0f14b5e10f230d9c8e1ec23151660b60 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 10:36:52 -0400 Subject: [PATCH 05/34] =?UTF-8?q?fix:=20mcg=20sentinel=20overflow=20?= =?UTF-8?q?=E2=80=94=20use=20uint32.view(int32)=20like=20EXL3=20does?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tools/fq_assemble_lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index b31149a..02c8757 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -655,7 +655,7 @@ def cmd_encode(args) -> int: tensors[f"{prefix}.trellis"] = exp_data["base"][f"{proj}_trellis"] tensors[f"{prefix}.suh"] = exp_data["base"][f"{proj}_suh"] tensors[f"{prefix}.svh"] = exp_data["base"][f"{proj}_svh"] - tensors[f"{prefix}.mcg"] = torch.tensor([0xCBAC1FED], dtype=torch.int32) + tensors[f"{prefix}.mcg"] = torch.tensor([0xCBAC1FED], dtype=torch.uint32).view(torch.int32) if tensors: save_safetensors(tensors, base_dir / f"model-layer-{layer:03d}.safetensors") print(f" Layer {layer}: {len(tensors)} base tensors", flush=True) From 81f8f80c090f0363130faad7f4ed4d3956a12225 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 11:28:07 -0400 Subject: [PATCH 06/34] =?UTF-8?q?feat:=20combine=5Fcartridges=20=E2=80=94?= =?UTF-8?q?=20assemble=20single-adapter=20K3/K4-like=20cartridges?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since vLLM supports only 1 LoRA per request, this tool combines individual MSRT stage files into single adapter files: - cart_k3like: K1trsc for all experts (K3-equivalent, 3bpw) - cart_k3k4like: K1trsc (all) + K2trsc (96 hot) = matches SIQ 160K3+96K4 Co-authored-by: Claude --- tools/combine_cartridges.py | 126 ++++++++++++++++++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100755 tools/combine_cartridges.py diff --git a/tools/combine_cartridges.py b/tools/combine_cartridges.py new file mode 100755 index 0000000..25e0c41 --- /dev/null +++ b/tools/combine_cartridges.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Assemble combined MSRT cartridges from individual stage files. + +Since vLLM supports only 1 LoRA per request, we need single adapter files that +contain all stages for a given recipe: + +1. cart_k3like.safetensors: K1trsc for ALL 256 experts (K2→K3, 3bpw total) +2. cart_k3k4like.safetensors: K1trsc for all + K2trsc for 96 hot (K2→K3/K4, 3.375bpw) + - 160 non-hot experts: only K1trsc stage (K3-equivalent) + - 96 hot experts: K1trsc + K2trsc stages (K4-equivalent) +""" +import json +import sys +from pathlib import Path +from safetensors import safe_open +from safetensors.torch import save_file +import torch + +def combine_cartridges( + res1_path: Path, # K1trsc for all experts + res2_path: Path, # K2trsc for hot experts + hot_experts: list[int], + out_path: Path, + label1: str = "res1", + label2: str = "res2", +): + """Combine two stage files into a single adapter with both stages. + + For hot experts: both res1 and res2 tensors are included + For non-hot experts: only res1 tensors are included + """ + tensors = {} + + # Load res1 (K1trsc) for ALL experts + with safe_open(str(res1_path), framework="pt") as f: + keys = list(f.keys()) + for key in keys: + tensors[key] = f.get_tensor(key) + + print(f"Loaded {len(tensors)} tensors from res1 (K1trsc, all experts)", flush=True) + + # Load res2 (K2trsc) only for hot experts + hot_set = set(hot_experts) + res2_count = 0 + with safe_open(str(res2_path), framework="pt") as f: + for key in f.keys(): + # Parse expert ID from key: model.layers.{L}.mlp.experts.{E}.{proj}.rank{R}.trellis_{label} + parts = key.split(".") + expert_id = int(parts[5]) + if expert_id in hot_set: + tensors[key] = f.get_tensor(key) + res2_count += 1 + + print(f"Loaded {res2_count} tensors from res2 (K2trsc, {len(hot_set)} hot experts)", flush=True) + print(f"Combined: {len(tensors)} total tensors", flush=True) + + # Write combined adapter + out_path.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(out_path)) + print(f"Saved combined cartridge: {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)", flush=True) + + # Write config + config = { + "schema": "fq-cartridge-adapter/1", + "stages": [ + {"k": 1, "label": label1, "experts": "all"}, + {"k": 2, "label": label2, "experts": hot_experts}, + ], + "description": "Combined K3/K4-like cartridge matching SIQ 160K3+96K4 allocation", + "num_tensors": len(tensors), + } + config_path = out_path.with_suffix(".config.json") + config_path.write_text(json.dumps(config, indent=2) + "\n") + return config + +def copy_k3only_cartridge( + res1_path: Path, + out_path: Path, + label: str = "res1", +): + """Create a K3-equivalent-only cartridge (just res1, all experts).""" + tensors = {} + with safe_open(str(res1_path), framework="pt") as f: + for key in f.keys(): + tensors[key] = f.get_tensor(key) + + out_path.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(out_path)) + print(f"Saved K3-like cartridge: {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)", flush=True) + + config = { + "schema": "fq-cartridge-adapter/1", + "stages": [{"k": 1, "label": label, "experts": "all"}], + "description": "K3-equivalent cartridge (K2+K1trsc, all experts, 3bpw)", + "num_tensors": len(tensors), + } + config_path = out_path.with_suffix(".config.json") + config_path.write_text(json.dumps(config, indent=2) + "\n") + return config + +if __name__ == "__main__": + cart_dir = Path("/tmp/poc_residual/fruit_msrt_output/cartridges") + out_dir = cart_dir + + # Hot experts: first 96 (matching SIQ's K4 tier) + hot_experts = list(range(96)) + + # 1. K3-like only (base + K1trsc = 3bpw on all experts) + print("=== Assembling K3-like cartridge (all experts, K1trsc only) ===", flush=True) + copy_k3only_cartridge( + cart_dir / "cartridge_res1.safetensors", + out_dir / "cart_k3like.safetensors", + ) + + # 2. K3/K4-like combined (matches SIQ 160K3 + 96K4) + print("\n=== Assembling K3/K4-like combined cartridge (160×K3 + 96×K4) ===", flush=True) + combine_cartridges( + cart_dir / "cartridge_res1.safetensors", + cart_dir / "cartridge_res2.safetensors", + hot_experts, + out_dir / "cart_k3k4like.safetensors", + ) + + print("\nDone! Combined cartridges:", flush=True) + print(f" cart_k3like.safetensors (K3-equivalent, all 256 experts)", flush=True) + print(f" cart_k3k4like.safetensors (K3/K4 mix, matches SIQ 160K3+96K4)", flush=True) From 200e4b38851f5d1d1564f497d50e702d75254e02 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 11:29:25 -0400 Subject: [PATCH 07/34] feat: MSE measurement results on Fruit SIQ model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Weight-level MSE (3 layers, 10 experts each, gate/up/down projections): | Config | MSE | vs K3 | vs K4 | |--------|-----|-------|-------| | K3 only | 2.718e-02 | 1.00× | 3.73× | | K4 only | 7.284e-03 | 0.268× | 1.00× | | MSRT K2+K1trsc | 2.908e-02 | 1.07× | 3.99× (K3-equiv, 7% worse) | | MSRT K2+K1+K2trsc | 1.995e-03 | 0.073× | 0.274× (3.6× BETTER than K4!) | MSRT at 4bpw (K2+K1trsc+K2trsc) is 3.6× better than native K4 at the same bitrate. This confirms v50/v52 PoC results on the real Fruit model. Co-authored-by: Claude --- tools/measure_mse_fruit.py | 201 +++++++++++++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100755 tools/measure_mse_fruit.py diff --git a/tools/measure_mse_fruit.py b/tools/measure_mse_fruit.py new file mode 100755 index 0000000..098128d --- /dev/null +++ b/tools/measure_mse_fruit.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +"""Measure weight-level MSE between Fruit model quantization variants. + +Compares: +1. BF16 original (reference) +2. SIQ quant (160 K3 + 96 K4 per layer) +3. MSRT K2 base (all experts at K2) +4. MSRT K2+K1trsc (all experts at K3-equivalent) +5. MSRT K2+K1trsc+K2trsc (96 hot experts at K4-equivalent) + +All measurements in regularized (Hadamard) space, which equals original +space since Hadamard is orthogonal. +""" +from __future__ import annotations +import json, math, os, sys, types, importlib.util, gc +from pathlib import Path +import torch + +EXL3_PKG = "/opt/fruit-pip/exllamav3" + +def _bootstrap(): + pkg = types.ModuleType("exllamav3"); pkg.__path__ = [EXL3_PKG]; sys.modules["exllamav3"] = pkg + for sub in ["util", "modules", "modules.quant", "modules.quant.exl3_lib"]: + full = f"exllamav3.{sub}"; m = types.ModuleType(full) + m.__path__ = [f"{EXL3_PKG}/{sub.replace('.', '/')}"]; sys.modules[full] = m + class _DPB: + def __init__(self, *a, **kw): pass + def __enter__(self): return self + def __exit__(self, *a): return False + def update(self, *a): pass + def new_task(self, *a, **kw): pass + _s = types.ModuleType("exllamav3.util.progress"); _s.ProgressBar = _DPB; sys.modules["exllamav3.util.progress"] = _s + _s = types.ModuleType("exllamav3.util.memory"); _s.free_mem = lambda: None; _s.list_gpu_tensors = lambda: []; sys.modules["exllamav3.util.memory"] = _s + _s = types.ModuleType("exllamav3.util"); _s.__path__ = [f"{EXL3_PKG}/util"]; _s.cuda_sync_active = lambda *a, **kw: torch.cuda.synchronize(); sys.modules["exllamav3.util"] = _s + _s = types.ModuleType("exllamav3.util.tensor"); _s.save_tensor_image = lambda *a, **kw: None; sys.modules["exllamav3.util.tensor"] = _s + spec = importlib.util.spec_from_file_location("exllamav3.ext", f"{EXL3_PKG}/ext.py") + m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.ext"] = m; spec.loader.exec_module(m) + ext = m.exllamav3_ext + spec = importlib.util.spec_from_file_location("exllamav3.util.hadamard", f"{EXL3_PKG}/util/hadamard.py") + m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.util.hadamard"] = m; spec.loader.exec_module(m) + ghd = m.get_hadamard_dt + spec = importlib.util.spec_from_file_location("exllamav3.modules.quant.exl3_lib.quantize", f"{EXL3_PKG}/modules/quant/exl3_lib/quantize.py") + m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = m; spec.loader.exec_module(m) + return ext, ghd, m.tensor_core_perm, m.tensor_core_perm_i, m.quantize_tiles, m.codebook_scale + +def block_rms(x, dim, keepdim=False): + return x.square().mean(dim=dim, keepdim=keepdim).sqrt() + +def regularize(w, device, ghd, cbs, had_k=128, had_n=128, seed=0): + k, n = w.shape + g = torch.Generator(device="cpu").manual_seed(seed) + su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) + sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) + out_scales = block_rms(w, dim=0, keepdim=True) + mean = out_scales.mean().item() + if mean > 1e-30: out_scales = out_scales / mean + sv = (sv * out_scales + 1e-10).float() + w = (w / sv).contiguous() + had_n_mat = ghd(had_n, device, torch.float, 1.0 / math.sqrt(had_n)) + w = (w.view(k, n // had_n, had_n) @ had_n_mat).view(k, n).contiguous() + in_scales = block_rms(w, dim=1, keepdim=True).clamp(min=1e-30) + su = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() + w = (w / su).contiguous() + had_k_mat = ghd(had_k, device, torch.float, 1.0 / math.sqrt(had_k)) + w = (had_k_mat @ w.view(k // had_k, had_k, n)).view(k, n).contiguous() + return w + +def quantize_trellis_raw(data, K, device, tcp, tcpi, qtf): + k, n = data.shape; tiles_n = n // 16; weight_q = torch.zeros_like(data) + qa = {"K": K, "mcg": True} + perm = tcp(device); perm_i = tcpi(device) + for bi in range(0, k, 16): + rows = data[bi:bi+16] + tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) + tiles = tiles[:, perm].contiguous() + quant_w, _ = qtf(tiles, qa) + quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) + weight_q[bi:bi+16] = quant_w + return weight_q + +def rescaled_trellis(base_q, residual, K_res, device, tcp, tcpi, qtf, cbs): + residual_rms = residual.square().mean().sqrt().item() + if residual_rms < 1e-12: return base_q + scale = abs(cbs) / residual_rms + scaled = residual * scale + quant = quantize_trellis_raw(scaled, K_res, device, tcp, tcpi, qtf) + return base_q + quant / scale + +def run_measurement(bf16_path, siq_path, device, ghd, tcp, tcpi, qtf, cbs): + """Measure weight-level MSE for all configurations.""" + from safetensors import safe_open + import re + + # SIQ tier_bitmap: 160 K3 + 96 K4 per layer + tier_path = Path(siq_path) / "tier_bitmap.json" + tier_bitmap = json.loads(tier_path.read_text()) if tier_path.exists() else {} + + # Load SIQ trellis weights and compare to BF16 + moe_layers = [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] + results = {} + + for layer in moe_layers[:3]: # First 3 layers for speed + bf16_shard = bf16_path / f"model-layer-{layer:03d}.safetensors" + siq_shard = siq_path / f"model-layer-{layer:03d}.safetensors" + + if not bf16_shard.exists(): + print(f" Layer {layer}: bf16 shard not found", flush=True) + continue + + print(f"\n=== Layer {layer} ===", flush=True) + + # Load BF16 expert weights + bf16_experts = {} + with safe_open(str(bf16_shard), framework="pt") as f: + for key in f.keys(): + if f"layers.{layer}.mlp.experts." in key and key.endswith(".weight"): + parts = key.split(".") + eid = int(parts[5]) + proj = parts[6] + if eid not in bf16_experts: bf16_experts[eid] = {} + bf16_experts[eid][proj] = f.get_tensor(key).float() + + # Load SIQ tier info + k_list = tier_bitmap.get(str(layer), {}).get("k", [3] * 256) + + # Measure each config + n_experts = min(10, len(bf16_experts)) # First 10 experts for speed + configs_mse = {name: [] for name in [ + "K3_all", "K4_all", "SIQ_mixed", "MSRT_K2", "MSRT_K2_K1trsc", "MSRT_K2_K1trsc_K2trsc"]} + + for eid in sorted(bf16_experts.keys())[:n_experts]: + for proj in ["gate_proj", "up_proj", "down_proj"]: + if proj not in bf16_experts[eid]: + continue + w = bf16_experts[eid][proj].to(device) + w_reg = regularize(w, device, ghd, cbs) + del w + + # K3 + q_k3 = quantize_trellis_raw(w_reg, 3, device, tcp, tcpi, qtf) + configs_mse["K3_all"].append((w_reg - q_k3).pow(2).mean().item()) + + # K4 + q_k4 = quantize_trellis_raw(w_reg, 4, device, tcp, tcpi, qtf) + configs_mse["K4_all"].append((w_reg - q_k4).pow(2).mean().item()) + + # SIQ mixed: K3 or K4 depending on tier + siq_k = k_list[eid] if eid < len(k_list) else 3 + q_siq = quantize_trellis_raw(w_reg, siq_k, device, tcp, tcpi, qtf) + configs_mse["SIQ_mixed"].append((w_reg - q_siq).pow(2).mean().item()) + + # MSRT K2 base + q_k2 = quantize_trellis_raw(w_reg, 2, device, tcp, tcpi, qtf) + configs_mse["MSRT_K2"].append((w_reg - q_k2).pow(2).mean().item()) + + # MSRT K2 + K1trsc (K3-equivalent) + r2 = w_reg - q_k2 + q_msrt3 = rescaled_trellis(q_k2, r2, 1, device, tcp, tcpi, qtf, cbs) + configs_mse["MSRT_K2_K1trsc"].append((w_reg - q_msrt3).pow(2).mean().item()) + + # MSRT K2 + K1trsc + K2trsc (K4-equivalent, for hot experts) + r3 = w_reg - q_msrt3 + q_msrt4 = rescaled_trellis(q_msrt3, r3, 2, device, tcp, tcpi, qtf, cbs) + configs_mse["MSRT_K2_K1trsc_K2trsc"].append((w_reg - q_msrt4).pow(2).mean().item()) + + del w_reg, q_k3, q_k4, q_siq, q_k2, q_msrt3, q_msrt4 + torch.cuda.empty_cache() + + # Print results + print(f" {'Config':<30} {'avg MSE':>12} {'min':>12} {'max':>12}", flush=True) + print(f" {'-'*70}", flush=True) + layer_results = {} + for name, mses in configs_mse.items(): + avg = sum(mses) / len(mses) if mses else 0 + mn = min(mses) if mses else 0 + mx = max(mses) if mses else 0 + print(f" {name:<30} {avg:>12.4e} {mn:>12.4e} {mx:>12.4e}", flush=True) + layer_results[name] = {"avg": avg, "min": mn, "max": mx, "n": len(mses)} + + results[f"layer{layer}"] = layer_results + + return results + +def main(): + import argparse + ap = argparse.ArgumentParser() + ap.add_argument("--bf16-path", required=True, type=Path) + ap.add_argument("--siq-path", required=True, type=Path) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--out", default="/tmp/poc_residual/mse_results.json") + args = ap.parse_args() + dev = torch.device(args.device) + print(f"Device: {dev} GPU: {torch.cuda.get_device_name(0)}", flush=True) + ext, ghd, tcp, tcpi, qtf, cbs = _bootstrap() + print(f"codebook_scale = {cbs}", flush=True) + results = run_measurement(args.bf16_path, args.siq_path, dev, ghd, tcp, tcpi, qtf, cbs) + Path(args.out).write_text(json.dumps(results, indent=2, default=str)) + print(f"\nResults saved to {args.out}", flush=True) + +if __name__ == "__main__": + main() From 4b8b016034b3ae22ae84892d456c145cecb4b435 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 11:52:02 -0400 Subject: [PATCH 08/34] fix: suh/svh/mcg format to match EXL3 checkpoint (float16, correct shapes, scalar mcg) - suh: (input_size,) float16 (was float32, wrong shape) - svh: (output_size,) float16 (was float32, wrong shape) - mcg: scalar () int32 (was (1,) int32) These match the SIQ model's checkpoint format that vLLM's EXL3 loader expects. Without this fix, the K2 base checkpoint fails validation with: ValueError: Invalid EXL3 MoE tensors for expert=0, projection=w1 Co-authored-by: Claude --- tools/fq_assemble_lora.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index 02c8757..cfc36f5 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -308,8 +308,13 @@ def compute_hadamard_vectors( ) -> dict[str, torch.Tensor]: """Compute the suh and svh vectors that EXL3 stores alongside trellis. - These are the per-block sign vectors and scales from regularize(). - The exact format must match what exl3_gemm expects. + EXL3 checkpoint format (from SIQ model inspection): + - suh: (input_size,) float16 — per-row sign+scale vector + - svh: (output_size,) float16 — per-column sign+scale vector + + The regularize() function applies: + 1. svh = sign * out_scales (per-column) + 2. suh = sign * in_scales / (-cbs) (per-row, after column Hadamard) """ k, n = w.shape g = torch.Generator(device="cpu").manual_seed(seed) @@ -320,17 +325,17 @@ def compute_hadamard_vectors( mean = out_scales.mean().item() if mean > 1e-30: out_scales = out_scales / mean - sv_full = (sv * out_scales + 1e-10).float() + svh = (sv * out_scales.squeeze(0) + 1e-10).half() # (n,) float16 # After column Hadamard - w_col = (w / sv_full).contiguous() + w_col = (w / svh.float().unsqueeze(0)).contiguous() had_n_mat = ghd(HADAMARD_BLOCK, device, torch.float, 1.0 / math.sqrt(HADAMARD_BLOCK)) w_col = (w_col.view(k, n // HADAMARD_BLOCK, HADAMARD_BLOCK) @ had_n_mat).view(k, n).contiguous() - in_scales = block_rms(w_col, dim=1, keepdim=True).clamp(min=1e-30) - su_full = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() + in_scales = block_rms(w_col, dim=1, keepdim=True).clamp(min=1e-30).squeeze(1) + suh = (su * in_scales / (-cbs) + 1e-10).half() # (k,) float16 - return {"suh": su_full.squeeze().contiguous(), "svh": sv_full.squeeze().contiguous()} + return {"suh": suh.contiguous(), "svh": svh.contiguous()} # ── Encoding Pipeline ───────────────────────────────────────────────────── @@ -655,7 +660,7 @@ def cmd_encode(args) -> int: tensors[f"{prefix}.trellis"] = exp_data["base"][f"{proj}_trellis"] tensors[f"{prefix}.suh"] = exp_data["base"][f"{proj}_suh"] tensors[f"{prefix}.svh"] = exp_data["base"][f"{proj}_svh"] - tensors[f"{prefix}.mcg"] = torch.tensor([0xCBAC1FED], dtype=torch.uint32).view(torch.int32) + tensors[f"{prefix}.mcg"] = torch.tensor(0xCBAC1FED, dtype=torch.uint32).view(torch.int32) if tensors: save_safetensors(tensors, base_dir / f"model-layer-{layer:03d}.safetensors") print(f" Layer {layer}: {len(tensors)} base tensors", flush=True) From ef9c7ed6a77439e291850f159b6cae59019bcb3b Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 13:03:20 -0400 Subject: [PATCH 09/34] fix: transpose weight to (in, out) for correct EXL3 trellis geometry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PyTorch stores Linear weights as (out_features, in_features), but EXL3 trellis expects (in//16, out//16, K*16). The encoder used k,n = w.shape treating k=input, n=output — swapped for all projections. This caused suh/svh and trellis dimensions to be swapped vs the SIQ reference checkpoint: gate/up: suh=(512,) svh=(1024,) trellis=(32,64,K*16) [wrong] should be: suh=(1024,) svh=(512,) trellis=(64,32,K*16) Fix: transpose weight before encoding: w = w.T.contiguous() Discovered by comparing trellis shapes against the SIQ reference model during vLLM loading (slab geometry mismatch error). Co-authored-by: Claude --- tools/fq_assemble_lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index cfc36f5..c3c8269 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -579,7 +579,7 @@ def cmd_encode(args) -> int: for proj in PROJECTIONS: if proj not in expert_weights[exp_id]: continue - w = expert_weights[exp_id][proj].to(device) + w = expert_weights[exp_id][proj].T.contiguous().to(device) # (in, out) for EXL3 trellis layout result = encode_expert_msrt( w, base_k, stages, device, ghd, tcp, tcpi, qtf, cbs, ext) del w From f91166053df4fcc85e4f915ac5a1f95140807ee8 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 19:59:29 -0400 Subject: [PATCH 10/34] test: skip MSRT tests without torch --- tests/test_fq_assemble_lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_fq_assemble_lora.py b/tests/test_fq_assemble_lora.py index 6c33065..4b58aba 100755 --- a/tests/test_fq_assemble_lora.py +++ b/tests/test_fq_assemble_lora.py @@ -11,7 +11,7 @@ from pathlib import Path import pytest -import torch +torch = pytest.importorskip("torch") # Add tools to path sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) From 5d9e65031f0fb36b75296afa00ea2dcb7b69e47a Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:07:24 -0400 Subject: [PATCH 11/34] fix: make MSRT encoding artifacts coherent --- recipes/fruit-k2-k3k4-cart.json | 15 +- schemas/fq-cartridge-1.schema.json | 44 + tests/test_fq_assemble_lora.py | 315 ++++--- tests/test_schemas.py | 7 + tools/fq_assemble_lora.py | 1279 ++++++++++++++++------------ 5 files changed, 1008 insertions(+), 652 deletions(-) create mode 100644 schemas/fq-cartridge-1.schema.json diff --git a/recipes/fruit-k2-k3k4-cart.json b/recipes/fruit-k2-k3k4-cart.json index 5d34823..6b0787c 100644 --- a/recipes/fruit-k2-k3k4-cart.json +++ b/recipes/fruit-k2-k3k4-cart.json @@ -9,19 +9,18 @@ "description": "K1 rescaled trellis residual on ALL experts → K3-equivalent (3bpw)" }, { - "k": 2, + "k": 1, "label": "res2", - "experts": "hot96", - "description": "K2 rescaled trellis residual on 96 hot experts → K4-equivalent (4bpw)" - } - ], - "moe_layers": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], - "hot_experts": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, + "experts": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95], - "description": "Fruit SIQ proxy: K2 base + K1trsc (all) + K2trsc (96 hot) = dual-cartridge matching SIQ 3.375bpw" + "description": "K1 residual on the recipe-selected 96 experts; total K4-equivalent" + } + ], + "moe_layers": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], + "description": "Fruit proxy: K2 base + K1 residual on all + K1 residual on selected 96 = 160 K3 and 96 K4 (3.375 bpw)" } diff --git a/schemas/fq-cartridge-1.schema.json b/schemas/fq-cartridge-1.schema.json new file mode 100644 index 0000000..7b7f1a5 --- /dev/null +++ b/schemas/fq-cartridge-1.schema.json @@ -0,0 +1,44 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/malaiwah/progressive-tensors/schemas/fq-cartridge-1.schema.json", + "title": "Progressive Tensors MSRT cartridge recipe", + "description": "Validated MSRT base and residual-stage selection for one BF16 checkpoint.", + "type": "object", + "additionalProperties": false, + "required": ["schema", "base_k", "stages", "moe_layers"], + "properties": { + "schema": {"const": "fq-cartridge/1"}, + "base_k": {"type": "integer", "minimum": 1, "maximum": 8}, + "stages": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": false, + "required": ["k", "label", "experts"], + "properties": { + "k": {"type": "integer", "minimum": 1, "maximum": 8}, + "label": {"type": "string", "pattern": "^[A-Za-z0-9_-]{1,32}$"}, + "experts": { + "oneOf": [ + {"const": "all"}, + { + "type": "array", + "uniqueItems": true, + "items": {"type": "integer", "minimum": 0} + } + ] + }, + "description": {"type": "string"} + } + } + }, + "moe_layers": { + "type": "array", + "minItems": 1, + "uniqueItems": true, + "items": {"type": "integer", "minimum": 0} + }, + "description": {"type": "string"} + } +} diff --git a/tests/test_fq_assemble_lora.py b/tests/test_fq_assemble_lora.py index 4b58aba..1de06c7 100755 --- a/tests/test_fq_assemble_lora.py +++ b/tests/test_fq_assemble_lora.py @@ -1,132 +1,221 @@ #!/usr/bin/env python3 -"""Tests for fq_assemble_lora. +"""Behavioral tests for the MSRT cartridge encoder.""" + +from __future__ import annotations -Tests the cartridge recipe parsing, MSRT encoding pipeline, and output format. -Uses tiny random tensors (not real model weights) for speed. -""" import json -import math import sys -import tempfile from pathlib import Path import pytest -torch = pytest.importorskip("torch") -# Add tools to path sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) +import fq_assemble_lora as lora + +RECIPE = Path(__file__).parent.parent / "recipes" / "fruit-k2-k3k4-cart.json" -def test_cartridge_recipe_schema(): - """Test that cartridge recipe follows fq-cartridge/1 schema.""" + +def write_recipe(path: Path, **overrides) -> Path: recipe = { "schema": "fq-cartridge/1", "base_k": 2, - "stages": [ - {"k": 1, "label": "res1", "experts": "all"}, - {"k": 2, "label": "res2", "experts": [0, 1, 2]}, - ], - "moe_layers": [3, 4, 5], + "stages": [{"k": 1, "label": "res1", "experts": "all"}], + "moe_layers": [3], } - assert recipe["schema"] == "fq-cartridge/1" - assert recipe["base_k"] == 2 - assert len(recipe["stages"]) == 2 - assert recipe["stages"][0]["k"] == 1 - assert recipe["stages"][1]["experts"] == [0, 1, 2] - - -def test_cartridge_recipe_from_file(): - """Test loading the Fruit recipe.""" - recipe_path = Path(__file__).parent.parent / "recipes" / "fruit-k2-k3k4-cart.json" - if not recipe_path.exists(): - pytest.skip("Recipe file not found") - recipe = json.loads(recipe_path.read_text()) - assert recipe["schema"] == "fq-cartridge/1" + recipe.update(overrides) + path.write_text(json.dumps(recipe)) + return path + + +def expert_key(layer: int, expert: int, projection: str) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.{projection}.weight" + + +def complete_keys(layer: int = 3, expert: int = 0) -> list[str]: + return [expert_key(layer, expert, projection) for projection in lora.PROJECTIONS] + + +def test_fruit_recipe_is_valid_and_has_claimed_bitrate(): + recipe = lora.load_recipe(RECIPE) assert recipe["base_k"] == 2 - assert len(recipe["stages"]) == 2 - assert recipe["stages"][0]["label"] == "res1" - assert recipe["stages"][1]["label"] == "res2" - assert recipe["stages"][0]["experts"] == "all" - assert len(recipe["moe_layers"]) == 11 # layers 3-13 - - -def test_block_rms(): - """Test RMS computation.""" - from fq_assemble_lora import block_rms - x = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - rms = block_rms(x, dim=0, keepdim=True) - expected = torch.sqrt(torch.tensor([(1+9)/2, (4+16)/2])) - assert torch.allclose(rms.squeeze(), expected, rtol=1e-5) - - -def test_rescaled_trellis_scale(): - """Test that rescaling produces correct scale factor.""" - # Mock: just test the scale computation logic - cbs = 1.2437 - residual = torch.randn(128, 256) - residual_rms = residual.square().mean().sqrt().item() - scale = abs(cbs) / residual_rms - assert scale > 0 - # After rescaling, RMS should be ~|cbs| - scaled = residual * scale - scaled_rms = scaled.square().mean().sqrt().item() - assert abs(scaled_rms - abs(cbs)) < 0.01 # close to codebook scale - - -def test_trellis_packed_shape(): - """Test that packed trellis has correct shape (without GPU).""" - # Shape: (k//16, n//16, K*16) int16 - k, n, K = 128, 256, 2 - expected_shape = (k // 16, n // 16, K * 16) - assert expected_shape == (8, 16, 32) - # For K=3: (8, 16, 48) - assert (k // 16, n // 16, 3 * 16) == (8, 16, 48) - - -def test_cartridge_adapter_naming(): - """Test that cartridge tensor names follow the expected pattern.""" - layer, exp, proj, rank, label = 3, 0, "gate_proj", 0, "res1" - prefix = f"model.layers.{layer}.mlp.experts.{exp}.{proj}.rank{rank}" - names = [ - f"{prefix}.trellis_{label}", - f"{prefix}.suh_{label}", - f"{prefix}.svh_{label}", - f"{prefix}.scale_{label}", + assert recipe["stages"][1]["k"] == 1 + assert len(recipe["stages"][1]["experts"]) == 96 + assert lora.effective_bpw(recipe, 256) == pytest.approx(3.375) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"schema": "fq-cartridge/2"}, "schema must be"), + ({"base_k": 0}, "base_k must be"), + ({"moe_layers": []}, "moe_layers"), + ({"stages": [{"k": 1, "label": "../../escape", "experts": "all"}]}, + "label must match"), + ({"stages": [{"k": 1, "label": "x", "experts": "hot96"}]}, + "experts must be"), + ({"stages": [ + {"k": 1, "label": "same", "experts": "all"}, + {"k": 1, "label": "same", "experts": [0]}, + ]}, "duplicate label"), + ], +) +def test_recipe_validation_rejects_unsafe_or_ambiguous_input( + tmp_path: Path, overrides: dict, message: str +): + path = write_recipe(tmp_path / "recipe.json", **overrides) + with pytest.raises((lora.CartridgeError, ValueError), match=message): + lora.load_recipe(path) + + +def test_stage_selection_supports_disjoint_expert_sets(): + stages = [ + {"k": 1, "label": "cold", "experts": [1]}, + {"k": 2, "label": "hot", "experts": [2]}, ] - for name in names: - assert f"trellis_{label}" in name or f"suh_{label}" in name \ - or f"svh_{label}" in name or f"scale_{label}" in name - - -def test_msrt_bpw_calculation(): - """Test effective bpw calculation for dual-cartridge configs.""" - n_experts = 256 - base_k = 2 - # Stage 1: K1 on all 256 experts - # Stage 2: K2 on 96 experts - total_bits = n_experts * base_k + n_experts * 1 + 96 * 2 - eff_bpw = total_bits / n_experts - assert eff_bpw == (512 + 256 + 192) / 256 # = 960/256 = 3.75 - - # willfalco comparison: 148 K3 + 108 K4 - willfalco_bpw = (148 * 3 + 108 * 4) / 256 - assert abs(willfalco_bpw - 3.422) < 0.01 - - -def test_encoding_summary_format(): - """Test that encoding summary has required fields.""" - summary = { - "tool": "fq_assemble_lora/1", - "base_k": 2, - "stages": [{"k": 1, "label": "res1", "experts": "all"}], - "moe_layers": [3, 4, 5], - "overall_mse": 0.001, - "n_experts_encoded": 768, - } - assert summary["tool"] == "fq_assemble_lora/1" - assert "overall_mse" in summary - assert summary["n_experts_encoded"] == 768 # 256 experts × 3 layers + assert [stage["label"] for stage in lora.selected_stages(stages, 1)] == ["cold"] + assert [stage["label"] for stage in lora.selected_stages(stages, 2)] == ["hot"] + assert lora.selected_stages(stages, 3) == [] -if __name__ == "__main__": - pytest.main([__file__, "-v"]) +def test_source_key_partition_preserves_router_norm_and_attention(): + keys = complete_keys() + [ + "model.layers.3.mlp.gate.weight", + "model.layers.3.self_attn.q_proj.weight", + "model.layers.3.input_layernorm.weight", + ] + experts = lora.inspect_source_layer(keys, 3) + assert set(experts[0]) == set(lora.PROJECTIONS) + assert lora.preserved_source_keys(keys) == keys[3:] + + +def test_missing_projection_fails_before_encoding(): + keys = complete_keys() + keys.remove(expert_key(3, 0, "gate_proj")) + with pytest.raises(lora.CartridgeError, match="projections"): + lora.inspect_source_layer(keys, 3) + + +def test_foreign_layer_in_per_layer_shard_is_rejected(): + keys = complete_keys(3, 0) + complete_keys(4, 0) + with pytest.raises(lora.CartridgeError, match="also contains"): + lora.inspect_source_layer(keys, 3) + + +def test_unsupported_source_layout_and_empty_layers_fail(tmp_path: Path): + source = tmp_path / "source" + source.mkdir() + (source / "config.json").write_text("{}") + (source / "model-00001-of-00002.safetensors").write_bytes(b"not-used") + with pytest.raises(lora.CartridgeError, match="not supported"): + lora.resolve_layer_shards(source, [3]) + with pytest.raises(lora.CartridgeError, match="moe_layers"): + lora.load_recipe(write_recipe(tmp_path / "empty.json", moe_layers=[])) + + +def test_checkpoint_copy_preserves_unselected_shards_and_drops_stale_metadata( + tmp_path: Path, +): + source = tmp_path / "source" + output = tmp_path / "base" + source.mkdir() + (source / "config.json").write_text("{}") + (source / "model-layer-002.safetensors").write_bytes(b"dense") + (source / "model-layer-003.safetensors").write_bytes(b"selected") + (source / "model.safetensors.index.json").write_text("stale") + (source / "MANIFEST.sha256").write_text("stale") + lora.copy_source_checkpoint(source, output, {"model-layer-003.safetensors"}) + assert (output / "model-layer-002.safetensors").read_bytes() == b"dense" + assert not (output / "model-layer-003.safetensors").exists() + assert not (output / "model.safetensors.index.json").exists() + assert not (output / "MANIFEST.sha256").exists() + + +def identity_hadamard(size, device, dtype, scale): + torch = pytest.importorskip("torch") + return torch.eye(size, device=device, dtype=dtype) + + +def test_zero_weight_serialized_scales_are_finite_and_round_trip(): + torch = pytest.importorskip("torch") + if lora.torch is None: + pytest.skip("fq_assemble_lora imported without torch") + weight = torch.zeros(128, 128) + regularized, suh, svh = lora.regularize_with_vectors( + weight, torch.device("cpu"), identity_hadamard, 1.2437) + assert torch.isfinite(suh).all() and torch.isfinite(svh).all() + assert (suh != 0).all() and (svh != 0).all() + restored = lora.inverse_regularize( + regularized, suh, svh, torch.device("cpu"), identity_hadamard) + assert torch.equal(restored, weight) + + +class FakeExtension: + def pack_trellis(self, packed, raw, bits): + packed.zero_() + + +def identity_permutation(device): + torch = pytest.importorskip("torch") + return torch.arange(256, device=device) + + +def identity_quantizer(tiles, options): + torch = pytest.importorskip("torch") + return tiles.clone(), torch.zeros_like(tiles, dtype=torch.int16) + + +def test_encoder_keys_stages_by_label_and_emits_packed_shape(): + torch = pytest.importorskip("torch") + if lora.torch is None: + pytest.skip("fq_assemble_lora imported without torch") + weight = torch.zeros(128, 128) + result = lora.encode_expert_msrt( + weight, 2, + [{"k": 1, "label": "only", "experts": "all"}], + torch.device("cpu"), identity_hadamard, + identity_permutation, identity_permutation, identity_quantizer, + 1.2437, FakeExtension(), + ) + assert set(result["stages"]) == {"only"} + assert result["base"]["trellis"].shape == (8, 8, 32) + assert result["stages"]["only"]["trellis"].shape == (8, 8, 16) + assert result["mse"] == 0.0 + + +def test_packed_quantization_refuses_missing_runtime_packer(): + torch = pytest.importorskip("torch") + if lora.torch is None: + pytest.skip("fq_assemble_lora imported without torch") + with pytest.raises(RuntimeError, match="pack_trellis"): + lora.quantize_trellis_packed( + torch.zeros(128, 128), 2, torch.device("cpu"), + identity_permutation, identity_permutation, identity_quantizer, None) + + +def test_base_metadata_is_regenerated_from_written_shards(tmp_path: Path): + torch = pytest.importorskip("torch") + pytest.importorskip("safetensors") + if lora.torch is None: + pytest.skip("fq_assemble_lora imported without torch") + from safetensors.torch import save_file + + base = tmp_path / "base" + base.mkdir() + (base / "config.json").write_text("{}") + save_file({"kept": torch.ones(1)}, str(base / "model-layer-002.safetensors")) + save_file({"quant": torch.ones(1)}, str(base / "model-layer-003.safetensors")) + (base / "MANIFEST.sha256").write_text("stale") + lora.write_base_metadata(base, 2, [3], {3: 1}) + + config = json.loads((base / "config.json").read_text()) + index = json.loads((base / "model.safetensors.index.json").read_text()) + manifest = (base / "MANIFEST.sha256").read_text() + assert config["quantization_config"]["quant_method"] == "exl3" + assert config["hybrid_tr3_tail"]["bits"] == 2.0 + assert index["weight_map"] == { + "kept": "model-layer-002.safetensors", + "quant": "model-layer-003.safetensors", + } + assert "stale" not in manifest + assert "model-layer-002.safetensors" in manifest diff --git a/tests/test_schemas.py b/tests/test_schemas.py index fcee94c..df3e3b1 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -167,6 +167,13 @@ def test_release_output_matches_the_schema(emitted): pointer="#/$defs/payload", label="release payload") +def test_fruit_cartridge_recipe_matches_schema(): + recipe = ( + Path(__file__).parent.parent / "recipes" / "fruit-k2-k3k4-cart.json" + ) + check("fq-cartridge-1", json.loads(recipe.read_text()), label=recipe.name) + + # ----------------------------------------------- documents that must NOT pass @pytest.mark.parametrize("doc,name,why", [ diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index c3c8269..2e628eb 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -1,17 +1,16 @@ #!/usr/bin/env python3 -"""fq_assemble_lora — Encode MSRT residual cartridges as LoRA-compatible adapters. +"""fq_assemble_lora — Encode BF16 weights as MSRT EXL3 cartridges. -Given a BF16 (or EXL3) source checkpoint and a cartridge recipe, this tool: - 1. Quantizes the base tier (K2 or K3) into standard EXL3 trellis format - 2. Computes residuals, rescales, and quantizes each residual stage - 3. Emits two outputs: - - A base checkpoint (standard EXL3 safetensors, loads normally in vLLM) - - One or more cartridge adapters (safetensors with per-stage trellis - tensors, loadable as LoRA adapters via vLLM's add_lora API) +Given a BF16 checkpoint and an ``fq-cartridge/1`` recipe, this tool: + 1. emits a complete, loadable EXL3 base checkpoint; + 2. quantizes selected residual stages with MSRT; and + 3. emits sharded custom cartridge adapters plus an explicit runtime contract. -The cartridge adapter is NOT a low-rank LoRA — it contains full-rank trellis- -quantized residual weights. The vLLM EXL3 LoRA wrapper (Exl3LoRAMoMethod) -applies them by running additional exl3_gemm passes and summing with rescaling. +Cartridges are full-rank additive trellis weights, not PEFT/LoRA matrices. +Their execution pattern is LoRA-like (base GEMM plus correction GEMMs), but +standard vLLM/SGLang ``add_lora`` APIs cannot load them without an EXL3 MSRT +runtime implementation. The emitted ``fq-cartridge-adapter/1`` config records +that custom contract instead of claiming standard LoRA compatibility. MSRT (Multi-Stage Rescaled Trellis) is described in: research/fungible-quant/poc/V50-LOW-BITRATE-MSRT.md @@ -42,102 +41,163 @@ import json import math import os +import re +import shutil import struct import sys import time +from contextlib import contextmanager from pathlib import Path from typing import Any -import torch +try: + import torch +except ModuleNotFoundError: # Base installs must still support --help. + torch = None + +sys.path.insert(0, str(Path(__file__).parent)) +from fq_assemble import ( # noqa: E402 + AssemblyError, + StagedOutput, + check_out_dir, + regenerate_manifest, + regenerate_shard_index, +) +from fq_repack import PROJ_ORDER # noqa: E402 # ── Constants ────────────────────────────────────────────────────────────── -TOOL_VERSION = "fq_assemble_lora/1" +TOOL_VERSION = "fq_assemble_lora/2" CARTRIDGE_SCHEMA = "fq-cartridge/1" ADAPTER_CONFIG_SCHEMA = "fq-cartridge-adapter/1" HADAMARD_BLOCK = 128 - -# Expert tensor name pattern — matches both BF16 (.weight) and EXL3 (.rank0.trellis) -EXPERT_RE_PATTERN = ( +MCG_SENTINEL_SIGNED = 0xCBAC1FED - (1 << 32) +LABEL_RE = re.compile(r"^[A-Za-z0-9_-]{1,32}$") +BF16_EXPERT_RE = re.compile( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." - r"(gate_proj|up_proj|down_proj)\.(?:rank\d+\.)?(?:weight|trellis)$" + r"(gate_proj|up_proj|down_proj)\.weight$" ) +PROJECTIONS = tuple(sorted(PROJ_ORDER, key=PROJ_ORDER.get)) -PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +class CartridgeError(RuntimeError): + """A recipe, source checkpoint, or encoded artifact is invalid.""" -# ── EXL3 Encoder Bootstrap ──────────────────────────────────────────────── -def bootstrap_encoder(encoder_source: str) -> tuple[Any, ...]: - """Load the EXL3 encoder without importing all of exllamav3. +def require_quant_dependencies() -> None: + if torch is None: + raise CartridgeError( + "MSRT encoding requires the 'quant' extra: " + "pip install 'progressive-tensors[quant]'") + try: + import safetensors # noqa: F401 + except ModuleNotFoundError as exc: + raise CartridgeError( + "MSRT encoding requires the 'quant' extra: " + "pip install 'progressive-tensors[quant]'") from exc + + +# ── EXL3 Encoder Bootstrap ──────────────────────────────────────────────── - Returns (ext, get_hadamard_dt, tensor_core_perm, tensor_core_perm_i, - quantize_tiles, codebook_scale). - """ +@contextmanager +def bootstrap_encoder(encoder_source: str): + """Load trusted EXL3 encoder modules temporarily and restore sys.modules.""" import importlib.util import types - pkg_root = Path(encoder_source) - pkg = types.ModuleType("exllamav3") - pkg.__path__ = [str(pkg_root)] - sys.modules["exllamav3"] = pkg - - for sub in ["util", "modules", "modules.quant", "modules.quant.exl3_lib"]: - full = f"exllamav3.{sub}" - m = types.ModuleType(full) - m.__path__ = [str(pkg_root / sub.replace(".", "/"))] - sys.modules[full] = m - - # Stub progress/memory to avoid flash_attn dependency - _stub = types.ModuleType("exllamav3.util.progress") - class _DPB: - def __init__(self, *a, **kw): pass - def __enter__(self): return self - def __exit__(self, *a): return False - def update(self, *a): pass - def new_task(self, *a, **kw): pass - _stub.ProgressBar = _DPB - sys.modules["exllamav3.util.progress"] = _stub - - _stub = types.ModuleType("exllamav3.util.memory") - _stub.free_mem = lambda: None - _stub.list_gpu_tensors = lambda: [] - sys.modules["exllamav3.util.memory"] = _stub - - _stub = types.ModuleType("exllamav3.util") - _stub.__path__ = [str(pkg_root / "util")] - _stub.cuda_sync_active = lambda *a, **kw: torch.cuda.synchronize() - sys.modules["exllamav3.util"] = _stub - - _stub = types.ModuleType("exllamav3.util.tensor") - _stub.save_tensor_image = lambda *a, **kw: None - sys.modules["exllamav3.util.tensor"] = _stub - - # Load the extension - spec = importlib.util.spec_from_file_location( - "exllamav3.ext", str(pkg_root / "ext.py")) - ext_mod = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.ext"] = ext_mod - spec.loader.exec_module(ext_mod) - - # Load Hadamard - spec = importlib.util.spec_from_file_location( - "exllamav3.util.hadamard", str(pkg_root / "util" / "hadamard.py")) - had_mod = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.util.hadamard"] = had_mod - spec.loader.exec_module(had_mod) - - # Load quantize module - quant_path = pkg_root / "modules" / "quant" / "exl3_lib" / "quantize.py" - spec = importlib.util.spec_from_file_location( - "exllamav3.modules.quant.exl3_lib.quantize", str(quant_path)) - quant_mod = importlib.util.module_from_spec(spec) - sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = quant_mod - spec.loader.exec_module(quant_mod) - - return (ext_mod.exllamav3_ext, had_mod.get_hadamard_dt, + pkg_root = Path(encoder_source).expanduser().resolve() + required = ( + pkg_root / "ext.py", + pkg_root / "util" / "hadamard.py", + pkg_root / "modules" / "quant" / "exl3_lib" / "quantize.py", + ) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise CartridgeError( + f"--encoder-source {pkg_root} is not an exllamav3 checkout; " + f"missing {missing}") + + names = [ + "exllamav3", + "exllamav3.util", + "exllamav3.modules", + "exllamav3.modules.quant", + "exllamav3.modules.quant.exl3_lib", + "exllamav3.util.progress", + "exllamav3.util.memory", + "exllamav3.util.tensor", + "exllamav3.ext", + "exllamav3.util.hadamard", + "exllamav3.modules.quant.exl3_lib.quantize", + ] + previous = {name: sys.modules.get(name) for name in names} + try: + pkg = types.ModuleType("exllamav3") + pkg.__path__ = [str(pkg_root)] + sys.modules["exllamav3"] = pkg + for sub in ["util", "modules", "modules.quant", "modules.quant.exl3_lib"]: + full = f"exllamav3.{sub}" + module = types.ModuleType(full) + module.__path__ = [str(pkg_root / sub.replace(".", "/"))] + sys.modules[full] = module + + progress = types.ModuleType("exllamav3.util.progress") + + class _DisabledProgress: + def __init__(self, *args, **kwargs): pass + def __enter__(self): return self + def __exit__(self, *args): return False + def update(self, *args): pass + def new_task(self, *args, **kwargs): pass + + progress.ProgressBar = _DisabledProgress + sys.modules["exllamav3.util.progress"] = progress + + memory = types.ModuleType("exllamav3.util.memory") + memory.free_mem = lambda: None + memory.list_gpu_tensors = lambda: [] + sys.modules["exllamav3.util.memory"] = memory + + util = types.ModuleType("exllamav3.util") + util.__path__ = [str(pkg_root / "util")] + util.cuda_sync_active = ( + lambda *args, **kwargs: + torch.cuda.synchronize() if torch.cuda.is_available() else None + ) + sys.modules["exllamav3.util"] = util + + tensor = types.ModuleType("exllamav3.util.tensor") + tensor.save_tensor_image = lambda *args, **kwargs: None + sys.modules["exllamav3.util.tensor"] = tensor + + def load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, str(path)) + if spec is None or spec.loader is None: + raise CartridgeError(f"cannot load encoder module {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + ext_mod = load("exllamav3.ext", required[0]) + had_mod = load("exllamav3.util.hadamard", required[1]) + quant_mod = load( + "exllamav3.modules.quant.exl3_lib.quantize", required[2]) + ext = getattr(ext_mod, "exllamav3_ext", None) + if not callable(getattr(ext, "pack_trellis", None)): + raise CartridgeError( + f"{pkg_root}: exllamav3 extension lacks pack_trellis") + yield ( + ext, had_mod.get_hadamard_dt, quant_mod.tensor_core_perm, quant_mod.tensor_core_perm_i, - quant_mod.quantize_tiles, quant_mod.codebook_scale) + quant_mod.quantize_tiles, quant_mod.codebook_scale, + ) + finally: + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # ── Quantization Primitives ──────────────────────────────────────────────── @@ -147,7 +207,31 @@ def block_rms(x: torch.Tensor, dim: int, keepdim: bool = False) -> torch.Tensor: return x.square().mean(dim=dim, keepdim=keepdim).sqrt() -def regularize( +def validate_quant_shape(w: torch.Tensor, *, who: str = "weight") -> tuple[int, int]: + """Return a valid EXL3 matrix shape or raise a useful error.""" + if w.ndim != 2: + raise ValueError(f"{who}: expected a 2-D BF16 weight, got shape {tuple(w.shape)}") + k, n = w.shape + if k % HADAMARD_BLOCK or n % HADAMARD_BLOCK: + raise ValueError( + f"{who}: shape {(k, n)} must be divisible by Hadamard block " + f"{HADAMARD_BLOCK} on both axes") + if k % 16 or n % 16: + raise ValueError(f"{who}: shape {(k, n)} must be divisible by trellis tile 16") + return k, n + + +def _finite_fp16_scale(x: torch.Tensor, sign: torch.Tensor) -> torch.Tensor: + """Round scales once while keeping every divisor finite and non-zero.""" + minimum = torch.finfo(torch.float16).tiny + safe = sign * x.abs().clamp_min(minimum) + rounded = safe.to(torch.float16) + if not torch.isfinite(rounded).all() or (rounded == 0).any(): + raise ValueError("Hadamard scale vector contains zero or non-finite values") + return rounded + + +def regularize_with_vectors( w: torch.Tensor, device: torch.device, ghd: Any, @@ -155,34 +239,87 @@ def regularize( had_k: int = HADAMARD_BLOCK, had_n: int = HADAMARD_BLOCK, seed: int = 0, -) -> torch.Tensor: - """Apply Hadamard regularization (in-place transform, returns new tensor). +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Regularize with the exact FP16 vectors that the checkpoint stores.""" + k, n = validate_quant_shape(w) + if not math.isfinite(float(cbs)) or float(cbs) == 0: + raise ValueError(f"codebook_scale must be finite and non-zero, got {cbs!r}") - This matches the EXL3 regularize() used in the PoC scripts (v35-v52). - The Hadamard is orthogonal, so MSE in regularized space = MSE in original. - """ - k, n = w.shape g = torch.Generator(device="cpu").manual_seed(seed) - su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) - sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) + su_sign = (torch.randn(k, generator=g).sign() + 1e-5).sign().to(device) + sv_sign = (torch.randn(n, generator=g).sign() + 1e-5).sign().to(device) - out_scales = block_rms(w, dim=0, keepdim=True) - mean = out_scales.mean().item() - if mean > 1e-30: + out_scales = block_rms(w, dim=0) + mean = out_scales.mean() + if torch.isfinite(mean) and mean.item() > 0: out_scales = out_scales / mean - sv = (sv * out_scales + 1e-10).float() - w = (w / sv).contiguous() + svh = _finite_fp16_scale(out_scales, sv_sign) + transformed = (w / svh.float().unsqueeze(0)).contiguous() had_n_mat = ghd(had_n, device, torch.float, 1.0 / math.sqrt(had_n)) - w = (w.view(k, n // had_n, had_n) @ had_n_mat).view(k, n).contiguous() + transformed = ( + transformed.view(k, n // had_n, had_n) @ had_n_mat + ).view(k, n).contiguous() + + in_scales = block_rms(transformed, dim=1) + suh_sign = su_sign * (-1.0 if cbs > 0 else 1.0) + suh = _finite_fp16_scale(in_scales / abs(float(cbs)), suh_sign) + transformed = (transformed / suh.float().unsqueeze(1)).contiguous() + + had_k_mat = ghd(had_k, device, torch.float, 1.0 / math.sqrt(had_k)) + transformed = ( + had_k_mat @ transformed.view(k // had_k, had_k, n) + ).view(k, n).contiguous() + if not torch.isfinite(transformed).all(): + raise ValueError("regularized weight contains non-finite values") + return transformed, suh.contiguous(), svh.contiguous() + + +def regularize( + w: torch.Tensor, + device: torch.device, + ghd: Any, + cbs: float, + had_k: int = HADAMARD_BLOCK, + had_n: int = HADAMARD_BLOCK, + seed: int = 0, +) -> torch.Tensor: + """Compatibility wrapper returning only the regularized weight.""" + return regularize_with_vectors( + w, device, ghd, cbs, had_k=had_k, had_n=had_n, seed=seed + )[0] - in_scales = block_rms(w, dim=1, keepdim=True).clamp(min=1e-30) - su = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() - w = (w / su).contiguous() +def inverse_regularize( + w_reg: torch.Tensor, + suh: torch.Tensor, + svh: torch.Tensor, + device: torch.device, + ghd: Any, + had_k: int = HADAMARD_BLOCK, + had_n: int = HADAMARD_BLOCK, +) -> torch.Tensor: + """Invert regularization using the exact serialized FP16 scale vectors.""" + k, n = validate_quant_shape(w_reg, who="reconstruction") had_k_mat = ghd(had_k, device, torch.float, 1.0 / math.sqrt(had_k)) - w = (had_k_mat @ w.view(k // had_k, had_k, n)).view(k, n).contiguous() - return w + restored = ( + had_k_mat.transpose(0, 1) + @ w_reg.view(k // had_k, had_k, n) + ).view(k, n) + restored = restored * suh.float().unsqueeze(1) + had_n_mat = ghd(had_n, device, torch.float, 1.0 / math.sqrt(had_n)) + restored = ( + restored.view(k, n // had_n, had_n) + @ had_n_mat.transpose(0, 1) + ).view(k, n) + restored = restored * svh.float().unsqueeze(0) + return restored.contiguous() + + +def _validate_k(K: int, *, who: str = "K") -> int: + if isinstance(K, bool) or not isinstance(K, int) or not 1 <= K <= 8: + raise ValueError(f"{who} must be an integer in 1..8, got {K!r}") + return K def quantize_trellis( @@ -193,12 +330,9 @@ def quantize_trellis( tcpi: Any, qtf: Any, ) -> torch.Tensor: - """Quantize a 2D tensor with EXL3 trellis at K bits. - - Returns the dequantized (reconstructed) tensor, NOT the packed indices. - The trellis tiles are 16×16, processed row-by-row in blocks of 16. - """ - k, n = data.shape + """Quantize a valid 2-D EXL3 matrix and return its reconstruction.""" + _validate_k(K) + k, n = validate_quant_shape(data) tiles_n = n // 16 weight_q = torch.zeros_like(data) qa = {"K": K, "mcg": True} @@ -206,13 +340,16 @@ def quantize_trellis( perm_i = tcpi(device) for bi in range(0, k, 16): - rows = data[bi:bi + 16] - tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) - tiles = tiles[:, perm].contiguous() - quant_w, _ = qtf(tiles, qa) - quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) + tiles = ( + data[bi:bi + 16].reshape(16, tiles_n, 16) + .permute(1, 0, 2).reshape(tiles_n, 256) + ) + quant_w, _ = qtf(tiles[:, perm].contiguous(), qa) + quant_w = ( + quant_w[:, perm_i].reshape(tiles_n, 16, 16) + .permute(1, 0, 2).reshape(16, n) + ) weight_q[bi:bi + 16] = quant_w - return weight_q @@ -223,48 +360,42 @@ def quantize_trellis_packed( tcp: Any, tcpi: Any, qtf: Any, - ext: Any = None, + ext: Any, ) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize and return BOTH reconstructed values and packed trellis indices. - - Returns (reconstructed_float, packed_trellis_int16) where packed_trellis - has shape (k // 16, n // 16, K * 16) dtype int16 — the EXL3 storage format. - - The quantize_tiles function returns raw Viterbi path indices of shape - (n_tiles, 256) — one int16 per weight element. ext.pack_trellis compresses - these to (n_tiles, K * 16) packed indices, which is what the EXL3 checkpoint - and vLLM loader expect (validated at exl3.py:2099-2102). - """ - k, n = data.shape + """Return reconstruction plus the runtime-compatible packed EXL3 trellis.""" + _validate_k(K) + if ext is None or not callable(getattr(ext, "pack_trellis", None)): + raise RuntimeError( + "the selected exllamav3 build lacks pack_trellis; refusing to " + "write raw Viterbi indices as an EXL3 checkpoint") + k, n = validate_quant_shape(data) tiles_n = n // 16 weight_q = torch.zeros_like(data) - # Raw (unpacked) indices: (tiles_k, tiles_n, 256) int16 - raw_indices = torch.zeros(k // 16, tiles_n, 256, dtype=torch.int16, device=device) + raw_indices = torch.zeros( + k // 16, tiles_n, 256, dtype=torch.int16, device=device + ) qa = {"K": K, "mcg": True} perm = tcp(device) perm_i = tcpi(device) for bi in range(0, k, 16): tk = bi // 16 - rows = data[bi:bi + 16] - tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) - tiles = tiles[:, perm].contiguous() - quant_w, quant_idx = qtf(tiles, qa) - # Store raw indices — quant_idx has shape (n_tiles, 256) + tiles = ( + data[bi:bi + 16].reshape(16, tiles_n, 16) + .permute(1, 0, 2).reshape(tiles_n, 256) + ) + quant_w, quant_idx = qtf(tiles[:, perm].contiguous(), qa) raw_indices[tk] = quant_idx - # Reconstruct - quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) + quant_w = ( + quant_w[:, perm_i].reshape(tiles_n, 16, 16) + .permute(1, 0, 2).reshape(16, n) + ) weight_q[bi:bi + 16] = quant_w - # Pack the raw indices to EXL3 format: (tiles_k, tiles_n, K * 16) int16 - if ext is not None and hasattr(ext, "pack_trellis"): - packed_shape = (k // 16, tiles_n, 256 * K // 16) - packed = torch.zeros(packed_shape, dtype=torch.int16, device=device) - ext.pack_trellis(packed, raw_indices.contiguous(), K) - else: - # Fallback: store raw indices (unpacked) — for testing without ext - packed = raw_indices - + packed = torch.zeros( + (k // 16, tiles_n, K * 16), dtype=torch.int16, device=device + ) + ext.pack_trellis(packed, raw_indices.contiguous(), K) return weight_q, packed @@ -275,26 +406,26 @@ def rescaled_trellis_quantize( device: torch.device, tcp: Any, tcpi: Any, qtf: Any, cbs: float, - ext: Any = None, + ext: Any, ) -> tuple[torch.Tensor, torch.Tensor, float]: - """Rescale residual to match codebook range, quantize, return (recon, packed, scale). - - The rescaling is the key MSRT innovation (v35 breakthrough): - scale = |codebook_scale| / RMS(residual) - quantized = trellis(residual * scale) / scale - """ + """Quantize one rescaled residual into runtime-compatible trellis form.""" + _validate_k(K_res, who="residual K") + k, n = validate_quant_shape(residual, who="residual") + if ext is None or not callable(getattr(ext, "pack_trellis", None)): + raise RuntimeError("exllamav3 pack_trellis is required") residual_rms = residual.square().mean().sqrt().item() + if not math.isfinite(residual_rms): + raise ValueError("residual RMS is non-finite") if residual_rms < 1e-12: - packed_w = 256 * K_res // 16 if ext else 256 return base_q, torch.zeros( - residual.shape[0] // 16, residual.shape[1] // 16, packed_w, + k // 16, n // 16, K_res * 16, dtype=torch.int16, device=device), 1.0 - scale = abs(cbs) / residual_rms - scaled = residual * scale - recon_packed, packed = quantize_trellis_packed(scaled, K_res, device, tcp, tcpi, qtf, ext) - recon = base_q + recon_packed / scale - return recon, packed, scale + scale = abs(float(cbs)) / residual_rms + recon_scaled, packed = quantize_trellis_packed( + residual * scale, K_res, device, tcp, tcpi, qtf, ext + ) + return base_q + recon_scaled / scale, packed, scale # ── Hadamard Vectors ────────────────────────────────────────────────────── @@ -306,36 +437,9 @@ def compute_hadamard_vectors( cbs: float, seed: int = 0, ) -> dict[str, torch.Tensor]: - """Compute the suh and svh vectors that EXL3 stores alongside trellis. - - EXL3 checkpoint format (from SIQ model inspection): - - suh: (input_size,) float16 — per-row sign+scale vector - - svh: (output_size,) float16 — per-column sign+scale vector - - The regularize() function applies: - 1. svh = sign * out_scales (per-column) - 2. suh = sign * in_scales / (-cbs) (per-row, after column Hadamard) - """ - k, n = w.shape - g = torch.Generator(device="cpu").manual_seed(seed) - su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) - sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) - - out_scales = block_rms(w, dim=0, keepdim=True) - mean = out_scales.mean().item() - if mean > 1e-30: - out_scales = out_scales / mean - svh = (sv * out_scales.squeeze(0) + 1e-10).half() # (n,) float16 - - # After column Hadamard - w_col = (w / svh.float().unsqueeze(0)).contiguous() - had_n_mat = ghd(HADAMARD_BLOCK, device, torch.float, 1.0 / math.sqrt(HADAMARD_BLOCK)) - w_col = (w_col.view(k, n // HADAMARD_BLOCK, HADAMARD_BLOCK) @ had_n_mat).view(k, n).contiguous() - - in_scales = block_rms(w_col, dim=1, keepdim=True).clamp(min=1e-30).squeeze(1) - suh = (su * in_scales / (-cbs) + 1e-10).half() # (k,) float16 - - return {"suh": suh.contiguous(), "svh": svh.contiguous()} + """Return the same finite FP16 vectors used by regularization.""" + _, suh, svh = regularize_with_vectors(w, device, ghd, cbs, seed=seed) + return {"suh": suh, "svh": svh} # ── Encoding Pipeline ───────────────────────────────────────────────────── @@ -347,388 +451,494 @@ def encode_expert_msrt( device: torch.device, ghd: Any, tcp: Any, tcpi: Any, qtf: Any, cbs: float, - ext: Any = None, + ext: Any, ) -> dict[str, Any]: - """Encode one expert weight matrix with MSRT. - - Returns dict with: - - base: {trellis, suh, svh} for the base tier - - stages: list of {trellis, suh, svh, scale} for each residual stage - - mse: reconstruction MSE vs original - """ - w_reg = regularize(w_bf16, device, ghd, cbs) - - # Base tier - base_recon, base_packed = quantize_trellis_packed(w_reg, base_k, device, tcp, tcpi, qtf, ext) - had_vectors = compute_hadamard_vectors(w_bf16, device, ghd, cbs) - + """Encode one matrix and measure the reconstruction actually emitted.""" + _validate_k(base_k, who="base_k") + w_reg, suh, svh = regularize_with_vectors(w_bf16, device, ghd, cbs) + base_recon, base_packed = quantize_trellis_packed( + w_reg, base_k, device, tcp, tcpi, qtf, ext + ) result = { "base": { "trellis": base_packed.cpu(), - "suh": had_vectors["suh"].cpu(), - "svh": had_vectors["svh"].cpu(), + "suh": suh.cpu(), + "svh": svh.cpu(), }, - "stages": [], + "stages": {}, } current_recon = base_recon - for stage in stages: + label = stage["label"] residual = w_reg - current_recon - recon, packed, scale = rescaled_trellis_quantize( - current_recon, residual, stage["k"], device, tcp, tcpi, qtf, cbs, ext) - result["stages"].append({ + current_recon, packed, scale = rescaled_trellis_quantize( + current_recon, residual, stage["k"], device, + tcp, tcpi, qtf, cbs, ext + ) + result["stages"][label] = { "trellis": packed.cpu(), - "suh": had_vectors["suh"].cpu(), # Same Hadamard vectors - "svh": had_vectors["svh"].cpu(), + "suh": suh.cpu(), + "svh": svh.cpu(), "scale": scale, - }) - current_recon = recon + } - result["mse"] = (w_reg - current_recon).pow(2).mean().item() + reconstructed = inverse_regularize( + current_recon, suh, svh, device, ghd + ) + result["mse"] = ( + w_bf16.float() - reconstructed.float() + ).square().mean().item() + result["regularized_mse"] = ( + w_reg - current_recon + ).square().mean().item() return result # ── Safetensors Output ──────────────────────────────────────────────────── def save_safetensors(tensors: dict[str, torch.Tensor], path: Path) -> None: - """Save tensors as safetensors file.""" + """Save an independent tensor mapping after validating finite metadata.""" from safetensors.torch import save_file + + if not tensors: + raise CartridgeError(f"refusing to write empty safetensors file {path}") + for name, tensor in tensors.items(): + if tensor.is_floating_point() and not torch.isfinite(tensor).all(): + raise CartridgeError(f"{path}: tensor {name} contains non-finite values") path.parent.mkdir(parents=True, exist_ok=True) save_file(tensors, str(path)) -def write_base_checkpoint( - source_dir: Path, - out_dir: Path, - layer_results: dict[int, dict[int, dict[str, Any]]], - moe_layers: list[int], - tp: int = 1, -) -> None: - """Write the base K checkpoint in standard EXL3 format. - - The base checkpoint has the same structure as a normal EXL3 quant: - - config.json with hybrid_tr3_tail - - tier_bitmap.json - - model-layer-*.safetensors with trellis, suh, svh, mcg tensors - """ - out_dir.mkdir(parents=True, exist_ok=True) - - # Copy non-layer files from source - import shutil - for f in source_dir.iterdir(): - if f.is_file() and not f.name.startswith("model-layer-"): - shutil.copy2(f, out_dir / f.name) - - # Write base tensors into layer shards - for layer in moe_layers: - if layer not in layer_results: +def load_recipe(path: Path) -> dict[str, Any]: + """Load and semantically validate one fq-cartridge/1 recipe.""" + try: + recipe = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise CartridgeError(f"{path}: cannot read cartridge recipe ({exc})") from exc + if not isinstance(recipe, dict): + raise CartridgeError(f"{path}: recipe must be a JSON object") + if recipe.get("schema") != CARTRIDGE_SCHEMA: + raise CartridgeError( + f"{path}: schema must be {CARTRIDGE_SCHEMA!r}, " + f"got {recipe.get('schema')!r}") + _validate_k(recipe.get("base_k"), who="base_k") + + layers = recipe.get("moe_layers") + if (not isinstance(layers, list) or not layers + or any(isinstance(v, bool) or not isinstance(v, int) or v < 0 + for v in layers) + or len(set(layers)) != len(layers)): + raise CartridgeError("moe_layers must be a non-empty list of unique integers") + recipe["moe_layers"] = sorted(layers) + + stages = recipe.get("stages") + if not isinstance(stages, list) or not stages: + raise CartridgeError("stages must be a non-empty list") + labels: set[str] = set() + for index, stage in enumerate(stages): + if not isinstance(stage, dict): + raise CartridgeError(f"stage {index}: must be an object") + _validate_k(stage.get("k"), who=f"stage {index} k") + label = stage.get("label") + if not isinstance(label, str) or not LABEL_RE.fullmatch(label): + raise CartridgeError( + f"stage {index}: label must match {LABEL_RE.pattern}") + if label in labels: + raise CartridgeError(f"stage {index}: duplicate label {label!r}") + labels.add(label) + experts = stage.get("experts") + if experts == "all": continue - tensors = {} - for exp_id, exp_data in layer_results[layer].items(): - base = exp_data["base"] - for proj_idx, proj in enumerate(PROJECTIONS): - rank = 0 # TP=1 for Fruit model - prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" - tensors[f"{prefix}.trellis"] = base["trellis"] - tensors[f"{prefix}.suh"] = base["suh"] - tensors[f"{prefix}.svh"] = base["svh"] - # mcg sentinel - import hashlib - mcg_val = 0xCBAC1FED - tensors[f"{prefix}.mcg"] = torch.tensor([mcg_val], dtype=torch.int32) - - shard_path = out_dir / f"model-layer-{layer:03d}.safetensors" - save_safetensors(tensors, shard_path) - print(f" Base layer {layer}: {len(tensors)} tensors -> {shard_path.name}", flush=True) - - -def write_cartridge_adapter( - out_dir: Path, - layer_results: dict[int, dict[int, dict[str, Any]]], - stage_idx: int, - stage_label: str, - moe_layers: list[int], - expert_filter: dict[int, list[int]] | None = None, -) -> Path: - """Write one cartridge stage as a LoRA-compatible safetensors adapter. - - If expert_filter is provided, only the specified experts per layer are included. - """ - out_dir.mkdir(parents=True, exist_ok=True) - tensors = {} - - for layer in moe_layers: - if layer not in layer_results: + if (not isinstance(experts, list) + or any(isinstance(v, bool) or not isinstance(v, int) or v < 0 + for v in experts) + or len(set(experts)) != len(experts)): + raise CartridgeError( + f"stage {label!r}: experts must be 'all' or unique non-negative IDs") + stage["experts"] = sorted(experts) + return recipe + + +def effective_bpw(recipe: dict[str, Any], expert_count: int) -> float: + """Nominal weight bits, excluding suh/svh metadata.""" + if expert_count <= 0: + raise ValueError("expert_count must be positive") + total = expert_count * recipe["base_k"] + for stage in recipe["stages"]: + selected = expert_count if stage["experts"] == "all" else len(stage["experts"]) + total += selected * stage["k"] + return total / expert_count + + +def selected_stages( + stages: list[dict[str, Any]], expert_id: int +) -> list[dict[str, Any]]: + """Resolve stage applicability before residual chaining begins.""" + return [ + stage for stage in stages + if stage["experts"] == "all" or expert_id in stage["experts"] + ] + + +def resolve_layer_shards(source: Path, layers: list[int]) -> dict[int, Path]: + """Require the supported one-safetensors-file-per-layer BF16 layout.""" + if not source.is_dir(): + raise CartridgeError(f"--source {source} is not a directory") + if not (source / "config.json").is_file(): + raise CartridgeError(f"--source {source} has no config.json") + resolved: dict[int, Path] = {} + for layer in layers: + candidates = [ + source / f"model-layer-{layer:03d}.safetensors", + source / f"model-layer-{layer:04d}.safetensors", + ] + matches = [path for path in candidates if path.is_file()] + if len(matches) != 1: + suffix = ( + "standard Hugging Face indexed shards are not supported by " + "this encoder; convert to per-layer BF16 shards first" + if not matches else f"ambiguous candidates: {matches}" + ) + raise CartridgeError(f"layer {layer}: source shard missing; {suffix}") + resolved[layer] = matches[0] + return resolved + + +def inspect_source_layer(keys: list[str], layer: int) -> dict[int, dict[str, str]]: + """Preflight exact BF16 expert coverage without loading weight payloads.""" + experts: dict[int, dict[str, str]] = {} + for key in keys: + match = BF16_EXPERT_RE.fullmatch(key) + if not match: + continue + key_layer, expert, projection = ( + int(match.group(1)), int(match.group(2)), match.group(3)) + if key_layer != layer: + raise CartridgeError( + f"layer {layer} shard also contains BF16 expert tensor {key}") + if projection in experts.setdefault(expert, {}): + raise CartridgeError( + f"layer {layer} expert {expert}: duplicate {projection}") + experts[expert][projection] = key + if not experts: + raise CartridgeError( + f"layer {layer}: no BF16 routed expert .weight tensors found") + expected = set(PROJECTIONS) + for expert, projections in experts.items(): + have = set(projections) + if have != expected: + raise CartridgeError( + f"layer {layer} expert {expert}: projections " + f"{sorted(have)} != {sorted(expected)}") + return experts + + +def preserved_source_keys(keys: list[str]) -> list[str]: + """Source tensors copied byte-for-value into a rewritten MoE shard.""" + return [key for key in keys if not BF16_EXPERT_RE.fullmatch(key)] + + +def load_preserved_tensors(source_file: Any, keys: list[str]) -> dict[str, Any]: + """Materialize only non-expert tensors from an open source shard.""" + return { + key: source_file.get_tensor(key) + for key in preserved_source_keys(keys) + } + + +def copy_source_checkpoint( + source: Path, base_dir: Path, selected_shards: set[str] +) -> None: + """Copy all untouched checkpoint content, excluding stale integrity files.""" + base_dir.mkdir(parents=True, exist_ok=True) + skip = selected_shards | { + "model.safetensors.index.json", + "MANIFEST.sha256", + } + for entry in source.iterdir(): + if entry.name in skip: continue - for exp_id, exp_data in layer_results[layer].items(): - if expert_filter and layer in expert_filter: - if exp_id not in expert_filter[layer]: - continue - if stage_idx >= len(exp_data["stages"]): - continue - stage = exp_data["stages"][stage_idx] - for proj in PROJECTIONS: - rank = 0 - prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" - tensors[f"{prefix}.trellis_{stage_label}"] = stage["trellis"] - tensors[f"{prefix}.suh_{stage_label}"] = stage["suh"] - tensors[f"{prefix}.svh_{stage_label}"] = stage["svh"] - tensors[f"{prefix}.scale_{stage_label}"] = torch.tensor([stage["scale"]], dtype=torch.float32) - - adapter_path = out_dir / f"cartridge_{stage_label}.safetensors" - save_safetensors(tensors, adapter_path) - - # Write adapter_config.json + destination = base_dir / entry.name + if entry.is_dir(): + shutil.copytree(entry, destination) + elif entry.is_file(): + shutil.copy2(entry, destination) + + +def write_base_metadata( + base_dir: Path, + base_k: int, + layers: list[int], + experts_per_layer: dict[int, int], +) -> None: + """Synchronize loader-visible EXL3 config, bitmap, index, and manifest.""" + config_path = base_dir / "config.json" + try: + config = json.loads(config_path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise CartridgeError(f"{config_path}: invalid source config ({exc})") from exc + if not isinstance(config, dict): + raise CartridgeError(f"{config_path}: config must be an object") + + counts = set(experts_per_layer.values()) + tail = config.setdefault("hybrid_tr3_tail", {}) + if not isinstance(tail, dict): + raise CartridgeError("config.json: hybrid_tr3_tail must be an object") + tail.update({ + "format": "exl3-trellis", + "bits": float(base_k), + "codebook": "mcg", + "moe_layers": [min(layers), max(layers)], + "tensor_schema": ( + "model.layers.{L}.mlp.experts.{E}.{proj}.rank{rank}.{component}"), + "tp": 1, + "mcg_multiplier": 0xCBAC1FED, + }) + if len(counts) == 1: + tail["experts_per_layer"] = next(iter(counts)) + else: + tail.pop("experts_per_layer", None) + tail.pop("k_values", None) + tail.pop("bits_per_expert", None) + config["quantization_config"] = { + "quant_method": "exl3", + "bits": float(base_k), + "codebook": "mcg", + "version": "rank-sliced", + } + config_path.write_text(json.dumps(config, indent=2) + "\n") + + bitmap_path = base_dir / "tier_bitmap.json" + try: + bitmap = json.loads(bitmap_path.read_text()) if bitmap_path.exists() else {} + except (OSError, json.JSONDecodeError) as exc: + raise CartridgeError(f"{bitmap_path}: invalid source bitmap ({exc})") from exc + if not isinstance(bitmap, dict): + raise CartridgeError(f"{bitmap_path}: bitmap must be an object") + for layer, count in experts_per_layer.items(): + entry = bitmap.setdefault(str(layer), {}) + if not isinstance(entry, dict): + entry = {} + bitmap[str(layer)] = entry + entry["k"] = [base_k] * count + entry["bits_per_expert"] = [base_k] * count + bitmap_path.write_text(json.dumps(bitmap, indent=2) + "\n") + + regenerate_shard_index(base_dir) + regenerate_manifest(base_dir) + + +def _stage_tensor_names( + layer: int, expert: int, projection: str, label: str, + result: dict[str, Any], +) -> dict[str, torch.Tensor]: + prefix = ( + f"model.layers.{layer}.mlp.experts.{expert}." + f"{projection}.rank0") + return { + f"{prefix}.trellis_{label}": result["trellis"], + f"{prefix}.suh_{label}": result["suh"], + f"{prefix}.svh_{label}": result["svh"], + f"{prefix}.scale_{label}": torch.tensor( + result["scale"], dtype=torch.float32), + } + + +def write_adapter_config( + directory: Path, + stages: list[dict[str, Any]], + shards: list[str], + tensor_count: int, +) -> None: + """Write the explicit custom MSRT runtime contract.""" config = { "schema": ADAPTER_CONFIG_SCHEMA, - "stage_label": stage_label, - "stage_k": stage_idx, - "num_tensors": len(tensors), + "format": "exl3-msrt-full-rank", + "standard_lora_compatible": False, + "runtime_operation": ( + "base_exl3_gemm + sum(stage_exl3_gemm / stage_scale)"), + "codebook": "mcg", + "mcg_multiplier": 0xCBAC1FED, + "mcg_ownership": "adapter-config", + "scale_shape": [], + "stages": [ + {"label": stage["label"], "k": stage["k"], + "experts": stage["experts"]} + for stage in stages + ], + "shards": shards, + "num_tensors": tensor_count, "tool_version": TOOL_VERSION, + "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } - (out_dir / f"cartridge_{stage_label}_config.json").write_text( + directory.mkdir(parents=True, exist_ok=True) + (directory / "adapter_config.json").write_text( json.dumps(config, indent=2) + "\n") - print(f" Cartridge '{stage_label}': {len(tensors)} tensors -> {adapter_path.name}", flush=True) - return adapter_path - - -# ── Main Encode Command ─────────────────────────────────────────────────── def cmd_encode(args) -> int: - """Encode a BF16 checkpoint into base K + cartridge adapters.""" - device = torch.device(args.device) - print(f"Device: {device} GPU: {torch.cuda.get_device_name(0)}", flush=True) - - # Bootstrap encoder - ext, ghd, tcp, tcpi, qtf, cbs = bootstrap_encoder(args.encoder_source) - print(f"codebook_scale = {cbs}", flush=True) - - # Load recipe - recipe = json.loads(args.recipe.read_text()) + """Encode a per-layer BF16 checkpoint into base EXL3 plus MSRT shards.""" + require_quant_dependencies() + recipe = load_recipe(args.recipe) base_k = recipe["base_k"] stages = recipe["stages"] - moe_layers = recipe.get("moe_layers", []) - if not moe_layers: - # Auto-detect from source config - cfg_path = args.source / "config.json" - if cfg_path.exists(): - cfg = json.loads(cfg_path.read_text()) - tail = cfg.get("hybrid_tr3_tail", {}) - if "moe_layers" in tail: - moe_layers = list(range(tail["moe_layers"][0], tail["moe_layers"][1] + 1)) - else: - moe_layers = list(range( - cfg.get("first_k_dense_replace", 3), - cfg.get("num_hidden_layers", 13))) - print(f"Base K={base_k}, {len(stages)} cartridge stages, " - f"MoE layers {moe_layers[0]}-{moe_layers[-1]} ({len(moe_layers)} layers)", flush=True) - - # Parse expert filters - expert_filters = [] - hot_experts = recipe.get("hot_experts", list(range(96))) # default: first 96 - for stage in stages: - exp_spec = stage["experts"] - if exp_spec == "all": - expert_filters.append(None) - elif isinstance(exp_spec, str) and exp_spec in ("hot96", "hot"): - expert_filters.append({l: hot_experts for l in moe_layers}) - elif isinstance(exp_spec, list): - expert_filters.append({l: exp_spec for l in moe_layers}) - else: - expert_filters.append(None) - - # Load source weights and encode - from safetensors import safe_open - layer_results: dict[int, dict[int, dict[str, Any]]] = {} + layers = recipe["moe_layers"] + layer_shards = resolve_layer_shards(args.source, layers) + out_dir = check_out_dir( + args.out, source=args.source, policy=args.recipe) + + device = torch.device(args.device) + if device.type == "cuda": + if not torch.cuda.is_available(): + raise CartridgeError(f"--device {device}: CUDA is unavailable") + device_name = torch.cuda.get_device_name(device) + else: + device_name = device.type + print(f"Device: {device} ({device_name})", flush=True) + print( + f"Base K={base_k}, {len(stages)} stages, " + f"MoE layers {layers[0]}-{layers[-1]} ({len(layers)} layers)", + flush=True) + + staged = StagedOutput(out_dir, args.force) + work = staged.begin() + base_dir = work / "base" + cartridge_dir = work / "cartridges" + stage_shards: dict[str, list[str]] = { + stage["label"]: [] for stage in stages} + stage_tensor_counts = {stage["label"]: 0 for stage in stages} + combined_shards: list[str] = [] + combined_tensor_count = 0 + experts_per_layer: dict[int, int] = {} total_mse = 0.0 - n_experts_total = 0 - - for layer in moe_layers: - shard_path = args.source / f"model-layer-{layer:03d}.safetensors" - if not shard_path.exists(): - # Try other shard naming patterns - shard_path = args.source / f"model-layer-{layer:04d}.safetensors" - if not shard_path.exists(): - print(f" Layer {layer}: shard not found, skipping", flush=True) - continue - - print(f"\nEncoding layer {layer}...", flush=True) - layer_results[layer] = {} - - # Load expert weights from shard - with safe_open(str(shard_path), framework="pt") as f: - keys = list(f.keys()) - # Find expert keys for this layer - import re - expert_pattern = re.compile(EXPERT_RE_PATTERN) - expert_weights: dict[int, dict[str, torch.Tensor]] = {} - for key in keys: - m = expert_pattern.match(key) - if m: - l, e, proj = int(m.group(1)), int(m.group(2)), m.group(3) - if l == layer: - if e not in expert_weights: - expert_weights[e] = {} - expert_weights[e][proj] = f.get_tensor(key).float() - - print(f" Found {len(expert_weights)} experts", flush=True) - - for exp_id in sorted(expert_weights.keys()): - for proj in PROJECTIONS: - if proj not in expert_weights[exp_id]: + total_regularized_mse = 0.0 + expert_count = 0 + projection_count = 0 + + try: + copy_source_checkpoint( + args.source, base_dir, + {path.name for path in layer_shards.values()}) + from safetensors import safe_open + + with bootstrap_encoder(args.encoder_source) as encoder: + ext, ghd, tcp, tcpi, qtf, cbs = encoder + print(f"codebook_scale = {cbs}", flush=True) + for layer in layers: + source_shard = layer_shards[layer] + print(f"\nEncoding layer {layer}...", flush=True) + with safe_open(str(source_shard), framework="pt") as source_file: + keys = list(source_file.keys()) + experts = inspect_source_layer(keys, layer) + experts_per_layer[layer] = len(experts) + base_tensors = load_preserved_tensors(source_file, keys) + stage_tensors: dict[str, dict[str, torch.Tensor]] = { + stage["label"]: {} for stage in stages} + layer_mses: list[float] = [] + + for expert in sorted(experts): + applicable = selected_stages(stages, expert) + for projection in PROJECTIONS: + source_name = experts[expert][projection] + source_weight = source_file.get_tensor(source_name) + if source_weight.ndim != 2: + raise CartridgeError( + f"{source_name}: expected 2-D BF16 weight, " + f"got {tuple(source_weight.shape)}") + weight = source_weight.float().T.contiguous().to(device) + result = encode_expert_msrt( + weight, base_k, applicable, device, + ghd, tcp, tcpi, qtf, cbs, ext) + prefix = ( + f"model.layers.{layer}.mlp.experts.{expert}." + f"{projection}.rank0") + base_tensors[f"{prefix}.trellis"] = ( + result["base"]["trellis"]) + base_tensors[f"{prefix}.suh"] = result["base"]["suh"] + base_tensors[f"{prefix}.svh"] = result["base"]["svh"] + base_tensors[f"{prefix}.mcg"] = torch.tensor( + MCG_SENTINEL_SIGNED, dtype=torch.int32) + for label, stage_result in result["stages"].items(): + stage_tensors[label].update(_stage_tensor_names( + layer, expert, projection, label, stage_result)) + total_mse += result["mse"] + total_regularized_mse += result["regularized_mse"] + layer_mses.append(result["mse"]) + projection_count += 1 + del weight, source_weight, result + expert_count += 1 + + save_safetensors(base_tensors, base_dir / source_shard.name) + combined: dict[str, torch.Tensor] = {} + for stage in stages: + label = stage["label"] + tensors = stage_tensors[label] + if not tensors: continue - w = expert_weights[exp_id][proj].T.contiguous().to(device) # (in, out) for EXL3 trellis layout - result = encode_expert_msrt( - w, base_k, stages, device, ghd, tcp, tcpi, qtf, cbs, ext) - del w + relative = f"{label}/{source_shard.name}" + save_safetensors(tensors, cartridge_dir / relative) + stage_shards[label].append(relative) + stage_tensor_counts[label] += len(tensors) + combined.update(tensors) + if combined: + relative = f"combined/{source_shard.name}" + save_safetensors(combined, cartridge_dir / relative) + combined_shards.append(relative) + combined_tensor_count += len(combined) + average = sum(layer_mses) / len(layer_mses) + print( + f" Layer {layer}: avg original-space MSE={average:.4e}; " + f"{len(experts)} experts, {len(layer_mses)} projections", + flush=True) + del base_tensors, stage_tensors, combined + if device.type == "cuda": torch.cuda.empty_cache() - # Store under (exp_id, proj) — flatten for writing - key = (exp_id, proj) - if exp_id not in layer_results[layer]: - layer_results[layer][exp_id] = {"base": {}, "stages": [[] for _ in stages], "mse": {}} - layer_results[layer][exp_id]["base"][proj] = result["base"] - for si, stage_result in enumerate(result["stages"]): - layer_results[layer][exp_id]["stages"][si].append(stage_result) - layer_results[layer][exp_id]["mse"][proj] = result["mse"] - total_mse += result["mse"] - n_experts_total += 1 - - # Print layer summary - layer_mses = [] - for exp_data in layer_results[layer].values(): - layer_mses.extend(exp_data["mse"].values()) - avg = sum(layer_mses) / len(layer_mses) if layer_mses else 0 - print(f" Layer {layer}: avg MSE = {avg:.4e} ({len(layer_mses)} projections)", flush=True) - - print(f"\nOverall avg MSE: {total_mse / max(n_experts_total, 1):.4e}", flush=True) - - # Write outputs - out_dir = Path(args.out) - print(f"\nWriting outputs to {out_dir}...", flush=True) - - # Restructure for writing: group by (layer, expert) -> {base: {trellis, suh, svh}, stages: [...]} - # The layer_results is already structured, but we need to reorganize for the writers - write_results: dict[int, dict[int, dict[str, Any]]] = {} - for layer, experts in layer_results.items(): - write_results[layer] = {} - for exp_id, exp_data in experts.items(): - # Merge projections into single base/stages - base_tensors = {} - for proj in PROJECTIONS: - if proj in exp_data["base"]: - for tname, tval in exp_data["base"][proj].items(): - base_tensors[f"{proj}_{tname}"] = tval - - stage_list = [] - for si in range(len(stages)): - stage_tensors = {} - for proj in PROJECTIONS: - if si < len(exp_data["stages"]) and proj == PROJECTIONS[0]: - # Only need one copy of suh/svh per expert - pass - if si < len(exp_data["stages"]): - for proj_idx, proj in enumerate(PROJECTIONS): - if proj_idx < len(exp_data["stages"][si]): - stage = exp_data["stages"][si][proj_idx] - stage_tensors[f"{proj}"] = stage - stage_list.append(stage_tensors) - - write_results[layer][exp_id] = { - "base": base_tensors, - "stages": stage_list, - "mse": exp_data["mse"], - } - - # Write base checkpoint - print("\nWriting base checkpoint...", flush=True) - base_dir = out_dir / "base" - base_dir.mkdir(parents=True, exist_ok=True) - # Simplified: write all experts for each layer into one shard - for layer in moe_layers: - if layer not in write_results: - continue - tensors = {} - for exp_id, exp_data in write_results[layer].items(): - for proj in PROJECTIONS: - if f"{proj}_trellis" not in exp_data["base"]: - continue - rank = 0 - prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" - tensors[f"{prefix}.trellis"] = exp_data["base"][f"{proj}_trellis"] - tensors[f"{prefix}.suh"] = exp_data["base"][f"{proj}_suh"] - tensors[f"{prefix}.svh"] = exp_data["base"][f"{proj}_svh"] - tensors[f"{prefix}.mcg"] = torch.tensor(0xCBAC1FED, dtype=torch.uint32).view(torch.int32) - if tensors: - save_safetensors(tensors, base_dir / f"model-layer-{layer:03d}.safetensors") - print(f" Layer {layer}: {len(tensors)} base tensors", flush=True) - - # Copy non-layer files from source - import shutil - for f in args.source.iterdir(): - if f.is_file() and not f.name.startswith("model-layer-"): - shutil.copy2(f, base_dir / f.name) - - # Write cartridge adapters - print("\nWriting cartridge adapters...", flush=True) - cart_dir = out_dir / "cartridges" - - for si, stage in enumerate(stages): - label = stage["label"] - expert_filter = expert_filters[si] - tensors = {} - - for layer in moe_layers: - if layer not in write_results: - continue - for exp_id, exp_data in write_results[layer].items(): - if expert_filter and layer in expert_filter: - if exp_id not in expert_filter[layer]: - continue - if si >= len(exp_data["stages"]): - continue - stage_data = exp_data["stages"][si] - for proj in PROJECTIONS: - if proj not in stage_data: - continue - rank = 0 - prefix = f"model.layers.{layer}.mlp.experts.{exp_id}.{proj}.rank{rank}" - s = stage_data[proj] - tensors[f"{prefix}.trellis_{label}"] = s["trellis"] - tensors[f"{prefix}.suh_{label}"] = s["suh"] - tensors[f"{prefix}.svh_{label}"] = s["svh"] - tensors[f"{prefix}.scale_{label}"] = torch.tensor([s["scale"]], dtype=torch.float32) - - adapter_path = cart_dir / f"cartridge_{label}.safetensors" - save_safetensors(tensors, adapter_path) - config = { - "schema": ADAPTER_CONFIG_SCHEMA, - "stage_label": label, - "stage_k": stage["k"], - "num_tensors": len(tensors), - "tool_version": TOOL_VERSION, + if expert_count == 0 or projection_count == 0: + raise CartridgeError("encoding produced no experts or projections") + for stage in stages: + label = stage["label"] + if not stage_shards[label]: + raise CartridgeError( + f"stage {label!r} selected no experts in the source") + write_adapter_config( + cartridge_dir / label, [stage], + stage_shards[label], stage_tensor_counts[label]) + write_adapter_config( + cartridge_dir / "combined", stages, + combined_shards, combined_tensor_count) + write_base_metadata(base_dir, base_k, layers, experts_per_layer) + + summary = { + "tool": TOOL_VERSION, + "base_k": base_k, + "effective_bpw_excluding_metadata": effective_bpw( + recipe, max(experts_per_layer.values())), + "stages": [ + {"k": stage["k"], "label": stage["label"], + "experts": stage["experts"]} + for stage in stages + ], + "moe_layers": layers, + "overall_mse_original_space": total_mse / projection_count, + "overall_mse_regularized_space": ( + total_regularized_mse / projection_count), + "n_experts_encoded": expert_count, + "n_projections_encoded": projection_count, "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } - (cart_dir / f"cartridge_{label}_config.json").write_text( - json.dumps(config, indent=2) + "\n") - print(f" Cartridge '{label}': {len(tensors)} tensors -> {adapter_path.name}", flush=True) - - # Write summary - summary = { - "tool": TOOL_VERSION, - "base_k": base_k, - "stages": [{"k": s["k"], "label": s["label"], - "experts": s["experts"]} for s in stages], - "moe_layers": moe_layers, - "overall_mse": total_mse / max(n_experts_total, 1), - "n_experts_encoded": n_experts_total, - "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), - } - (out_dir / "encoding_summary.json").write_text(json.dumps(summary, indent=2) + "\n") + (work / "encoding_summary.json").write_text( + json.dumps(summary, indent=2) + "\n") + staged.commit() + except Exception: + staged.abort() + raise + print(f"\nDone! Output: {out_dir}", flush=True) - print(f" Overall MSE: {summary['overall_mse']:.4e}", flush=True) + print( + f" Original-space MSE: " + f"{total_mse / projection_count:.4e}", flush=True) return 0 @@ -746,10 +956,17 @@ def main(argv=None) -> int: enc.add_argument("--encoder-source", required=True, type=Path, help="Path to exllamav3 Python package") enc.add_argument("--device", default="cuda:0") + enc.add_argument( + "--force", action="store_true", + help="Replace a previous fq_assemble_lora output carrying its sentinel") args = p.parse_args(argv) - if args.command == "encode": - return cmd_encode(args) + try: + if args.command == "encode": + return cmd_encode(args) + except (AssemblyError, CartridgeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 return 1 From 2116bbddb8a41ba9be2c13f54f6cc4e84c850308 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:12:41 -0400 Subject: [PATCH 12/34] feat: validate and package MSRT cartridge tools --- CHANGELOG.md | 12 +- pyproject.toml | 6 + tests/test_fq_combine_cartridges.py | 73 +++++++++ tests/test_fq_measure_mse_fruit.py | 35 +++++ tools/fq_assemble_lora.py | 12 +- tools/fq_combine_cartridges.py | 203 ++++++++++++++++++++++++ tools/fq_measure_mse_fruit.py | 232 ++++++++++++++++++++++++++++ 7 files changed, 567 insertions(+), 6 deletions(-) create mode 100644 tests/test_fq_combine_cartridges.py create mode 100644 tests/test_fq_measure_mse_fruit.py create mode 100755 tools/fq_combine_cartridges.py create mode 100755 tools/fq_measure_mse_fruit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 73ca175..1a9933a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -59,14 +59,20 @@ measured, what is implemented, and what is not. does *not*, and what an attacker with full control of the artifact repository can and cannot do under fingerprint pinning. - **JSON Schemas** in [`schemas/`](schemas/) for `fq-segment/1` (segment - metadata and index), `fq-attestation/1`, `fq-manifest/1`, `fq-policy/2` - and `fq-release/1` — derived from real emitted artifacts and re-validated - against freshly emitted documents on every CI run. + metadata and index), `fq-attestation/1`, `fq-manifest/1`, `fq-policy/2`, + `fq-cartridge/1`, and `fq-release/1` — derived from real emitted artifacts + and re-validated against freshly emitted documents on every CI run. - **Packaging** — `pyproject.toml` with console entry points (`fq-repack`, `fq-assemble`, `fq-fetch`, `fq-prime`, `fq-verify`, `fq-release`, `fq-eps`), a hashed universal dev lock (`requirements-dev.txt`), and GitHub Actions CI running the suite on ubuntu-latest and macos-latest for Python 3.11 / 3.12 / 3.13, plus wheel-build and trust-root jobs. +- **MSRT cartridge tools** — `fq-assemble-lora` creates a complete EXL3 base + checkpoint plus validated, sharded full-rank residual cartridges; + `fq-combine-cartridges` validates and combines separately encoded stages; + `fq-measure-mse-fruit` compares the actual SIQ checkpoint and MSRT variants + in original weight space. These custom cartridges are explicitly not + standard PEFT/LoRA adapters and require an EXL3 MSRT-aware runtime. - **[docs/PRIOR-ART.md](docs/PRIOR-ART.md)** — commissioned independent prior-art review, and the single narrow claim this project makes. diff --git a/pyproject.toml b/pyproject.toml index 5d88d6d..369be06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ hub = ["huggingface_hub>=0.23"] # the bare wheel and runs `--help` on all of them to keep that true. numeric = ["numpy>=1.24"] charts = ["numpy>=1.24", "matplotlib>=3.7"] +quant = ["torch>=2.3", "safetensors>=0.4"] [project.urls] Homepage = "https://github.com/malaiwah/progressive-tensors" @@ -49,6 +50,9 @@ fq-prime = "fq_prime:main" fq-verify = "fq_verify:main" fq-release = "fq_release:main" fq-eps = "fq_eps:main" +fq-assemble-lora = "fq_assemble_lora:main" +fq-combine-cartridges = "fq_combine_cartridges:main" +fq-measure-mse-fruit = "fq_measure_mse_fruit:main" # The tools are flat, mutually-importing single-file modules (they add their # own directory to sys.path when run as scripts). Ship them as top-level @@ -65,12 +69,14 @@ sources = ["tools"] [tool.hatch.build.targets.wheel.force-include] "keys/FINGERPRINTS" = "fq_data/keys/FINGERPRINTS" "schemas" = "fq_data/schemas" +"recipes/fruit-k2-k3k4-cart.json" = "fq_data/recipes/fruit-k2-k3k4-cart.json" [tool.hatch.build.targets.sdist] include = [ "/tools", "/tests", "/schemas", + "/recipes", "/keys", "/docs", "/README.md", diff --git a/tests/test_fq_combine_cartridges.py b/tests/test_fq_combine_cartridges.py new file mode 100644 index 0000000..e81aa19 --- /dev/null +++ b/tests/test_fq_combine_cartridges.py @@ -0,0 +1,73 @@ +"""Validation tests for fq_combine_cartridges.""" + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) +import fq_assemble_lora as lora +import fq_combine_cartridges as combine + + +class TensorShape: + def __init__(self, shape): + self.shape = shape + self.ndim = len(shape) + + +def stage_tensors(label="res1", k=1): + prefix = "model.layers.3.mlp.experts.0" + tensors = {} + for projection in lora.PROJECTIONS: + root = f"{prefix}.{projection}.rank0" + tensors[f"{root}.trellis_{label}"] = TensorShape((8, 8, k * 16)) + tensors[f"{root}.suh_{label}"] = TensorShape((128,)) + tensors[f"{root}.svh_{label}"] = TensorShape((128,)) + tensors[f"{root}.scale_{label}"] = TensorShape(()) + return tensors + + +def test_stage_tensor_validation_requires_complete_components_and_k_geometry(): + stage = {"label": "res1", "k": 1, "experts": [0]} + combine.validate_stage_tensors(stage_tensors(), stage, "stage.safetensors") + + missing = stage_tensors() + missing.pop(next(key for key in missing if ".scale_res1" in key)) + with pytest.raises(lora.CartridgeError, match="components"): + combine.validate_stage_tensors(missing, stage, "stage.safetensors") + + wrong_k = stage_tensors(k=2) + with pytest.raises(lora.CartridgeError, match="expected last dimension 16"): + combine.validate_stage_tensors(wrong_k, stage, "stage.safetensors") + + +def write_config(path: Path, stage: dict, **overrides): + config = { + "schema": "fq-cartridge-adapter/1", + "format": "exl3-msrt-full-rank", + "standard_lora_compatible": False, + "base_k": 2, + "base_manifest_sha256": "a" * 64, + "stages": [stage], + "shards": [f"{stage['label']}/model-layer-003.safetensors"], + } + config.update(overrides) + path.write_text(json.dumps(config)) + + +def test_stage_config_rejects_path_traversal_and_base_mismatch(tmp_path: Path): + stage = {"label": "res1", "k": 1, "experts": [0]} + path = tmp_path / "adapter_config.json" + write_config(path, stage) + _, identity = combine.load_stage_config(path, stage, None) + assert identity == (2, "a" * 64) + + write_config(path, stage, shards=["../escape.safetensors"]) + with pytest.raises(lora.CartridgeError, match="safe relative paths"): + combine.load_stage_config(path, stage, None) + + write_config(path, stage, base_manifest_sha256="b" * 64) + with pytest.raises(lora.CartridgeError, match="base checkpoint identity"): + combine.load_stage_config(path, stage, (2, "a" * 64)) diff --git a/tests/test_fq_measure_mse_fruit.py b/tests/test_fq_measure_mse_fruit.py new file mode 100644 index 0000000..fc985c0 --- /dev/null +++ b/tests/test_fq_measure_mse_fruit.py @@ -0,0 +1,35 @@ +"""Pure orchestration tests for fq_measure_mse_fruit.""" + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) +import fq_measure_mse_fruit as measure + + +def test_stratified_sampling_includes_each_stage_chain(): + stages = [ + {"label": "res1", "experts": "all"}, + {"label": "res2", "experts": [0, 1]}, + ] + tiers = measure.select_stratified_experts([0, 1, 2, 3], stages, 1) + assert tiers == {"res1+res2": [0], "res1": [2]} + + +def test_parse_ids_rejects_duplicates_and_negative_values(): + assert measure.parse_ids("3,1") == [3, 1] + with pytest.raises(Exception, match="unique non-negative"): + measure.parse_ids("1,1") + with pytest.raises(Exception, match="unique non-negative"): + measure.parse_ids("-1") + + +def test_atomic_output_creates_parent_and_replaces_file(tmp_path: Path): + output = tmp_path / "nested" / "results.json" + measure.write_json_atomic(output, {"version": 1}) + measure.write_json_atomic(output, {"schema": "fq-msrt-mse/1"}) + assert json.loads(output.read_text()) == {"schema": "fq-msrt-mse/1"} + assert not (output.parent / ".results.json.tmp").exists() diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index 2e628eb..30c9365 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -62,6 +62,7 @@ check_out_dir, regenerate_manifest, regenerate_shard_index, + sha256_file, ) from fq_repack import PROJ_ORDER # noqa: E402 @@ -747,6 +748,8 @@ def _stage_tensor_names( def write_adapter_config( directory: Path, + base_k: int, + base_manifest_sha256: str, stages: list[dict[str, Any]], shards: list[str], tensor_count: int, @@ -754,6 +757,8 @@ def write_adapter_config( """Write the explicit custom MSRT runtime contract.""" config = { "schema": ADAPTER_CONFIG_SCHEMA, + "base_k": base_k, + "base_manifest_sha256": base_manifest_sha256, "format": "exl3-msrt-full-rank", "standard_lora_compatible": False, "runtime_operation": ( @@ -897,18 +902,19 @@ def cmd_encode(args) -> int: if expert_count == 0 or projection_count == 0: raise CartridgeError("encoding produced no experts or projections") + write_base_metadata(base_dir, base_k, layers, experts_per_layer) + base_manifest_sha256 = sha256_file(base_dir / "MANIFEST.sha256") for stage in stages: label = stage["label"] if not stage_shards[label]: raise CartridgeError( f"stage {label!r} selected no experts in the source") write_adapter_config( - cartridge_dir / label, [stage], + cartridge_dir / label, base_k, base_manifest_sha256, [stage], stage_shards[label], stage_tensor_counts[label]) write_adapter_config( - cartridge_dir / "combined", stages, + cartridge_dir / "combined", base_k, base_manifest_sha256, stages, combined_shards, combined_tensor_count) - write_base_metadata(base_dir, base_k, layers, experts_per_layer) summary = { "tool": TOOL_VERSION, diff --git a/tools/fq_combine_cartridges.py b/tools/fq_combine_cartridges.py new file mode 100755 index 0000000..391d333 --- /dev/null +++ b/tools/fq_combine_cartridges.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Combine validated MSRT stage shards into one sharded custom adapter.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + +sys.path.insert(0, str(Path(__file__).parent)) +import fq_assemble_lora as lora +from fq_assemble import AssemblyError, StagedOutput, check_out_dir + +STAGE_KEY_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.rank(\d+)\." + r"(trellis|suh|svh|scale)_([A-Za-z0-9_-]{1,32})$" +) +COMPONENTS = {"trellis", "suh", "svh", "scale"} + + +def load_stage_config( + path: Path, expected_stage: dict[str, Any], identity: tuple[int, str] | None +) -> tuple[dict[str, Any], tuple[int, str]]: + """Validate one stage config and its base-checkpoint identity.""" + try: + config = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise lora.CartridgeError(f"{path}: invalid adapter config ({exc})") from exc + if not isinstance(config, dict) or config.get("schema") != lora.ADAPTER_CONFIG_SCHEMA: + raise lora.CartridgeError( + f"{path}: schema must be {lora.ADAPTER_CONFIG_SCHEMA!r}") + if config.get("format") != "exl3-msrt-full-rank": + raise lora.CartridgeError(f"{path}: unsupported format {config.get('format')!r}") + if config.get("standard_lora_compatible") is not False: + raise lora.CartridgeError(f"{path}: standard_lora_compatible must be false") + stages = config.get("stages") + expected = { + "label": expected_stage["label"], + "k": expected_stage["k"], + "experts": expected_stage["experts"], + } + if stages != [expected]: + raise lora.CartridgeError( + f"{path}: stage metadata {stages!r} != recipe {expected!r}") + if (not isinstance(config.get("shards"), list) + or not config["shards"] + or any(not isinstance(value, str) or Path(value).is_absolute() + or ".." in Path(value).parts for value in config["shards"])): + raise lora.CartridgeError(f"{path}: shards must be safe relative paths") + current = (config.get("base_k"), config.get("base_manifest_sha256")) + if (isinstance(current[0], bool) or not isinstance(current[0], int) + or not isinstance(current[1], str) or len(current[1]) != 64): + raise lora.CartridgeError(f"{path}: invalid base checkpoint identity") + if identity is not None and current != identity: + raise lora.CartridgeError( + f"{path}: base checkpoint identity {current!r} != {identity!r}") + return config, current + + +def validate_stage_tensors( + tensors: dict[str, Any], stage: dict[str, Any], shard: str +) -> None: + """Require exact components, K geometry, projection, and expert coverage.""" + label = stage["label"] + groups: dict[tuple[int, int, str, int], set[str]] = {} + observed_experts: dict[int, set[int]] = {} + for key, tensor in tensors.items(): + match = STAGE_KEY_RE.fullmatch(key) + if not match: + raise lora.CartridgeError(f"{shard}: unexpected tensor key {key!r}") + layer, expert, projection, rank, component, key_label = ( + int(match.group(1)), int(match.group(2)), match.group(3), + int(match.group(4)), match.group(5), match.group(6)) + if key_label != label: + raise lora.CartridgeError( + f"{shard}: tensor {key!r} carries label {key_label!r}, expected {label!r}") + groups.setdefault((layer, expert, projection, rank), set()).add(component) + observed_experts.setdefault(layer, set()).add(expert) + if component == "trellis": + if tensor.ndim != 3 or tensor.shape[-1] != stage["k"] * 16: + raise lora.CartridgeError( + f"{shard}: {key} has shape {tuple(tensor.shape)}, expected last " + f"dimension {stage['k'] * 16}") + elif component in {"suh", "svh"} and tensor.ndim != 1: + raise lora.CartridgeError(f"{shard}: {key} must be a vector") + elif component == "scale" and tensor.ndim != 0: + raise lora.CartridgeError(f"{shard}: {key} must be a scalar") + if not groups: + raise lora.CartridgeError(f"{shard}: stage contains no tensors") + for group, components in groups.items(): + if components != COMPONENTS: + raise lora.CartridgeError( + f"{shard}: {group} components {sorted(components)} != " + f"{sorted(COMPONENTS)}") + for layer, experts in observed_experts.items(): + expected_experts = ( + experts if stage["experts"] == "all" else set(stage["experts"]) + ) + if experts != expected_experts: + raise lora.CartridgeError( + f"{shard}: layer {layer} experts {sorted(experts)} != " + f"recipe {sorted(expected_experts)}") + expected_groups = { + (layer, expert, projection, 0) + for expert in expected_experts for projection in lora.PROJECTIONS + } + actual_groups = {group for group in groups if group[0] == layer} + if actual_groups != expected_groups: + raise lora.CartridgeError( + f"{shard}: incomplete expert/projection/rank coverage") + + +def combine(args) -> int: + lora.require_quant_dependencies() + recipe = lora.load_recipe(args.recipe) + root = args.cartridges.expanduser().resolve() + if not root.is_dir(): + raise lora.CartridgeError(f"--cartridges {root} is not a directory") + out = check_out_dir(args.out, source=root, policy=args.recipe) + + configs: list[tuple[dict[str, Any], dict[str, Any], Path]] = [] + identity = None + basenames = None + for stage in recipe["stages"]: + stage_dir = root / stage["label"] + config, identity = load_stage_config( + stage_dir / "adapter_config.json", stage, identity) + names = {Path(value).name for value in config["shards"]} + if len(names) != len(config["shards"]): + raise lora.CartridgeError( + f"{stage_dir}: duplicate shard basenames are ambiguous") + if basenames is None: + basenames = names + elif names != basenames: + raise lora.CartridgeError( + f"{stage_dir}: shard coverage {sorted(names)} != {sorted(basenames)}") + configs.append((stage, config, stage_dir)) + + from safetensors import safe_open + + staged = StagedOutput(out, args.force) + work = staged.begin() + output_shards: list[str] = [] + tensor_count = 0 + try: + for basename in sorted(basenames or []): + combined: dict[str, Any] = {} + for stage, config, stage_dir in configs: + relative = next( + value for value in config["shards"] + if Path(value).name == basename) + source = stage_dir / relative + if not source.is_file(): + # Encoder configs store paths relative to cartridge root. + source = root / relative + if not source.is_file(): + raise lora.CartridgeError(f"missing stage shard {relative}") + with safe_open(str(source), framework="pt") as handle: + tensors = { + key: handle.get_tensor(key) for key in handle + } + validate_stage_tensors(tensors, stage, str(source)) + collisions = set(combined) & set(tensors) + if collisions: + raise lora.CartridgeError( + f"{source}: duplicate tensor keys {sorted(collisions)[:3]}") + combined.update(tensors) + lora.save_safetensors(combined, work / basename) + output_shards.append(basename) + tensor_count += len(combined) + lora.write_adapter_config( + work, identity[0], identity[1], recipe["stages"], + output_shards, tensor_count) + staged.commit() + except Exception: + staged.abort() + raise + print(f"Combined {len(output_shards)} shards into {out}", flush=True) + return 0 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--recipe", required=True, type=Path) + parser.add_argument( + "--cartridges", required=True, type=Path, + help="Directory containing one validated stage directory per recipe label") + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--force", action="store_true") + args = parser.parse_args(argv) + try: + return combine(args) + except (AssemblyError, lora.CartridgeError, ValueError, OSError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/fq_measure_mse_fruit.py b/tools/fq_measure_mse_fruit.py new file mode 100755 index 0000000..6e386c4 --- /dev/null +++ b/tools/fq_measure_mse_fruit.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +"""Measure original-space MSE for real SIQ and simulated MSRT weights.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from pathlib import Path +from typing import Any + +sys.path.insert(0, str(Path(__file__).parent)) +import fq_assemble_lora as lora # noqa: E402 +from fq_verify import decode_proj # noqa: E402 + +EXL3_KEY_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.rank(\d+)\." + r"(trellis|suh|svh|mcg)$" +) + + +def parse_ids(value: str | None) -> list[int] | None: + if value is None: + return None + try: + values = [int(token) for token in value.split(",") if token] + except ValueError as exc: + raise lora.CartridgeError(f"invalid comma-separated integer list {value!r}") from exc + if not values or any(value < 0 for value in values) or len(set(values)) != len(values): + raise lora.CartridgeError("ID lists must contain unique non-negative integers") + return values + + +def select_stratified_experts( + expert_ids: list[int], stages: list[dict[str, Any]], per_tier: int +) -> dict[str, list[int]]: + """Group by emitted stage chain and sample each tier independently.""" + tiers: dict[str, list[int]] = {} + for expert in expert_ids: + labels = [stage["label"] for stage in lora.selected_stages(stages, expert)] + tier = "+".join(labels) if labels else "base" + tiers.setdefault(tier, []).append(expert) + if per_tier > 0: + tiers = {tier: ids[:per_tier] for tier, ids in tiers.items()} + return tiers + + +def load_siq_slots(path: Path, layer: int, expert: int, projection: str): + """Load every rank slice for one actual SIQ expert projection.""" + from safetensors import safe_open + + ranks: dict[int, dict[str, Any]] = {} + with safe_open(str(path), framework="pt") as source: + for key in source.keys(): + match = EXL3_KEY_RE.fullmatch(key) + if not match: + continue + key_layer, key_expert, key_projection, rank, component = ( + int(match.group(1)), int(match.group(2)), match.group(3), + int(match.group(4)), match.group(5)) + if (key_layer, key_expert, key_projection) != ( + layer, expert, projection): + continue + ranks.setdefault(rank, {})[component] = source.get_tensor(key) + if not ranks: + raise lora.CartridgeError( + f"{path}: no SIQ tensors for layer {layer} expert {expert} {projection}") + required = {"trellis", "suh", "svh", "mcg"} + for rank, slot in ranks.items(): + missing = required - set(slot) + if missing: + raise lora.CartridgeError( + f"{path}: rank {rank} is missing SIQ components {sorted(missing)}") + return [ranks[rank] for rank in sorted(ranks)] + + +def mse(reference, reconstructed) -> float: + return (reference.float() - reconstructed.float()).square().mean().item() + + +def simulate_msrt( + weight, recipe, device, ghd, tcp, tcpi, qtf, cbs, expert: int +) -> dict[str, float]: + """Use production regularization/quantization primitives for all variants.""" + regularized, suh, svh = lora.regularize_with_vectors( + weight, device, ghd, cbs) + current = lora.quantize_trellis( + regularized, recipe["base_k"], device, tcp, tcpi, qtf) + reconstructed = lora.inverse_regularize(current, suh, svh, device, ghd) + values = {"msrt_base": mse(weight, reconstructed)} + for stage in lora.selected_stages(recipe["stages"], expert): + residual = regularized - current + residual_rms = residual.square().mean().sqrt().item() + if residual_rms >= 1e-12: + scale = abs(float(cbs)) / residual_rms + correction = lora.quantize_trellis( + residual * scale, stage["k"], device, tcp, tcpi, qtf) + current = current + correction / scale + reconstructed = lora.inverse_regularize(current, suh, svh, device, ghd) + values[f"msrt_through_{stage['label']}"] = mse(weight, reconstructed) + return values + + +def run_measurement(args) -> dict[str, Any]: + lora.require_quant_dependencies() + recipe = lora.load_recipe(args.recipe) + layers = parse_ids(args.layers) or recipe["moe_layers"] + requested_experts = parse_ids(args.experts) + bf16_shards = lora.resolve_layer_shards(args.bf16, layers) + siq_shards = lora.resolve_layer_shards(args.siq, layers) + device = lora.torch.device(args.device) + if device.type == "cuda" and not lora.torch.cuda.is_available(): + raise lora.CartridgeError(f"--device {device}: CUDA is unavailable") + + sums: dict[str, float] = {} + counts: dict[str, int] = {} + samples: dict[str, Any] = {} + encoder_parent = str(args.encoder_source.expanduser().resolve().parent) + sys.path.insert(0, encoder_parent) + try: + from safetensors import safe_open + + with lora.bootstrap_encoder(args.encoder_source) as encoder: + _, ghd, tcp, tcpi, qtf, cbs = encoder + for layer in layers: + with safe_open(str(bf16_shards[layer]), framework="pt") as source: + keys = list(source.keys()) + experts = lora.inspect_source_layer(keys, layer) + available = sorted(experts) + if requested_experts is not None: + missing = set(requested_experts) - set(available) + if missing: + raise lora.CartridgeError( + f"layer {layer}: requested experts missing: {sorted(missing)}") + tiers = {"explicit": requested_experts} + else: + tiers = select_stratified_experts( + available, recipe["stages"], args.experts_per_tier) + selected = [expert for ids in tiers.values() for expert in ids] + samples[str(layer)] = {"tiers": tiers, "experts": selected} + + for expert in selected: + for projection in lora.PROJECTIONS: + source_weight = source.get_tensor( + experts[expert][projection]).float() + internal = source_weight.T.contiguous().to(device) + for name, value in simulate_msrt( + internal, recipe, device, ghd, tcp, tcpi, + qtf, cbs, expert).items(): + sums[name] = sums.get(name, 0.0) + value + counts[name] = counts.get(name, 0) + 1 + + slots = load_siq_slots( + siq_shards[layer], layer, expert, projection) + decoded = decode_proj(slots, projection, str(device)) + if decoded.shape != source_weight.shape: + raise lora.CartridgeError( + f"SIQ shape {tuple(decoded.shape)} != BF16 shape " + f"{tuple(source_weight.shape)} for layer {layer} " + f"expert {expert} {projection}") + value = mse(source_weight, decoded.cpu()) + sums["siq_actual"] = sums.get("siq_actual", 0.0) + value + counts["siq_actual"] = counts.get("siq_actual", 0) + 1 + del source_weight, internal, decoded, slots + if device.type == "cuda": + lora.torch.cuda.empty_cache() + finally: + if sys.path and sys.path[0] == encoder_parent: + sys.path.pop(0) + + if not counts: + raise lora.CartridgeError("measurement selected no projections") + return { + "schema": "fq-msrt-mse/1", + "recipe": str(args.recipe), + "bf16": str(args.bf16), + "siq": str(args.siq), + "layers": layers, + "sampling": { + "experts_explicit": requested_experts, + "experts_per_tier": args.experts_per_tier, + "selected": samples, + }, + "metrics": { + name: {"mse": sums[name] / counts[name], "projections": counts[name]} + for name in sorted(counts) + }, + } + + +def write_json_atomic(path: Path, document: dict[str, Any]) -> None: + path = path.expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp") + temporary.write_text(json.dumps(document, indent=2) + "\n") + os.replace(temporary, path) + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bf16", required=True, type=Path) + parser.add_argument("--siq", required=True, type=Path) + parser.add_argument("--recipe", required=True, type=Path) + parser.add_argument("--encoder-source", required=True, type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--layers", help="Comma-separated layer IDs; defaults to the recipe") + parser.add_argument( + "--experts", help="Explicit comma-separated expert IDs for every layer") + parser.add_argument( + "--experts-per-tier", type=int, default=10, + help="Stratified experts per emitted stage chain; 0 means all") + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args(argv) + if args.experts_per_tier < 0: + parser.error("--experts-per-tier must be non-negative") + try: + results = run_measurement(args) + write_json_atomic(args.out, results) + except (lora.CartridgeError, ValueError, OSError, ImportError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(json.dumps(results["metrics"], indent=2)) + print(f"Results saved to {args.out}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From a4749734ac6fa1f7fcca4d3b16821ffa89b188f3 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:15:58 -0400 Subject: [PATCH 13/34] fix: finalize sharded cartridge runtime contract --- CHANGELOG.md | 5 +- pyproject.toml | 2 +- schemas/fq-cartridge-adapter-1.schema.json | 61 +++++++ tests/test_fq_combine_cartridges.py | 50 +++++ tests/test_schemas.py | 17 ++ tools/combine_cartridges.py | 126 ------------- tools/fq_assemble_lora.py | 8 +- tools/fq_combine_cartridges.py | 3 +- tools/fq_measure_mse_fruit.py | 7 +- tools/measure_mse_fruit.py | 201 --------------------- 10 files changed, 141 insertions(+), 339 deletions(-) create mode 100644 schemas/fq-cartridge-adapter-1.schema.json delete mode 100755 tools/combine_cartridges.py delete mode 100755 tools/measure_mse_fruit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a9933a..5fd8e79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,8 +60,9 @@ measured, what is implemented, and what is not. repository can and cannot do under fingerprint pinning. - **JSON Schemas** in [`schemas/`](schemas/) for `fq-segment/1` (segment metadata and index), `fq-attestation/1`, `fq-manifest/1`, `fq-policy/2`, - `fq-cartridge/1`, and `fq-release/1` — derived from real emitted artifacts - and re-validated against freshly emitted documents on every CI run. + `fq-cartridge/1`, `fq-cartridge-adapter/1`, and `fq-release/1` — derived + from real emitted artifacts and re-validated against freshly emitted + documents on every CI run. - **Packaging** — `pyproject.toml` with console entry points (`fq-repack`, `fq-assemble`, `fq-fetch`, `fq-prime`, `fq-verify`, `fq-release`, `fq-eps`), a hashed universal dev lock (`requirements-dev.txt`), and diff --git a/pyproject.toml b/pyproject.toml index 369be06..64125b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ hub = ["huggingface_hub>=0.23"] # the bare wheel and runs `--help` on all of them to keep that true. numeric = ["numpy>=1.24"] charts = ["numpy>=1.24", "matplotlib>=3.7"] -quant = ["torch>=2.3", "safetensors>=0.4"] +quant = ["torch>=2.3", "safetensors>=0.4", "numpy>=1.24"] [project.urls] Homepage = "https://github.com/malaiwah/progressive-tensors" diff --git a/schemas/fq-cartridge-adapter-1.schema.json b/schemas/fq-cartridge-adapter-1.schema.json new file mode 100644 index 0000000..4796aa4 --- /dev/null +++ b/schemas/fq-cartridge-adapter-1.schema.json @@ -0,0 +1,61 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/malaiwah/progressive-tensors/schemas/fq-cartridge-adapter-1.schema.json", + "title": "Progressive Tensors MSRT cartridge adapter", + "description": "Runtime contract for sharded full-rank EXL3 MSRT residual weights.", + "type": "object", + "additionalProperties": false, + "required": [ + "schema", "base_k", "base_manifest_sha256", "format", + "standard_lora_compatible", "runtime_operation", "codebook", + "mcg_multiplier", "mcg_ownership", "scale_shape", "stages", + "shards", "num_tensors", "tool_version", "created_utc" + ], + "properties": { + "schema": {"const": "fq-cartridge-adapter/1"}, + "base_k": {"type": "integer", "minimum": 1, "maximum": 8}, + "base_manifest_sha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"}, + "format": {"const": "exl3-msrt-full-rank"}, + "standard_lora_compatible": {"const": false}, + "runtime_operation": {"type": "string", "minLength": 1}, + "codebook": {"const": "mcg"}, + "mcg_multiplier": {"const": 3417055213}, + "mcg_ownership": {"const": "adapter-config"}, + "scale_shape": {"const": []}, + "stages": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": false, + "required": ["label", "k", "experts"], + "properties": { + "label": {"type": "string", "pattern": "^[A-Za-z0-9_-]{1,32}$"}, + "k": {"type": "integer", "minimum": 1, "maximum": 8}, + "experts": { + "oneOf": [ + {"const": "all"}, + { + "type": "array", + "uniqueItems": true, + "items": {"type": "integer", "minimum": 0} + } + ] + } + } + } + }, + "shards": { + "type": "array", + "minItems": 1, + "uniqueItems": true, + "items": { + "type": "string", + "pattern": "^(?!/)(?!.*(?:^|/)\\.\\.(?:/|$)).+\\.safetensors$" + } + }, + "num_tensors": {"type": "integer", "minimum": 1}, + "tool_version": {"type": "string", "pattern": "^fq_assemble_lora/[0-9]+$"}, + "created_utc": {"type": "string", "format": "date-time"} + } +} diff --git a/tests/test_fq_combine_cartridges.py b/tests/test_fq_combine_cartridges.py index e81aa19..122c798 100644 --- a/tests/test_fq_combine_cartridges.py +++ b/tests/test_fq_combine_cartridges.py @@ -3,6 +3,7 @@ import json import sys from pathlib import Path +from types import SimpleNamespace import pytest @@ -71,3 +72,52 @@ def test_stage_config_rejects_path_traversal_and_base_mismatch(tmp_path: Path): write_config(path, stage, base_manifest_sha256="b" * 64) with pytest.raises(lora.CartridgeError, match="base checkpoint identity"): combine.load_stage_config(path, stage, (2, "a" * 64)) + + +def test_combiner_merges_validated_stage_shards(tmp_path: Path): + torch = pytest.importorskip("torch") + pytest.importorskip("safetensors") + if lora.torch is None: + pytest.skip("fq_assemble_lora imported without torch") + + stages = [ + {"label": "res1", "k": 1, "experts": "all"}, + {"label": "res2", "k": 1, "experts": [0]}, + ] + recipe = tmp_path / "recipe.json" + recipe.write_text(json.dumps({ + "schema": "fq-cartridge/1", + "base_k": 2, + "stages": stages, + "moe_layers": [3], + })) + root = tmp_path / "stages" + shard_name = "model-layer-003.safetensors" + for stage in stages: + directory = root / stage["label"] + tensors = {} + for key, shape in stage_tensors(stage["label"], stage["k"]).items(): + tensors[key] = torch.zeros(shape.shape) + if ".trellis_" in key: + tensors[key] = tensors[key].to(torch.int16) + lora.save_safetensors(tensors, directory / shard_name) + lora.write_adapter_config( + directory, + 2, + "a" * 64, + [stage], + [f"{stage['label']}/{shard_name}"], + len(tensors), + ) + + output = tmp_path / "combined" + assert combine.combine(SimpleNamespace( + recipe=recipe, + cartridges=root, + out=output, + force=False, + )) == 0 + config = json.loads((output / "adapter_config.json").read_text()) + assert config["stages"] == stages + assert config["num_tensors"] == 24 + assert (output / shard_name).is_file() diff --git a/tests/test_schemas.py b/tests/test_schemas.py index df3e3b1..46ea20b 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -23,6 +23,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) import fq_fetch # noqa: E402,F401 (imported for its emitted documents) import fq_release # noqa: E402 +import fq_assemble_lora # noqa: E402 import fq_repack # noqa: E402 from test_fq_fetch import (REV, build_source, served, trust_root, # noqa: E402,F401 write_policy) @@ -174,6 +175,22 @@ def test_fruit_cartridge_recipe_matches_schema(): check("fq-cartridge-1", json.loads(recipe.read_text()), label=recipe.name) +def test_emitted_cartridge_adapter_config_matches_schema(tmp_path): + fq_assemble_lora.write_adapter_config( + tmp_path, + 2, + "a" * 64, + [{"label": "res1", "k": 1, "experts": [0]}], + ["res1/model-layer-003.safetensors"], + 12, + ) + check( + "fq-cartridge-adapter-1", + json.loads((tmp_path / "adapter_config.json").read_text()), + label="adapter_config.json", + ) + + # ----------------------------------------------- documents that must NOT pass @pytest.mark.parametrize("doc,name,why", [ diff --git a/tools/combine_cartridges.py b/tools/combine_cartridges.py deleted file mode 100755 index 25e0c41..0000000 --- a/tools/combine_cartridges.py +++ /dev/null @@ -1,126 +0,0 @@ -#!/usr/bin/env python3 -"""Assemble combined MSRT cartridges from individual stage files. - -Since vLLM supports only 1 LoRA per request, we need single adapter files that -contain all stages for a given recipe: - -1. cart_k3like.safetensors: K1trsc for ALL 256 experts (K2→K3, 3bpw total) -2. cart_k3k4like.safetensors: K1trsc for all + K2trsc for 96 hot (K2→K3/K4, 3.375bpw) - - 160 non-hot experts: only K1trsc stage (K3-equivalent) - - 96 hot experts: K1trsc + K2trsc stages (K4-equivalent) -""" -import json -import sys -from pathlib import Path -from safetensors import safe_open -from safetensors.torch import save_file -import torch - -def combine_cartridges( - res1_path: Path, # K1trsc for all experts - res2_path: Path, # K2trsc for hot experts - hot_experts: list[int], - out_path: Path, - label1: str = "res1", - label2: str = "res2", -): - """Combine two stage files into a single adapter with both stages. - - For hot experts: both res1 and res2 tensors are included - For non-hot experts: only res1 tensors are included - """ - tensors = {} - - # Load res1 (K1trsc) for ALL experts - with safe_open(str(res1_path), framework="pt") as f: - keys = list(f.keys()) - for key in keys: - tensors[key] = f.get_tensor(key) - - print(f"Loaded {len(tensors)} tensors from res1 (K1trsc, all experts)", flush=True) - - # Load res2 (K2trsc) only for hot experts - hot_set = set(hot_experts) - res2_count = 0 - with safe_open(str(res2_path), framework="pt") as f: - for key in f.keys(): - # Parse expert ID from key: model.layers.{L}.mlp.experts.{E}.{proj}.rank{R}.trellis_{label} - parts = key.split(".") - expert_id = int(parts[5]) - if expert_id in hot_set: - tensors[key] = f.get_tensor(key) - res2_count += 1 - - print(f"Loaded {res2_count} tensors from res2 (K2trsc, {len(hot_set)} hot experts)", flush=True) - print(f"Combined: {len(tensors)} total tensors", flush=True) - - # Write combined adapter - out_path.parent.mkdir(parents=True, exist_ok=True) - save_file(tensors, str(out_path)) - print(f"Saved combined cartridge: {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)", flush=True) - - # Write config - config = { - "schema": "fq-cartridge-adapter/1", - "stages": [ - {"k": 1, "label": label1, "experts": "all"}, - {"k": 2, "label": label2, "experts": hot_experts}, - ], - "description": "Combined K3/K4-like cartridge matching SIQ 160K3+96K4 allocation", - "num_tensors": len(tensors), - } - config_path = out_path.with_suffix(".config.json") - config_path.write_text(json.dumps(config, indent=2) + "\n") - return config - -def copy_k3only_cartridge( - res1_path: Path, - out_path: Path, - label: str = "res1", -): - """Create a K3-equivalent-only cartridge (just res1, all experts).""" - tensors = {} - with safe_open(str(res1_path), framework="pt") as f: - for key in f.keys(): - tensors[key] = f.get_tensor(key) - - out_path.parent.mkdir(parents=True, exist_ok=True) - save_file(tensors, str(out_path)) - print(f"Saved K3-like cartridge: {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)", flush=True) - - config = { - "schema": "fq-cartridge-adapter/1", - "stages": [{"k": 1, "label": label, "experts": "all"}], - "description": "K3-equivalent cartridge (K2+K1trsc, all experts, 3bpw)", - "num_tensors": len(tensors), - } - config_path = out_path.with_suffix(".config.json") - config_path.write_text(json.dumps(config, indent=2) + "\n") - return config - -if __name__ == "__main__": - cart_dir = Path("/tmp/poc_residual/fruit_msrt_output/cartridges") - out_dir = cart_dir - - # Hot experts: first 96 (matching SIQ's K4 tier) - hot_experts = list(range(96)) - - # 1. K3-like only (base + K1trsc = 3bpw on all experts) - print("=== Assembling K3-like cartridge (all experts, K1trsc only) ===", flush=True) - copy_k3only_cartridge( - cart_dir / "cartridge_res1.safetensors", - out_dir / "cart_k3like.safetensors", - ) - - # 2. K3/K4-like combined (matches SIQ 160K3 + 96K4) - print("\n=== Assembling K3/K4-like combined cartridge (160×K3 + 96×K4) ===", flush=True) - combine_cartridges( - cart_dir / "cartridge_res1.safetensors", - cart_dir / "cartridge_res2.safetensors", - hot_experts, - out_dir / "cart_k3k4like.safetensors", - ) - - print("\nDone! Combined cartridges:", flush=True) - print(f" cart_k3like.safetensors (K3-equivalent, all 256 experts)", flush=True) - print(f" cart_k3k4like.safetensors (K3/K4 mix, matches SIQ 160K3+96K4)", flush=True) diff --git a/tools/fq_assemble_lora.py b/tools/fq_assemble_lora.py index 30c9365..46084bf 100755 --- a/tools/fq_assemble_lora.py +++ b/tools/fq_assemble_lora.py @@ -40,10 +40,8 @@ import argparse import json import math -import os import re import shutil -import struct import sys import time from contextlib import contextmanager @@ -56,7 +54,7 @@ torch = None sys.path.insert(0, str(Path(__file__).parent)) -from fq_assemble import ( # noqa: E402 +from fq_assemble import ( AssemblyError, StagedOutput, check_out_dir, @@ -64,7 +62,7 @@ regenerate_shard_index, sha256_file, ) -from fq_repack import PROJ_ORDER # noqa: E402 +from fq_repack import PROJ_ORDER # ── Constants ────────────────────────────────────────────────────────────── @@ -156,7 +154,7 @@ def new_task(self, *args, **kwargs): pass memory = types.ModuleType("exllamav3.util.memory") memory.free_mem = lambda: None - memory.list_gpu_tensors = lambda: [] + memory.list_gpu_tensors = list sys.modules["exllamav3.util.memory"] = memory util = types.ModuleType("exllamav3.util") diff --git a/tools/fq_combine_cartridges.py b/tools/fq_combine_cartridges.py index 391d333..61cd013 100755 --- a/tools/fq_combine_cartridges.py +++ b/tools/fq_combine_cartridges.py @@ -160,8 +160,9 @@ def combine(args) -> int: if not source.is_file(): raise lora.CartridgeError(f"missing stage shard {relative}") with safe_open(str(source), framework="pt") as handle: + keys = list(handle.keys()) tensors = { - key: handle.get_tensor(key) for key in handle + key: handle.get_tensor(key) for key in keys } validate_stage_tensors(tensors, stage, str(source)) collisions = set(combined) & set(tensors) diff --git a/tools/fq_measure_mse_fruit.py b/tools/fq_measure_mse_fruit.py index 6e386c4..ec2b247 100755 --- a/tools/fq_measure_mse_fruit.py +++ b/tools/fq_measure_mse_fruit.py @@ -12,8 +12,8 @@ from typing import Any sys.path.insert(0, str(Path(__file__).parent)) -import fq_assemble_lora as lora # noqa: E402 -from fq_verify import decode_proj # noqa: E402 +import fq_assemble_lora as lora +from fq_verify import decode_proj EXL3_KEY_RE = re.compile( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." @@ -54,7 +54,8 @@ def load_siq_slots(path: Path, layer: int, expert: int, projection: str): ranks: dict[int, dict[str, Any]] = {} with safe_open(str(path), framework="pt") as source: - for key in source.keys(): + keys = list(source.keys()) + for key in keys: match = EXL3_KEY_RE.fullmatch(key) if not match: continue diff --git a/tools/measure_mse_fruit.py b/tools/measure_mse_fruit.py deleted file mode 100755 index 098128d..0000000 --- a/tools/measure_mse_fruit.py +++ /dev/null @@ -1,201 +0,0 @@ -#!/usr/bin/env python3 -"""Measure weight-level MSE between Fruit model quantization variants. - -Compares: -1. BF16 original (reference) -2. SIQ quant (160 K3 + 96 K4 per layer) -3. MSRT K2 base (all experts at K2) -4. MSRT K2+K1trsc (all experts at K3-equivalent) -5. MSRT K2+K1trsc+K2trsc (96 hot experts at K4-equivalent) - -All measurements in regularized (Hadamard) space, which equals original -space since Hadamard is orthogonal. -""" -from __future__ import annotations -import json, math, os, sys, types, importlib.util, gc -from pathlib import Path -import torch - -EXL3_PKG = "/opt/fruit-pip/exllamav3" - -def _bootstrap(): - pkg = types.ModuleType("exllamav3"); pkg.__path__ = [EXL3_PKG]; sys.modules["exllamav3"] = pkg - for sub in ["util", "modules", "modules.quant", "modules.quant.exl3_lib"]: - full = f"exllamav3.{sub}"; m = types.ModuleType(full) - m.__path__ = [f"{EXL3_PKG}/{sub.replace('.', '/')}"]; sys.modules[full] = m - class _DPB: - def __init__(self, *a, **kw): pass - def __enter__(self): return self - def __exit__(self, *a): return False - def update(self, *a): pass - def new_task(self, *a, **kw): pass - _s = types.ModuleType("exllamav3.util.progress"); _s.ProgressBar = _DPB; sys.modules["exllamav3.util.progress"] = _s - _s = types.ModuleType("exllamav3.util.memory"); _s.free_mem = lambda: None; _s.list_gpu_tensors = lambda: []; sys.modules["exllamav3.util.memory"] = _s - _s = types.ModuleType("exllamav3.util"); _s.__path__ = [f"{EXL3_PKG}/util"]; _s.cuda_sync_active = lambda *a, **kw: torch.cuda.synchronize(); sys.modules["exllamav3.util"] = _s - _s = types.ModuleType("exllamav3.util.tensor"); _s.save_tensor_image = lambda *a, **kw: None; sys.modules["exllamav3.util.tensor"] = _s - spec = importlib.util.spec_from_file_location("exllamav3.ext", f"{EXL3_PKG}/ext.py") - m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.ext"] = m; spec.loader.exec_module(m) - ext = m.exllamav3_ext - spec = importlib.util.spec_from_file_location("exllamav3.util.hadamard", f"{EXL3_PKG}/util/hadamard.py") - m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.util.hadamard"] = m; spec.loader.exec_module(m) - ghd = m.get_hadamard_dt - spec = importlib.util.spec_from_file_location("exllamav3.modules.quant.exl3_lib.quantize", f"{EXL3_PKG}/modules/quant/exl3_lib/quantize.py") - m = importlib.util.module_from_spec(spec); sys.modules["exllamav3.modules.quant.exl3_lib.quantize"] = m; spec.loader.exec_module(m) - return ext, ghd, m.tensor_core_perm, m.tensor_core_perm_i, m.quantize_tiles, m.codebook_scale - -def block_rms(x, dim, keepdim=False): - return x.square().mean(dim=dim, keepdim=keepdim).sqrt() - -def regularize(w, device, ghd, cbs, had_k=128, had_n=128, seed=0): - k, n = w.shape - g = torch.Generator(device="cpu").manual_seed(seed) - su = (torch.randn(k, generator=g).sign() + 1e-5).sign().float().to(device) - sv = (torch.randn(n, generator=g).sign() + 1e-5).sign().float().to(device) - out_scales = block_rms(w, dim=0, keepdim=True) - mean = out_scales.mean().item() - if mean > 1e-30: out_scales = out_scales / mean - sv = (sv * out_scales + 1e-10).float() - w = (w / sv).contiguous() - had_n_mat = ghd(had_n, device, torch.float, 1.0 / math.sqrt(had_n)) - w = (w.view(k, n // had_n, had_n) @ had_n_mat).view(k, n).contiguous() - in_scales = block_rms(w, dim=1, keepdim=True).clamp(min=1e-30) - su = (su.unsqueeze(1) * in_scales / (-cbs) + 1e-10).float() - w = (w / su).contiguous() - had_k_mat = ghd(had_k, device, torch.float, 1.0 / math.sqrt(had_k)) - w = (had_k_mat @ w.view(k // had_k, had_k, n)).view(k, n).contiguous() - return w - -def quantize_trellis_raw(data, K, device, tcp, tcpi, qtf): - k, n = data.shape; tiles_n = n // 16; weight_q = torch.zeros_like(data) - qa = {"K": K, "mcg": True} - perm = tcp(device); perm_i = tcpi(device) - for bi in range(0, k, 16): - rows = data[bi:bi+16] - tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256) - tiles = tiles[:, perm].contiguous() - quant_w, _ = qtf(tiles, qa) - quant_w = quant_w[:, perm_i].reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, n) - weight_q[bi:bi+16] = quant_w - return weight_q - -def rescaled_trellis(base_q, residual, K_res, device, tcp, tcpi, qtf, cbs): - residual_rms = residual.square().mean().sqrt().item() - if residual_rms < 1e-12: return base_q - scale = abs(cbs) / residual_rms - scaled = residual * scale - quant = quantize_trellis_raw(scaled, K_res, device, tcp, tcpi, qtf) - return base_q + quant / scale - -def run_measurement(bf16_path, siq_path, device, ghd, tcp, tcpi, qtf, cbs): - """Measure weight-level MSE for all configurations.""" - from safetensors import safe_open - import re - - # SIQ tier_bitmap: 160 K3 + 96 K4 per layer - tier_path = Path(siq_path) / "tier_bitmap.json" - tier_bitmap = json.loads(tier_path.read_text()) if tier_path.exists() else {} - - # Load SIQ trellis weights and compare to BF16 - moe_layers = [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] - results = {} - - for layer in moe_layers[:3]: # First 3 layers for speed - bf16_shard = bf16_path / f"model-layer-{layer:03d}.safetensors" - siq_shard = siq_path / f"model-layer-{layer:03d}.safetensors" - - if not bf16_shard.exists(): - print(f" Layer {layer}: bf16 shard not found", flush=True) - continue - - print(f"\n=== Layer {layer} ===", flush=True) - - # Load BF16 expert weights - bf16_experts = {} - with safe_open(str(bf16_shard), framework="pt") as f: - for key in f.keys(): - if f"layers.{layer}.mlp.experts." in key and key.endswith(".weight"): - parts = key.split(".") - eid = int(parts[5]) - proj = parts[6] - if eid not in bf16_experts: bf16_experts[eid] = {} - bf16_experts[eid][proj] = f.get_tensor(key).float() - - # Load SIQ tier info - k_list = tier_bitmap.get(str(layer), {}).get("k", [3] * 256) - - # Measure each config - n_experts = min(10, len(bf16_experts)) # First 10 experts for speed - configs_mse = {name: [] for name in [ - "K3_all", "K4_all", "SIQ_mixed", "MSRT_K2", "MSRT_K2_K1trsc", "MSRT_K2_K1trsc_K2trsc"]} - - for eid in sorted(bf16_experts.keys())[:n_experts]: - for proj in ["gate_proj", "up_proj", "down_proj"]: - if proj not in bf16_experts[eid]: - continue - w = bf16_experts[eid][proj].to(device) - w_reg = regularize(w, device, ghd, cbs) - del w - - # K3 - q_k3 = quantize_trellis_raw(w_reg, 3, device, tcp, tcpi, qtf) - configs_mse["K3_all"].append((w_reg - q_k3).pow(2).mean().item()) - - # K4 - q_k4 = quantize_trellis_raw(w_reg, 4, device, tcp, tcpi, qtf) - configs_mse["K4_all"].append((w_reg - q_k4).pow(2).mean().item()) - - # SIQ mixed: K3 or K4 depending on tier - siq_k = k_list[eid] if eid < len(k_list) else 3 - q_siq = quantize_trellis_raw(w_reg, siq_k, device, tcp, tcpi, qtf) - configs_mse["SIQ_mixed"].append((w_reg - q_siq).pow(2).mean().item()) - - # MSRT K2 base - q_k2 = quantize_trellis_raw(w_reg, 2, device, tcp, tcpi, qtf) - configs_mse["MSRT_K2"].append((w_reg - q_k2).pow(2).mean().item()) - - # MSRT K2 + K1trsc (K3-equivalent) - r2 = w_reg - q_k2 - q_msrt3 = rescaled_trellis(q_k2, r2, 1, device, tcp, tcpi, qtf, cbs) - configs_mse["MSRT_K2_K1trsc"].append((w_reg - q_msrt3).pow(2).mean().item()) - - # MSRT K2 + K1trsc + K2trsc (K4-equivalent, for hot experts) - r3 = w_reg - q_msrt3 - q_msrt4 = rescaled_trellis(q_msrt3, r3, 2, device, tcp, tcpi, qtf, cbs) - configs_mse["MSRT_K2_K1trsc_K2trsc"].append((w_reg - q_msrt4).pow(2).mean().item()) - - del w_reg, q_k3, q_k4, q_siq, q_k2, q_msrt3, q_msrt4 - torch.cuda.empty_cache() - - # Print results - print(f" {'Config':<30} {'avg MSE':>12} {'min':>12} {'max':>12}", flush=True) - print(f" {'-'*70}", flush=True) - layer_results = {} - for name, mses in configs_mse.items(): - avg = sum(mses) / len(mses) if mses else 0 - mn = min(mses) if mses else 0 - mx = max(mses) if mses else 0 - print(f" {name:<30} {avg:>12.4e} {mn:>12.4e} {mx:>12.4e}", flush=True) - layer_results[name] = {"avg": avg, "min": mn, "max": mx, "n": len(mses)} - - results[f"layer{layer}"] = layer_results - - return results - -def main(): - import argparse - ap = argparse.ArgumentParser() - ap.add_argument("--bf16-path", required=True, type=Path) - ap.add_argument("--siq-path", required=True, type=Path) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--out", default="/tmp/poc_residual/mse_results.json") - args = ap.parse_args() - dev = torch.device(args.device) - print(f"Device: {dev} GPU: {torch.cuda.get_device_name(0)}", flush=True) - ext, ghd, tcp, tcpi, qtf, cbs = _bootstrap() - print(f"codebook_scale = {cbs}", flush=True) - results = run_measurement(args.bf16_path, args.siq_path, dev, ghd, tcp, tcpi, qtf, cbs) - Path(args.out).write_text(json.dumps(results, indent=2, default=str)) - print(f"\nResults saved to {args.out}", flush=True) - -if __name__ == "__main__": - main() From 66f4b4ae4c8acd06dd79e2ad172863b5a6b94748 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:19:03 -0400 Subject: [PATCH 14/34] ci: smoke MSRT console entry points --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0fc2e83..15c3741 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -71,7 +71,7 @@ jobs: VIRTUAL_ENV=.venv-base uv pip install dist/*.whl VIRTUAL_ENV=.venv-base uv pip list | grep -qi '^numpy ' \ && { echo "numpy leaked into the base install"; exit 1; } || true - for cmd in fq-repack fq-assemble fq-fetch fq-prime fq-verify fq-release fq-eps; do + for cmd in fq-repack fq-assemble fq-assemble-lora fq-combine-cartridges fq-measure-mse-fruit fq-fetch fq-prime fq-verify fq-release fq-eps; do echo "--- $cmd --help" ./.venv-base/bin/$cmd --help > /dev/null done From ce64c4ac98c650ee6b7109a0bd064692be3bf89a Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:21:54 -0400 Subject: [PATCH 15/34] test: finalize MSRT regression collection --- tests/test_fq_assemble_lora.py | 1 - tests/test_schemas.py | 20 ++++++++++++-------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/test_fq_assemble_lora.py b/tests/test_fq_assemble_lora.py index 1de06c7..cd258e5 100755 --- a/tests/test_fq_assemble_lora.py +++ b/tests/test_fq_assemble_lora.py @@ -12,7 +12,6 @@ sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) import fq_assemble_lora as lora - RECIPE = Path(__file__).parent.parent / "recipes" / "fruit-k2-k3k4-cart.json" diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 46ea20b..f7822ee 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -21,14 +21,18 @@ jsonschema = pytest.importorskip("jsonschema") sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) -import fq_fetch # noqa: E402,F401 (imported for its emitted documents) -import fq_release # noqa: E402 -import fq_assemble_lora # noqa: E402 -import fq_repack # noqa: E402 -from test_fq_fetch import (REV, build_source, served, trust_root, # noqa: E402,F401 - write_policy) -from test_fq_repack import LAYERS # noqa: E402 - +import fq_assemble_lora +import fq_fetch # imported for its emitted documents +import fq_release +from test_fq_fetch import ( + REV, + build_source, + trust_root, + write_policy, +) +from test_fq_repack import LAYERS + +pytest_plugins = ("test_fq_fetch",) def _schemas_dir() -> Path | None: From 22bf818b956bedfe1d7a81753c17ad1e8c87bb1e Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 12 Aug 2026 20:28:15 -0400 Subject: [PATCH 16/34] test: keep cache authentication fixtures satisfiable --- tests/test_fq_fetch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_fq_fetch.py b/tests/test_fq_fetch.py index 268b3c7..a60e5cc 100644 --- a/tests/test_fq_fetch.py +++ b/tests/test_fq_fetch.py @@ -1386,7 +1386,7 @@ def test_full_header_fallback_reauthenticates_cache_and_reuses_pass_bytes( "build", "--dir", str(repo), "--release", "test 0.1.0", "--repo", "test/pub", "--revision", REV, "--sign-key", str(key)]) == 0 served["mount"]("test/pub", repo) - policy = write_policy(tmp_path / "recipe.json", {LAYERS[0]: [3, 4, 4, 4]}) + policy = write_policy(tmp_path / "recipe.json", {LAYERS[0]: [3] * E}) out = tmp_path / "fetched" argv = ["--policy", str(policy), "--out", str(out), "--source", f"test/pub@{REV}", "--trust-signer", pub, @@ -1443,7 +1443,7 @@ def test_signed_header_cache_reduces_restart_authentication_cost(tmp_path, serve "build", "--dir", str(repo), "--release", "test 0.1.0", "--repo", "test/pub", "--revision", REV, "--sign-key", str(key)]) == 0 served["mount"]("test/pub", repo) - policy = write_policy(tmp_path / "recipe.json", {LAYERS[0]: [3, 4, 4, 4]}) + policy = write_policy(tmp_path / "recipe.json", {LAYERS[0]: [3] * E}) out = tmp_path / "fetched" argv = ["--policy", str(policy), "--out", str(out), "--source", f"test/pub@{REV}", "--trust-signer", pub, From cd5970bbf4f73581f10301fa8dd5f3d2cc268ed3 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 13 Aug 2026 00:34:27 -0400 Subject: [PATCH 17/34] feat: MSRT campaign encoder with a DAG recipe and signed provenance Rewrite the cartridge encoder around the graph the products actually form, and make every emitted fragment carry checkable provenance. fq-cartridge/2 replaces the linear recipe: bases plus stages that name the parent reconstruction they correct. Nine products spanning 35 nominal trellis bpw now cost nine quantization passes emitting 14, measured at 1.83x less trellis kernel time than encoding each product separately on real GLM-5.2 experts. Campaign mechanics: plan/skeleton/encode/finalize; reads standard indexed HF shards or per-layer shards without loading a whole shard; work addressed as (layer, 32-expert block), claimed with an O_EXCL lock and committed as one atomic unit so a resumed block can never pin a residual to a parent it only recomputed in memory; --devices runs one worker per GPU over disjoint blocks. TILE_BATCH=128 is the measured tiling optimum at both model scales. Provenance: each shard ships a signed fq-attestation/1 line naming the sha256 of every expert's contiguous byte range, the encoder bundle including the compiled extension, the determinism scope, the quant args, and the exact parent shard digest a residual corrects. Shard payloads carry no timestamp, so re-encoding inside the declared scope reproduces them byte for byte. finalize re-hashes every fragment before publishing and refuses two signers, two encoder builds, unbound chains or shards the recipe does not describe. Consumers pin a key: fq_combine_cartridges verifies the signed assembly plan, requires the campaign identity to match every fragment, checks each chain edge, re-derives the runtime constants instead of copying them, validates tensor geometry, and narrows a product to chosen experts from the signed plan before reading any payload. Fixes found while building this: a hardlinked config.json made both bases claim one K and mutated the skeleton; finalize wiped a base's commit markers; a zero-RMS residual shipped an all-zero trellis that decodes to a nonzero codebook value; save_file's dtype ordering meant no expert occupied one byte range, so expert_sha256 could not mean what it says. Evidence: 13 GPU tests on an RTX 5090 with the real EXL3 kernels, including a published campaign decoded end to end through ext.reconstruct; 420 local tests. --- CHANGELOG.md | 47 +- README.md | 43 + docs/MSRT-CAMPAIGN.md | 381 +++ pyproject.toml | 2 +- recipes/fruit-k2-k3k4-cart.json | 22 +- recipes/fruit-k2k3-dag.json | 36 + recipes/glm52-k2k3-dag.json | 42 + schemas/fq-attestation-1.schema.json | 287 +- schemas/fq-cartridge-1.schema.json | 44 - schemas/fq-cartridge-2.schema.json | 86 + schemas/fq-cartridge-adapter-1.schema.json | 61 - schemas/fq-cartridge-adapter-2.schema.json | 283 ++ schemas/fq-cartridge-assembly-1.schema.json | 282 ++ tests/test_fq_assemble_lora.py | 728 ++++- tests/test_fq_combine_cartridges.py | 421 ++- tests/test_msrt_decode_parity.py | 268 ++ tests/test_schemas.py | 54 +- tools/fq_assemble_lora.py | 2666 +++++++++++++++---- tools/fq_combine_cartridges.py | 617 ++++- tools/fq_measure_mse_fruit.py | 52 +- tools/fq_verify.py | 19 +- 21 files changed, 5385 insertions(+), 1056 deletions(-) create mode 100644 docs/MSRT-CAMPAIGN.md create mode 100644 recipes/fruit-k2k3-dag.json create mode 100644 recipes/glm52-k2k3-dag.json delete mode 100644 schemas/fq-cartridge-1.schema.json create mode 100644 schemas/fq-cartridge-2.schema.json delete mode 100644 schemas/fq-cartridge-adapter-1.schema.json create mode 100644 schemas/fq-cartridge-adapter-2.schema.json create mode 100644 schemas/fq-cartridge-assembly-1.schema.json create mode 100755 tests/test_msrt_decode_parity.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fd8e79..0d5f238 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,7 +60,8 @@ measured, what is implemented, and what is not. repository can and cannot do under fingerprint pinning. - **JSON Schemas** in [`schemas/`](schemas/) for `fq-segment/1` (segment metadata and index), `fq-attestation/1`, `fq-manifest/1`, `fq-policy/2`, - `fq-cartridge/1`, `fq-cartridge-adapter/1`, and `fq-release/1` — derived + `fq-cartridge/2`, `fq-cartridge-adapter/2`, `fq-cartridge-assembly/1`, and + `fq-release/1` — derived from real emitted artifacts and re-validated against freshly emitted documents on every CI run. - **Packaging** — `pyproject.toml` with console entry points (`fq-repack`, @@ -68,12 +69,44 @@ measured, what is implemented, and what is not. `fq-eps`), a hashed universal dev lock (`requirements-dev.txt`), and GitHub Actions CI running the suite on ubuntu-latest and macos-latest for Python 3.11 / 3.12 / 3.13, plus wheel-build and trust-root jobs. -- **MSRT cartridge tools** — `fq-assemble-lora` creates a complete EXL3 base - checkpoint plus validated, sharded full-rank residual cartridges; - `fq-combine-cartridges` validates and combines separately encoded stages; - `fq-measure-mse-fruit` compares the actual SIQ checkpoint and MSRT variants - in original weight space. These custom cartridges are explicitly not - standard PEFT/LoRA adapters and require an EXL3 MSRT-aware runtime. +- **MSRT cartridge campaign tools** — `fq-assemble-lora` encodes a whole + `fq-cartridge/2` graph in one pass over the weights: every declared base + tier becomes a complete EXL3 checkpoint, and every stage is a rescaled + trellis residual against the reconstruction of the `parent` it names, so + nine loadable products spanning 35 bits per weight cost nine quantization + passes emitting 14 (measured on real GLM-5.2 experts: 1.83x less trellis + kernel time and 2.5x fewer bytes than encoding each product separately). + Subcommands + `plan` / `skeleton` / `encode` / `finalize`; reads standard indexed Hugging + Face shards or per-layer shards without ever loading a whole shard; work is + addressed as (layer, 32-expert block), claimed with an `O_EXCL` lock and + committed as one atomic unit, so an interrupted or preempted run resumes at + block granularity and `--devices` runs one worker per GPU over disjoint + blocks with no shared state. `fq-combine-cartridges` turns one published + assembly plan into a self-contained `fq-cartridge-adapter/2` cartridge under + a pinned signer, narrowing a full-expert stage to the experts a consumer + actually wants — decided from the signed plan before any payload is read. + `fq-measure-mse-fruit` compares the actual SIQ checkpoint against every + graph node through the production encoder itself. These custom cartridges are + explicitly not standard PEFT/LoRA adapters and require an EXL3 MSRT-aware + runtime. +- **Signed provenance for every encoded fragment.** Each shard ships a + `fq-attestation/1` line beside it: `encode-of` for expert shards, naming the + sha256 of each expert's contiguous byte range, the encoder bundle (Python + modules *and* the compiled extension), the determinism scope, the effective + quant arguments, and the exact parent shard digest the residual corrects; + `repack-of` for skeleton shards, naming per-tensor digests and the source + file the bytes were copied from. Shard payloads carry no timestamp, so + re-encoding inside the declared scope reproduces them byte for byte. + `finalize` re-hashes all of it before publishing anything and refuses a + campaign that spans two signers or two encoder builds, whose stages do not + name the parents this campaign published, or that holds shards the recipe + does not describe. +- **[docs/MSRT-CAMPAIGN.md](docs/MSRT-CAMPAIGN.md)** — the GLM-5.2 campaign + runbook: per-K trellis cost and a full 32-expert block measured on real + GLM-5.2 weights, the resulting 201 GPU-hour / 1.332 TB projection with its + measured/derived/unmeasured labels, the gates to run before renting a fleet, + and the resume, publish and verification procedure. - **[docs/PRIOR-ART.md](docs/PRIOR-ART.md)** — commissioned independent prior-art review, and the single narrow claim this project makes. diff --git a/README.md b/README.md index c5ce058..70cb019 100644 --- a/README.md +++ b/README.md @@ -543,6 +543,47 @@ The encoder driver and capture tooling are documented at the immutable [research revision `69fbef710e558e9cf8e2ad634eccc774f9a806fb`](https://github.com/malaiwah/vllm-voipmonitor/tree/69fbef710e558e9cf8e2ad634eccc774f9a806fb/research/fungible-quant); they are not a supported runtime component of this repository. +## Encode a progressive cartridge graph (`fq_assemble_lora`) + +A segment tree ships one bit-width per expert. An **MSRT cartridge graph** +ships several, from one pass over the weights: a `fq-cartridge/2` recipe +declares complete base tiers and *rescaled trellis residual stages*, each +naming the `parent` reconstruction it corrects, so the products share their +ancestors' work instead of re-encoding it. + +```bash +# nine loadable products (K2, K3, and K3/K4/K5-like cartridges over both +# bases) in nine quantization passes emitting 14 bits per weight +uv run tools/fq_assemble_lora.py plan --source \ + --recipe recipes/glm52-k2k3-dag.json --block-size 32 +uv run tools/fq_assemble_lora.py skeleton --source \ + --recipe recipes/glm52-k2k3-dag.json --out ./campaign --sign-key ~/.fq_keys/c.key +uv run tools/fq_assemble_lora.py encode --source \ + --recipe recipes/glm52-k2k3-dag.json --out ./campaign \ + --encoder-source --sign-key ~/.fq_keys/c.key \ + --devices cuda:0,cuda:1 +uv run tools/fq_assemble_lora.py finalize --source \ + --recipe recipes/glm52-k2k3-dag.json --out ./campaign --sign-key ~/.fq_keys/c.key + +# consumer: one product, narrowed to the experts you actually want upgraded, +# under a pinned signer +uv run tools/fq_combine_cartridges.py --root ./campaign \ + --assembly k2-k4like-direct --out ./k4like-hot96 --experts 0-95 \ + --trust-key <64-hex campaign signer> +``` + +Work is addressed as (layer, 32-expert block) and committed as one atomic unit, +so a preempted run resumes at block granularity and `--devices` runs one worker +per GPU over disjoint blocks with no shared state. Every shard ships a signed +`fq-attestation/1` line naming the sha256 of each expert's byte range, the +encoder bundle that produced it and — for a residual — the exact parent shard +digest it corrects; `finalize` re-hashes all of it before publishing and the +combiner re-checks it under a key you pin. Measured on real GLM-5.2 experts, the +shared-parent graph costs **1.83x less trellis kernel time and 2.5x fewer +bytes** than encoding the same nine products separately; +[docs/MSRT-CAMPAIGN.md](docs/MSRT-CAMPAIGN.md) carries the per-K cost table, the +full-campaign projection and the runbook. + ## Status & roadmap | Piece | Status | @@ -558,6 +599,8 @@ they are not a supported runtime component of this repository. | Mixed-size (true mixed-K) assembly + loader metadata | offline assembly is working and tested; serving an output remains subject to the runtime's TP4-only / EP-and-DP refusal and hardware constraints | | Four tiers in the artifact tree (K2/K3/K4/K5) | root K3 is complete (layers 3–78); root K2/K4/K5 are `encode-of` tiers; nested `sources/willfalco-*` contains community-primed material for layers 3–10. For current coverage, use `per_k[K].layer_coverage.layers` (`fq-layer-coverage/1`) or signed index keys for older manifests; `per_k[K].layers` is legacy extrema only | | Runtime progressive loader + live bit-width reallocation (vLLM/GG) | separate experimental research, TP-only and not wired as an end-to-end supported workflow; no live-reallocation claim is made by these tools | +| MSRT cartridge graph (`fq_assemble_lora`, `fq_combine_cartridges`) | working, tested; encode→decode parity proven against the runtime's own `ext.reconstruct` on real GLM-5.2 experts. Publishing a full GLM-5.2 graph is costed and gated in [docs/MSRT-CAMPAIGN.md](docs/MSRT-CAMPAIGN.md) but **not yet run** | +| Serving an MSRT cartridge | **blocked on runtime**: the reference EXL3 MSRT implementation ([local-inference-lab/vllm#299](https://github.com/local-inference-lab/vllm/pull/299)) is draft, TP=1, one model-wide slot, and materializes dense FP16 shadow weights, which GLM-5.2's 734 G routed weights do not fit on one node | | Packaging, CI (ubuntu + macOS, py3.11–3.13), JSON Schemas | landed this release | ## Prior art and positioning diff --git a/docs/MSRT-CAMPAIGN.md b/docs/MSRT-CAMPAIGN.md new file mode 100644 index 0000000..196f5f9 --- /dev/null +++ b/docs/MSRT-CAMPAIGN.md @@ -0,0 +1,381 @@ +# GLM-5.2 MSRT campaign runbook + +Numbers below are measured on real GLM-5.2 weights with this encoder, on an +RTX 5090 (Blackwell, 170 SMs, 1.79 TB/s), because that is the local card. The +rental target is an RTX PRO 6000 Blackwell (188 SMs, but 1.60 TB/s on the +Server Edition), so per-GPU throughput is **not** assumed equal: §2 states a +range and §4.0 makes a one-block probe on the rented card a gate. + +Every claim is labelled **measured**, **derived** or **unmeasured**. + +## 1. What the campaign produces + +One `fq-cartridge/2` recipe (`recipes/glm52-k2k3-dag.json`) declares a +quantization graph, not a list of checkpoints: + +``` +W ─┬─ k2 (2 bpw base) ─┬─ k2r1 (+1) ─── k2r1r1 (+1) + │ └─ k2r2 (+2) ─── k2r2r1 (+1) + └─ k3 (3 bpw base) ─┬─ k3r1 (+1) ─── k3r1r1 (+1) + └─ k3r2 (+2) +``` + +Nine product contracts for an MSRT-aware runtime (§6), over all 256 routed +experts of all 76 MoE layers (3–78, including the MTP layer): + +| Assembly | bpw | Base + chain | +|---|---:|---| +| `k2` | 2.0 | k2 | +| `k2-k3like` | 3.0 | k2 + k2r1 | +| `k2-k4like-stepped` | 4.0 | k2 + k2r1 + k2r1r1 | +| `k2-k4like-direct` | 4.0 | k2 + k2r2 | +| `k2-k5like` | 5.0 | k2 + k2r2 + k2r2r1 | +| `k3` | 3.0 | k3 | +| `k3-k4like` | 4.0 | k3 + k3r1 | +| `k3-k5like-stepped` | 5.0 | k3 + k3r1 + k3r1r1 | +| `k3-k5like-direct` | 5.0 | k3 + k3r2 | + +Sharing parent reconstructions is the point: nine products spanning 35 *nominal +trellis* bits per weight are encoded in **9 passes emitting 14 nominal trellis +bits per weight**. Nominal means trellis payload only; §3 has the actual bytes, +which come to 14.09 bpw for the routed artifacts (the `suh`/`svh` vectors and +scales) and 14.51 bpw with the skeleton amortised in. + +Every emitted fragment carries a signed `fq-attestation/1` line beside it: +`encode-of` for expert shards, with the sha256 of each expert's contiguous byte +range and the exact parent shard digest the residual corrects; `repack-of` for +skeleton shards, with per-tensor digests and the source file they were copied +from. `finalize` refuses to publish anything that does not hash to its recorded +digest, name its own bytes, agree on one signer and one encoder build, and name +the parent fragment this campaign actually published. + +## 2. Cost + +**Measured** per-matrix trellis cost, real GLM-5.2 expert (layer 40, expert +128), `tile_batch=128`, median of 2, foreign GPU utilisation 0%: + +| K | s / 6144x2048 matrix | trellis states | +|---|---:|---:| +| K1 | 1.618 | 32768 | +| K2 | 0.806 | 16384 | +| K3 | 0.674 | 8192 | + +Lower K is **more** expensive: the DP table is `65536 >> K` wide. The graph uses +five K1 stages, so **K1 is 72% of campaign GPU time**. + +**Measured** end-to-end, current code, `tile_batch=128`, clean GPU at start: +one real 32-expert block (layer 40, experts 128–159, all 9 nodes) committed in +**1188.4 s = 12.379 s/matrix**, of which 12.333 s/matrix is the matrix loop. +Everything the block needs to be *done* — the nine grouped safetensors writes, +per-expert span hashing, nine ed25519 signatures and the digest sidecars — costs +**0.046 s/matrix, 0.37%**. Two earlier blocks of the same work measured 12.06 s +(older unbounded tiling) and 12.30 s/matrix, both with a foreign job resident +for part of the run, so the three runs cluster at 12.06–12.38 s/matrix. + +| Quantity | Value | +|---|---:| +| Expert-projection matrices | 58,368 (76 layers x 256 experts x 3) | +| Routed weights | 734.44 G | +| Passes per matrix | 9 | +| Committed rate, measured | **12.379 s/matrix** | +| **Whole campaign, one RTX 5090** | **200.7 GPU-hours** | +| Same nine products, kernel time only | 19 passes, 20.47 s/matrix | +| **DAG saving, trellis kernels** | **1.83x** (measured per-K mix) | +| DAG saving, end to end | ~1.80x (inferred: independent encoding's own write and metadata cost is unmeasured) | +| Bytes saved | 2.5x (14 vs 35 nominal trellis bpw) | + +The isolated kernel sum for one expert is 34.08 s (11.36 s/matrix), so the real +loop costs ~8% more than the kernels alone: regularization, nine inverse +transforms, eighteen MSE reductions and the device-to-host copies serialize with +the trellis search. That gap is why the campaign figure comes from a block, not +from the sweep. + +Fleet projection at the live JarvisLabs spot price of $0.99/GPU-hour. A roofline +weighting of the measured per-K mix (K1 is 72.4% of kernel time and +memory-bound at 1.79 -> 1.60 TB/s; K2/K3 lean on SM count, 170 -> 188) puts the +parity-card centre at `0.724 x 1.122 + 0.276 x 0.904 = 1.062`: + +| Assumption | GPU-h | 8-GPU compute wall | Compute | +|---|---:|---:|---:| +| RTX PRO 6000 10% faster | 181 | 22.6 h | $179 | +| **roofline centre (+6.2%)** | **213** | **26.6 h** | **$211** | +| RTX PRO 6000 25% slower | 251 | 31.4 h | $248 | + +**Compute wall is not elapsed wall.** Nothing overlaps the first source window +or the final output drain, and finalize (§4.6) needs 0.4–1.9 h with GPUs paused. +Realistic elapsed time at the roofline centre and 350 Mbps is **28–31 h**. + +Budget: **$240 planning centre, $300 authorization cap** — engineering +contingencies, not statistical bounds. Compute $179–$248; 2 TB for two days +~$13; gates, CPU-side finalize/upload tails and one all-GPU in-flight-block redo +$3–$10; the cap additionally covers one full 8-layer-window redo (~$21). +A routine spot preemption with persistent storage loses at most the in-flight +block per GPU: 2.64 GPU-h / $2.61 for all eight, not a window. + +Resolve the card range with the §4.0 probe. One block gives the committed rate; +one complete layer across all eight GPUs (2.64 GPU-h, ~$2.61) gives the fleet +rate under real storage and multi-GPU contention. The campaign figure is +**extrapolated linearly from that gate**, not proven by it: one block samples +kernel work on identical geometry, not shard locality, eight-process I/O, or +sustained clocks over a day. + +## 3. Storage + +| Artifact | Size | +|---|---:| +| k2 base trellis | 184.6 GB | +| k3 base trellis | 276.4 GB | +| K1 stages (k2r1, k2r1r1, k2r2r1, k3r1, k3r1r1) | 92.8 GB each | +| K2 stages (k2r2, k3r2) | 184.6 GB each | +| suh/svh metadata, all 9 nodes | 8.6 GB | +| skeleton (everything that is not a routed expert) | 37.8 GB | +| **campaign on disk** | **1.332 TB** | +| upload (the second base re-ships the skeleton) | 1.370 TB | +| source download | 1.507 TB | +| **total WAN** | **2.876 TB** | + +Per layer, all nine nodes emit 17.0 GB in 72 block files (9 nodes x 8 blocks of +32 experts). The campaign is 5,472 block files plus 282 skeleton shards, each +with a digest sidecar and a signed attestation beside it. + +Actual routed bitrate is **14.09 bpw** (nominal 14 plus 8.6 GB of `suh`/`svh` +and scales), or **14.51 bpw** with the skeleton amortised over routed weights. +The upload figure is *logical repository bytes*: the second base re-ships the +same skeleton content, which Hugging Face's Xet content-addressed storage should +deduplicate, so physical WAN is likely ~2.838 TB. 2.876 TB is the conservative +number. + +**Rent 2 TB of persistent storage.** Campaign payload 1.332 TB plus one 150 GB +source window is already 1.482 TB, and `finalize` re-reads and re-hashes every +fragment before publishing a manifest, so the whole campaign has to be local at +that moment. 1.6 TB leaves ~118 GB for partial files, HF/Xet caches, logs and +retries; the extra 0.4 TB costs ~$2.67 for two days. Rolling publication deletes +*source* shards, never campaign output. A 1 TB volume cannot run this plan. + +Aggregate throughput of **280 Mbps** is sufficient to hide transfer under +overlap; require **≥350 Mbps measured** before committing the fleet. Below that, +stage bytes on a regional filesystem with a CPU VM instead of paying GPUs to +wait. + +## 4. Rental procedure + +Set once: + +```bash +REPO_ID=zai-org/GLM-5.2 +REV=<40-hex-commit> # the revision every attestation pins +SRC=~/.cache/huggingface/hub/models--zai-org--GLM-5.2/snapshots/$REV +CAMPAIGN=/data/glm52-msrt # on the 2 TB volume +KEY=~/.fq_keys/glm52-campaign.key # NEVER inside $CAMPAIGN or $SRC +RECIPE=recipes/glm52-k2k3-dag.json +ENC=/opt/exllamav3-python/exllamav3 +``` + +Every download must pass `--revision "$REV"`. Without it `hf download` populates +the `main` snapshot instead of `$SRC`, which either leaves the pinned source +incomplete or silently stages a different revision. + +`--base-model` and `--base-revision` are inferred from a Hugging Face snapshot +path like the one above. Any other layout must state them explicitly; the +revision must be an immutable 40-hex commit, because that is what every +attestation pins. + +### 4.0 Gates before the fleet + +1. `pytest tests/ -q` — full suite green. +2. On the rented card: + `FQ_ENCODER_SOURCE=$ENC FQ_PARITY_SOURCE= + pytest tests/test_msrt_decode_parity.py -q`. Thirteen tests prove, with the + runtime's own `ext.reconstruct`, that a *published* campaign decodes to the + MSE it attested — encode, finalize, combine under a pinned key, then decode + the combined adapter against the source weights. +3. `fq-assemble-lora plan --source $SRC --recipe $RECIPE` — confirms layout, + block count and which source shards are absent. +4. One block on one GPU (below). Compare the **committed** rate — the encoder + prints both `s/matrix` committed and the quantizing part of it — with §2. + Record the SKU, power limit and sustained clocks while it runs. +5. One complete layer across all eight GPUs, timed with external wall clock: + 2.44 GPU-h, ~$2.41. This is the only sample that includes eight-process I/O + on one volume and multi-GPU contention, and it is what the campaign figure is + extrapolated from. +6. Measure throughput to the hub (`hf download` one 5 GB shard): need ≥350 Mbps. + +```bash +fq-assemble-lora encode --source $SRC --recipe $RECIPE --out $CAMPAIGN \ + --encoder-source $ENC --sign-key $KEY --block-size 32 \ + --layers 40 --shard-index 4 --shard-count 8 --device cuda:0 +``` + +### 4.1 Stage source bytes per window + +`plan` reports the exact shard files each layer needs: + +```bash +fq-assemble-lora plan --source $SRC --recipe $RECIPE \ + --layers 3-10 --block-size 32 --out-plan window.json +python - <<'PY' +import json +plan = json.load(open("window.json")) +shards = sorted({s for L in plan["layers"] for s in L["shards"]}) +print(" ".join(f"--include {s}" for s in shards)) +PY +``` + +Feed those to +`hf download "$REPO_ID" --revision "$REV" --include ...`. A GLM MoE layer spans +~4 shards (~21.4 GB); an 8-layer window is ~150 GB. Stage the first window +*before* starting the GPUs. + +Optional but recommended: fetch the hub's own LFS digests once and pass them as +`--source-digests`, so the skeleton pass never re-hashes source payloads to +attest what it copied: + +```bash +curl -s "https://huggingface.co/api/models/$REPO_ID/tree/$REV?recursive=1" \ + > source-tree.json # lfs.oid is the sha256 of each file +``` + +### 4.2 Skeleton (once per window) + +```bash +fq-assemble-lora skeleton --source $SRC --recipe $RECIPE --out $CAMPAIGN \ + --sign-key $KEY --block-size 32 --source-digests source-tree.json +``` + +Copies non-expert tensors out of whatever shards are present and reports how +many were absent; re-run after every window, completed shards are skipped. It +also creates the signing key, so run it before any parallel encode. + +**Measured shape:** 85 of GLM-5.2's 282 shards hold non-expert tensors, 1,217 +tensors totalling 37.8 GB. Without `--source-digests`, those 85 shards +(~454 GB) are hashed once each to pin `repack-of` materials; digests are cached +in `$CAMPAIGN/source-digests.json` so windows never repeat the work. + +### 4.3 Encode + +```bash +fq-assemble-lora encode --source $SRC --recipe $RECIPE --out $CAMPAIGN \ + --encoder-source $ENC --sign-key $KEY --block-size 32 \ + --layers 3-10 --devices cuda:0,cuda:1,cuda:2,cuda:3,cuda:4,cuda:5,cuda:6,cuda:7 +``` + +`--devices` runs one single-device worker per GPU over disjoint +`(layer, block)` work, logging to `$CAMPAIGN/logs/`. Workers have disjoint work +and output paths — no NCCL, no shared writes — but they do share the source +shards, the page cache and the output volume's I/O bandwidth. Each block is +claimed with an `O_EXCL` lock, so two launchers pointed at one campaign cannot +burn GPU hours on the same block; the launcher clears stale claims once, before +it forks. + +**Run one launcher at a time.** The work split is disjoint within a launcher, +not across launchers. + +### 4.4 Resume after preemption + +Re-run the identical command. A block counts as done only when every node of it +has a committed digest, and a block is re-encoded as a whole: a residual is only +valid against the parent bytes published beside it, so the tool never rewrites +one stage against a parent it merely recomputed in memory. Commit markers are +retracted before a payload is rewritten, so an interrupted rewrite can never +leave a digest that agrees with stale bytes. + +Nothing is deleted. A recipe, block-size or source-revision change is refused +outright — with or without `--force` — because converting a campaign in place +would leave the previous layout's shards to be published alongside the new ones. +`--force` re-encodes blocks of *this* campaign. + +### 4.5 Publish per window + +Every block file is final and self-describing on arrival, with its digest and +signed attestation beside it, so a window can be staged to the repository as +soon as it finishes — but a **base tier is not loadable until `finalize` links +the skeleton into it**, so per-window uploads are staging only (use a private +branch or a staging prefix), and the authoritative upload happens after §4.6. + +```bash +# per window: stage what is final, then delete that window's SOURCE shards +for path in "$CAMPAIGN"/stages/*; do + hf upload "$path" "stages/$(basename "$path")" # + digests, attestations +done + +# after finalize: publish the complete base trees, including the skeleton +# hardlinks, metadata, digests and attestations that make them loadable +for label in k2 k3; do + hf upload "$CAMPAIGN/base/$label" "base/$label" +done +hf upload "$CAMPAIGN/assemblies" assemblies # signed plans +hf upload "$CAMPAIGN/campaign_summary.json" campaign_summary.json +``` + +Xet deduplication means the second base's identical skeleton content should cost +almost no extra WAN. Upload from a CPU process while the GPUs work, and pause the +fleet before the final drain. + +### 4.6 Finalize once, at the end + +```bash +fq-assemble-lora finalize --source $SRC --recipe $RECIPE --out $CAMPAIGN \ + --sign-key $KEY --block-size 32 +``` + +Needs only `config.json` and the source index — deleted source payloads do not +have to be restaged — but it **re-reads and re-hashes all 1.332 TB of campaign +output**: a digest computed by the process that wrote a file proves nothing +about the file that survived. Budget 0.4–1.9 h depending on volume throughput +and **pause the GPU fleet first**; this pass needs no GPU. + +It fails, naming what is wrong, unless every expected block and skeleton shard +is present, hashes to its recorded digest, carries a signed line naming those +exact bytes and span digests, agrees on one signer and one encoder build, and +(for stages) names the parent fragment this campaign published. It also refuses +shards the recipe does not describe. On success it writes, per base: +`config.json` with `hybrid_tr3_tail` and `quantization_config`, +`tier_bitmap.json`, `model.safetensors.index.json`, `MANIFEST.sha256` (built +from the verified digests), and hardlinks the skeleton payload in. Then one +signed `assemblies/