[RWKV] Final fix RWMV 4bit

[younesbelkada]

What does this PR do?

Fixes #23848

Double quantization was not working properly for RWKV models as stated in the PR above - leading to an error. This PR proposed a global fix for RWKV models so that they can be ran in 4bit bitsandbytes without any problem.

The followed approach here is the following:

• For each target layer, de-quantize the 4bit weights using bnb.nn.functional.dequantize_4bit
• Perform the weights scaling
• Requantize the weights again.

That way it is possible to make sure to cover both the double quantization and classic 4bit quantization and match the results together

import torch
from transformers import RwkvForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "RWKV/rwkv-4-169m-pile"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True
)

model = RwkvForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)
tok = AutoTokenizer.from_pretrained(model_id)

text = "Hello my name is"
input_ids = tok.encode(text, return_tensors="pt").to(0)

out = model.generate(input_ids, max_new_tokens=30)
print(tok.decode(out[0], skip_special_tokens=True))

model_non_dequant = RwkvForCausalLM.from_pretrained(model_id, load_in_4bit=True)

text = "Hello my name is"
input_ids = tok.encode(text, return_tensors="pt").to(0)

out = model.generate(input_ids, max_new_tokens=30)
print(tok.decode(out[0], skip_special_tokens=True))

cc @amyeroberts and @SunMarc for your information!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[amyeroberts]

Thanks for adding this fix!

Could you add a test which would have failed before this change and passes now? Ideally it should be applied to all eligible models

[amyeroberts]

Why are we setting requires_grad=False here?

[younesbelkada]

This seems to be a requirement from bnb, all quantized parameters need to have that value set to False whereas the default is True :/ I can open an issue on bnb if this is a bug

Otherwise you get

  File "/home/younes_huggingface_co/miniconda3/envs/fix-test/lib/python3.9/site-packages/bitsandbytes/nn/modules.py", line 179, in to
    return self.cuda(device)
  File "/home/younes_huggingface_co/miniconda3/envs/fix-test/lib/python3.9/site-packages/bitsandbytes/nn/modules.py", line 158, in cuda
    self.data = w_4bit
RuntimeError: data set to a tensor that requires gradients must be floating point or complex dtype

[amyeroberts]

Yes, even if it's not a bug, then it would be good to get clarification about this.

Will this affectively set these layers to non-trainable even if they were trainable before?

[younesbelkada]

Hm I am not sure here, in any case, quantized layers cannot be trained in any case as this is not supported
I have added more clarifications here: ba1b10f

[younesbelkada]

Thanks for the review, I added a test that should be applicable to other checkpoints as well, as they use the same arch

[amyeroberts]

Thanks for fixing and iterating!