ThereforeGames
afc8d7ff01
Add W6A8 checkpoint loading, conversion, inference, LoRA requantization, compile compatibility, and export, plus architecture-aware w4a8_mixed selection. Add MOSS-Image detection and quantization presets, update documentation, tests, and version to 2.4.0.
1004 lines
37 KiB
Python
1004 lines
37 KiB
Python
import importlib
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
COMFY_ROOT = Path(__file__).resolve().parents[3]
|
|
CUSTOM_NODES_ROOT = COMFY_ROOT / "custom_nodes"
|
|
PACKAGE_NAME = Path(__file__).resolve().parents[1].name
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
sys.path.insert(0, str(CUSTOM_NODES_ROOT))
|
|
|
|
quant = importlib.import_module(f"{PACKAGE_NAME}.int8_quant")
|
|
loader = importlib.import_module(f"{PACKAGE_NAME}.int8_unet_loader")
|
|
model_adapter = importlib.import_module(f"{PACKAGE_NAME}.int8_model_adapter")
|
|
quantization_policy = importlib.import_module(f"{PACKAGE_NAME}.quantization_policy")
|
|
|
|
|
|
class QuantizationModeTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.quantization_state = {
|
|
"dynamic_quantize": quant.Int8TensorwiseOps.dynamic_quantize,
|
|
"_is_prequantized": quant.Int8TensorwiseOps._is_prequantized,
|
|
"quantization_mode": quant.Int8TensorwiseOps.quantization_mode,
|
|
"keep_float_names": list(quant.Int8TensorwiseOps.keep_float_names),
|
|
"int4_sensitive_names": list(quant.Int8TensorwiseOps.int4_sensitive_names),
|
|
"int4_mixed_selected_names": set(quant.Int8TensorwiseOps.int4_mixed_selected_names),
|
|
"int4_mixed_ratio": quant.Int8TensorwiseOps.int4_mixed_ratio,
|
|
}
|
|
quant.Int8TensorwiseOps.reset_otf_progress()
|
|
|
|
def tearDown(self):
|
|
quant.Int8TensorwiseOps.reset_otf_progress()
|
|
for name, value in self.quantization_state.items():
|
|
setattr(quant.Int8TensorwiseOps, name, value)
|
|
|
|
def test_int_mm_dispatch_preserves_aligned_and_unaligned_results(self):
|
|
for rows, inner, columns in ((17, 8, 8), (3, 7, 5)):
|
|
left = torch.randint(-8, 9, (rows, inner), dtype=torch.int8)
|
|
right = torch.randint(-8, 9, (inner, columns), dtype=torch.int8)
|
|
|
|
actual = quant._torch_int_mm_dispatch(left, right, output_columns=columns)
|
|
expected = left.to(torch.int32) @ right.to(torch.int32)
|
|
|
|
self.assertTrue(torch.equal(actual, expected))
|
|
|
|
def test_int8_convrot_uses_native_runtime_and_preserves_dynamic_lora(self):
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = torch.nn.Parameter(torch.randint(-8, 9, (16, 256), dtype=torch.int8), requires_grad=False)
|
|
module.weight_scale = torch.ones((16, 1), dtype=torch.float32)
|
|
module._is_quantized = True
|
|
module._quant_format = "int8_tensorwise"
|
|
module._is_per_row = True
|
|
module._outlier_method = quant.OUTLIER_METHOD_CONVROT
|
|
module._convrot_groupsize = 256
|
|
module.quarot_hadamard = quant._build_outlier_hadamard(
|
|
quant.OUTLIER_METHOD_CONVROT,
|
|
256,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
x = torch.randn(3, 256, dtype=torch.float32)
|
|
lora_a = torch.randn(4, 256, dtype=torch.float32)
|
|
lora_b = torch.randn(16, 4, dtype=torch.float32)
|
|
rotated_lora_a = quant._transform_weight_like_for_outlier_method(
|
|
lora_a,
|
|
quant.OUTLIER_METHOD_CONVROT,
|
|
torch.device("cpu"),
|
|
)
|
|
module.dynamic_lora_entries = [{"A": rotated_lora_a, "B": lora_b, "offset": None}]
|
|
base_output = torch.randn(3, 16, dtype=torch.bfloat16)
|
|
expected_base = base_output.clone()
|
|
|
|
with mock.patch.object(quant, "_native_int8_convrot_linear", return_value=base_output) as native_linear:
|
|
actual = module(x)
|
|
|
|
native_linear.assert_called_once()
|
|
expected = expected_base + torch.nn.functional.linear(
|
|
torch.nn.functional.linear(x, lora_a),
|
|
lora_b,
|
|
).to(base_output.dtype)
|
|
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.02)
|
|
|
|
def test_linear_cast_cleanup_runs_when_execution_raises(self):
|
|
for quant_format, kernel_owner, kernel_name in (
|
|
(None, quant.F, "linear"),
|
|
("convrot_w4a4", quant, "native_int4_linear"),
|
|
(quant.W4A8_FORMAT, quant, "native_w4a8_linear"),
|
|
):
|
|
with self.subTest(quant_format=quant_format or "float"):
|
|
module = quant.Int8TensorwiseOps.Linear(4, 3, bias=False, device="cpu", dtype=torch.float32)
|
|
module._is_quantized = quant_format is not None
|
|
module._quant_format = quant_format
|
|
offload_stream = object()
|
|
|
|
with mock.patch(
|
|
"comfy.ops.cast_bias_weight",
|
|
return_value=(module.weight, module.bias, offload_stream),
|
|
), mock.patch("comfy.ops.uncast_bias_weight") as uncast_bias_weight, mock.patch.object(
|
|
kernel_owner,
|
|
kernel_name,
|
|
side_effect=RuntimeError("execution failed"),
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "execution failed"):
|
|
module(torch.randn(2, 4, dtype=torch.float32))
|
|
|
|
uncast_bias_weight.assert_called_once_with(
|
|
module,
|
|
module.weight,
|
|
module.bias,
|
|
offload_stream,
|
|
)
|
|
|
|
def test_native_int8_convrot_metadata_round_trip_preserves_layer_group_size(self):
|
|
state_dict = {
|
|
"weight": torch.randint(-8, 9, (16, 2688), dtype=torch.int8),
|
|
"weight_scale": torch.ones((16, 1), dtype=torch.float32),
|
|
"comfy_quant": quant._encode_comfy_quant_config({
|
|
"format": "int8_tensorwise",
|
|
"convrot": True,
|
|
"convrot_groupsize": 64,
|
|
}),
|
|
}
|
|
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
|
|
module.load_state_dict(state_dict, strict=False)
|
|
|
|
self.assertEqual(module._convrot_groupsize, 64)
|
|
self.assertEqual(tuple(module.quarot_hadamard.shape), (64, 64))
|
|
self.assertEqual(module._outlier_method, quant.OUTLIER_METHOD_CONVROT)
|
|
with mock.patch.object(
|
|
quant,
|
|
"_native_int8_convrot_linear",
|
|
return_value=torch.zeros((2, 16), dtype=torch.bfloat16),
|
|
) as native_linear:
|
|
module(torch.randn(2, 2688, dtype=torch.float32))
|
|
native_linear.assert_called_once()
|
|
self.assertEqual(native_linear.call_args.args[-1], 64)
|
|
|
|
saved_state_dict = module.state_dict()
|
|
saved_config = quant._decode_comfy_quant_config(saved_state_dict["comfy_quant"])
|
|
self.assertEqual(saved_config["convrot_groupsize"], 64)
|
|
self.assertNotIn("quarot_hadamard", saved_state_dict)
|
|
|
|
reloaded = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
reloaded.load_state_dict(saved_state_dict, strict=False)
|
|
self.assertEqual(reloaded._convrot_groupsize, 64)
|
|
self.assertEqual(tuple(reloaded.quarot_hadamard.shape), (64, 64))
|
|
|
|
def test_native_int8_convrot_load_rejects_invalid_group_metadata(self):
|
|
for group_size, message in (
|
|
(0, "must be positive"),
|
|
(256, "2688 is not divisible"),
|
|
(32, "must be a power of 4"),
|
|
("invalid", "invalid ConvRot group size"),
|
|
):
|
|
with self.subTest(group_size=group_size):
|
|
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
state_dict = {
|
|
"weight": torch.randint(-8, 9, (16, 2688), dtype=torch.int8),
|
|
"weight_scale": torch.ones((16, 1), dtype=torch.float32),
|
|
"comfy_quant": quant._encode_comfy_quant_config({
|
|
"format": "int8_tensorwise",
|
|
"convrot": True,
|
|
"convrot_groupsize": group_size,
|
|
}),
|
|
}
|
|
|
|
with self.assertRaisesRegex(ValueError, message):
|
|
module.load_state_dict(state_dict, strict=False)
|
|
|
|
def test_native_int8_convrot_save_rejects_inconsistent_hadamard_shape(self):
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = nn.Parameter(torch.randint(-8, 9, (16, 256), dtype=torch.int8), requires_grad=False)
|
|
module.weight_scale = torch.ones((16, 1), dtype=torch.float32)
|
|
module._is_quantized = True
|
|
module._quant_format = "int8_tensorwise"
|
|
module._outlier_method = quant.OUTLIER_METHOD_CONVROT
|
|
module._convrot_groupsize = 256
|
|
module.quarot_hadamard = quant._build_outlier_hadamard(
|
|
quant.OUTLIER_METHOD_CONVROT,
|
|
64,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Hadamard shape must be"):
|
|
module.state_dict()
|
|
|
|
def test_native_int8_convrot_save_rejects_nondivisible_group_size(self):
|
|
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = nn.Parameter(torch.randint(-8, 9, (16, 2688), dtype=torch.int8), requires_grad=False)
|
|
module.weight_scale = torch.ones((16, 1), dtype=torch.float32)
|
|
module._is_quantized = True
|
|
module._quant_format = "int8_tensorwise"
|
|
module._outlier_method = quant.OUTLIER_METHOD_CONVROT
|
|
module._convrot_groupsize = 256
|
|
module.quarot_hadamard = quant._build_outlier_hadamard(
|
|
quant.OUTLIER_METHOD_CONVROT,
|
|
256,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "2688 is not divisible"):
|
|
module.state_dict()
|
|
|
|
def test_dynamic_lora_hook_supports_baked_native_int8_linear(self):
|
|
class DiffusionModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layer = nn.Linear(4, 3, bias=False)
|
|
self.layer.quant_format = "int8_tensorwise"
|
|
|
|
diffusion_model = DiffusionModel()
|
|
input_tensor = torch.randn(2, 4)
|
|
base_output = diffusion_model.layer(input_tensor).detach().clone()
|
|
lora_up = torch.randn(3, 2)
|
|
lora_down = torch.randn(2, 4)
|
|
adapter = quant.LoRAAdapter(
|
|
[],
|
|
(lora_up, lora_down, None, None, None, None),
|
|
)
|
|
dynamic_loras = [{
|
|
"name": "native.safetensors",
|
|
"strength": 0.75,
|
|
"patches": {"diffusion_model.layer.weight": adapter},
|
|
"patch_uuid": "native-test",
|
|
}]
|
|
|
|
quant.DynamicLoRAHook().apply_composition(diffusion_model, dynamic_loras)
|
|
actual = diffusion_model.layer(input_tensor)
|
|
expected_delta = torch.nn.functional.linear(
|
|
torch.nn.functional.linear(input_tensor, lora_down * 0.75),
|
|
lora_up,
|
|
)
|
|
|
|
self.assertTrue(hasattr(diffusion_model.layer, quant._NATIVE_DYNAMIC_LORA_HOOK_ATTRIBUTE))
|
|
torch.testing.assert_close(actual, base_output + expected_delta)
|
|
|
|
def test_native_dynamic_lora_combines_compatible_factors_without_changing_delta(self):
|
|
class DiffusionModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layer = nn.Linear(4, 3, bias=False)
|
|
self.layer.quant_format = "int8_tensorwise"
|
|
|
|
diffusion_model = DiffusionModel()
|
|
input_tensor = torch.randn(2, 4)
|
|
base_output = diffusion_model.layer(input_tensor).detach().clone()
|
|
adapters = []
|
|
expected_delta = torch.zeros_like(base_output)
|
|
for index, (rank, strength, alpha) in enumerate(((2, 0.75, 1.0), (3, -0.4, 6.0))):
|
|
lora_up = torch.randn(3, rank)
|
|
lora_down = torch.randn(rank, 4)
|
|
adapters.append({
|
|
"name": f"native-{index}.safetensors",
|
|
"strength": strength,
|
|
"patches": {
|
|
"diffusion_model.layer.weight": quant.LoRAAdapter(
|
|
[],
|
|
(lora_up, lora_down, alpha, None, None, None),
|
|
),
|
|
},
|
|
"patch_uuid": f"native-combined-{index}",
|
|
})
|
|
expected_delta += torch.nn.functional.linear(
|
|
torch.nn.functional.linear(input_tensor, lora_down * ((alpha / rank) * strength)),
|
|
lora_up,
|
|
)
|
|
|
|
quant.DynamicLoRAHook().apply_composition(diffusion_model, adapters)
|
|
actual = diffusion_model.layer(input_tensor)
|
|
|
|
self.assertEqual(len(diffusion_model.layer.dynamic_lora_entries), 1)
|
|
torch.testing.assert_close(actual, base_output + expected_delta, rtol=1e-5, atol=1e-5)
|
|
|
|
def test_dynamic_lora_hook_matches_linear_inside_compiled_module_wrapper(self):
|
|
class CompiledModule(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._orig_mod = nn.Module()
|
|
self._orig_mod.layer = nn.Linear(4, 3, bias=False)
|
|
self._orig_mod.layer.quant_format = "int8_tensorwise"
|
|
|
|
class DiffusionModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.blocks = nn.ModuleList([CompiledModule()])
|
|
|
|
diffusion_model = DiffusionModel()
|
|
lora_up = torch.randn(3, 2)
|
|
lora_down = torch.randn(2, 4)
|
|
adapter = quant.LoRAAdapter(
|
|
[],
|
|
(lora_up, lora_down, None, None, None, None),
|
|
)
|
|
dynamic_loras = [{
|
|
"name": "compiled.safetensors",
|
|
"strength": 1.0,
|
|
"patches": {"diffusion_model.blocks.0.layer.weight": adapter},
|
|
"patch_uuid": "compiled-test",
|
|
}]
|
|
|
|
quant.DynamicLoRAHook().apply_composition(diffusion_model, dynamic_loras)
|
|
|
|
self.assertEqual(
|
|
len(diffusion_model.blocks[0]._orig_mod.layer.dynamic_lora_entries),
|
|
1,
|
|
)
|
|
|
|
def test_dynamic_lora_hook_unwraps_hidden_lazy_compile_source(self):
|
|
class SourceBlock(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layer = nn.Linear(4, 3, bias=False)
|
|
self.layer.quant_format = "int8_tensorwise"
|
|
|
|
class HiddenCompiledModule(nn.Module):
|
|
def __init__(self, source_module):
|
|
super().__init__()
|
|
object.__setattr__(self, "_source_module", source_module)
|
|
|
|
class DiffusionModel(nn.Module):
|
|
def __init__(self, source_block):
|
|
super().__init__()
|
|
self.blocks = nn.ModuleList([HiddenCompiledModule(source_block)])
|
|
|
|
source_block = SourceBlock()
|
|
diffusion_model = DiffusionModel(source_block)
|
|
lora_up = torch.randn(3, 2)
|
|
lora_down = torch.randn(2, 4)
|
|
adapter = quant.LoRAAdapter(
|
|
[],
|
|
(lora_up, lora_down, None, None, None, None),
|
|
)
|
|
dynamic_loras = [{
|
|
"name": "hidden-compiled.safetensors",
|
|
"strength": 1.0,
|
|
"patches": {"diffusion_model.blocks.0.layer.weight": adapter},
|
|
"patch_uuid": "hidden-compiled-test",
|
|
}]
|
|
|
|
quant.DynamicLoRAHook().apply_composition(diffusion_model, dynamic_loras)
|
|
|
|
self.assertEqual(len(source_block.layer.dynamic_lora_entries), 1)
|
|
|
|
def test_legacy_modes_are_not_exposed_or_normalized(self):
|
|
for mode in ("none", "convrot", "quarot", "hadanorm"):
|
|
self.assertNotIn(mode, quant.QUANTIZATION_MODE_CHOICES)
|
|
self.assertEqual(quant.normalize_quantization_mode(mode), "int8")
|
|
|
|
def test_grouped_int8_modes_are_exposed_as_distinct_low_bit_modes(self):
|
|
for quantization_mode in (
|
|
quant.QUANTIZATION_MODE_W4A8,
|
|
quant.QUANTIZATION_MODE_W4A8_MIXED,
|
|
quant.QUANTIZATION_MODE_W6A8,
|
|
):
|
|
with self.subTest(quantization_mode=quantization_mode):
|
|
self.assertIn(quantization_mode, quant.QUANTIZATION_MODE_CHOICES)
|
|
self.assertTrue(quant.quantization_mode_is_int4(quantization_mode))
|
|
self.assertEqual(
|
|
quant.quantization_mode_outlier_method(quantization_mode),
|
|
quant.OUTLIER_METHOD_NONE,
|
|
)
|
|
|
|
self.assertEqual(quant.quantization_mode_grouped_bits(quant.QUANTIZATION_MODE_W4A8), 4)
|
|
self.assertEqual(quant.quantization_mode_grouped_bits(quant.QUANTIZATION_MODE_W4A8_MIXED), 4)
|
|
self.assertEqual(quant.quantization_mode_grouped_bits(quant.QUANTIZATION_MODE_W6A8), 6)
|
|
self.assertEqual(quant.quantization_mode_grouped_format(quant.QUANTIZATION_MODE_W6A8), quant.W6A8_FORMAT)
|
|
self.assertTrue(quant.quantization_mode_uses_mixed_int8(quant.QUANTIZATION_MODE_W4A8_MIXED))
|
|
|
|
def test_model_adapter_rejects_w4a8_without_native_layout(self):
|
|
with mock.patch.object(model_adapter, "native_w4a8_available", return_value=False):
|
|
with self.assertRaisesRegex(RuntimeError, "ComfyUI 0.32.0 or newer"):
|
|
model_adapter.INT8ModelAdapter().apply_quantization(
|
|
object(),
|
|
enable_quantization=model_adapter.QUANTIZATION_CONTROL_ALWAYS,
|
|
model_type=model_adapter.AUTO_MODEL_TYPE,
|
|
quantization_mode=quant.QUANTIZATION_MODE_W4A8,
|
|
log_progress=False,
|
|
)
|
|
|
|
@unittest.skipUnless(quant._NATIVE_W4A8_AVAILABLE, "Native ComfyUI W4A8 layout is unavailable")
|
|
def test_native_w4a8_state_dict_round_trip_and_requantization(self):
|
|
q_weight = quant.quantize_native_w4a8(torch.randn(16, 256, dtype=torch.float32))
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = torch.nn.Parameter(q_weight, requires_grad=False)
|
|
module.weight_scale = None
|
|
module._is_quantized = True
|
|
module._quant_format = quant.W4A8_FORMAT
|
|
module._convrot_groupsize = q_weight._params.convrot_groupsize
|
|
module._w4a8_group_size = q_weight._params.group_size
|
|
|
|
state_dict = module.state_dict()
|
|
self.assertIn("weight_s_rel", state_dict)
|
|
self.assertIn("weight_s_channel", state_dict)
|
|
self.assertEqual(
|
|
quant._decode_comfy_quant_config(state_dict["comfy_quant"])["format"],
|
|
quant.W4A8_FORMAT,
|
|
)
|
|
|
|
loaded = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
loaded.load_state_dict(state_dict, strict=False)
|
|
self.assertEqual(loaded._quant_format, quant.W4A8_FORMAT)
|
|
self.assertEqual(loaded.weight._layout_cls, "AsymW4A8Int8Layout")
|
|
|
|
requantized = loaded.set_weight(
|
|
loaded.convert_weight(loaded.weight) + 0.01,
|
|
seed=17,
|
|
return_weight=True,
|
|
)
|
|
self.assertEqual(requantized._layout_cls, "AsymW4A8Int8Layout")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_W6A8_AVAILABLE, "Native ComfyUI W6A8 support is unavailable")
|
|
def test_native_w6a8_state_dict_round_trip_and_requantization(self):
|
|
q_weight = quant.quantize_native_w6a8(torch.randn(8, 256, dtype=torch.float32))
|
|
module = quant.Int8TensorwiseOps.Linear(256, 8, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = torch.nn.Parameter(q_weight, requires_grad=False)
|
|
module.weight_scale = None
|
|
module._is_quantized = True
|
|
module._quant_format = quant.W6A8_FORMAT
|
|
module._convrot_groupsize = q_weight._params.convrot_groupsize
|
|
module._w4a8_group_size = q_weight._params.group_size
|
|
module._grouped_int8_bits = 6
|
|
|
|
state_dict = module.state_dict()
|
|
self.assertEqual(tuple(state_dict["weight"].shape), (8, 192))
|
|
self.assertEqual(state_dict["weight_s_rel"].shape, (8, 8))
|
|
self.assertNotIn("weight_codebook", state_dict)
|
|
quant_config = quant._decode_comfy_quant_config(state_dict["comfy_quant"])
|
|
self.assertEqual(quant_config["format"], quant.W6A8_FORMAT)
|
|
self.assertEqual(quant_config["group_size"], 32)
|
|
|
|
loaded = quant.Int8TensorwiseOps.Linear(256, 8, bias=False, device="cpu", dtype=torch.float32)
|
|
loaded.load_state_dict(state_dict, strict=False)
|
|
self.assertEqual(loaded._quant_format, quant.W6A8_FORMAT)
|
|
self.assertEqual(loaded._grouped_int8_bits, 6)
|
|
self.assertEqual(loaded.weight._layout_cls, "AsymW4A8Int8Layout")
|
|
|
|
requantized = loaded.set_weight(
|
|
loaded.convert_weight(loaded.weight) + 0.01,
|
|
seed=17,
|
|
return_weight=True,
|
|
)
|
|
self.assertEqual(requantized._layout_cls, "AsymW4A8Int8Layout")
|
|
self.assertEqual(tuple(requantized._qdata.shape), (8, 192))
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_native_int4_state_dict_round_trip(self):
|
|
weight = torch.randn(16, 256, dtype=torch.float32)
|
|
q_weight = quant.quantize_native_int4(weight)
|
|
|
|
module = quant.Int8TensorwiseOps.Linear(
|
|
256,
|
|
16,
|
|
bias=True,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
)
|
|
module.weight = torch.nn.Parameter(q_weight, requires_grad=False)
|
|
module.weight_scale = q_weight._params.scale
|
|
module._is_quantized = True
|
|
module._quant_format = "convrot_w4a4"
|
|
module._convrot_groupsize = 256
|
|
module._quant_group_size = 64
|
|
module._linear_dtype = "int4"
|
|
|
|
state_dict = module.state_dict()
|
|
self.assertEqual(tuple(state_dict["weight"].shape), (16, 128))
|
|
|
|
loaded = quant.Int8TensorwiseOps.Linear(
|
|
256,
|
|
16,
|
|
bias=True,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
)
|
|
loaded.load_state_dict(state_dict, strict=False)
|
|
|
|
self.assertEqual(loaded._quant_format, "convrot_w4a4")
|
|
self.assertEqual(loaded.weight._layout_cls, "TensorCoreConvRotW4A4Layout")
|
|
output = loaded(torch.randn(3, 256, dtype=torch.float32))
|
|
self.assertEqual(tuple(output.shape), (3, 16))
|
|
self.assertTrue(torch.isfinite(output).all())
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_native_int4_int8_compute_metadata_round_trip(self):
|
|
q_weight = quant.quantize_native_int4(
|
|
torch.randn(16, 256, dtype=torch.float32),
|
|
linear_dtype="int8",
|
|
)
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = torch.nn.Parameter(q_weight, requires_grad=False)
|
|
module.weight_scale = q_weight._params.scale
|
|
module._is_quantized = True
|
|
module._quant_format = "convrot_w4a4"
|
|
|
|
loaded = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
loaded.load_state_dict(module.state_dict(), strict=False)
|
|
self.assertEqual(loaded._linear_dtype, "int8")
|
|
self.assertEqual(loaded.weight._params.linear_dtype, "int8")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_native_int4_dynamic_lora_preserves_runtime_delta(self):
|
|
q_weight = quant.quantize_native_int4(torch.randn(16, 256, dtype=torch.float32))
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module.weight = torch.nn.Parameter(q_weight, requires_grad=False)
|
|
module.weight_scale = q_weight._params.scale
|
|
module._is_quantized = True
|
|
module._quant_format = "convrot_w4a4"
|
|
|
|
x = torch.randn(3, 256, dtype=torch.float32)
|
|
base_output = module(x)
|
|
lora_a = torch.randn(4, 256, dtype=torch.float32)
|
|
lora_b = torch.randn(16, 4, dtype=torch.float32)
|
|
module.dynamic_lora_entries = [{"A": lora_a, "B": lora_b, "offset": None}]
|
|
|
|
actual = module(x)
|
|
expected = base_output + torch.nn.functional.linear(torch.nn.functional.linear(x, lora_a), lora_b)
|
|
torch.testing.assert_close(actual, expected)
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_int4_mixed_keeps_sensitive_layer_in_int8(self):
|
|
quant.Int8TensorwiseOps.dynamic_quantize = True
|
|
quant.Int8TensorwiseOps._is_prequantized = False
|
|
quant.Int8TensorwiseOps.quantization_mode = quant.QUANTIZATION_MODE_INT4_MIXED
|
|
quant.Int8TensorwiseOps.keep_float_names = []
|
|
quant.Int8TensorwiseOps.int4_sensitive_names = ["sensitive"]
|
|
quant.Int8TensorwiseOps._init_otf_progress({
|
|
"bulk.weight": torch.randn(16, 256),
|
|
"sensitive.weight": torch.randn(16, 256),
|
|
})
|
|
|
|
bulk = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
bulk._load_from_state_dict(
|
|
{"bulk.weight": torch.randn(16, 256)},
|
|
"bulk.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
sensitive = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
sensitive._load_from_state_dict(
|
|
{"sensitive.weight": torch.randn(16, 256)},
|
|
"sensitive.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
self.assertEqual(bulk._quant_format, "convrot_w4a4")
|
|
self.assertEqual(sensitive._quant_format, "int8_tensorwise")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_W4A8_AVAILABLE, "Native ComfyUI W4A8 layout is unavailable")
|
|
def test_w4a8_mixed_keeps_sensitive_layer_in_int8(self):
|
|
quant.Int8TensorwiseOps.dynamic_quantize = True
|
|
quant.Int8TensorwiseOps._is_prequantized = False
|
|
quant.Int8TensorwiseOps.quantization_mode = quant.QUANTIZATION_MODE_W4A8_MIXED
|
|
quant.Int8TensorwiseOps.keep_float_names = []
|
|
quant.Int8TensorwiseOps.int4_sensitive_names = ["sensitive"]
|
|
quant.Int8TensorwiseOps.int4_mixed_ratio = 0.5
|
|
state_dict = {
|
|
"bulk.weight": torch.randn(16, 256),
|
|
"sensitive.weight": torch.randn(16, 256),
|
|
}
|
|
quant.Int8TensorwiseOps._init_otf_progress(state_dict)
|
|
|
|
bulk = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
bulk._load_from_state_dict(
|
|
{"bulk.weight": state_dict["bulk.weight"]},
|
|
"bulk.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
sensitive = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
sensitive._load_from_state_dict(
|
|
{"sensitive.weight": state_dict["sensitive.weight"]},
|
|
"sensitive.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
self.assertEqual(bulk._quant_format, quant.W4A8_FORMAT)
|
|
self.assertEqual(sensitive._quant_format, "int8_tensorwise")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_W6A8_AVAILABLE, "Native ComfyUI W6A8 support is unavailable")
|
|
def test_w6a8_on_the_fly_uses_native_layout(self):
|
|
quant.Int8TensorwiseOps.dynamic_quantize = True
|
|
quant.Int8TensorwiseOps._is_prequantized = False
|
|
quant.Int8TensorwiseOps.quantization_mode = quant.QUANTIZATION_MODE_W6A8
|
|
quant.Int8TensorwiseOps.keep_float_names = []
|
|
|
|
module = quant.Int8TensorwiseOps.Linear(256, 8, bias=False, device="cpu", dtype=torch.float32)
|
|
module._load_from_state_dict(
|
|
{"bulk.weight": torch.randn(8, 256)},
|
|
"bulk.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
self.assertEqual(module._quant_format, quant.W6A8_FORMAT)
|
|
self.assertEqual(module._grouped_int8_bits, 6)
|
|
self.assertEqual(module.weight._layout_cls, "AsymW4A8Int8Layout")
|
|
|
|
def test_model_adapter_quantizes_w6a8_and_w4a8_mixed(self):
|
|
def quantize(mode, int4_sensitive=False):
|
|
module = nn.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
with mock.patch.object(model_adapter, "_get_int8_compute_device", return_value=torch.device("cpu")):
|
|
return model_adapter._quantize_linear_module(
|
|
"blocks.0.proj",
|
|
module,
|
|
module.weight.detach(),
|
|
mode,
|
|
int4_sensitive=int4_sensitive,
|
|
)[0]
|
|
|
|
if quant._NATIVE_W6A8_AVAILABLE:
|
|
w6_module = quantize(quant.QUANTIZATION_MODE_W6A8)
|
|
self.assertEqual(w6_module._quant_format, quant.W6A8_FORMAT)
|
|
self.assertEqual(w6_module._grouped_int8_bits, 6)
|
|
self.assertEqual(w6_module._w4a8_group_size, 32)
|
|
|
|
mixed_w4a8_module = quantize(quant.QUANTIZATION_MODE_W4A8_MIXED)
|
|
self.assertEqual(mixed_w4a8_module._quant_format, quant.W4A8_FORMAT)
|
|
|
|
mixed_int8_module = quantize(quant.QUANTIZATION_MODE_W4A8_MIXED, int4_sensitive=True)
|
|
self.assertEqual(mixed_int8_module._quant_format, "int8_tensorwise")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_int4_full_promotes_sensitive_layer_to_int4(self):
|
|
quant.Int8TensorwiseOps.dynamic_quantize = True
|
|
quant.Int8TensorwiseOps._is_prequantized = False
|
|
quant.Int8TensorwiseOps.quantization_mode = quant.QUANTIZATION_MODE_INT4_FULL
|
|
quant.Int8TensorwiseOps.keep_float_names = []
|
|
quant.Int8TensorwiseOps.int4_sensitive_names = ["sensitive"]
|
|
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module._load_from_state_dict(
|
|
{"sensitive.weight": torch.randn(16, 256)},
|
|
"sensitive.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
self.assertEqual(module._quant_format, "convrot_w4a4")
|
|
|
|
@unittest.skipUnless(quant._NATIVE_INT4_AVAILABLE, "Native ComfyUI INT4 layout is unavailable")
|
|
def test_int4_modes_preserve_keep_float_layers(self):
|
|
quant.Int8TensorwiseOps.dynamic_quantize = True
|
|
quant.Int8TensorwiseOps.keep_float_names = ["protected"]
|
|
quant.Int8TensorwiseOps.int4_sensitive_names = []
|
|
|
|
for quantization_mode in (
|
|
quant.QUANTIZATION_MODE_INT4_MIXED,
|
|
quant.QUANTIZATION_MODE_INT4_FULL,
|
|
):
|
|
with self.subTest(quantization_mode=quantization_mode):
|
|
quant.Int8TensorwiseOps._is_prequantized = False
|
|
quant.Int8TensorwiseOps.quantization_mode = quantization_mode
|
|
weight = torch.randn(16, 256)
|
|
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
|
|
module._load_from_state_dict(
|
|
{"protected.weight": weight.clone()},
|
|
"protected.",
|
|
{},
|
|
False,
|
|
[],
|
|
[],
|
|
[],
|
|
)
|
|
|
|
self.assertFalse(module._is_quantized)
|
|
self.assertIsNone(module._quant_format)
|
|
torch.testing.assert_close(module.weight, weight)
|
|
|
|
def test_architecture_presets_protect_residual_writeback_paths(self):
|
|
krea_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_KREA2)
|
|
anima_preset = loader.get_model_type_quantization_preset("anima")
|
|
minimax_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_MINIMAX_H3)
|
|
moss_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_MOSS)
|
|
|
|
self.assertIn("first", krea_preset["keep_float"])
|
|
self.assertEqual(krea_preset["int4_sensitive"], ("attn.wo", "mlp.down"))
|
|
self.assertIn("final_layer", anima_preset["keep_float"])
|
|
self.assertEqual(
|
|
anima_preset["int4_sensitive"],
|
|
("self_attn.output_proj", "cross_attn.output_proj", "mlp.layer2"),
|
|
)
|
|
self.assertEqual(
|
|
minimax_preset["keep_float"],
|
|
(
|
|
"video_patch_proj", "audio_patch_proj", "condition_proj", "time_embedder",
|
|
"token_refiner", "adaln_proj", "final_layer",
|
|
),
|
|
)
|
|
self.assertEqual(minimax_preset["int4_sensitive"], ("attn.out_proj", "mlp.fc2"))
|
|
self.assertEqual(
|
|
moss_preset["keep_float"],
|
|
(
|
|
"video_embedder", "text_embedder", "time_embedding", "shared_block_modulation",
|
|
"video_head_modulation", "final_linear_video", "final_norm_video",
|
|
),
|
|
)
|
|
self.assertEqual(moss_preset["int4_sensitive"], ("proj.o", "proj.ffn_down"))
|
|
|
|
def test_model_adapter_auto_detects_minimax_h3(self):
|
|
class MiniMaxH3(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.video_patch_proj = nn.Linear(4, 4)
|
|
self.audio_patch_proj = nn.Linear(4, 4)
|
|
self.condition_proj = nn.Linear(4, 4)
|
|
self.token_refiner = nn.Sequential(nn.Linear(4, 4))
|
|
self.blocks = nn.ModuleList([nn.Module()])
|
|
self.blocks[0].adaln_proj = nn.Sequential(nn.Linear(4, 4))
|
|
|
|
self.assertEqual(
|
|
model_adapter._infer_model_type_from_modules(MiniMaxH3()),
|
|
loader.MODEL_TYPE_MINIMAX_H3,
|
|
)
|
|
|
|
def test_model_adapter_auto_detects_moss_image(self):
|
|
class MossImage(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.video_embedder = nn.Linear(4, 4)
|
|
self.text_embedder = nn.Module()
|
|
self.text_embedder.blocks = nn.ModuleList([nn.Module()])
|
|
self.time_embedding = nn.Sequential(nn.Linear(4, 4))
|
|
self.shared_block_modulation = nn.Linear(4, 4)
|
|
self.video_head_modulation = nn.Linear(4, 4)
|
|
self.final_linear_video = nn.Linear(4, 4)
|
|
self.layers = nn.ModuleList([nn.Module()])
|
|
self.layers[0].proj = nn.ModuleDict({"ffn_up_gate": nn.Linear(4, 4)})
|
|
|
|
self.assertEqual(
|
|
model_adapter._infer_model_type_from_modules(MossImage()),
|
|
loader.MODEL_TYPE_MOSS,
|
|
)
|
|
|
|
def test_moss_preset_matches_the_official_quantized_layer_boundary(self):
|
|
class MossImage(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.video_embedder = nn.Linear(256, 16, bias=False)
|
|
self.text_embedder = nn.Sequential(nn.Linear(256, 16, bias=False))
|
|
self.time_embedding = nn.Sequential(nn.Linear(256, 16, bias=False))
|
|
self.shared_block_modulation = nn.Linear(256, 16, bias=False)
|
|
self.video_head_modulation = nn.Linear(256, 16, bias=False)
|
|
self.final_linear_video = nn.Linear(256, 16, bias=False)
|
|
self.layers = nn.ModuleList([nn.Module()])
|
|
self.layers[0].proj = nn.ModuleDict({
|
|
"qkv": nn.Linear(256, 16, bias=False),
|
|
"o": nn.Linear(256, 16, bias=False),
|
|
"ffn_up_gate": nn.Linear(256, 16, bias=False),
|
|
"ffn_down": nn.Linear(256, 16, bias=False),
|
|
})
|
|
|
|
preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_MOSS)
|
|
candidate_names = {
|
|
name
|
|
for name, _module in model_adapter._collect_int8_candidates(
|
|
MossImage(),
|
|
preset["keep_float"],
|
|
)
|
|
}
|
|
|
|
self.assertEqual(
|
|
candidate_names,
|
|
{
|
|
"layers.0.proj.qkv",
|
|
"layers.0.proj.o",
|
|
"layers.0.proj.ffn_up_gate",
|
|
"layers.0.proj.ffn_down",
|
|
},
|
|
)
|
|
|
|
def test_architecture_ratios_cover_residual_writeback_paths(self):
|
|
krea_names = [
|
|
f"blocks.{block}.{path}"
|
|
for block in range(28)
|
|
for path in (
|
|
"attn.wq", "attn.wk", "attn.wv", "attn.gate", "attn.wo",
|
|
"mlp.gate", "mlp.up", "mlp.down",
|
|
)
|
|
]
|
|
anima_names = [
|
|
f"blocks.{block}.{path}"
|
|
for block in range(28)
|
|
for path in (
|
|
"self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.output_proj",
|
|
"cross_attn.q_proj", "cross_attn.k_proj", "cross_attn.v_proj", "cross_attn.output_proj",
|
|
"mlp.layer1", "mlp.layer2",
|
|
)
|
|
]
|
|
krea_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_KREA2)
|
|
anima_preset = loader.get_model_type_quantization_preset("anima")
|
|
|
|
krea_selected = quantization_policy.select_int4_mixed_layers(
|
|
krea_names,
|
|
0.25,
|
|
krea_preset["int4_sensitive"],
|
|
)
|
|
anima_selected = quantization_policy.select_int4_mixed_layers(
|
|
anima_names,
|
|
0.3,
|
|
anima_preset["int4_sensitive"],
|
|
)
|
|
|
|
self.assertEqual(len(krea_selected), 56)
|
|
self.assertTrue(all(name.endswith(("attn.wo", "mlp.down")) for name in krea_selected))
|
|
self.assertEqual(len(anima_selected), 84)
|
|
self.assertTrue(all(name.endswith(("output_proj", "mlp.layer2")) for name in anima_selected))
|
|
|
|
def test_mixed_ratio_selects_a_stable_model_agnostic_budget(self):
|
|
layer_names = [f"blocks.{index}.linear" for index in range(10)]
|
|
|
|
first = quantization_policy.select_int4_mixed_layers(
|
|
layer_names,
|
|
0.2,
|
|
priority_patterns=("blocks.9.",),
|
|
)
|
|
second = quantization_policy.select_int4_mixed_layers(
|
|
layer_names,
|
|
0.2,
|
|
priority_patterns=("blocks.9.",),
|
|
)
|
|
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(len(first), 2)
|
|
self.assertIn("blocks.9.linear", first)
|
|
|
|
def test_mixed_ratio_boundaries_select_none_or_all(self):
|
|
layer_names = [f"layer.{index}" for index in range(5)]
|
|
|
|
self.assertEqual(
|
|
quantization_policy.select_int4_mixed_layers(layer_names, 0.0),
|
|
(),
|
|
)
|
|
self.assertEqual(
|
|
quantization_policy.select_int4_mixed_layers(layer_names, 1.0),
|
|
tuple(layer_names),
|
|
)
|
|
|
|
def test_mixed_ratio_profiles_are_nested_and_preserve_the_default_selection(self):
|
|
layer_names = [f"layer.{index}" for index in range(100)]
|
|
lower_profile = set(quantization_policy.select_int4_mixed_layers(layer_names, 0.15))
|
|
default_profile = set(quantization_policy.select_int4_mixed_layers(layer_names, 0.2))
|
|
higher_profile = set(quantization_policy.select_int4_mixed_layers(layer_names, 0.3))
|
|
|
|
self.assertLess(lower_profile, default_profile)
|
|
self.assertLess(default_profile, higher_profile)
|
|
self.assertEqual(
|
|
set(quantization_policy.select_int4_mixed_layers(layer_names[:10], 0.2)),
|
|
{"layer.2", "layer.7"},
|
|
)
|
|
|
|
def test_otf_mixed_profile_selects_the_requested_fraction(self):
|
|
quant.Int8TensorwiseOps.quantization_mode = quant.QUANTIZATION_MODE_INT4_MIXED
|
|
quant.Int8TensorwiseOps.keep_float_names = []
|
|
quant.Int8TensorwiseOps.int4_sensitive_names = ["layer.4"]
|
|
quant.Int8TensorwiseOps.int4_mixed_ratio = 0.4
|
|
state_dict = {
|
|
f"layer.{index}.weight": torch.randn(16, 256)
|
|
for index in range(5)
|
|
}
|
|
|
|
quant.Int8TensorwiseOps._init_otf_progress(state_dict)
|
|
|
|
self.assertEqual(len(quant.Int8TensorwiseOps.int4_mixed_selected_names), 2)
|
|
self.assertIn("layer.4", quant.Int8TensorwiseOps.int4_mixed_selected_names)
|
|
|
|
def test_loader_preselects_mixed_profile_from_safetensors_shapes(self):
|
|
from safetensors.torch import save_file
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path = Path(temp_dir) / "mixed.safetensors"
|
|
save_file(
|
|
{
|
|
"model.diffusion_model.layers.0.proj.qkv.weight": torch.empty(16, 256),
|
|
"model.diffusion_model.layers.0.proj.o.weight": torch.empty(16, 256),
|
|
"model.diffusion_model.text_embedder.blocks.0.qkv.weight": torch.empty(16, 256),
|
|
"model.diffusion_model.layers.0.norm.weight": torch.empty(16),
|
|
},
|
|
str(path),
|
|
)
|
|
|
|
selected = loader._select_int4_mixed_layers_from_safetensors(
|
|
str(path),
|
|
keep_float_names=("text_embedder",),
|
|
priority_patterns=("proj.qkv",),
|
|
ratio=0.5,
|
|
)
|
|
|
|
self.assertEqual(selected, {"layers.0.proj.qkv"})
|
|
|
|
def test_public_quantization_inputs_do_not_use_legacy_int8_names(self):
|
|
adapter_inputs = model_adapter.INT8ModelAdapter.INPUT_TYPES()["required"]
|
|
loader_inputs = loader.UNetLoaderINTW8A8.INPUT_TYPES()["required"]
|
|
|
|
self.assertIn("enable_quantization", adapter_inputs)
|
|
self.assertNotIn("enable_int8", adapter_inputs)
|
|
self.assertIn("prepack_weights", adapter_inputs)
|
|
self.assertNotIn("prepack_int8_weights", adapter_inputs)
|
|
self.assertIn("prepack_weights", loader_inputs)
|
|
self.assertNotIn("prepack_int8_weights", loader_inputs)
|
|
self.assertEqual(
|
|
adapter_inputs["enable_quantization"],
|
|
(
|
|
model_adapter.QUANTIZATION_CONTROL_CHOICES,
|
|
{
|
|
"default": model_adapter.QUANTIZATION_CONTROL_AS_NEEDED,
|
|
"tooltip": "as_needed converts FP8 and floating-point inputs, but leaves MODEL inputs containing Toolkit-supported INT8, W4A4, W4A8, or W6A8 layers unchanged. always converts remaining eligible layers; bypass returns the MODEL unchanged.",
|
|
},
|
|
),
|
|
)
|
|
self.assertEqual(adapter_inputs["int4_mixed_ratio"][1]["default"], 0.2)
|
|
self.assertEqual(loader_inputs["int4_mixed_ratio"][1]["default"], 0.2)
|
|
self.assertEqual(model_adapter.INT8ModelAdapter.FUNCTION, "apply_quantization")
|
|
|
|
def test_as_needed_returns_a_supported_model_unchanged(self):
|
|
shared_model = nn.Module()
|
|
shared_model.diffusion_model = nn.Sequential(nn.Linear(4, 4))
|
|
shared_model.diffusion_model[0]._is_quantized = True
|
|
model = SimpleNamespace(model=shared_model, object_patches={}, object_patches_backup={})
|
|
|
|
output = model_adapter.INT8ModelAdapter().apply_quantization(
|
|
model,
|
|
enable_quantization=model_adapter.QUANTIZATION_CONTROL_AS_NEEDED,
|
|
model_type=model_adapter.AUTO_MODEL_TYPE,
|
|
quantization_mode=quant.QUANTIZATION_MODE_INT8,
|
|
log_progress=False,
|
|
)[0]
|
|
|
|
self.assertIs(output, model)
|
|
|
|
def test_as_needed_detects_supported_object_patches(self):
|
|
quantized_module = nn.Linear(4, 4)
|
|
quantized_module._is_quantized = True
|
|
model = SimpleNamespace(
|
|
model=SimpleNamespace(diffusion_model=nn.Module()),
|
|
object_patches={"diffusion_model.block": quantized_module},
|
|
object_patches_backup={},
|
|
)
|
|
|
|
self.assertTrue(model_adapter._model_has_quantized_modules(model))
|
|
|
|
def test_as_needed_recognizes_native_int8_but_not_fp8(self):
|
|
native_int8_module = nn.Linear(4, 4)
|
|
native_int8_module.quant_format = "int8_tensorwise"
|
|
native_w4a8_module = nn.Linear(4, 4)
|
|
native_w4a8_module.quant_format = "asym_w4a8_int8"
|
|
native_w6a8_module = nn.Linear(4, 4)
|
|
native_w6a8_module.quant_format = "w6a8_int8"
|
|
native_fp8_module = nn.Linear(4, 4)
|
|
native_fp8_module.quant_format = "float8_e4m3fn"
|
|
|
|
self.assertTrue(model_adapter._module_has_target_quantization(native_int8_module))
|
|
self.assertTrue(model_adapter._module_has_target_quantization(native_w4a8_module))
|
|
self.assertTrue(model_adapter._module_has_target_quantization(native_w6a8_module))
|
|
self.assertFalse(model_adapter._module_has_target_quantization(native_fp8_module))
|
|
|
|
def test_bypass_returns_model_without_inspection(self):
|
|
model = object()
|
|
|
|
output = model_adapter.INT8ModelAdapter().apply_quantization(
|
|
model,
|
|
enable_quantization=model_adapter.QUANTIZATION_CONTROL_BYPASS,
|
|
model_type=model_adapter.AUTO_MODEL_TYPE,
|
|
quantization_mode=quant.QUANTIZATION_MODE_INT8,
|
|
)[0]
|
|
|
|
self.assertIs(output, model)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|