nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/quantization/qtensor/fp8_tensor.py

@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Implements FP8 quantization for efficient tensor storage and computation."""
+
+import math
+
+import torch
+
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+from ..utils import (
+    convert_quantization_axis_to_reduce_axis,
+    reduce_amax,
+    reduce_block_amax,
+    reduce_block_padding,
+)
+
+__all__ = ["FP8QTensor"]
+
+
+class FP8QTensor(BaseQuantizedTensor):
+    """Implements the FP8 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as a packed fp8 tensor.
+    """
+
+    @classmethod
+    def quantize(
+        cls,
+        input: torch.Tensor,
+        scales: torch.Tensor = None,
+        axis: tuple | int | None = None,
+        block_sizes: dict | None = None,
+    ) -> tuple:
+        """Converting a tensor to a quantized format based on FP8 quantization. Only E4M3 is supported.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            scales (torch.Tensor): The scales for quantization.
+            axis: The dimensions to reduce for quantization. None or int or tuple of ints.
+            block_sizes (dict): A dictionary specifying the block size for each dimension.
+
+        Note: One can only provide axis or block_sizes for FP8 quantization.
+
+        Returns:
+            tuple: FP8QTensor, scales
+        """
+        original_input = input
+
+        # If block_sizes is provided, pad the input so that each dimension is divisible by the block size.
+        if block_sizes:
+            input = reduce_block_padding(input, block_sizes)
+
+        # Compute scales if not provided
+        if scales is None:
+            if block_sizes:
+                amax = reduce_block_amax(input, block_sizes)
+            else:
+                reduce_axis = convert_quantization_axis_to_reduce_axis(input, axis)
+                amax = reduce_amax(input, axis=reduce_axis)
+            scales = amax / 448.0  # Consider parameterizing the divisor if needed
+
+        # Determine the expected scales shape from the (possibly padded) input
+        expected_shape = list(input.shape)
+        expanded_scales = scales.clone()
+        if block_sizes:
+            for dim, block_size in block_sizes.items():
+                # Convert negative indices to positive ones.
+                dim = dim if dim >= 0 else len(input.shape) + dim
+                # After padding, this should always hold.
+                assert input.shape[dim] % block_size == 0, (
+                    f"Tensor dimension {dim}, {input.shape[dim]} is not divisible by {block_size} even after padding."
+                )
+                # The scales tensor is expected to have size equal to input.shape[dim] // block_size.
+                expected_shape[dim] = input.shape[dim] // block_size
+
+            # If we get amax from tensor_quantizer, it might not be the expected shape but has the
+            # same number of elements, reshape it.
+            if scales.shape != tuple(expected_shape) and scales.numel() == math.prod(
+                expected_shape
+            ):
+                scales = scales.reshape(expected_shape)
+                expanded_scales = expanded_scales.reshape(expected_shape)
+
+            assert scales.shape == tuple(expected_shape), (
+                f"Mismatch in expected scale shape: {scales.shape} vs {tuple(expected_shape)}"
+            )
+
+            # Expand scales along each block dimension for broadcasting.
+            for dim, block_size in block_sizes.items():
+                expanded_scales = expanded_scales.repeat_interleave(block_size, dim=dim)
+
+        # Perform quantization using FP8 (E4M3) format.
+        quantized_data = (input / expanded_scales).to(torch.float8_e4m3fn)
+
+        # Crop quantized_data back to the original shape (if padding was added).
+        slices = tuple(slice(0, dim) for dim in original_input.shape)
+        quantized_data_cropped = quantized_data[slices]
+
+        return cls(original_input.shape, original_input.dtype, quantized_data_cropped), scales
+
+    def dequantize(self, dtype: torch.dtype = None, **kwarg):
+        """Dequantze FP8 packed tensor to a target dtype."""
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+        assert "scale" in kwarg, "Require scale for FP8 dequantization."
+
+        # get args
+        scales = kwarg["scale"]
+        block_sizes = kwarg.get("block_sizes")
+
+        quantized_data = self._quantized_data
+        if block_sizes:
+            # pad the weight if needed
+            quantized_data = reduce_block_padding(quantized_data, block_sizes)
+            shape = quantized_data.shape
+
+            # Compute expanded shape for broadcasting scales
+            expanded_shape = list(shape)
+            for dim, block_size in block_sizes.items():
+                assert shape[dim] % block_size == 0, (
+                    f"Dimension {shape[dim]} is not divisible by {block_size}."
+                )
+                expanded_shape[dim] //= block_size  # Reduce the dimension size for blocks
+
+            assert tuple(expanded_shape) == scales.shape, (
+                f"Scales shape {scales.shape} must match expected {tuple(expanded_shape)}."
+            )
+
+            # Expand scales for broadcasting
+            for dim, block_size in block_sizes.items():
+                scales = scales.repeat_interleave(block_size, dim=dim)
+
+        # handle padded tensors
+        slices = tuple(slice(0, dim) for dim in self.metadata["shape"])
+        scales = scales.to(self._quantized_data.device)
+
+        return (quantized_data.to(dtype) * scales.to(dtype))[slices]

modelopt/torch/quantization/qtensor/int4_tensor.py

@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Implements INT4 quantization for efficient tensor storage and computation."""
+
+import numpy as np
+import torch
+
+from ..extensions import get_cuda_ext
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+from ..utils import reduce_amax, reduce_block_padding
+
+__all__ = ["INT4QTensor"]
+
+
+class INT4QTensor(BaseQuantizedTensor):
+    """Implements the INT4 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as a packed uint8 tensor.
+    """
+
+    @staticmethod
+    def _get_quant_maxbound(num_bits):
+        return 2 ** (num_bits - 1) - 1
+
+    @classmethod
+    def quantize(cls, input: torch.Tensor, block_size: int) -> tuple:
+        """Converting a tensor to a quantized format based on INT4 (AWQ) quantization.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            block_size (int): The size of each block for quantization.
+
+        Returns:
+            tuple: Contains quantized data, input quantization config, and scale quantization config.
+        """
+        cuda_ext = get_cuda_ext()
+
+        scale_quant_maxbound = cls._get_quant_maxbound(num_bits=4)
+
+        assert input.shape[-1] % 2 == 0, "Input tensor must have even number on last dimension."
+
+        # pad the input if needed
+        original_input = input
+        input = reduce_block_padding(input.view(-1), block_sizes={-1: block_size})
+
+        # get scales for each block
+        block_input = input.view(-1, block_size)
+        scales = scale_quant_maxbound / reduce_amax(block_input, -1)
+        # expand scalers to match shape of input
+        scales = scales.view(block_input.shape[0], -1)  # shape: (block_input.shape[0], 1)
+
+        if cuda_ext and input.is_cuda:
+            # use a custom cuda kernel if available
+            packed_output_uint8 = cuda_ext.INT4_quantize(input, scales, block_size)
+        else:
+            scaled_blocks = block_input * scales
+            flattened = scaled_blocks.flatten()
+            # uint4: 0 - 15
+            flattened = flattened.round().clamp(
+                -(scale_quant_maxbound + 1), scale_quant_maxbound
+            ) + (scale_quant_maxbound + 1)
+            flattened = flattened.to(torch.uint8)
+
+            packed_output_uint8 = torch.empty(
+                input.numel() // 2, dtype=torch.uint8, device=input.device
+            )
+            # pack the int4 weights into a uint8 tensor
+            # packing format: w0, w1, w2, w3, w4, w5, ...
+            #               | byte  | byte  | byte  |
+            packed_output_uint8 = flattened[::2] << 4 | flattened[1::2]
+
+        # reshape the packed_output_uint8 to the original dimensions for sharding
+        packed_output_uint8 = packed_output_uint8.reshape(*original_input.shape[:-1], -1)
+        return cls(original_input.shape, input.dtype, packed_output_uint8), scales
+
+    def dequantize(self, dtype: torch.dtype = None, **kwarg):
+        """Dequantze INT4 packed tensor to a target dtype."""
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+        cuda_ext = get_cuda_ext()
+
+        # get kwargs
+        scales = kwarg["scale"]
+        block_sizes = kwarg["block_sizes"]
+        scale_quant_maxbound = self._get_quant_maxbound(num_bits=4)
+
+        if cuda_ext and self._quantized_data.is_cuda:
+            # use a custom cuda kernel if available
+            scales = scales.to(self._quantized_data.device)
+            output = cuda_ext.INT4_dequantize(
+                self._quantized_data.view(-1), scales, block_sizes[-1]
+            )
+            return (
+                output.view(-1)[: np.prod(self.metadata["shape"])]  # handle padding
+                .view(self.metadata["shape"])
+                .to(dtype)
+            )
+        else:
+            # indexing in torch required long dtype, we may need to optimize this with customized kernels
+            # convert (0, 15) -> (-8, 7)
+            first_half = (self._quantized_data >> 4).to(torch.long) - (scale_quant_maxbound + 1)
+            second_half = (self._quantized_data & 0x0F).to(torch.long) - (scale_quant_maxbound + 1)
+
+            # de-quantize tensor
+            first_half = first_half.view(-1, block_sizes[-1] // 2) / scales.view(-1, 1)
+            second_half = second_half.view(-1, block_sizes[-1] // 2) / scales.view(-1, 1)
+
+            # merge the interleaving elements
+            first_half = first_half.flatten().unsqueeze(-1).transpose(0, 1)
+            second_half = second_half.flatten().unsqueeze(-1).transpose(0, 1)
+            return (
+                torch.stack([first_half, second_half], dim=-1)
+                .view(-1)[: np.prod(self.metadata["shape"])]  # handle padding
+                .reshape(self.metadata["shape"])
+                .to(dtype)
+            )

modelopt/torch/quantization/qtensor/int8_tensor.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""Implements INT8 quantization for efficient tensor storage and computation."""
+
+import torch
+
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+from ..utils import (
+    convert_quantization_axis_to_reduce_axis,
+    reduce_amax,
+    reduce_block_amax,
+    reduce_block_padding,
+)
+
+
+class INT8QTensor(BaseQuantizedTensor):
+    """Implements the INT8 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as an INT8 tensor.
+    """
+
+    @classmethod
+    def quantize(
+        cls,
+        input: torch.Tensor,
+        scales: torch.Tensor | None = None,
+        axis: tuple | int | None = None,
+        block_sizes: dict | None = None,
+    ) -> tuple:
+        """Converting a tensor to a quantized format based on INT8 quantization.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            scales (torch.Tensor): The scales for quantization.
+            axis: The dimensions to reduce for quantization. None or int or tuple of ints.
+            block_sizes (dict): A dictionary specifying the block size for each dimension.
+        Note: One can only provide axis or block_sizes for INT8 quantization.
+
+        Returns:
+            tuple: INT8QTensor, scales
+        """
+        original_input = input
+        if scales is None:
+            if block_sizes:
+                input = reduce_block_padding(input, block_sizes)
+                amax = reduce_block_amax(input, block_sizes)
+            else:
+                reduce_axis = convert_quantization_axis_to_reduce_axis(input, axis)
+                amax = reduce_amax(input, axis=reduce_axis)
+            scales = amax / 127.0
+
+        # Calculate the scale shape and make sure it aligns with input and block_sizes
+        expected_shape = list(input.shape)
+        expanded_scales = scales.clone()
+        if block_sizes:
+            for dim, block_size in block_sizes.items():
+                dim = dim if dim >= 0 else len(input.shape) + dim  # Convert negative index
+                assert input.shape[dim] % block_size == 0, (
+                    f"Tensor dimension {dim}, {input.shape[dim]} is not divisible by {block_size}."
+                )
+                expected_shape[dim] = (
+                    input.shape[dim] // block_size
+                )  # Adjust expected shape for blocks
+
+            # Assert the shape of `scales` matches expected reduced dimensions
+            assert scales.shape == tuple(expected_shape), (
+                f"Mismatch in expected scale shape: {scales.shape} vs {tuple(expected_shape)}"
+            )
+
+            # Expand scales for broadcasting
+            for dim, block_size in block_sizes.items():
+                expanded_scales = expanded_scales.repeat_interleave(block_size, dim=dim)
+
+        # Quantization
+        quantized_data = (input / expanded_scales).round().clamp(-128, 127).to(torch.int8)
+
+        return cls(original_input.shape, original_input.dtype, quantized_data), scales
+
+    def dequantize(self, dtype: torch.dtype = None, **kwarg):
+        """Dequantize INT8 packed tensor to a target dtype."""
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+        assert "scale" in kwarg, "Require scale for INT8 dequantization."
+
+        # Get args
+        scales = kwarg["scale"]
+        block_sizes = kwarg.get("block_sizes")
+
+        shape = self._quantized_data.shape
+        if block_sizes:
+            # Compute expanded shape for broadcasting scales
+            expanded_shape = list(shape)
+            for dim, block_size in block_sizes.items():
+                assert shape[dim] % block_size == 0, (
+                    f"Dimension {shape[dim]} is not divisible by {block_size}."
+                )
+                expanded_shape[dim] //= block_size  # Reduce the dimension size for blocks
+
+            assert tuple(expanded_shape) == scales.shape, (
+                f"Scales shape {scales.shape} must match expected {tuple(expanded_shape)}."
+            )
+
+            # Expand scales for broadcasting
+            for dim, block_size in block_sizes.items():
+                scales = scales.repeat_interleave(block_size, dim=dim)
+
+        # Handle padded tensors
+        slices = tuple(slice(0, dim) for dim in self.metadata["shape"])
+        scales = scales.to(self._quantized_data.device)
+
+        return (self._quantized_data.view(torch.int8).to(dtype) * scales.to(dtype))[slices]

modelopt/torch/quantization/qtensor/mxfp4_tensor.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Implements MXFP4 quantization for efficient tensor storage and computation."""
+
+import torch
+
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+
+__all__ = ["MXFP4QTensor"]
+
+
+class MXFP4QTensor(BaseQuantizedTensor):
+    """Implements the MXFP4 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as a packed fp8 tensor.
+    """
+
+    E2M1_max = 6.0
+
+    E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
+    E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
+
+    @classmethod
+    def quantize(cls, input: torch.Tensor, block_size: int | None) -> tuple:
+        """Converting a tensor to a quantized format based on MXFP4 quantization. Only E4M3 is supported.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            block_sizes (dict | None): The block sizes for quantization.
+        """
+
+        def cast_fp4(x):
+            sign = torch.sign(x)
+            sign_bit = (2 - sign) // 2
+            ord_ = torch.sum(
+                (x.abs().unsqueeze(-1) - MXFP4QTensor.E2M1_bounds.to(x.device)) > 0, dim=-1
+            )
+            fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
+            return fp4_val
+
+        def fuse_uint4_to_uint8(x):
+            # If the last dimension is odd, pad with zeros
+            # If this behavior is not desired, please modify the code accordingly
+            left_side = x[..., 0::2]  # Even indices (0, 2, 4...)
+            right_side = x[..., 1::2]  # Odd indices (1, 3, 5...)
+            new_data = right_side.clone() << 4  # Put odd indices (higher addresses) in high bits
+            new_data[..., : left_side.shape[-1]] += left_side  # Put even indices in low bits
+            return new_data
+
+        if block_size is None:
+            block_size = 32
+
+        original_shape = input.shape
+        original_dtype = input.dtype
+        input = input.view(-1, block_size)
+        # get scales
+        # Casting scales to float yields better match with fakequant kernel (needs further investigation)
+        input_amax = input.float().abs().max(dim=-1, keepdim=True).values
+        descale = input_amax / cls.E2M1_max
+        min_value = torch.tensor(-127.0, device=descale.device)
+        e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
+
+        input = (input / torch.exp2(e8m0_scale)).view(original_shape)
+        input_q = cast_fp4(input)
+        input_q = fuse_uint4_to_uint8(input_q)
+        e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
+        return cls(original_shape, original_dtype, input_q), e8m0_scale
+
+    def dequantize(self, dtype: torch.dtype = None, **kwarg):
+        """Dequantze MXFP4 packed tensor to a target dtype."""
+
+        def unfuse_uint8_to_uint4(x):
+            """Unfuse uint8 values back to uint4 values.
+
+            This is the inverse operation of fuse_uint4_to_uint8.
+            """
+            # Extract the lower 4 bits (even indices)
+            left_side = x & 0x0F
+
+            # Extract the upper 4 bits (odd indices)
+            right_side = (x >> 4) & 0x0F
+
+            # Create a new tensor with alternating values
+            shape = list(x.shape)
+            shape[-1] = shape[-1] * 2
+            result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
+
+            # Fill in the values - even indices get low bits, odd indices get high bits
+            result[..., 0::2] = left_side  # Even indices from low bits
+            result[..., 1::2] = right_side  # Odd indices from high bits
+
+            return result
+
+        e8m0_scale = kwarg["scale"]
+        block_size = kwarg["block_sizes"][-1]
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+
+        # Unfuse the uint8 values back to uint4
+        x_unfused = unfuse_uint8_to_uint4(self._quantized_data)
+        # Extract sign and magnitude
+        sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(
+            torch.float32
+        )  # Extract sign bit and convert to +1/-1
+        magnitude = x_unfused & 0b0111  # Extract magnitude bits
+        magnitude = magnitude.to(torch.long)
+
+        # Create a tensor with the E2M1 values
+        values = torch.tensor(MXFP4QTensor.E2M1_values, device=self._quantized_data.device)
+
+        # Use gather to index the values tensor properly
+        # We need to reshape magnitude to match the dimensions we want to gather along
+        original_shape = magnitude.shape
+        x_float = values[magnitude.reshape(-1)].reshape(original_shape)
+
+        # Apply sign and scale
+        x_float = sign.float() * x_float
+
+        # Reshape to apply block-wise scaling
+        x_float = x_float.reshape(-1, block_size)
+
+        # Apply the E8M0 scale
+        scale_factor = torch.exp2(e8m0_scale.float() - 127)
+        scale_factor = scale_factor.reshape(-1, 1)  # Reshape for proper broadcasting
+
+        # Apply scaling and reshape back to original shape
+        x_float = x_float * scale_factor
+
+        # Reshape back to the original shape
+        return x_float.reshape(original_shape).to(dtype)