chore: import upstream snapshot with attribution

tests/backend/quantization/gguf/test_ggml_tensor.py

@@ -0,0 +1,128 @@
+import gguf
+import pytest
+import torch
+
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+def quantize_tensor(data: torch.Tensor, ggml_quantization_type: gguf.GGMLQuantizationType) -> GGMLTensor:
+    """Quantize a torch.Tensor to a GGMLTensor.
+
+    Uses the gguf library's numpy implementation to quantize the tensor.
+    """
+    data_np = data.detach().cpu().numpy()
+    quantized_np = gguf.quantize(data_np, ggml_quantization_type)
+    return GGMLTensor(
+        data=torch.from_numpy(quantized_np),
+        ggml_quantization_type=ggml_quantization_type,
+        tensor_shape=data.shape,
+        compute_dtype=data.dtype,
+    ).to(device=data.device)  # type: ignore
+
+
+@pytest.mark.parametrize(
+    ["device", "x1_quant_type", "x2_quant_type"],
+    [
+        # Test with no quantization.
+        ("cpu", None, None),
+        # Test with Q8_0 quantization.
+        ("cpu", gguf.GGMLQuantizationType.Q8_0, gguf.GGMLQuantizationType.Q8_0),
+        ("cpu", None, gguf.GGMLQuantizationType.Q8_0),
+        ("cpu", gguf.GGMLQuantizationType.Q8_0, None),
+        # Test with F16 quantization (i.e. torch-compmatible quantization).
+        ("cpu", gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.F16),
+        ("cpu", None, gguf.GGMLQuantizationType.F16),
+        ("cpu", gguf.GGMLQuantizationType.F16, None),
+        # Test all of above cases on CUDA.
+        ("cuda", None, None),
+        # Test with Q8_0 quantization.
+        ("cuda", gguf.GGMLQuantizationType.Q8_0, gguf.GGMLQuantizationType.Q8_0),
+        ("cuda", None, gguf.GGMLQuantizationType.Q8_0),
+        ("cuda", gguf.GGMLQuantizationType.Q8_0, None),
+        # Test with F16 quantization (i.e. torch-compmatible quantization).
+        ("cuda", gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.F16),
+        ("cuda", None, gguf.GGMLQuantizationType.F16),
+        ("cuda", gguf.GGMLQuantizationType.F16, None),
+    ],
+)
+def test_ggml_tensor_multiply(
+    device: str, x1_quant_type: gguf.GGMLQuantizationType | None, x2_quant_type: gguf.GGMLQuantizationType | None
+):
+    # Skip test if CUDA is not available.
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available.")
+
+    generator = torch.Generator().manual_seed(123)
+
+    x1 = torch.randn(32, 64, generator=generator).to(device=device)
+    x2 = torch.randn(32, 64, generator=generator).to(device=device)
+
+    # Quantize the tensors.
+    x1_quantized = quantize_tensor(x1, x1_quant_type) if x1_quant_type is not None else x1
+    x2_quantized = quantize_tensor(x2, x2_quant_type) if x2_quant_type is not None else x2
+
+    # Check devices.
+    for x in [x1, x2, x1_quantized, x2_quantized]:
+        assert x.device.type == device
+
+    # Perform the multiplication.
+    result = x1 * x2
+    result_quantized = x1_quantized * x2_quantized
+
+    assert result.shape == result_quantized.shape
+    assert result.dtype == result_quantized.dtype
+    assert torch.allclose(result, result_quantized, atol=1e-1)
+
+
+def test_ggml_tensor_to_dtype_raises_error():
+    x = torch.randn(32, 64)
+    x_quantized = quantize_tensor(x, gguf.GGMLQuantizationType.Q8_0)
+
+    with pytest.raises(ValueError):
+        x_quantized.to(dtype=torch.float32)
+
+    with pytest.raises(ValueError):
+        x_quantized.float()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device")
+def test_ggml_tensor_to_device():
+    x = torch.randn(32, 64)
+    x_cpu = quantize_tensor(x, gguf.GGMLQuantizationType.Q8_0)
+
+    x_gpu = x_cpu.to(device=torch.device("cuda"))
+
+    assert x_cpu.device.type == "cpu"
+    assert x_gpu.device.type == "cuda"
+
+    assert torch.allclose(x_cpu.quantized_data, x_gpu.quantized_data.cpu(), atol=1e-5)
+
+
+def test_ggml_tensor_shape():
+    x = torch.randn(32, 64)
+    x_quantized = quantize_tensor(x, gguf.GGMLQuantizationType.Q8_0)
+
+    assert x_quantized.shape == x.shape
+    assert x_quantized.size() == x.size()
+
+
+def test_ggml_tensor_quantized_shape():
+    x = torch.randn(32, 64)
+    x_quantized = quantize_tensor(x, gguf.GGMLQuantizationType.Q8_0)
+
+    # This is mainly just a smoke test to confirm that .quantized_shape can be accesses and doesn't hit any weird
+    # dispatch errors.
+    assert x_quantized.quantized_shape != x.shape
+
+
+def test_ggml_tensor_calc_size():
+    """Test that the calc_tensor_size(...) utility function correctly uses the underlying quantized tensor to calculate
+    size rather than the unquantized tensor.
+    """
+    x = torch.randn(32, 64)
+    x_quantized = quantize_tensor(x, gguf.GGMLQuantizationType.Q8_0)
+
+    compression_ratio = calc_tensor_size(x) / calc_tensor_size(x_quantized)
+    # Assert that the compression ratio is approximately 4x.
+    assert abs(compression_ratio - 4) < 0.5

