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,14 +30,17 @@ from transformers import (
pipeline,
)
from transformers.testing_utils import (
apply_skip_if_not_implemented,
is_accelerate_available,
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,
)
@@ -77,10 +80,14 @@ if is_torch_available():
return self.module(input, *args, **kwargs) + self.adapter(input)
if is_bitsandbytes_available():
import bitsandbytes as bnb
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
class BaseMixedInt8Test(unittest.TestCase):
# We keep the constants inside the init function and model loading inside setUp function
@@ -108,6 +115,7 @@ class BaseMixedInt8Test(unittest.TestCase):
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
@apply_skip_if_not_implemented
class MixedInt8Test(BaseMixedInt8Test):
def setUp(self):
super().setUp()
@@ -240,7 +248,6 @@ class MixedInt8Test(BaseMixedInt8Test):
r"""
A simple test to check if `llm_int8_skip_modules` works as expected
"""
import bitsandbytes as bnb
quantization_config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["classifier"])
seq_classification_model = AutoModelForSequenceClassification.from_pretrained(
@@ -263,7 +270,9 @@ class MixedInt8Test(BaseMixedInt8Test):
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_8bit.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = self.model_8bit.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)
@@ -280,7 +289,7 @@ class MixedInt8Test(BaseMixedInt8Test):
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_8bit_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)
@@ -298,7 +307,9 @@ class MixedInt8Test(BaseMixedInt8Test):
model_8bit.dequantize()
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_8bit.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model_8bit.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)
@@ -319,8 +330,10 @@ class MixedInt8Test(BaseMixedInt8Test):
def test_device_and_dtype_assignment(self):
r"""
Test whether trying to cast (or assigning a device to) a model after converting it in 8-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 8-bit precision will raise an appropriate error.
The test ensures that such operations are prohibited on 8-bit models
to prevent invalid conversions.
"""
with self.assertRaises(ValueError):
# Tries with `str`
@@ -332,21 +345,21 @@ class MixedInt8Test(BaseMixedInt8Test):
with self.assertRaises(ValueError):
# Tries with a `device`
self.model_8bit.to(torch.device("cuda:0"))
self.model_8bit.to(torch.device(torch_device))
with self.assertRaises(ValueError):
# Tries with a `device`
# Tries to cast the 8-bit model to float32 using `float()`
self.model_8bit.float()
with self.assertRaises(ValueError):
# Tries with a `device`
# Tries to cast the 4-bit model to float16 using `half()`
self.model_8bit.half()
# Test if we did not break anything
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 this does not throw an error
_ = self.model_fp16.to("cpu")
@@ -385,7 +398,9 @@ class MixedInt8Test(BaseMixedInt8Test):
# generate
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_from_saved.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model_from_saved.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)
@@ -410,7 +425,9 @@ class MixedInt8Test(BaseMixedInt8Test):
# generate
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_from_saved.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model_from_saved.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)
@@ -435,7 +452,9 @@ class MixedInt8Test(BaseMixedInt8Test):
# generate
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model_from_saved.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model_from_saved.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)
@@ -455,7 +474,7 @@ class MixedInt8Test(BaseMixedInt8Test):
# generate
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model.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)
@@ -463,7 +482,7 @@ class MixedInt8Test(BaseMixedInt8Test):
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@require_torch_gpu_if_bnb_not_multi_backend_enabled
@slow
class MixedInt8T5Test(unittest.TestCase):
@classmethod
@@ -494,14 +513,14 @@ class MixedInt8T5Test(unittest.TestCase):
# test with `google-t5/t5-small`
model = T5ForConditionalGeneration.from_pretrained(self.model_name, load_in_8bit=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_8bit=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
@@ -511,7 +530,6 @@ class MixedInt8T5Test(unittest.TestCase):
`flan-t5-small` uses `T5DenseGatedActDense` whereas `google-t5/t5-small` uses `T5DenseReluDense`. We need to test
both cases.
"""
import bitsandbytes as bnb
from transformers import T5ForConditionalGeneration
@@ -521,14 +539,14 @@ class MixedInt8T5Test(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.Linear8bitLt))
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_8bit=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)
def test_inference_with_keep_in_fp32_serialized(self):
@@ -538,7 +556,6 @@ class MixedInt8T5Test(unittest.TestCase):
`flan-t5-small` uses `T5DenseGatedActDense` whereas `google-t5/t5-small` uses `T5DenseReluDense`. We need to test
both cases.
"""
import bitsandbytes as bnb
from transformers import T5ForConditionalGeneration
@@ -553,14 +570,14 @@ class MixedInt8T5Test(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.Linear8bitLt))
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_8bit=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)
@@ -614,6 +631,7 @@ class MixedInt8ModelClassesTest(BaseMixedInt8Test):
self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter)
@apply_skip_if_not_implemented
class MixedInt8TestPipeline(BaseMixedInt8Test):
def setUp(self):
super().setUp()
@@ -623,7 +641,8 @@ class MixedInt8TestPipeline(BaseMixedInt8Test):
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()
@@ -648,6 +667,7 @@ class MixedInt8TestPipeline(BaseMixedInt8Test):
@require_torch_multi_gpu
@apply_skip_if_not_implemented
class MixedInt8TestMultiGpu(BaseMixedInt8Test):
def setUp(self):
super().setUp()
@@ -669,11 +689,14 @@ class MixedInt8TestMultiGpu(BaseMixedInt8Test):
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)
@require_torch_multi_gpu
@apply_skip_if_not_implemented
class MixedInt8TestCpuGpu(BaseMixedInt8Test):
def setUp(self):
super().setUp()
@@ -683,7 +706,7 @@ class MixedInt8TestCpuGpu(BaseMixedInt8Test):
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
# Check the exactness of the results
output_parallel = model.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_parallel = model.generate(input_ids=encoded_input["input_ids"].to(torch_device), max_new_tokens=10)
# Get the generation
output_text = self.tokenizer.decode(output_parallel[0], skip_special_tokens=True)
@@ -819,6 +842,7 @@ class MixedInt8TestCpuGpu(BaseMixedInt8Test):
self.check_inference_correctness(model_8bit)
@apply_skip_if_not_implemented
class MixedInt8TestTraining(BaseMixedInt8Test):
def setUp(self):
self.model_name = "facebook/opt-350m"
@@ -831,7 +855,10 @@ class MixedInt8TestTraining(BaseMixedInt8Test):
# Step 1: freeze all parameters
model = AutoModelForCausalLM.from_pretrained(self.model_name, load_in_8bit=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
@@ -847,10 +874,10 @@ class MixedInt8TestTraining(BaseMixedInt8Test):
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()
@@ -862,6 +889,7 @@ class MixedInt8TestTraining(BaseMixedInt8Test):
self.assertTrue(module.weight.grad is None)
@apply_skip_if_not_implemented
class MixedInt8GPT2Test(MixedInt8Test):
model_name = "openai-community/gpt2-xl"
EXPECTED_RELATIVE_DIFFERENCE = 1.8720077507258357
@@ -870,6 +898,9 @@ class MixedInt8GPT2Test(MixedInt8Test):
EXPECTED_OUTPUTS.add("Hello my name is John Doe, and I'm a fan of the")
# Expected values on a A10
EXPECTED_OUTPUTS.add("Hello my name is John Doe, and I am a member of the")
# Expected values on Intel CPU
EXPECTED_OUTPUTS.add("Hello my name is John Doe. I am a man. I am")
EXPECTED_OUTPUTS.add("Hello my name is John, and I'm a writer. I'm")
def test_int8_from_pretrained(self):
r"""
@@ -887,6 +918,6 @@ class MixedInt8GPT2Test(MixedInt8Test):
# generate
encoded_input = self.tokenizer(self.input_text, return_tensors="pt")
output_sequences = model.generate(input_ids=encoded_input["input_ids"].to(0), max_new_tokens=10)
output_sequences = model.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)