AttributeError: 'Parameter' object has no attribute 'CB' from bitsandbytes library when running the Mistral model on spark dataframe

[nagarjuna1995]

Platform: Databricks
LLM Model: MistralAI-7B

I am trying to run a quantized LLM(4bit) model on a spark dataframe consisting of blocks of text and questions. As the model is quantized, some of the layers are on CPU and some of them are on GPU. I have created a Spark UDF to run the model inference on spark dataframe(20K rows). I am getting an error from BitsAndBytes library

An exception was thrown from the Python worker. Please see the stack trace below.
AttributeError: 'Parameter' object has no attribute 'CB'
Traceback (most recent call last):
File "/databricks/python/lib/python3.10/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/generation/utils.py", line 1479, in generate
return self.greedy_search(
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/generation/utils.py", line 2340, in greedy_search
outputs = self(
File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 1154, in forward
outputs = self.model(
File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 1039, in forward
layer_outputs = decoder_layer(
File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 754, in forward
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 255, in forward
query_states = self.q_proj(hidden_states)
File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-f335e7c4-adba-47d3-8a10-4fc922549da8/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 443, in forward
if self.weight.CB is not None:
AttributeError: 'Parameter' object has no attribute 'CB'

I tried to run on the sample of the rows but was getting the same error. I am suspecting it is due to quantization and the distribution of layers on GPU and CPU

[younesbelkada]

Hi @nagarjuna1995
Hmm, can you try to use latest bitsandbytes, transformers and accelerate?

pip install -U transformers bitsandbytes accelerate

[nagarjuna1995]

Hello @younesbelkada , I have tried installing the latest version but still the same issue