Claude
b108d4f4f5
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
import json
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from .. import mxfp8 as m
|
|
|
|
|
|
def _sd(blocks=8):
|
|
g = torch.Generator().manual_seed(0)
|
|
r = lambda *s: torch.randn(*s, generator=g, dtype=torch.bfloat16)
|
|
sd = {"token_refiner.blocks.0.attn.qkv_proj.weight": r(96, 64), "video_patch_proj.weight": r(32, 64)}
|
|
for i in range(blocks):
|
|
sd[f"blocks.{i}.attn.qkv_proj.weight"] = r(96, 64)
|
|
sd[f"blocks.{i}.mlp.fc2.weight"] = r(32, 100) # 100 is not a multiple of 32
|
|
sd[f"blocks.{i}.adaln_proj.linear.weight"] = r(64, 128)
|
|
return sd
|
|
|
|
|
|
def _quantizable():
|
|
pytest.importorskip("comfy.quant_ops")
|
|
from comfy.quant_ops import QUANT_ALGOS
|
|
if "mxfp8" not in QUANT_ALGOS:
|
|
pytest.skip("this comfy_kitchen has no MXFP8")
|
|
|
|
|
|
def test_edges_and_non_projections_stay_as_they_are():
|
|
_quantizable()
|
|
sd = _sd()
|
|
keep = {k: v.clone() for k, v in sd.items()}
|
|
n, note = m.quantize_state_dict(sd)
|
|
assert n == 3 and "blocks 2-4" in note and "3 left" in note # 8 blocks: 0,1 and 5,6,7 stay
|
|
for i in (2, 3, 4):
|
|
k = f"blocks.{i}.attn.qkv_proj"
|
|
assert sd[k + ".weight"].dtype == torch.float8_e4m3fn
|
|
assert json.loads(sd[k + ".comfy_quant"].numpy().tobytes()) == {"format": "mxfp8"}
|
|
assert sd[k + ".weight_scale"].dtype == torch.float8_e8m0fnu
|
|
for k, v in keep.items():
|
|
if ".qkv_proj" in k and k.startswith("blocks.") and int(k.split(".")[1]) in (2, 3, 4):
|
|
continue
|
|
assert torch.equal(sd[k], v), k
|
|
|
|
|
|
def test_second_pass_and_foreign_files_are_left_alone():
|
|
_quantizable()
|
|
sd = _sd()
|
|
m.quantize_state_dict(sd)
|
|
again = {k: v.clone() for k, v in sd.items()}
|
|
assert m.quantize_state_dict(sd)[0] == 0 and all(torch.equal(sd[k].view(torch.uint8), again[k].view(torch.uint8)) for k in sd)
|
|
n, note = m.quantize_state_dict({"unet.weight": torch.zeros(4, 4)})
|
|
assert n == 0 and "loaded as is" in note
|