Output wrong result when using W2A8

[HaisongDing]

test.zip

The W2A8 fails the assert_close test. Bitblas version is 0.1.0.post

Code:

import torch
import bitblas
import torch.nn as nn

# enabling debug output

N = 4096
K = 2560

import pickle
with open('log.bin', 'rb') as f:
 loaded_tensors = pickle.load(f)
 input_tensor, weight_tensor = loaded_tensors

 print(input_tensor.size())
 print(input_tensor.max())
 print(input_tensor.min())

 print(weight_tensor.size())
 print(weight_tensor.max())
 print(weight_tensor.min())

bitblas.set_log_level("Info")
matmul_config = bitblas.MatmulConfig(
 # M=1, # M dimension
 N=N, # N dimension
 K=K, # K dimension
 A_dtype="int8", # activation A dtype
 W_dtype="int2", # weight W dtype
 accum_dtype="int32", # accumulation dtype
 out_dtype="float32", # output dtype
 layout="nt", # matrix layout, "nt" indicates the layout of A is non-transpose and the layout of W is transpose
 with_bias=False, # bias
 # configs for weight only quantization
 group_size=None, # setting for grouped quantization
 with_scaling=False, # setting for scaling factor
 with_zeros=False, # setting for zeros
 zeros_mode=None, # setting for how to calculating zeros
)

matmul = bitblas.Matmul(config=matmul_config)

# Transform weight tensor to int4 data type
weight_tensor_int8 = matmul.transform_weight(weight_tensor)
# Perform mixed-precision matrix multiplication
output_tensor = matmul(input_tensor, weight_tensor_int8)

# Reference result using PyTorch matmul for comparison
# ref_result = torch.matmul(input_tensor.to(torch.float32), weight_tensor.t().to(torch.float32))

ref_result = nn.functional.linear(input_tensor.to(torch.float32), weight_tensor.to(torch.float32))

print("Ref output:", ref_result)
print("BitBLAS output:", output_tensor)
torch.testing.assert_close(output_tensor, ref_result, rtol=1e-2, atol=1e-0)

The output of one go is:

torch.Size([1, 10, 2560])
tensor(127, device='cuda:0', dtype=torch.int8)
tensor(-128, device='cuda:0', dtype=torch.int8)
torch.Size([4096, 2560])
tensor(1, device='cuda:0', dtype=torch.int8)
tensor(-1, device='cuda:0', dtype=torch.int8)
Ref output: tensor([[[ -297., 476., 115., ..., 374., 27., 261.],
 [-1267., 83., -95., ..., 700., 69., 2158.],
 [ 515., -2524., 572., ..., 398., -636., -1122.],
 ...,
 [ -297., 476., 115., ..., 374., 27., 261.],
 [ 621., -1438., -826., ..., -850., -388., 464.],
 [ 515., -2524., 572., ..., 398., -636., -1122.]]],
 device='cuda:0')
BitBLAS output: tensor([[[ -297., 476., 115., ..., 374., 27., 181.],
 [-1267., 83., -95., ..., 700., 69., 1391.],
 [ 515., -2524., 572., ..., 398., -636., -587.],
 ...,
 [ -297., 476., 115., ..., 374., 27., 181.],
 [ 621., -1438., -826., ..., -850., -388., 788.],
 [ 515., -2524., 572., ..., 398., -636., -587.]]],
 device='cuda:0')
Traceback (most recent call last):
AssertionError: Tensor-likes are not close!

Mismatched elements: 2554 / 40960 (6.2%)
Greatest absolute difference: 19645.0 at index (0, 2, 463) (up to 1.0 allowed)
Greatest relative difference: 303.5 at index (0, 0, 1087) (up to 0.01 allowed)

Oddly, when I run another go, the results match! My guess is different kernel is selected. See attached (see test.zip above) for the complete log and the data I use.

[tanvoontao]

Haisong Ding (@HaisongDing) Have u figure out the issue?

[HaisongDing]

No, I switched to llama.cpp for low-bit LMM inference.