bitsandbytes is widely integrated with many of the libraries in the Hugging Face and wider PyTorch ecosystem. This guide provides a brief overview of the integrations and how to use bitsandbytes with them. For more details, you should refer to the linked documentation for each library.
Tip
Learn more in the bitsandbytes Transformers integration guide.
With Transformers, it's very easy to load any model in 4 or 8-bit and quantize them on the fly. To configure the quantization parameters, specify them in the [~transformers.BitsAndBytesConfig] class.
For example, to load and quantize a model to 4-bits and use the bfloat16 data type for compute:
Warning
bfloat16 is the ideal compute_dtype if your hardware supports it. While the default compute_dtype, float32, ensures backward compatibility (due to wide-ranging hardware support) and numerical stability, it is large and slows down computations. In contrast, float16 is smaller and faster but can lead to numerical instabilities. bfloat16 combines the best aspects of both; it offers the numerical stability of float32 and the reduced memory footprint and speed of a 16-bit data type. Check if your hardware supports bfloat16 and configure it using the bnb_4bit_compute_dtype parameter in [~transformers.BitsAndBytesConfig]!
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
model_4bit = AutoModelForCausalLM.from_pretrained(
"bigscience/bloom-1b7",
device_map=device_map,
quantization_config=quantization_config,
)You can use any of the 8-bit or paged optimizers with Transformers by passing them to the [~transformers.Trainer] class on initialization. All bitsandbytes optimizers are supported by passing the correct string in the [~transformers.TrainingArguments] optim parameter. For example, to load a [~bitsandbytes.optim.PagedAdamW32bit] optimizer:
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(
...,
optim="paged_adamw_32bit",
)
trainer = Trainer(model, training_args, ...)
trainer.train()Tip
Learn more in the bitsandbytes PEFT integration guide.
PEFT builds on the bitsandbytes Transformers integration, and extends it for training with a few more steps. Let's prepare the 4-bit model from the section above for training.
Call the [~peft.prepare_model_for_kbit_training] method to prepare the model for training. This only works for Transformers models!
from peft import prepare_model_for_kbit_training
model_4bit = prepare_model_for_kbit_training(model_4bit)Setup a [~peft.LoraConfig] to use QLoRA:
from peft import LoraConfig
config = LoraConfig(
r=16,
lora_alpha=8,
target_modules="all-linear",
lora_dropout=0.05
bias="none",
task_type="CAUSAL_LM"
)Now call the [~peft.get_peft_model] function on your model and config to create a trainable [PeftModel].
from peft import get_peft_model
model = get_peft_model(model_4bit, config)Tip
Learn more in the bitsandbytes Accelerate integration guide.
bitsandbytes is also easily usable from Accelerate and you can quantize any PyTorch model by passing a [~accelerate.utils.BnbQuantizationConfig] with your desired settings, and then calling the [~accelerate.utils.load_and_quantize_model] function to quantize it.
from accelerate import init_empty_weights
from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model
from mingpt.model import GPT
model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024
with init_empty_weights():
empty_model = GPT(model_config)
bnb_quantization_config = BnbQuantizationConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16, # optional
bnb_4bit_use_double_quant=True, # optional
bnb_4bit_quant_type="nf4" # optional
)
quantized_model = load_and_quantize_model(
empty_model,
weights_location=weights_location,
bnb_quantization_config=bnb_quantization_config,
device_map = "auto"
)bitsandbytes is available from:
- PyTorch Lightning, a deep learning framework for professional AI researchers and machine learning engineers who need maximal flexibility without sacrificing performance at scale.
- Lightning Fabric, a fast and lightweight way to scale PyTorch models without boilerplate.
Learn more in the bitsandbytes PyTorch Lightning integration guide.
bitsandbytes is integrated with Lit-GPT, a hackable implementation of state-of-the-art open-source large language models. Lit-GPT is based on Lightning Fabric, and it can be used for quantization during training, finetuning, and inference.
Learn more in the bitsandbytes Lit-GPT integration guide.
To learn in more detail about some of bitsandbytes integrations, take a look at the following blog posts: