#
Model quantization

##
bitsandbytes Integration

Accelerate brings `bitsandbytes` quantization to your model. You can now load any pytorch model in 8-bit or 4-bit with a few lines of code.

If you want to use Transformers models with `bitsandbytes`, you should follow this documentation.

To learn more about how the `bitsandbytes` quantization works, check out the blog posts on 8-bit quantization and 4-bit quantization.

###
Pre-Requisites

You will need to install the following requirements:

- Install `bitsandbytes` library

    pip install bitsandbytes

For non-cuda devices, you can refer to the bitsandbytes installation guide here.

- Install latest `accelerate` from source

    pip install git+https://github.com/huggingface/accelerate.git

- Install `minGPT` and `huggingface_hub` to run examples

    git clone https://github.com/karpathy/minGPT.git
    pip install minGPT/
    pip install huggingface_hub

###
How it works

First, we need to initialize our model. To save memory, we can initialize an empty model using the context manager init_empty_weights().

Let’s take the GPT2 model from minGPT library.

    from accelerate import init_empty_weights
    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)

Then, we need to get the path to the weights of your model. The path can be the state_dict file (e.g. “pytorch_model.bin”) or a folder containing the sharded checkpoints.

    from huggingface_hub import snapshot_download
    weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")

Finally, you need to set your quantization configuration with BnbQuantizationConfig.

Here’s an example for 8-bit quantization:

    from accelerate.utils import BnbQuantizationConfig
    bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)

Here’s an example for 4-bit quantization:

    from accelerate.utils import BnbQuantizationConfig
    bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")

To quantize your empty model with the selected configuration, you need to use load_and_quantize_model().

    from accelerate.utils import load_and_quantize_model
    quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config)

###
Saving and loading 8-bit model

You can save your 8-bit model with accelerate using save_model().

    from accelerate import Accelerator
    accelerate = Accelerator()
    new_weights_location = "path/to/save_directory"
    accelerate.save_model(quantized_model, new_weights_location)

    quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")

Note that 4-bit model serialization is currently not supported.

###
Offload modules to cpu and disk

You can offload some modules to cpu/disk if you don’t have enough space on the GPU to store the entire model on your GPUs. This uses big model inference under the hood. Check this documentation for more details.

For 8-bit quantization, the selected modules will be converted to 8-bit precision.

For 4-bit quantization, the selected modules will be kept in `torch_dtype` that the user passed in `BnbQuantizationConfig`. We will add support to convert these offloaded modules in 4-bit when 4-bit serialization will be possible.

You just need to pass a custom `device_map` in order to offload modules on cpu/disk. The offload modules will be dispatched on the GPU when needed. Here’s an example :

    device_map = {
        "transformer.wte": 0,
        "transformer.wpe": 0,
        "transformer.drop": 0,
        "transformer.h": "cpu",
        "transformer.ln_f": "disk",
        "lm_head": "disk",
    }

###
Fine-tune a quantized model

It is not possible to perform pure 8bit or 4bit training on these models. However, you can train these models by leveraging parameter efficient fine tuning methods (PEFT) and train for example adapters on top of them. Please have a look at peft library for more details.

Currently, you can’t add adapters on top of any quantized model. However, with the official support of adapters with Transformers models, you can fine-tune quantized models. If you want to finetune a Transformers model , follow this documentation instead. Check out this demo on how to fine-tune a 4-bit Transformers model.

Note that you don’t need to pass `device_map` when loading the model for training. It will automatically load your model on your GPU. Please note that `device_map=auto` should be used for inference only.

###
Example demo - running GPT2 1.5b on a Google Colab

Check out the Google Colab demo for running quantized models on a GPT2 model. The GPT2-1.5B model checkpoint is in FP32 which uses 6GB of memory. After quantization, it uses 1.6GB with 8-bit modules and 1.2GB with 4-bit modules.

Update on GitHub