tests/backend/quantization/test_bnb_llm_int8.py

@@ -0,0 +1,85 @@
+import pytest
+import torch
+
+try:
+    from invokeai.backend.quantization.bnb_llm_int8 import InvokeLinear8bitLt
+except ImportError:
+    pass
+
+
+def test_invoke_linear_8bit_lt_quantization():
+    """Test quantization with InvokeLinear8bitLt."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    # Set the seed for reproducibility since we are using a pretty tight atol.
+    torch.manual_seed(3)
+
+    orig_layer = torch.nn.Linear(32, 64)
+    orig_layer_state_dict = orig_layer.state_dict()
+
+    # Initialize a InvokeLinear8bitLt layer (it is not quantized yet).
+    quantized_layer = InvokeLinear8bitLt(input_features=32, output_features=64, has_fp16_weights=False)
+
+    # Load the non-quantized layer's state dict into the quantized layer.
+    quantized_layer.load_state_dict(orig_layer_state_dict)
+
+    # Move the InvokeLinear8bitLt layer to the GPU. This triggers quantization.
+    quantized_layer.to("cuda")
+
+    # Assert that the InvokeLinear8bitLt layer is quantized.
+    assert quantized_layer.weight.CB is not None
+    assert quantized_layer.weight.SCB is not None
+    assert quantized_layer.weight.CB.dtype == torch.int8
+
+    # Run inference on both the original and quantized layers.
+    x = torch.randn(1, 32)
+    y = orig_layer(x)
+    y_quantized = quantized_layer(x.to("cuda"))
+    assert y.shape == y_quantized.shape
+    # All within ~20% of each other.
+    assert torch.allclose(y, y_quantized.to("cpu"), atol=0.05)
+
+
+def test_invoke_linear_8bit_lt_state_dict_roundtrip():
+    """Test that we can roundtrip the state dict of a quantized InvokeLinear8bitLt layer."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    # Set the seed for reproducibility since we are using a pretty tight atol.
+    torch.manual_seed(3)
+
+    orig_layer = torch.nn.Linear(32, 64)
+    orig_layer_state_dict = orig_layer.state_dict()
+
+    # Run inference on the original layer.
+    x = torch.randn(1, 32)
+    y = orig_layer(x)
+
+    # Prepare a quantized InvokeLinear8bitLt layer.
+    quantized_layer_1 = InvokeLinear8bitLt(input_features=32, output_features=64, has_fp16_weights=False)
+    quantized_layer_1.load_state_dict(orig_layer_state_dict)
+    quantized_layer_1.to("cuda")
+
+    # Assert that the InvokeLinear8bitLt layer is quantized.
+    assert quantized_layer_1.weight.CB is not None
+    assert quantized_layer_1.weight.SCB is not None
+    assert quantized_layer_1.weight.CB.dtype == torch.int8
+
+    # Run inference on the quantized layer.
+    y_quantized_1 = quantized_layer_1(x.to("cuda"))
+
+    # Save the state dict of the quantized layer.
+    quantized_layer_1_state_dict = quantized_layer_1.state_dict()
+
+    # Load the state dict of the quantized layer into a new quantized layer.
+    quantized_layer_2 = InvokeLinear8bitLt(input_features=32, output_features=64, has_fp16_weights=False)
+    quantized_layer_2.load_state_dict(quantized_layer_1_state_dict)
+    quantized_layer_2.to("cuda")
+
+    # Run inference on the new quantized layer.
+    y_quantized_2 = quantized_layer_2(x.to("cuda"))
+
+    # Assert that the inference results are the same.
+    assert torch.allclose(y, y_quantized_1.to("cpu"), atol=0.05)
+    assert torch.allclose(y_quantized_1, y_quantized_2, atol=1e-5)