Class TestLLMInt8Functional

tests/test_functional.py:350–572

class TestLLMInt8Functional:
    @staticmethod
    def vectorwise_mm_dequant(xq, S1, S2, dtype=torch.half):
        """Reference implementation for the F.int8_mm_dequant function."""
        C = 127.0

        x = xq.float()
        if len(S1.shape) == 3 and len(x.shape) == 2:
            S1 = S1.squeeze(0)
        if len(S2.shape) == 3 and len(x.shape) == 2:
            S2 = S2.squeeze(0)
        if len(S1.shape) == 2:
            x *= S1 / C
        else:
            x *= S1 / C
        x *= S2 / C
        return x.to(dtype)

    @staticmethod
    def vectorwise_quant(x, dim=1):
        """Reference implementation"""
        max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True)
        xq = torch.round(x * (127.0 / max1)).to(torch.int8)
        return xq, max1

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [128], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [256], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("dim3", [499, 512], ids=id_formatter("dim3"))
    @pytest.mark.parametrize("dim4", [512], ids=id_formatter("dim4"))
    @pytest.mark.parametrize("dims", (2, 3), ids=id_formatter("dims"))
    @pytest.mark.parametrize("ldb", (0,), ids=id_formatter("ldb"))
    def test_int8_linear_matmul(self, device, dim1, dim2, dim3, dim4, dims, ldb):
        for i in range(k):
            if dims == 2:
                A = torch.randint(-128, 127, size=(dim1, dim3), dtype=torch.int8, device=device)
            elif dims == 3:
                A = torch.randint(-128, 127, size=(dim1, dim2, dim3), dtype=torch.int8, device=device)
            B = torch.randint(-128, 127, size=(dim4, dim3), dtype=torch.int8, device=device)
            C1 = torch.matmul(A.float(), B.t().float())

            C2 = F.int8_linear_matmul(A, B)
            torch.testing.assert_close(C1, C2.float())

    @pytest.mark.parametrize("device", get_available_devices())
    @pytest.mark.parametrize("dim1", [32], ids=id_formatter("dim1"))
    @pytest.mark.parametrize("dim2", [32], ids=id_formatter("dim2"))
    @pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3"))
    @pytest.mark.parametrize("dim4", [32], ids=id_formatter("dim4"))
    @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims"))
    def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims):
        for i in range(k):
            if dims == 2:
                A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).half()
            elif dims == 3:
                A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).half()
            B = torch.randn((dim4, dim3), device=device).half()
            torch.nn.init.xavier_uniform_(B)