"""4-bit weight quantization study: element format x scale precision x extras. CISM's ``hybrid-int4`` packs MLP weights as symmetric INT4 (levels -7..7, blocks of 32, per-block FP32 scale = amax/7, round-half-away-from-zero) with everything else 2D per-row INT8. This study holds that pipeline fixed and swaps one knob at a time: * element format: uniform INT4 vs E2M1 (MXFP4) vs NF4 (QLoRA) * scale precision: FP32 vs FP8 E4M3 vs power-of-two E8M0 (OCP MXFP4) * block size: 32 vs 16 (NVFP4 granularity) * error feedback: per-block residual diffusion on top of the winner The NumPy simulation mirrors the native quantizer bit-for-bit (validated against the real engine below), rebuilds FP32 engines from the dequantized weights, and scores perplexity on the AutoBench corpus over small and not-so- small checkpoints. Quality only; speed is a kernel question, decided later. """ import argparse import json import numpy as np import torch from cism import Engine from cism.autobench import CORPUS from cism.loader import import_model E2M1 = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=np.float64) NF4 = np.array([-1.0, -0.6962, -0.5239, -0.3650, -0.1848, -0.1274, -0.0400, 0.0, 0.0397, 0.1292, 0.2826, 0.5626, 0.5824, 0.8400, 1.0770, 1.2712]) TINY = float(np.finfo(np.float32).tiny) def round_half_away(values: np.ndarray) -> np.ndarray: # std::round semantics: halfway cases go away from zero. np.round does # banker's rounding and would NOT mirror the native quantizer. return np.floor(np.abs(values) + 0.5) * np.sign(values) def _blocked(matrix: np.ndarray, block: int) -> np.ndarray: rows, cols = matrix.shape pad = (-cols) % block padded = np.pad(matrix, ((0, 0), (0, pad))) return padded.reshape(rows, -1, block) def _scale_for(amax: np.ndarray, mode: str, qmax: float) -> np.ndarray: if mode == "fp32": return np.maximum(amax / qmax, TINY) if mode == "e8m0": # OCP shared exponent: power of two, floor(log2(amax)) - (2 or 3) so # the largest element value can reach but not exceed amax. return np.exp2(np.floor(np.log2(np.maximum(amax, TINY))) - np.log2(qmax)) if mode == "e4m3": base = torch.from_numpy(np.maximum(amax / qmax, TINY)) return np.maximum(base.to(torch.float8_e4m3fn).to(torch.float32).numpy(), TINY) raise ValueError(scale_mode := mode) def _lut_round(magnitude: np.ndarray, lut: np.ndarray) -> np.ndarray: index = np.argmin(np.abs(magnitude[..., None] - lut), axis=-1) return lut[index] def _deblock(dequantized: np.ndarray, rows: int, cols: int) -> np.ndarray: return dequantized.reshape(rows, -1)[:, :cols].astype(np.float32) def quantize_uniform(matrix: np.ndarray, *, levels: int = 7, scale_mode: str = "fp32", block: int = 32) -> np.ndarray: """Symmetric integer grid (-levels..levels), mirrors native int4.""" rows, cols = matrix.shape blocks = _blocked(matrix, block) amax = np.abs(blocks).max(axis=-1, keepdims=True) scale = _scale_for(amax, scale_mode, float(levels)) quantized = np.clip(round_half_away(blocks / scale), -levels, levels) return _deblock(quantized * scale, rows, cols) def quantize_fp_e2m1(matrix: np.ndarray, *, scale_mode: str = "e4m3", block: int = 32) -> np.ndarray: rows, cols = matrix.shape blocks = _blocked(matrix, block) amax = np.abs(blocks).max(axis=-1, keepdims=True) scale = _scale_for(amax, scale_mode, 6.0) dequantized = np.sign(blocks) * _lut_round(np.abs(blocks) / scale, E2M1) * scale return _deblock(dequantized, rows, cols) def quantize_nf4(matrix: np.ndarray, *, scale_mode: str = "fp32", block: int = 32) -> np.ndarray: rows, cols = matrix.shape blocks = _blocked(matrix, block) amax = np.abs(blocks).max(axis=-1, keepdims=True) scale = _scale_for(amax, scale_mode, 1.0) dequantized = _lut_round(blocks / scale, NF4) * scale return _deblock(dequantized, rows, cols) def quantize_uniform_feedback(matrix: np.ndarray, *, levels: int = 7) -> np.ndarray: """Fixed-scale RTN plus per-block residual diffusion along the row.""" rows, cols = matrix.shape blocks = _blocked(matrix, 32) amax = np.abs(blocks).max(axis=-1) scale = np.maximum(amax / float(levels), TINY) # (rows, nblocks) quantized = np.zeros_like(blocks) residual = np.zeros((rows, blocks.shape[1])) for position in range(blocks.shape[-1]): target = blocks[:, :, position] + residual quantized[:, :, position] = np.clip(round_half_away(target / scale), -levels, levels) residual = target - quantized[:, :, position] * scale return _deblock(quantized * scale[..., None], rows, cols) def quantize_e2m1_feedback(matrix: np.ndarray, *, scale_mode: str = "e4m3") -> np.ndarray: rows, cols = matrix.shape blocks = _blocked(matrix, 32) amax = np.abs(blocks).max(axis=-1) scale = _scale_for(amax, scale_mode, 6.0) # (rows, nblocks) quantized = np.zeros_like(blocks) residual = np.zeros((rows, blocks.shape[1])) for position in range(blocks.shape[-1]): target = blocks[:, :, position] + residual quantized[:, :, position] = np.sign(target) * _lut_round(np.abs(target) / scale, E2M1) residual = target - quantized[:, :, position] * scale return _deblock(quantized * scale[..., None], rows, cols) def quantize_int8_row(matrix: np.ndarray) -> np.ndarray: """Mirror of the native Storage::int8 constructor (per-row scale).""" amax = np.abs(matrix).max(axis=1, keepdims=True) scale = np.maximum(amax / 127.0, TINY) quantized = np.clip(round_half_away(matrix / scale), -127, 127) return (quantized * scale).astype(np.float32) def simulate(weights: dict, scheme) -> dict: out = {} for name, array in weights.items(): if array.ndim == 1: out[name] = array.astype(np.float32) elif ".mlp." in name: out[name] = scheme(array) else: out[name] = quantize_int8_row(array) return out def perplexity(config: dict, weights: dict, tokenizer) -> float: engine = Engine.from_weights(config, weights, tokenizer, model_id="study") return engine.nll(CORPUS, window=128)["perplexity"] SCHEMES = { "int4 fp32 b32 (ours)": lambda m: quantize_uniform(m), "int4 e4m3 b32": lambda m: quantize_uniform(m, scale_mode="e4m3"), "int4 e8m0 b32": lambda m: quantize_uniform(m, scale_mode="e8m0"), "E2M1 e8m0 b32 (OCP MXFP4)": lambda m: quantize_fp_e2m1(m, scale_mode="e8m0"), "E2M1 e4m3 b32 (NVFP4-ish)": lambda m: quantize_fp_e2m1(m, scale_mode="e4m3"), "E2M1 e4m3 b16": lambda m: quantize_fp_e2m1(m, scale_mode="e4m3", block=16), "int4 e4m3 b16": lambda m: quantize_uniform(m, scale_mode="e4m3", block=16), "E2M1 fp32 b32": lambda m: quantize_fp_e2m1(m, scale_mode="fp32"), "NF4 fp32 b32 (QLoRA)": lambda m: quantize_nf4(m, scale_mode="fp32"), "NF4 e4m3 b32": lambda m: quantize_nf4(m, scale_mode="e4m3"), "int4 fb fp32 b32": lambda m: quantize_uniform_feedback(m), "E2M1 fb e4m3 b32": lambda m: quantize_e2m1_feedback(m), } def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("models", nargs="+") args = parser.parse_args() results = {} for model_id in args.models: loaded = import_model(model_id, local_files_only=True) config, weights, tokenizer = dict(loaded.config), loaded.weights, loaded.tokenizer row = {} row["fp32 (reference)"] = perplexity(config, weights, tokenizer) row["hybrid-int4 (native engine)"] = Engine.from_pretrained( model_id, precision="hybrid-int4", local_files_only=True).nll(CORPUS, window=128)["perplexity"] for label, scheme in SCHEMES.items(): row[label] = perplexity(config, simulate(weights, scheme), tokenizer) results[model_id] = row reference = row["fp32 (reference)"] print(f"\n=== {model_id} (fp32 PPL {reference:.2f}) ===") for label, value in row.items(): delta = 100 * (value - reference) / reference print(f" {label:32s} PPL {value:10.3f} ({delta:+.3f}%)") print("\n" + json.dumps(results, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())