rwkv-quant 0.1.0__py3-none-any.whl

rwkv_quant/__init__.py

@@ -0,0 +1,4 @@
+from .api import quantize, calibrate
+from .calibration import QuantConfig, GROUPS
+
+__all__ = ["quantize", "calibrate", "QuantConfig", "GROUPS"]

rwkv_quant/api.py

@@ -0,0 +1,111 @@
+"""
+Высокоуровневый API. Два входа:
+  - quantize(ckpt, out, preset=...)   -- быстрый старт, готовые пресеты
+  - quantize(ckpt, out, config=...)   -- полный контроль через QuantConfig
+  - calibrate(ckpt, corpus)           -- автоматически подобрать QuantConfig
+                                          под конкретный чекпоинт вместо
+                                          пресета "с потолка" (см. README:
+                                          чувствительность к квантованию НЕ
+                                          переносится между масштабами модели)
+"""
+import torch
+
+from .presets import PRESETS
+from .calibration import GROUPS, QuantConfig
+from .calibration.ablation import single_group_ablation, perplexity, combined_sanity_check
+from .models.rwkv7_ref import RWKV7Ref
+from .formats import save
+
+
+def quantize(checkpoint_path: str, output_path: str, preset: str = "reduction",
+             config: QuantConfig = None):
+    """
+    Quick-start: quantize(ckpt, out, preset="compression")
+    Advanced:    quantize(ckpt, out, config=QuantConfig(proj=4, ...))
+
+    preset игнорируется, если передан config.
+    """
+    if config is None:
+        if preset not in PRESETS:
+            raise ValueError(f"unknown preset {preset!r}, choose from {list(PRESETS)}")
+        config = PRESETS[preset]
+
+    # метаданные (naming/n_layer/...) дешевле всего взять из уже написанного
+    # загрузчика RWKV7Ref, вместо повторной ручной детекции формата
+    model = RWKV7Ref(checkpoint_path, device="cpu", dtype=torch.bfloat16)
+    naming, n_layer, n_embd = model.naming, model.n_layer, model.n_embd
+    head_size, vocab_size = model.head_size, model.vocab_size
+    del model
+
+    if checkpoint_path.endswith(".pth"):
+        sd = torch.load(checkpoint_path, map_location="cpu")
+    else:
+        from safetensors.torch import load_file
+        sd = load_file(f"{checkpoint_path}/model.safetensors")
+
+    return save(sd, config, output_path, naming, n_layer, n_embd, head_size, vocab_size)
+
+
+def calibrate(checkpoint_path: str, eval_corpus_path: str, device: str = "mps",
+              ppl_threshold_pct: float = 5.0, spqr_outlier_frac: float = 0.02,
+              small_clip_percentile: float = 99.9, verbose: bool = True) -> QuantConfig:
+    """
+    Прогоняет single-group ablation на конкретном чекпоинте и подбирает
+    QuantConfig вместо использования пресета "с потолка".
+
+    Эвристика (по находкам из README, не универсальный поиск):
+      1. Для каждой группы -- минимальная битность из {8,6,4,2} с Δppl% <=
+         ppl_threshold_pct без какой-либо доп. обработки.
+      2. Если группа из {proj, cmix, emb_head} не прошла порог даже на INT8 --
+         пробуем INT4 + SpQR sparse outlier (percentile-clipping для них НЕ
+         пробуем: он делает эти группы хуже, см. README).
+      3. Если группа "small" не прошла порог -- пробуем INT6 + percentile
+         clip (SpQR для неё не нужен, clip уже полностью решает проблему).
+    """
+    model = RWKV7Ref(checkpoint_path, device=device, dtype=torch.bfloat16)
+    data = torch.load(eval_corpus_path).to(device)
+
+    baseline_ppl, results = single_group_ablation(
+        model, data, groups=GROUPS, bits_grid=(8, 6, 4, 2), verbose=verbose)
+
+    best_bits = {}
+    for group in GROUPS:
+        group_results = [r for r in results if r[0] == group and r[3] <= ppl_threshold_pct]
+        best_bits[group] = min((r[1] for r in group_results), default=16)
+
+    outlier_fracs, clip_percentiles = {}, {}
+
+    for group in ("proj", "cmix", "emb_head"):
+        if best_bits[group] > 4:
+            cfg_try = QuantConfig(**{group: 4}, outlier_fracs={group: spqr_outlier_frac})
+            ppl_try = perplexity(model, data, cfg_try)
+            delta = 100 * (ppl_try - baseline_ppl) / baseline_ppl
+            if verbose:
+                print(f"  retry {group} @ INT4+SpQR({spqr_outlier_frac*100:.0f}%): Δ={delta:+.2f}%")
+            if delta <= ppl_threshold_pct * 3:  # мягче: это уже "сжатие ценой качества"
+                best_bits[group] = 4
+                outlier_fracs[group] = spqr_outlier_frac
+
+    if best_bits["small"] > 6:
+        cfg_try = QuantConfig(small=6, clip_percentiles={"small": small_clip_percentile})
+        ppl_try = perplexity(model, data, cfg_try)
+        delta = 100 * (ppl_try - baseline_ppl) / baseline_ppl
+        if verbose:
+            print(f"  retry small @ INT6+clip(p{small_clip_percentile}): Δ={delta:+.2f}%")
+        if delta <= ppl_threshold_pct:
+            best_bits["small"] = 6
+            clip_percentiles["small"] = small_clip_percentile
+
+    # single-group ablation выше оценивает каждую группу НЕЗАВИСИМО и не
+    # может поймать эффекты взаимодействия при одновременном квантовании
+    # (см. presets.py: все 4 LoRA-ветки на INT4 порознь безобидны, вместе --
+    # ~150x взрыв ppl). Обязательная проверка целиком перед выдачей конфига.
+    best_bits, outlier_fracs, final_ppl, final_delta = combined_sanity_check(
+        model, data, best_bits, outlier_fracs, clip_percentiles,
+        baseline_ppl=baseline_ppl, ppl_threshold_pct=ppl_threshold_pct, verbose=verbose)
+
+    config = QuantConfig(clip_percentiles=clip_percentiles, outlier_fracs=outlier_fracs, **best_bits)
+    if verbose:
+        print(f"\nCalibrated config: {config}\noutlier_fracs={outlier_fracs}  clip={clip_percentiles}")
+        print(f"Финальный Δppl(целиком) = {final_delta:+.2f}%")
+    return config

