saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

FIX / Quantization: Add extra validation for bnb config (#31135)

Browse source 
add validation for bnb config
This commit is contained in:
Younes Belkada 2024-05-30 11:45:03 +02:00 committed by GitHub
parent 2b9e252b16
commit 5e5c4d629d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 11 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/transformers/utils/quantization_config.py
View file

@ -383,6 +383,10 @@ class BitsAndBytesConfig(QuantizationConfigMixin):
if bnb_4bit_quant_storage is None:
self.bnb_4bit_quant_storage = torch.uint8
elif isinstance(bnb_4bit_quant_storage, str):
if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]:
raise ValueError(
"`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') "
)
self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage)
elif isinstance(bnb_4bit_quant_storage, torch.dtype):
self.bnb_4bit_quant_storage = bnb_4bit_quant_storage

7
tests/quantization/bnb/test_4bit.py
View file

@ -303,6 +303,13 @@ class Bnb4BitTest(Base4bitTest):
model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-small", load_in_4bit=True, device_map="auto")
self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.float32)
def test_bnb_4bit_wrong_config(self):
r"""
Test whether creating a bnb config with unsupported values leads to errors.
"""
with self.assertRaises(ValueError):
_ = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_storage="add")
@require_bitsandbytes
@require_accelerate