Enable BNB multi-backend support (#31098)

* enable cpu bnb path

* fix style

* fix code style

* fix 4 bit path

* Update src/transformers/utils/import_utils.py

Co-authored-by: Aarni Koskela <akx@iki.fi>

* add multi backend refactor tests

* fix style

* tweak 4bit quantizer + fix corresponding tests

* tweak 8bit quantizer + *try* fixing corresponding tests

* fix dequant bnb 8bit

* account for Intel CPU in variability of expected outputs

* enable cpu and xpu device map

* further tweaks to account for Intel CPU

* fix autocast to work with both cpu + cuda

* fix comments

* fix comments

* switch to testing_utils.torch_device

* allow for xpu in multi-gpu tests

* fix tests 4bit for CPU NF4

* fix bug with is_torch_xpu_available needing to be called as func

* avoid issue where test reports attr err due to other failure

* fix formatting

* fix typo from resolving of merge conflict

* polish based on last PR review

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

* fix CI

* Update src/transformers/integrations/integration_utils.py

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

* Update src/transformers/integrations/integration_utils.py

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

* fix error log

* fix error msg

* add \n in error log

* make quality

* rm bnb cuda restriction in doc

* cpu model don't need dispatch

* fix doc

* fix style

* check cuda avaliable in testing

* fix tests

* Update docs/source/en/model_doc/chameleon.md

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

* Update docs/source/en/model_doc/llava_next.md

Co-authored-by: Aarni Koskela <akx@iki.fi>

* Update tests/quantization/bnb/test_4bit.py

Co-authored-by: Aarni Koskela <akx@iki.fi>

* Update tests/quantization/bnb/test_4bit.py

Co-authored-by: Aarni Koskela <akx@iki.fi>

* fix doc

* fix check multibackends

* fix import sort

* remove check torch in bnb

* docs: update bitsandbytes references with multi-backend info

* docs: fix small mistakes in bnb paragraph

* run formatting

* reveret bnb check

* move bnb multi-backend check to import_utils

* Update src/transformers/utils/import_utils.py

Co-authored-by: Aarni Koskela <akx@iki.fi>

* fix bnb check

* minor fix for bnb

* check lib first

* fix code style

* Revert "run formatting"

This reverts commit ac108c6d6b34f45a5745a736ba57282405cfaa61.

* fix format

* give warning when bnb version is low and no cuda found]

* fix device assignment check to be multi-device capable

* address akx feedback on get_avlbl_dev fn

* revert partially, as we don't want the function that public, as docs would be too much (enforced)

---------

Co-authored-by: Aarni Koskela <akx@iki.fi>
Co-authored-by: Titus von Koeller <9048635+Titus-von-Koeller@users.noreply.github.com>
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
jiqing-feng
2024-09-24 17:40:56 +08:00
committed by GitHub
parent e15687fffe
commit 11c27dd331
20 changed files with 436 additions and 100 deletions

View File

