feat(qtensor): add MarlinQBitsTensor

Adding more tests revealed a bug in the Marlin int4 kernel when the
weights and inputs are large enough.
Failing configurations are marked as xfail.

optimum/quanto/tensor/weights/marlin/int4/__init__.py

@@ -1 +1,2 @@
 from .packed import *
+from .qbits import *

optimum/quanto/tensor/weights/marlin/int4/qbits.py

@@ -0,0 +1,168 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ast
+
+import torch
+from torch.autograd import Function
+
+from ....function import QuantizedLinearFunction
+from ....grouped import group, ungroup
+from ....qtype import qtypes
+from ...qbits import WeightQBitsTensor
+from ..permutations import marlin_permute
+from .packed import MarlinInt4PackedTensor
+
+
+__all__ = ["MarlinInt4WeightQBitsTensor"]
+
+
+class MarlinQBitsDequantizer(Function):
+ @staticmethod
+ def forward(ctx, t):
+ unpacked = t._data.unpack()
+ scale = t._scale
+ shift = t._shift
+ unpacked = group(unpacked, axis=0, group_size=t._group_size)
+ # Apply inverted permutations
+ scale = marlin_permute(scale, reverse=True)
+ shift = marlin_permute(shift, reverse=True)
+ n_scales = scale.numel()
+ scale = scale.t().reshape((n_scales, 1))
+ shift = shift.t().reshape((n_scales, 1))
+ # Shift is already scaled and negated
+ dqt = scale * unpacked + shift
+ return ungroup(dqt, axis=t.axis, orig_shape=t.shape)
+
+ @staticmethod
+ def backward(ctx, gO):
+ return gO
+
+
+class MarlinQBitsLinearFunction(QuantizedLinearFunction):
+ @staticmethod
+ def forward(ctx, input, other, bias):
+ ctx.save_for_backward(input, other)
+ if type(input) is not torch.Tensor:
+ input = input.dequantize()
+ out_features, in_features = other.shape
+ output = torch.ops.quanto.gemm_f16i4_marlin(
+ input,
+ other._data._data,
+ other._scale,
+ other._shift,
+ other._workspace,
+ )
+ if bias is not None:
+ output = output + bias
+ return output
+
+
+class MarlinInt4WeightQBitsTensor(WeightQBitsTensor):
+ @staticmethod
+ def __new__(cls, qtype, axis, group_size, size, stride, data, scale, shift, requires_grad=False):
+ assert data.device.type == "cuda"
+ assert data.device == scale.device
+ assert data.device == shift.device
+ return torch.Tensor._make_wrapper_subclass(
+ cls, size, strides=stride, dtype=scale.dtype, device=data.device, requires_grad=requires_grad
+ )
+
+ def __init__(self, qtype, axis, group_size, size, stride, data, scale, shift, requires_grad=False):
+ assert axis == 0
+ out_features, in_features = size
+ if not isinstance(data, MarlinInt4PackedTensor):
+ assert type(data) is torch.Tensor
+ # Format data, scale and shift for optimized CUDA gemm
+ ungrouped = ungroup(data, axis=0, orig_shape=size)
+ data = MarlinInt4PackedTensor.pack(ungrouped)
+ scale = scale.reshape(out_features, in_features // group_size).t().contiguous()
+ shift = shift.reshape(out_features, in_features // group_size).t()
+ if not shift.dtype.is_floating_point:
+ # Integer shift must be scaled
+ shift = scale * shift
+ # Shift must be negated
+ shift = -shift.contiguous()
+ # Finally, apply scale and shift permutations
+ scale = marlin_permute(scale)
+ shift = marlin_permute(shift)
+ super().__init__(qtype, axis, group_size, size, stride, data, scale, shift)
+ self._workspace = torch.zeros(out_features // 128 * 16, dtype=torch.int, device=data.device)
+
+ def dequantize(self):
+ return MarlinQBitsDequantizer.apply(self)
+
+ def weight_qbits_tensor(self):
+ """Convert back to a WeightQBitsTensor
+
+ This is required to make sure only standard packing is used when serializing.
+ """
+ data = group(self._data.unpack(), axis=self.axis, group_size=self._group_size)
+ scale = marlin_permute(self._scale, reverse=True)
+ shift = marlin_permute(self._shift, reverse=True)
+ n_scales = scale.numel()
+ scale = scale.t().reshape((n_scales, 1))
+ shift = -shift.t().reshape((n_scales, 1))
+ return WeightQBitsTensor(
+ self._qtype, self._axis, self._group_size, self.size(), self.stride(), data, scale, shift
+ )
+
+ def __tensor_flatten__(self):
+ inner_tensors = ["_data", "_scale", "_shift"]
+ # Since meta can be used for serialization, use only strings
+ meta = {
+ "qtype": self._qtype.name,
+ "axis": str(self._axis),
+ "group_size": str(self._group_size),
+ "size": str(list(self.size())),
+ "stride": str(list(self.stride())),
+ }
+ return inner_tensors, meta
+
+ @staticmethod
+ def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
+ assert len(inner_tensors) == 3
+ assert len(meta) == 5
+ data, scale, shift = inner_tensors["_data"], inner_tensors["_scale"], inner_tensors["_shift"]
+ # Meta should only contain strings, AST compatible except qtype
+ qtype = qtypes[meta["qtype"]]
+ axis = ast.literal_eval(meta["axis"])
+ group_size = ast.literal_eval(meta["group_size"])
+ size = ast.literal_eval(meta["size"])
+ stride = ast.literal_eval(meta["stride"])
+ return MarlinInt4WeightQBitsTensor(qtype, axis, group_size, size, stride, data, scale, shift)
+
+ @classmethod
+ def __torch_function__(cls, func, types, args=(), kwargs=None):
+ """Dispatch torch functions applied on this subtensor
+
+ This method is called whenever a torch function (such as `torch.nn.functional.linear`)
+ is called with at least one parameter coresponding to this subtensor:
+
+ - if a quantized implementation exists for the selected function, it is called,
+ - otherwise, the original implementation is called, deactivating further functional dispatch.
+
+ During the execution of the standard torch function, a second-level of dispatch will
+ happen, but this time directly on individual torch Tensor operations (mainly ATEN).
+ """
+ kwargs = kwargs or {}
+ if func is torch.nn.functional.linear:
+
+ def qlinear(input, other, bias=None):
+ return MarlinQBitsLinearFunction.apply(input, other, bias)
+
+ return qlinear(*args, **kwargs)
+ # Defer to operations dispatcher
+ with torch._C.DisableTorchFunctionSubclass():
+ return func(*args, **kwargs)

test/tensor/weights/optimized/test_marlin_int4_weight_qbits_tensor.py

@@ -0,0 +1,150 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from helpers import device_eq, random_qweight
+from tensor.weights.weight_helpers import check_weight_qtensor_linear
+
+from optimum.quanto import qint4
+from optimum.quanto.library.extensions import is_extension_available
+from optimum.quanto.tensor.weights import WeightQBitsTensor
+from optimum.quanto.tensor.weights.marlin.int4 import MarlinInt4WeightQBitsTensor
+
+
+@pytest.mark.skipif(
+ not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, reason="CUDA >= sm80 not available"
+)
+@pytest.mark.parametrize("in_features", [128, 256, 512, 1024])
+@pytest.mark.parametrize("out_features", [128, 256, 512, 1024])
+def test_marlin_int4_weight_qbits_tensor_from_qbits_tensor(in_features, out_features):
+ qtype = qint4
+ group_size = 128
+ dtype = torch.float16
+ shape = (out_features, in_features)
+ device = torch.device("cuda")
+ qbt = random_qweight(shape, qtype, dtype, group_size=group_size, device=device)
+ # Create a MarlinInt4WeightQBitsTensor from the WeightQBitsTensor members
+ marlinqbt = MarlinInt4WeightQBitsTensor(
+ qtype=qbt.qtype,
+ axis=qbt.axis,
+ group_size=qbt._group_size,
+ size=qbt.size(),
+ stride=qbt.stride(),
+ data=qbt._data.unpack(),
+ scale=qbt._scale,
+ shift=qbt._shift,
+ )
+ assert marlinqbt.dtype == dtype
+ assert marlinqbt.qtype == qtype
+ assert marlinqbt.shape == shape
+ assert device_eq(marlinqbt.device, device)
+ # Verify the dequantized tensors are identical
+ assert torch.equal(marlinqbt.dequantize(), qbt.dequantize())
+ # Now verify that we can reconstruct the WeightQBitsTensor
+ new_qbt = marlinqbt.weight_qbits_tensor()
+ assert type(new_qbt) is WeightQBitsTensor
+ assert new_qbt.dtype == dtype
+ assert new_qbt.qtype == qtype
+ assert new_qbt.shape == shape
+ assert torch.equal(new_qbt._data, qbt._data)
+ assert torch.equal(new_qbt._scale, qbt._scale)
+ assert torch.equal(new_qbt._shift, qbt._shift)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_marlin_int4_weight_qbits_tensor_move(device):
+ qtype = qint4
+ group_size = 128
+ dtype = torch.float16
+ shape = (1024, 1024)
+ device = torch.device("cuda")
+ # Create an MarlinInt4WeightQBitsTensor from a QBitsTensor on CUDA
+ qbt = random_qweight(shape, qtype, dtype, group_size=group_size, device=torch.device("cuda"))
+ marlinqbt = MarlinInt4WeightQBitsTensor(
+ qtype=qbt.qtype,
+ axis=qbt.axis,
+ group_size=qbt._group_size,
+ size=qbt.size(),
+ stride=qbt.stride(),
+ data=qbt._data.unpack(),
+ scale=qbt._scale,
+ shift=qbt._shift,
+ )
+ # Move to device, dequantize and compare
+ moved_qbt = marlinqbt.to(device)
+ assert isinstance(moved_qbt, WeightQBitsTensor)
+ if device.type != "cuda":
+ assert type(moved_qbt) is not MarlinInt4WeightQBitsTensor
+ assert marlinqbt.dtype == moved_qbt.dtype
+ assert marlinqbt.qtype == moved_qbt.qtype
+ assert marlinqbt.shape == moved_qbt.shape
+ assert torch.equal(marlinqbt.dequantize().to(device), moved_qbt.dequantize())
+
+
+def _test_marlin_int4_weight_qbits_tensor_linear(
+ dtype, weight_qtype, group_size, batch_size, tokens, in_features, out_features, use_bias
+):
+ # Create an MarlinInt4WeightQBitsTensor from a QBitsTensor on CUDA
+ qbt = random_qweight(
+ (out_features, in_features), weight_qtype, dtype, group_size=group_size, device=torch.device("cuda")
+ )
+ marlin_qweight = MarlinInt4WeightQBitsTensor(
+ qtype=qbt.qtype,
+ axis=qbt.axis,
+ group_size=qbt._group_size,
+ size=qbt.size(),
+ stride=qbt.stride(),
+ data=qbt._data.unpack(),
+ scale=qbt._scale,
+ shift=qbt._shift,
+ )
+ check_weight_qtensor_linear(marlin_qweight, batch_size, tokens, use_bias)
+
+
+@pytest.mark.skipif(
+ not is_extension_available("quanto_cuda") or torch.cuda.get_device_capability()[0] < 8,
+ reason="CUDA >= sm80 not available",
+)
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("tokens", [16, 32])
+@pytest.mark.parametrize("in_features", [1024])
+@pytest.mark.parametrize("out_features", [1024, 2048, 4096])
+@pytest.mark.parametrize("use_bias", [True, False], ids=["bias", "no-bias"])
+def test_marlin_int4_weight_qbits_tensor_linear(batch_size, tokens, in_features, out_features, use_bias):
+ dtype = torch.float16
+ weight_qtype = qint4
+ group_size = 128
+ _test_marlin_int4_weight_qbits_tensor_linear(
+ dtype, weight_qtype, group_size, batch_size, tokens, in_features, out_features, use_bias
+ )
+
+
+@pytest.mark.xfail(reason="Bug in Marlin kernel", strict=False)
+@pytest.mark.skipif(
+ not is_extension_available("quanto_cuda") or torch.cuda.get_device_capability()[0] < 8,
+ reason="CUDA >= sm80 not available",
+)
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("tokens", [48, 64])
+# @pytest.mark.parametrize("in_features", [1024, 2048, 4096, 16384])
+@pytest.mark.parametrize("in_features", [4096, 16384])
+@pytest.mark.parametrize("out_features", [2048, 4096])
+def test_marlin_int4_weight_qbits_tensor_linear_failing(batch_size, tokens, in_features, out_features):
+ dtype = torch.float16
+ weight_qtype = qint4
+ group_size = 128
+ _test_marlin_int4_weight_qbits_tensor_linear(
+ dtype, weight_qtype, group_size, batch_size, tokens, in_features, out_features, use_bias=False
+ )