CUDA error: an illegal memory access was encountered when using BitBlas on multiple GPUs

[mobicham]

There seems to be an issue with BitBlas when using multi-gpu, even allocating a new array causes this problem:

In [8]: model_loaded.model.layers[0].self_attn.q_proj.device
Out[8]: 0

In [11]: model_loaded.model.layers[16].self_attn.q_proj.device
Out[11]: 1

In [12]: l = model_loaded.model.layers[16].self_attn.q_proj
In [13]: x = torch.randn((1, 4096), device=l.device, dtype=torch.float16)
In [14]: out = l(x) #Runs OK

In [15]: l = model_loaded.model.layers[0].self_attn.q_proj
In [16]: x = torch.randn((1, 4096), device=l.device, dtype=torch.float16)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[16], line 1
----> 1 x = torch.randn((1, 4096), device=l.device, dtype=torch.float16)

RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

### Tasks
- [x] revert the transform weight function implementation of BitBLAS Matmul

[LeiWang1999]

Hi @mobicham , which framework are you using for this issue? would you mind provide the scripts for us to reproduce :)

[mobicham]

Using BitBlas via hqq on transformers models:

• pip install git+https://github.com/mobiusml/hqq.git
• pip install bitblas

# Code tested with 2 x 24 GB gpus
import torch, gc
from transformers import AutoModelForCausalLM, AutoTokenizer, HqqConfig

device    = 'auto'
dtype     = torch.float16
model_id  = 'meta-llama/Meta-Llama-3-8B-Instruct'
cache_dir = '.' 

quant_config  = HqqConfig(nbits=4, group_size=64, quant_zero=False, quant_scale=False, axis=1)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=dtype, 
    cache_dir=cache_dir,
    device_map=device, 
    quantization_config=quant_config
)

#Patching
from hqq.utils.patching import *
from hqq.core.quantize import *
HQQLinear.set_backend(HQQBackend.PYTORCH)
prepare_for_inference(model, backend='bitblas', verbose=True) 

# Check if the layers have parameters on the right device
for layer_id in range(0, 14):
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.device) #cuda:0
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.W_q.device) #cuda:0
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.scale.device) #cuda:0
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.zero.device) #cuda:0
    print('--------------------------------------------------------')

for layer_id in range(15, len(model.model.layers)):
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.device) #cuda:1
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.W_q.device) #cuda:1
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.scale.device) #cuda:1
    print(layer_id, model.model.layers[layer_id].self_attn.q_proj.zero.device) #cuda:1
    print('--------------------------------------------------------')


# Allocate example
layer_id = 0
l        = model.model.layers[layer_id].self_attn.q_proj
x        = torch.randn((1, 4096), device=l.device, dtype=torch.float16)
out      = l(x) #Runs OK


layer_id = 16
l        = model.model.layers[layer_id].self_attn.q_proj
x        = torch.randn((1, 4096), device=l.device, dtype=torch.float16)
out      = l(x) #Runs OK

print(out) #RuntimeError: CUDA error: an illegal memory access was encountered


# #Test 
# input_tensor = torch.zeros((1, 8), dtype=torch.int32, device='cuda:0')
# with torch.no_grad():
# 	out_ref = model.forward(input_tensor)

Thank you!

[LeiWang1999]

hi @mobicham , amazing project and thanks for reporting.

I suspect there’s a reason behind transform_weight, given that we refactored this API to include weight compression in 0.0.1.dev.14. However, I’m not certain if this was the intended rationale. Let me take a look.

In the meantime, I believe we should consider reverting the change to ensure the bitblas-based project functions correctly.

[mobicham]

Thank you @LeiWang1999
I noticed that transform_weights on a cuda:1 device tensor returns a cuda:0 device tensor for some reason. So I put it back on the right device https://github.com/mobiusml/hqq/blob/master/hqq/backends/bitblas.py#L115 but even when all the parameters (quantized weights, scale, zero) are on the right device, it produces that error

[LeiWang1999]

Hi @mobicham , looks like it works on my single gpu environment. btw, I'm using the latest release 0.0.1.dev15. I'll soon test it on my multi-gpu env.