@@ -30,12 +30,13 @@ from transformers import (
pipeline,
)
from transformers.testing_utils import (
apply_skip_if_not_implemented,
is_bitsandbytes_available,
is_torch_available,
require_accelerate,
require_bitsandbytes,
require_torch,
require_torch_gpu,
require_torch_gpu_if_bnb_not_multi_backend_enabled,
require_torch_multi_gpu,
slow,
torch_device,
@@ -85,7 +86,7 @@ if is_bitsandbytes_available():
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
class Base4bitTest(unittest.TestCase):
# We keep the constants inside the init function and model loading inside setUp function
@@ -111,6 +112,7 @@ class Base4bitTest(unittest.TestCase):
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
@apply_skip_if_not_implemented
class Bnb4BitTest(Base4bitTest):
def setUp(self):
super().setUp()
@@ -206,7 +208,7 @@ class Bnb4BitTest(Base4bitTest):
tok = AutoTokenizer.from_pretrained(model_id)
text = "Hello my name is"
input_ids = tok.encode(text, return_tensors="pt").to(0)
input_ids = tok.encode(text, return_tensors="pt").to(torch_device)
_ = model.generate(input_ids, max_new_tokens=30)
@@ -217,7 +219,9 @@ class Bnb4BitTest(Base4bitTest):
the same output across GPUs. So we'll generate few tokens (5-10) and check their output.
"""
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = self.model_4bit.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = self.model_4bit.generate(
input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10
)
self.assertIn(self.tokenizer.decode(output_sequences[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS)
@@ -234,7 +238,7 @@ class Bnb4BitTest(Base4bitTest):
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_4bit_from_config.generate(
input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10
input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10
)
self.assertIn(self.tokenizer.decode(output_sequences[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS)
@@ -252,7 +256,9 @@ class Bnb4BitTest(Base4bitTest):
model_4bit.dequantize()
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_4bit.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model_4bit.generate(
input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10
)
self.assertIn(self.tokenizer.decode(output_sequences[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS)
@@ -267,15 +273,18 @@ class Bnb4BitTest(Base4bitTest):
self.assertEqual(self.model_4bit.device.type, "cpu")
self.assertAlmostEqual(self.model_4bit.get_memory_footprint(), mem_before)
# Move back to CUDA device
self.model_4bit.to(0)
self.assertEqual(self.model_4bit.device, torch.device(0))
self.assertAlmostEqual(self.model_4bit.get_memory_footprint(), mem_before)
if torch.cuda.is_available():
# Move back to CUDA device
self.model_4bit.to("cuda")
self.assertEqual(self.model_4bit.device.type, "cuda")
self.assertAlmostEqual(self.model_4bit.get_memory_footprint(), mem_before)
def test_device_and_dtype_assignment(self):
r"""
Test whether trying to cast (or assigning a device to) a model after converting it in 4-bit will throw an error.
Checks also if other models are casted correctly.
Test whether attempting to change the device or cast the dtype of a model
after converting it to 4-bit precision will raise an appropriate error.
The test ensures that such operations are prohibited on 4-bit models
to prevent invalid conversions.
"""
# Moving with `to` or `cuda` is not supported with versions < 0.43.2.
@@ -297,25 +306,24 @@ class Bnb4BitTest(Base4bitTest):
self.model_4bit.to(torch.float16)
with self.assertRaises(ValueError):
# Tries with a `dtype` and `device`
self.model_4bit.to(device="cuda:0", dtype=torch.float16)
with self.assertRaises(ValueError):
# Tries with a cast
# Tries to cast the 4-bit model to float32 using `float()`
self.model_4bit.float()
with self.assertRaises(ValueError):
# Tries with a cast
# Tries to cast the 4-bit model to float16 using `half()`
self.model_4bit.half()
# Test if we did not break anything
self.model_4bit.to(torch.device(torch_device))
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
self.model_fp16 = self.model_fp16.to(torch.float32)
_ = self.model_fp16.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
_ = self.model_fp16.generate(input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10)
# Check that this does not throw an error
_ = self.model_fp16.cuda()
if torch.cuda.is_available():
# Check that this does not throw an error
_ = self.model_fp16.cuda()
# Check this does not throw an error
_ = self.model_fp16.to("cpu")
@@ -344,8 +352,9 @@ class Bnb4BitTest(Base4bitTest):
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
@apply_skip_if_not_implemented
class Bnb4BitT5Test(unittest.TestCase):
@classmethod
def setUpClass(cls):
@@ -375,14 +384,14 @@ class Bnb4BitT5Test(unittest.TestCase):
# test with `google-t5/t5-small`
model = T5ForConditionalGeneration.from_pretrained(self.model_name, load_in_4bit=True, device_map="auto")
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(0)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(torch_device)
_ = model.generate(**encoded_input)
# test with `flan-t5-small`
model = T5ForConditionalGeneration.from_pretrained(
self.dense_act_model_name, load_in_4bit=True, device_map="auto"
)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(0)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(torch_device)
_ = model.generate(**encoded_input)
T5ForConditionalGeneration._keep_in_fp32_modules = modules
@@ -400,17 +409,18 @@ class Bnb4BitT5Test(unittest.TestCase):
# there was a bug with decoders - this test checks that it is fixed
self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q, bnb.nn.Linear4bit))
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(0)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(torch_device)
_ = model.generate(**encoded_input)
# test with `flan-t5-small`
model = T5ForConditionalGeneration.from_pretrained(
self.dense_act_model_name, load_in_4bit=True, device_map="auto"
)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(0)
encoded_input = self.tokenizer(self.input_text, return_tensors="pt").to(torch_device)
_ = model.generate(**encoded_input)
@apply_skip_if_not_implemented
class Classes4BitModelTest(Base4bitTest):
def setUp(self):
super().setUp()
@@ -460,6 +470,7 @@ class Classes4BitModelTest(Base4bitTest):
self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter)
@apply_skip_if_not_implemented
class Pipeline4BitTest(Base4bitTest):
def setUp(self):
super().setUp()
@@ -469,7 +480,8 @@ class Pipeline4BitTest(Base4bitTest):
TearDown function needs to be called at the end of each test to free the GPU memory and cache, also to
avoid unexpected behaviors. Please see: https://discuss.pytorch.org/t/how-can-we-release-gpu-memory-cache/14530/27
"""
del self.pipe
if hasattr(self, "pipe"):
del self.pipe
gc.collect()
torch.cuda.empty_cache()
@@ -484,7 +496,12 @@ class Pipeline4BitTest(Base4bitTest):
self.pipe = pipeline(
"text-generation",
model=self.model_name,
model_kwargs={"device_map": "auto", "load_in_4bit": True, "torch_dtype": torch.float16},
model_kwargs={
"device_map": "auto",
"load_in_4bit": True,
# float16 isn't supported on CPU, use bfloat16 instead
"torch_dtype": torch.bfloat16 if torch_device == "cpu" else torch.float16,
},
max_new_tokens=self.MAX_NEW_TOKENS,
)
@@ -494,6 +511,7 @@ class Pipeline4BitTest(Base4bitTest):
@require_torch_multi_gpu
@apply_skip_if_not_implemented
class Bnb4bitTestMultiGpu(Base4bitTest):
def setUp(self):
super().setUp()
@@ -515,10 +533,13 @@ class Bnb4bitTestMultiGpu(Base4bitTest):
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
# Second real batch
output_parallel = model_parallel.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_parallel = model_parallel.generate(
input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10
)
self.assertIn(self.tokenizer.decode(output_parallel[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS)
@apply_skip_if_not_implemented
class Bnb4BitTestTraining(Base4bitTest):
def setUp(self):
self.model_name = "facebook/opt-350m"
@@ -531,7 +552,10 @@ class Bnb4BitTestTraining(Base4bitTest):
# Step 1: freeze all parameters
model = AutoModelForCausalLM.from_pretrained(self.model_name, load_in_4bit=True)
self.assertEqual(set(model.hf_device_map.values()), {torch.cuda.current_device()})
if torch.cuda.is_available():
self.assertEqual(set(model.hf_device_map.values()), {torch.cuda.current_device()})
else:
self.assertTrue(all(param.device.type == "cpu" for param in model.parameters()))
for param in model.parameters():
param.requires_grad = False # freeze the model - train adapters later
@@ -547,10 +571,10 @@ class Bnb4BitTestTraining(Base4bitTest):
module.v_proj = LoRALayer(module.v_proj, rank=16)
# Step 3: dummy batch
batch = self.tokenizer("Test batch ", return_tensors="pt").to(0)
batch = self.tokenizer("Test batch ", return_tensors="pt").to(torch_device)
# Step 4: Check if the gradient is not None
with torch.cuda.amp.autocast():
with torch.autocast(torch_device):
out = model.forward(**batch)
out.logits.norm().backward()
@@ -562,6 +586,7 @@ class Bnb4BitTestTraining(Base4bitTest):
self.assertTrue(module.weight.grad is None)
@apply_skip_if_not_implemented
class Bnb4BitGPT2Test(Bnb4BitTest):
model_name = "openai-community/gpt2-xl"
EXPECTED_RELATIVE_DIFFERENCE = 3.3191854854152187
@@ -570,8 +595,9 @@ class Bnb4BitGPT2Test(Bnb4BitTest):
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
@apply_skip_if_not_implemented
class BaseSerializationTest(unittest.TestCase):
model_name = "facebook/opt-125m"
input_text = "Mars colonists' favorite meals are"
@@ -635,7 +661,9 @@ class BaseSerializationTest(unittest.TestCase):
d1[k].quant_state.as_dict().values(),
):
if isinstance(v0, torch.Tensor):
self.assertTrue(torch.equal(v0, v1.to(v0.device)))
# The absmax will not be saved in the quant_state when using NF4 in CPU
if v0.numel() != 0:
self.assertTrue(torch.equal(v0, v1.to(v0.device)))
else:
self.assertTrue(v0 == v1)
@@ -659,6 +687,7 @@ class BaseSerializationTest(unittest.TestCase):
)
@apply_skip_if_not_implemented
class ExtendedSerializationTest(BaseSerializationTest):
"""
tests more combinations of parameters
@@ -706,8 +735,9 @@ class GPTSerializationTest(BaseSerializationTest):
@require_bitsandbytes
@require_accelerate
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
@apply_skip_if_not_implemented
class Bnb4BitTestBasicConfigTest(unittest.TestCase):
def test_load_in_4_and_8_bit_fails(self):
with self.assertRaisesRegex(ValueError, "load_in_4bit and load_in_8bit are both True"):