rwkv_quant/backends/metal/quant_linear.py

@@ -0,0 +1,106 @@
+"""
+backends/metal/quant_linear.py — INT8 per-row scale + SpQR sparse-outlier
+Linear на Metal (MLX custom kernel), без промежуточного bf16-разворачивания.
+
+Заменяет formats.reader.load_dequantized() (весь тензор int8 -> bf16 в
+памяти) на прямой матмул по codes/scale + точечную поправку по
+разреженным outlier'ам. codes уже нулевые в outlier-позициях (см.
+writer._real_quantize_sparse_outlier), поэтому outlier-коррекция — это
+чистое дополнение, а не overwrite.
+"""
+import numpy as np
+import mlx.core as mx
+
+from ...formats.schema import int8_codes
+
+
+def _build_outlier_csr(outlier_indices, outlier_values, n_rows):
+    """torch [n_outliers,2] (row,col) + [n_outliers] значения ->
+    row_offsets[n_rows+1], cols[n_outliers], vals[n_outliers] — отсортировано
+    по row, для CSR-обхода внутри kernel'я. Выполняется один раз при загрузке."""
+    if outlier_indices is None or outlier_indices.numel() == 0:
+        return (mx.array(np.zeros(n_rows + 1, dtype=np.int32)),
+                mx.array(np.zeros(1, dtype=np.int32)),
+                mx.array(np.zeros(1, dtype=np.float32)))
+
+    rows = outlier_indices[:, 0].numpy()
+    cols = outlier_indices[:, 1].numpy()
+    vals = outlier_values.float().numpy()
+
+    order = np.argsort(rows, kind="stable")
+    rows, cols, vals = rows[order], cols[order], vals[order]
+
+    row_offsets = np.zeros(n_rows + 1, dtype=np.int32)
+    row_offsets[1:] = np.cumsum(np.bincount(rows, minlength=n_rows))
+
+    return (mx.array(row_offsets), mx.array(cols.astype(np.int32)), mx.array(vals.astype(np.float32)))
+
+
+_kernel_cache = {}
+
+
+def _get_kernel(IN: int, OUT: int, has_outliers: bool):
+    key = (IN, OUT, has_outliers)
+    if key in _kernel_cache:
+        return _kernel_cache[key]
+
+    hdr = f"""
+constant uint IN_C  = {IN};
+constant uint OUT_C = {OUT};
+"""
+    outlier_body = """
+    uint start = row_offsets[oi];
+    uint end   = row_offsets[oi+1];
+    for (uint idx = start; idx < end; idx++) {
+        uint c = (uint)outlier_cols[idx];
+        acc += x[n*IN_C + c] * outlier_vals[idx];
+    }
+""" if has_outliers else ""
+
+    body = """
+    uint oi = thread_position_in_grid.x;   // output feature
+    uint n  = thread_position_in_grid.y;   // batch row
+
+    float acc = 0.0f;
+    for (uint c = 0; c < IN_C; c++)
+        acc += x[n*IN_C + c] * (float)codes[oi*IN_C + c];
+    acc *= (float)scale[oi];
+""" + outlier_body + """
+    out[n*OUT_C + oi] = acc;
+"""
+    kern = mx.fast.metal_kernel(
+        name=f"quant_linear_{'spqr' if has_outliers else 'plain'}_{IN}_{OUT}",
+        input_names=["x", "codes", "scale", "row_offsets", "outlier_cols", "outlier_vals"],
+        output_names=["out"],
+        header=hdr, source=body,
+    )
+    _kernel_cache[key] = kern
+    return kern
+
+
+class QuantLinear:
+    """Один Linear на int8 codes + per-row scale (+ опц. SpQR), веса живут
+    на GPU в упакованном виде — никогда не разворачиваются в полный bf16."""
+
+    def __init__(self, qt):
+        # qt: rwkv_quant.formats.schema.QuantizedTensor, bits < 16
+        assert qt.bits < 16, "QuantLinear только для квантованных тензоров"
+        self.out_features, self.in_features = qt.shape
+        self.codes = mx.array(int8_codes(qt).numpy())    # int8 [out, in] (packed распаковывается: v1 -- референс)
+        self.scale = mx.array(qt.scale.float().numpy().reshape(-1))  # fp32 [out]
+        self.row_offsets, self.outlier_cols, self.outlier_vals = _build_outlier_csr(
+            qt.outlier_indices, qt.outlier_values, self.out_features)
+        self.has_outliers = qt.outlier_indices is not None and qt.outlier_indices.numel() > 0
+
+    def __call__(self, x):
+        lead_shape = x.shape[:-1]
+        x2d = x.reshape(-1, self.in_features).astype(mx.float32)
+        N = x2d.shape[0]
+        kern = _get_kernel(self.in_features, self.out_features, self.has_outliers)
+        out = kern(
+            inputs=[x2d, self.codes, self.scale, self.row_offsets, self.outlier_cols, self.outlier_vals],
+            grid=(self.out_features, N, 1), threadgroup=(1, 1, 1),
+            output_shapes=[(N, self.out_features)],
+            output_dtypes=[mx.float32],
+        )[0]
+        return out.reshape(*lead_shape, self.out_features)