File size: 1,797 Bytes
f2c0505 | 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 | import torch
from orbitquant_wan_a2.nibbles import unpack_uint4
from orbitquant_wan_a2.source_quant import TARGET_SUFFIXES, is_target_key, _quantize_rows_to_packed
def test_exact_480_target_rule():
keys = [f"blocks.{b}{s}" for b in range(40) for s in TARGET_SUFFIXES]
assert len(keys) == 480
assert len(set(keys)) == 480
assert all(is_target_key(k) for k in keys)
assert not is_target_key("blocks.0.block.modulation")
assert not is_target_key("patch_embedding.weight")
def test_direct_quant_packs_codes_and_bf16_scale():
torch.manual_seed(123)
d = 8
cb = torch.linspace(-0.8, 0.8, 16, dtype=torch.float32)
bank_item = {
"perm": torch.arange(d, dtype=torch.int64),
"signs": torch.ones(d, dtype=torch.int8),
"codebook": cb,
"block_size": torch.tensor([8], dtype=torch.int32),
}
w = torch.randn(3, d, dtype=torch.bfloat16)
packed, scale, stats = _quantize_rows_to_packed(w, bank_item, device=torch.device("cpu"), row_chunk=2)
assert packed.dtype == torch.uint8
assert packed.shape == (3, d // 2)
assert scale.dtype == torch.bfloat16
assert scale.shape == (3,)
codes = unpack_uint4(packed, d)
assert int(codes.min()) >= 0 and int(codes.max()) <= 15
assert stats["row_scale_dtype"] == "bfloat16"
def test_safetensors_shape_probe_uses_slice_metadata(tmp_path):
import json
from safetensors.torch import save_file
from orbitquant_wan_a2.source_quant import _shape_of
shard = tmp_path / 'model-00001-of-00001.safetensors'
save_file({'blocks.0.block.self_attn.q.weight': torch.zeros(5, 8, dtype=torch.bfloat16)}, str(shard))
wm = {'blocks.0.block.self_attn.q.weight': shard.name}
assert _shape_of(tmp_path, wm, 'blocks.0.block.self_attn.q.weight') == (5, 8)
|