| 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) |
|
|