[mobicham]

Thanks for testing! Yeah it works fine on single gpu, the issue is only with multi-gpu. Let me know how it goes!

[LeiWang1999]

@mobicham , yeah I've identified the problem, this issue is similar to #105.

BitBLAS was initialized with the shape n=4096, k=4096, while the shard model dispatched n=4096 and k=2048.

One possible solution could be to make the initialization of the Linear Layer aware of the shard information, similar to how it’s done in vLLM.

[mobicham]

Hm I am not sure how you got that, but in my example all seems correct. The code works perfectly fine on a single gpu for 7 models I tested.

    layer = model.model.layers[layer_id].self_attn.q_proj
    print(layer_id, layer.name, layer.device) #cuda:0
    print(layer_id, layer.name, layer.W_q.device, layer.W_q.shape, 'N', layer.out_features, 'K', layer.in_features) #cuda:0
    print(layer_id, layer.name, layer.scale.device, layer.scale.shape) #cuda:0
    print(layer_id, layer.name, layer.zero.device, layer.zero.shape) #cuda:0
    print(layer.eng_tag, layer.matmul_eng)
    print('--------------------------------------------------------')

prints

16 model.layers.16.self_attn.q_proj 1
16 model.layers.16.self_attn.q_proj cuda:1 torch.Size([4096, 2048]) N 4096 K 4096
16 model.layers.16.self_attn.q_proj cuda:1 torch.Size([4096, 64])
16 model.layers.16.self_attn.q_proj cuda:1 torch.Size([4096, 64])
torch.Size([4096, 4096])_4_64 <bitblas.ops.general_matmul.Matmul object at 0x7f904b238eb0>

W_q is int4 packed in int8, so its size should be 4096 x 2048 as a packed tensor. The unpacked size is 4096 x 4096, and it should be calling the 4096 x 4096 kernel, so it's all correct.

[LeiWang1999]

@mobicham , oh yeah thanks, sorry, that's my mistake, let me dig further.

[mobicham]

Here's a minimal example:

from hqq.core.quantize import *
from hqq.backends.bitblas import patch_hqq_to_bitblas
HQQLinear.set_backend(HQQBackend.PYTORCH)

def init_layers(K, N, device, quant_config, compute_dtype=torch.float16):
    layer_0   = torch.nn.Linear(in_features=N, out_features=K, bias=False).to(device=device, dtype=compute_dtype)
    hqq_0     = HQQLinear(layer_0, quant_config, compute_dtype=compute_dtype, device=device, del_orig=False)
    bitblas_0 = patch_hqq_to_bitblas(hqq_0, None)
    layer_0.device = device
    layer_0.compute_dtype = compute_dtype
    return [layer_0, bitblas_0]

def eval(layers):
    for layer in layers:
        x        = torch.randn((1, layer.out_features), device=layer.device, dtype=layer.compute_dtype)
        with torch.no_grad():
            out = layer(x)
        print(layer, layer.device, 'OK')


quant_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_scale=False, quant_zero=False, axis=1)
K, N         = 4096, 4096

layers = []
layers += init_layers(K, N, 'cuda:0', quant_config)
eval(layers)
# Linear(in_features=4096, out_features=4096, bias=False) cuda:0 OK
# HQQLinearBitBlas() cuda:0 OK
# Linear(in_features=4096, out_features=4096, bias=False) cuda:0 OK


layers += init_layers(K, N, 'cuda:1', quant_config)
eval(layers)

for i in range(5):
    eval(layers)
# RuntimeError: CUDA error: an illegal memory access was encountered
# CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
# For debugging consider passing CUDA_LAUNCH_BLOCKING=1
# Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

[LeiWang1999]

Hello @mobicham, I've identified the bug related to kernels on multi-GPUs; we need to call torch.cuda.set_device prior to invoking the kernel.

The device set should ultimately be integrated within the BitBLAS operator's forward function. However, a hotfix can be implemented prior to the HQQ forward execution :)

[mobicham]

Thank you! I just tried it dropbox/hqq@f8c4519 and it works!
Closing this issue as a result.