462 lines
17 KiB
Python
462 lines
17 KiB
Python
"""Tests for OrbitQuant loader support in ComfyUI-INT-Crush.
|
|
|
|
Test strategy: manually construct state dicts with the OrbitQuant format
|
|
(packed indices, bfloat16 norms, global RPBH/codebook tensors,
|
|
format_version=4 metadata) — no converter dependency needed.
|
|
"""
|
|
|
|
import math
|
|
import torch
|
|
import pytest
|
|
|
|
from _intcrush import load
|
|
|
|
_oq = load("_orbitquant")
|
|
_qu = load("_quant_utils")
|
|
|
|
|
|
# ── Helpers ──
|
|
|
|
def _make_codebook(d: int, bits: int) -> torch.Tensor:
|
|
"""Build a simple uniform codebook for testing.
|
|
|
|
In production, this comes from Lloyd-Max on the f_d marginal.
|
|
For tests, a uniform grid suffices to verify the dequant path.
|
|
"""
|
|
n = 2 ** bits
|
|
centroids = torch.linspace(-1.0, 1.0, n, dtype=torch.float32)
|
|
return centroids
|
|
|
|
|
|
def _make_orbitquant_state_dict(
|
|
d: int = 256,
|
|
out_features: int = 32,
|
|
bits: int = 4,
|
|
seed: int = 42,
|
|
include_adaln: bool = False,
|
|
) -> tuple[dict, dict]:
|
|
"""Build a synthetic OrbitQuant model's state_dict and metadata.
|
|
|
|
Returns (state_dict, metadata) ready for save_file().
|
|
"""
|
|
h = _oq.largest_power_of_two_dividing(d)
|
|
perm, signs, h = _oq.make_rpbh_params(d, seed=seed)
|
|
d_padded = perm.shape[0]
|
|
|
|
codebook = _make_codebook(d, bits)
|
|
codebook_key = f"_orbitquant_codebook_{d}_{bits}"
|
|
|
|
# Build a fake weight: random unit directions * random norms
|
|
W = torch.randn(out_features, d)
|
|
norms = W.norm(dim=1, keepdim=True)
|
|
W_dir = W / (norms + 1e-10)
|
|
|
|
# Quantize directions to codebook
|
|
# (simplified: use nearest centroid)
|
|
centroids = codebook
|
|
boundaries = (centroids[:-1] + centroids[1:]) / 2.0
|
|
indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
|
|
indices = indices.clamp(0, len(centroids) - 1).reshape(out_features, d)
|
|
|
|
# Offset to signed int8 range
|
|
offset = 1 << (bits - 1)
|
|
indices_int8 = (indices - offset).to(torch.int8)
|
|
|
|
if bits == 4:
|
|
packed = _qu.pack_int4(indices_int8).to(torch.uint8)
|
|
else:
|
|
packed = indices_int8
|
|
|
|
scale = norms.to(torch.bfloat16)
|
|
|
|
state_dict = {
|
|
"layers.0.q_proj.weight": packed,
|
|
"layers.0.q_proj.weight_scale": scale,
|
|
"_rpbh_perm_256": perm,
|
|
"_rpbh_signs_256": signs,
|
|
codebook_key: codebook,
|
|
}
|
|
|
|
metadata = {
|
|
"int_crush.format_version": "4",
|
|
"int_crush.orbitquant": "true",
|
|
"int_crush.rpbh_seed": str(seed),
|
|
"int_crush.codebook_bits": str(bits),
|
|
"int_crush.orbitquant_dims": str(d),
|
|
"int_crush.orbitquant_d_padded": f"{d}:{d_padded}",
|
|
"int_crush.layers.0.q_proj.weight.quant_scheme": "orbitquant_codebook",
|
|
}
|
|
|
|
return state_dict, metadata
|
|
|
|
|
|
# ── Tests ──
|
|
|
|
class TestRPBHRotate:
|
|
"""Test RPBH rotation matches converter implementation."""
|
|
|
|
def test_norm_preservation(self):
|
|
"""‖x @ Π^T‖ = ‖x‖ (orthogonal)."""
|
|
d = 256
|
|
perm, signs, h = _oq.make_rpbh_params(d, seed=42)
|
|
x = torch.randn(8, d)
|
|
y = _oq.rpbh_rotate(x, perm, signs, h)
|
|
assert torch.allclose(x.norm(), y.norm(), rtol=1e-3)
|
|
|
|
def test_cancellation(self):
|
|
"""(W @ Π^T) @ (x @ Π^T)^T = W @ x^T."""
|
|
for d in [256, 512, 1024]:
|
|
perm, signs, h = _oq.make_rpbh_params(d, seed=42)
|
|
W = torch.randn(16, d)
|
|
x = torch.randn(8, d)
|
|
W_rot = _oq.rpbh_rotate(W, perm, signs, h)
|
|
x_rot = _oq.rpbh_rotate(x, perm, signs, h)
|
|
y_rot = W_rot @ x_rot.T
|
|
y_orig = W @ x.T
|
|
rel_err = (y_rot - y_orig).norm() / y_orig.norm()
|
|
assert rel_err < 1e-4
|
|
|
|
def test_fwht_matches_sylvester(self):
|
|
"""FWHT output matches Sylvester Hadamard @ x."""
|
|
h = 16
|
|
H = torch.tensor([[1.0]])
|
|
while H.shape[0] < h:
|
|
H = torch.kron(torch.tensor([[1., 1.], [1., -1.]]), H)
|
|
H = H / math.sqrt(h)
|
|
x = torch.randn(5, h)
|
|
y_fwht = _oq.fast_walsh_hadamard(x)
|
|
y_matrix = x @ H.T
|
|
assert torch.allclose(y_fwht, y_matrix, atol=1e-4)
|
|
|
|
def test_mass_balancing(self):
|
|
"""Spiked input → uniform variance after RPBH (CV < 0.3)."""
|
|
d = 1024
|
|
h = _oq.largest_power_of_two_dividing(d)
|
|
k = d // h
|
|
if k < 2:
|
|
pytest.skip("single-block dimension")
|
|
perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
|
|
n_samples = 5000
|
|
x = torch.zeros(n_samples, d)
|
|
x[:, :h] = torch.randn(n_samples, h)
|
|
x = x / x.norm(dim=1, keepdim=True)
|
|
z = _oq.rpbh_rotate(x, perm, signs, h)
|
|
var_per_coord = z.double().pow(2).mean(0)
|
|
cv = var_per_coord.std().item() / var_per_coord.mean().item()
|
|
assert cv < 0.3
|
|
|
|
|
|
class TestCodebookDequant:
|
|
"""Test codebook dequantization."""
|
|
|
|
def test_int4_roundtrip(self):
|
|
"""Quantize → pack → dequant ≈ original (within codebook MSE)."""
|
|
d = 256
|
|
bits = 4
|
|
codebook = _make_codebook(d, bits)
|
|
W = torch.randn(32, d) * 0.1
|
|
norms = W.norm(dim=1, keepdim=True)
|
|
W_dir = W / (norms + 1e-10)
|
|
boundaries = (codebook[:-1] + codebook[1:]) / 2.0
|
|
indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
|
|
indices = indices.clamp(0, 15).reshape(32, d)
|
|
offset = 1 << (bits - 1)
|
|
indices_int8 = (indices - offset).to(torch.int8)
|
|
packed = _qu.pack_int4(indices_int8).to(torch.uint8)
|
|
|
|
W_dequant = _oq.codebook_dequant_int4(
|
|
packed, norms.to(torch.bfloat16), codebook, d, d)
|
|
|
|
# Should be close to original (within codebook quantization error)
|
|
mse = (W - W_dequant.float()).pow(2).mean().item()
|
|
assert mse < 0.01
|
|
|
|
def test_int8_roundtrip(self):
|
|
"""INT8 indices → codebook lookup → norms multiply."""
|
|
d = 256
|
|
bits = 8
|
|
codebook = _make_codebook(d, bits)
|
|
W = torch.randn(16, d) * 0.1
|
|
norms = W.norm(dim=1, keepdim=True)
|
|
W_dir = W / (norms + 1e-10)
|
|
boundaries = (codebook[:-1] + codebook[1:]) / 2.0
|
|
indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
|
|
indices = indices.clamp(0, 255).reshape(16, d)
|
|
offset = 1 << (bits - 1)
|
|
indices_int8 = (indices - offset).to(torch.int8)
|
|
|
|
W_dequant = _oq.codebook_dequant_int8(
|
|
indices_int8, norms.to(torch.bfloat16), codebook, d, d)
|
|
|
|
mse = (W - W_dequant.float()).pow(2).mean().item()
|
|
assert mse < 0.001 # INT8 should be more accurate
|
|
|
|
def test_no_padding_trim(self):
|
|
"""d_padded > d → output is [out, d_padded] (no trim at inference)."""
|
|
d = 200 # not divisible by most block sizes
|
|
h = 128 # explicit block size
|
|
d_padded = 256 # next multiple of 128
|
|
bits = 4
|
|
codebook = _make_codebook(d_padded, bits)
|
|
norms = torch.ones(8, 1, dtype=torch.bfloat16)
|
|
indices_int8 = torch.zeros(8, d_padded, dtype=torch.int8)
|
|
|
|
W_dequant = _oq.codebook_dequant_int4(
|
|
_qu.pack_int4(indices_int8).to(torch.uint8),
|
|
norms, codebook, d, d_padded)
|
|
|
|
assert W_dequant.shape == (8, d_padded) # NOT (8, d)
|
|
|
|
|
|
class TestAdaLNDequant:
|
|
"""Test AdaLN RTN dequantization with group scales."""
|
|
|
|
def test_group_scale_shape(self):
|
|
"""AdaLN scales [out, d/64] → dequant produces [out, d]."""
|
|
d = 256
|
|
group_size = 64
|
|
out = 32
|
|
num_groups = d // group_size
|
|
scales = torch.randn(out, num_groups).abs().clamp(min=1e-4).to(torch.float16)
|
|
packed = _qu.pack_int4(
|
|
torch.randint(-8, 7, (out, d), dtype=torch.int8)
|
|
).to(torch.uint8)
|
|
|
|
W_dequant = _oq.adaln_rtn_dequant(packed, scales, d, bits=4)
|
|
assert W_dequant.shape == (out, d)
|
|
|
|
def test_group_scale_correctness(self):
|
|
"""Each group of 64 values is multiplied by its scale."""
|
|
d = 128
|
|
group_size = 64
|
|
out = 4
|
|
num_groups = d // group_size
|
|
scales = torch.tensor([[2.0, 3.0]] * out, dtype=torch.float16)
|
|
# All indices = 0 → dequant value = 0 * scale = 0
|
|
# Use index 8 (offset 0) → value = -8 * scale
|
|
values = torch.full((out, d), -8, dtype=torch.int8)
|
|
packed = _qu.pack_int4(values).to(torch.uint8)
|
|
|
|
W_dequant = _oq.adaln_rtn_dequant(packed, scales, d, bits=4)
|
|
# Group 0: -8 * 2.0 = -16.0
|
|
# Group 1: -8 * 3.0 = -24.0
|
|
assert torch.allclose(W_dequant[:, :64], torch.full((out, 64), -16.0), atol=0.1)
|
|
assert torch.allclose(W_dequant[:, 64:], torch.full((out, 64), -24.0), atol=0.1)
|
|
|
|
|
|
class TestOrbitQuantLayoutDequantize:
|
|
"""Test OrbitQuantLayout.dequantize() works for both paths."""
|
|
|
|
def test_codebook_path(self):
|
|
"""dequantize with codebook produces correct output."""
|
|
_ql = load("quant_layout")
|
|
OrbitQuantLayout = _ql.OrbitQuantLayout
|
|
d = 256
|
|
bits = 4
|
|
codebook = _make_codebook(d, bits)
|
|
norms = torch.ones(8, 1, dtype=torch.bfloat16)
|
|
packed = _qu.pack_int4(
|
|
torch.zeros(8, d, dtype=torch.int8)
|
|
).to(torch.uint8)
|
|
|
|
params = OrbitQuantLayout.Params(
|
|
scale=norms, orig_dtype=torch.float16, orig_shape=(8, d),
|
|
quant_scheme="orbitquant_codebook", d=d, d_padded=d,
|
|
bits=bits, codebook=codebook,
|
|
)
|
|
W = OrbitQuantLayout.dequantize(packed, params)
|
|
assert W.shape == (8, d)
|
|
|
|
def test_adaln_path(self):
|
|
"""dequantize with adaln_rtn uses group scales."""
|
|
_ql = load("quant_layout")
|
|
OrbitQuantLayout = _ql.OrbitQuantLayout
|
|
d = 256
|
|
out = 8
|
|
scales = torch.ones(out, d // 64, dtype=torch.float16)
|
|
packed = _qu.pack_int4(
|
|
torch.zeros(out, d, dtype=torch.int8)
|
|
).to(torch.uint8)
|
|
|
|
params = OrbitQuantLayout.Params(
|
|
scale=scales, orig_dtype=torch.float16, orig_shape=(out, d),
|
|
quant_scheme="adaln_rtn", d=d, d_padded=d, bits=4,
|
|
)
|
|
W = OrbitQuantLayout.dequantize(packed, params)
|
|
assert W.shape == (out, d)
|
|
|
|
|
|
class TestMetadataParsing:
|
|
"""Test OrbitQuant metadata parsing."""
|
|
|
|
def test_parse_orbitquant_metadata(self):
|
|
"""Valid metadata is parsed correctly."""
|
|
metadata = {
|
|
"int_crush.format_version": "4",
|
|
"int_crush.orbitquant": "true",
|
|
"int_crush.rpbh_seed": "42",
|
|
"int_crush.codebook_bits": "4",
|
|
"int_crush.orbitquant_dims": "256,512",
|
|
"int_crush.orbitquant_d_padded": "256:256,512:512",
|
|
}
|
|
result = _oq.parse_orbitquant_metadata(metadata)
|
|
assert result is not None
|
|
assert result["rpbh_seed"] == 42
|
|
assert result["codebook_bits"] == 4
|
|
assert result["dims"] == [256, 512]
|
|
assert result["d_padded_map"] == {256: 256, 512: 512}
|
|
|
|
def test_parse_non_orbitquant(self):
|
|
"""Non-OrbitQuant metadata returns None."""
|
|
assert _oq.parse_orbitquant_metadata({}) is None
|
|
assert _oq.parse_orbitquant_metadata(None) is None
|
|
assert _oq.parse_orbitquant_metadata({
|
|
"int_crush.format_version": "3"
|
|
}) is None
|
|
|
|
def test_get_quant_scheme_from_metadata(self):
|
|
"""Per-layer quant_scheme is read correctly."""
|
|
metadata = {
|
|
"int_crush.layers.0.q_proj.weight.quant_scheme": "orbitquant_codebook",
|
|
"int_crush.layers.0.k_proj.weight.quant_scheme": "adaln_rtn",
|
|
}
|
|
assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.q_proj.") == "orbitquant_codebook"
|
|
assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.k_proj.") == "adaln_rtn"
|
|
assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.v_proj.") is None
|
|
|
|
def test_get_quant_scheme_strips_model_prefix(self):
|
|
"""Metadata lookup strips model-level prefix from ComfyUI."""
|
|
metadata = {
|
|
"int_crush.single_blocks.0.linear1.weight.quant_scheme": "orbitquant_codebook",
|
|
"int_crush.double_blocks.0.img_attn.q_proj.weight.quant_scheme": "adaln_rtn",
|
|
}
|
|
# ComfyUI passes prefix with model.diffusion_model. prefix
|
|
assert _oq.get_quant_scheme_from_metadata(
|
|
metadata, "model.diffusion_model.single_blocks.0.linear1.") == "orbitquant_codebook"
|
|
assert _oq.get_quant_scheme_from_metadata(
|
|
metadata, "model.diffusion_model.double_blocks.0.img_attn.q_proj.") == "adaln_rtn"
|
|
# Non-existent layer still returns None
|
|
assert _oq.get_quant_scheme_from_metadata(
|
|
metadata, "model.diffusion_model.single_blocks.0.linear2.") is None
|
|
|
|
def test_infer_metadata_from_state_dict(self):
|
|
"""Metadata is inferred from global tensor keys in state_dict."""
|
|
state_dict = {
|
|
"_rpbh_perm_256": torch.randperm(256),
|
|
"_rpbh_signs_256": torch.ones(256),
|
|
"_orbitquant_codebook_256_4": torch.linspace(-1, 1, 16),
|
|
"layers.0.weight": torch.zeros(1),
|
|
}
|
|
result = _oq.infer_orbitquant_metadata_from_state_dict(state_dict)
|
|
assert result is not None
|
|
assert result["int_crush.format_version"] == "4"
|
|
assert result["int_crush.orbitquant"] == "true"
|
|
assert "256" in result["int_crush.orbitquant_dims"]
|
|
assert result["int_crush.codebook_bits"] == "4"
|
|
|
|
def test_infer_metadata_returns_none_without_globals(self):
|
|
"""Returns None when no OrbitQuant globals in state_dict."""
|
|
state_dict = {"layers.0.weight": torch.zeros(1)}
|
|
assert _oq.infer_orbitquant_metadata_from_state_dict(state_dict) is None
|
|
|
|
def test_infer_quant_scheme_from_scale(self):
|
|
"""Scale shape inference works correctly."""
|
|
# [out, 1] → orbitquant_codebook
|
|
scale_cb = torch.ones(8, 1)
|
|
assert _oq.infer_quant_scheme_from_scale(scale_cb) == "orbitquant_codebook"
|
|
|
|
# [out, 4] → adaln_rtn (d=256, group_size=64 → 4 groups)
|
|
scale_adaln = torch.ones(8, 4)
|
|
assert _oq.infer_quant_scheme_from_scale(scale_adaln) == "adaln_rtn"
|
|
|
|
|
|
class TestGlobalTensorCache:
|
|
"""Test OrbitQuant global tensor caching."""
|
|
|
|
def test_load_and_cache_globals(self):
|
|
"""Global tensors are loaded and cached correctly."""
|
|
state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
|
|
model_path = "test_model"
|
|
|
|
globals_cache = _oq.load_orbitquant_globals(state_dict, metadata, model_path)
|
|
|
|
assert 256 in globals_cache["rpbh"]
|
|
assert "perm" in globals_cache["rpbh"][256]
|
|
assert "signs" in globals_cache["rpbh"][256]
|
|
assert (256, 4) in globals_cache["codebooks"]
|
|
|
|
# Verify global tensors were popped from state_dict
|
|
assert "_rpbh_perm_256" not in state_dict
|
|
assert "_rpbh_signs_256" not in state_dict
|
|
assert "_orbitquant_codebook_256_4" not in state_dict
|
|
|
|
# Verify per-layer tensors remain
|
|
assert "layers.0.q_proj.weight" in state_dict
|
|
|
|
# Cleanup
|
|
_oq.clear_orbitquant_globals(model_path)
|
|
|
|
def test_cache_idempotent(self):
|
|
"""Second call returns cached result without popping again."""
|
|
state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
|
|
model_path = "test_model_idempotent"
|
|
|
|
cache1 = _oq.load_orbitquant_globals(state_dict, metadata, model_path)
|
|
# State dict now has globals popped — reconstruct for second call
|
|
# But the function should return cached result without touching state_dict
|
|
cache2 = _oq.load_orbitquant_globals(state_dict, metadata, model_path)
|
|
assert cache1 is cache2
|
|
|
|
_oq.clear_orbitquant_globals(model_path)
|
|
|
|
def test_clear_globals(self):
|
|
"""clear_orbitquant_globals removes cached data."""
|
|
state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
|
|
model_path = "test_model_clear"
|
|
|
|
_oq.load_orbitquant_globals(state_dict, metadata, model_path)
|
|
assert model_path in _oq._orbitquant_globals
|
|
|
|
_oq.clear_orbitquant_globals(model_path)
|
|
assert model_path not in _oq._orbitquant_globals
|
|
|
|
def test_clear_all_globals(self):
|
|
"""clear_orbitquant_globals(None) clears everything."""
|
|
sd1, meta1 = _make_orbitquant_state_dict(d=256, bits=4)
|
|
sd2, meta2 = _make_orbitquant_state_dict(d=256, bits=4)
|
|
|
|
_oq.load_orbitquant_globals(sd1, meta1, "model_a")
|
|
_oq.load_orbitquant_globals(sd2, meta2, "model_b")
|
|
|
|
_oq.clear_orbitquant_globals()
|
|
assert len(_oq._orbitquant_globals) == 0
|
|
|
|
|
|
class TestLoadersDetection:
|
|
"""Test OrbitQuant detection helpers in _loaders."""
|
|
|
|
def test_detect_orbitquant(self):
|
|
"""_detect_orbitquant returns True for valid metadata."""
|
|
_ld = load("_loaders")
|
|
_detect_orbitquant = _ld._detect_orbitquant
|
|
|
|
assert _detect_orbitquant({
|
|
"int_crush.format_version": "4",
|
|
"int_crush.orbitquant": "true",
|
|
}) is True
|
|
|
|
def test_detect_non_orbitquant(self):
|
|
"""_detect_orbitquant returns False for non-OrbitQuant metadata."""
|
|
_ld = load("_loaders")
|
|
_detect_orbitquant = _ld._detect_orbitquant
|
|
|
|
assert _detect_orbitquant({}) is False
|
|
assert _detect_orbitquant(None) is False
|
|
assert _detect_orbitquant({
|
|
"int_crush.format_version": "3",
|
|
}) is False
|
|
assert _detect_orbitquant({
|
|
"int_crush.format_version": "4",
|
|
# missing orbitquant=true
|
|
}) is False
|