Quantizing a 7B Language Model to 4-bit NF4

This project demonstrates how to load and run the 7-billion-parameter language model mistralai/Mistral-7B-v0.1 using 4-bit NormalFloat (NF4) quantization with bitsandbytes. It also includes a benchmark for measuring model loading time, GPU memory usage, and text-generation speed.

Quantization reduces the amount of memory required to run a model. While a 7B model typically needs around 14 GB just for its FP16 weights, a 4-bit version can use roughly 4 GB for the quantized weights and related data. Actual GPU memory usage will be higher and depends on the hardware, input length, and generation settings.

Overview

The notebook focuses on three tasks:

1. Configure 4-bit NF4 quantization with BitsAndBytesConfig.
2. Load the tokenizer and quantized model, then generate text from a prompt.
3. Benchmark a quantized model and, optionally, an FP16 model by measuring:
   • Model loading time
   • Peak GPU memory usage
   • Text-generation speed in tokens per second

The benchmark runs the 4-bit configuration by default to reduce the chance of running out of GPU memory. Enable the FP16 comparison only if your GPU has enough available memory. Around 16 GB of VRAM may be a reasonable starting point for this model, but the exact requirement depends on your workload and setup.

Features

Feature	Description
4-bit NF4 quantization	Uses the NF4 data type, designed for quantizing model weights that are approximately normally distributed.
Double quantization	Quantizes the quantization constants to reduce memory use further.
Higher-precision computation	Uses bfloat16 for computation when supported by the GPU. float16 can be used on compatible hardware that does not support bfloat16.
Automatic device mapping	Uses device_map="auto" to place model layers on available devices when supported by the installed libraries and hardware.
Performance benchmarking	Measures model loading time, peak allocated GPU memory, and generation speed.

Requirements

• Python 3.9 or newer
• A CUDA-capable GPU is recommended
• PyTorch installed with CUDA support
• Enough available GPU memory for the model and the chosen generation settings
• Internet access to download the model from Hugging Face

A GPU with 6–8 GB of VRAM may be able to run the quantized model with modest settings, but this is not guaranteed. Longer inputs, longer outputs, and other GPU workloads require additional memory.

Installation

Install the dependencies in your Python environment:

pip install -U torch transformers accelerate "bitsandbytes>=0.46.1"

If you are using a notebook, run the command in a code cell. Make sure that your PyTorch installation matches your CUDA environment.

To check whether bitsandbytes is installed, run:

python -c "import bitsandbytes as bnb; print(bnb.__version__)"

If the model repository requires authentication in your environment, sign in to Hugging Face before downloading it.

Usage

1. Load the model with 4-bit quantization

The following example loads Mistral-7B-v0.1 in 4-bit NF4 and generates a short response.

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)

model_id = "mistralai/Mistral-7B-v0.1"

# Use bfloat16 only when your GPU supports it; otherwise use float16.
compute_dtype = (
    torch.bfloat16
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    else torch.float16
)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
)
model.eval()

prompt = "Explain quantum computing in one short sentence."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=False,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

The decoded output includes the original prompt as well as the generated text. This is expected for a standard causal language model.

2. Run the benchmark

The notebook's measure_model_performance() function measures the model's loading time, peak GPU memory use, and generation speed. Use the following code after defining that function in the notebook:

int4_results = measure_model_performance(is_quantized=True)
print(int4_results)

Example output might look like this:

Quantized (4-bit) Model
Model Load Time: 12.34 seconds
Peak VRAM Used: 4.12 GB
Generation Speed: 18.45 tokens/sec

These values are illustrative only. Your results will vary with the GPU, drivers, software versions, input length, and system load.

To benchmark the FP16 version, uncomment the following lines only if you have enough GPU memory:

# Run only if your GPU has sufficient available VRAM.
# fp16_results = measure_model_performance(is_quantized=False)
# print(fp16_results)

How the quantization settings work

Parameter	Purpose
load_in_4bit=True	Loads the model weights using 4-bit quantization.
bnb_4bit_quant_type="nf4"	Selects NF4, a quantization data type designed for normally distributed weights.
bnb_4bit_compute_dtype=torch.bfloat16	Sets the compute data type used for operations after weights are dequantized as needed. Use a supported data type for your GPU.
bnb_4bit_use_double_quant=True	Quantizes the quantization constants as well, which can reduce memory use.

Approximate model-weight memory

Precision	Approximate memory for weights	Notes
FP32	28 GB	Full-precision weights
FP16 or BF16	14 GB	Half-precision weights
4-bit NF4	Around 4 GB	Quantized weights plus some quantization overhead; total runtime memory is higher

These are estimates based on the model's parameter count. They do not include all memory used by activations, the key-value cache, temporary tensors, or other runtime components.

Benchmark methodology

The measure_model_performance() function is intended to perform the following steps:

1. Clear unused GPU cache and reset peak-memory statistics with PyTorch.
2. Load the model and record the loading time.
3. Record peak GPU memory using torch.cuda.max_memory_allocated().
4. Run a short warm-up generation that is excluded from the timing.
5. Generate a fixed number of new tokens and calculate tokens per second.
6. Release model objects and clear the CUDA cache where possible.

For a fair comparison, use the same prompt, generation settings, software environment, and GPU for each run. Note that torch.cuda.max_memory_allocated() measures memory allocated by PyTorch, so it may not include every type of GPU memory use reported by system-monitoring tools.

Troubleshooting

Issue	What to try
CUDA out-of-memory error	Close other GPU workloads, reduce the input or output length, and avoid running the FP16 benchmark.
bitsandbytes installation or CUDA error	Check the installed PyTorch and CUDA environment, then update bitsandbytes using the installation command above.
device_map="auto" does not work	Confirm that accelerate is installed and that your hardware and library versions support the selected device mapping.
Generation is slow	Check GPU utilization, reduce the input length, and make sure the model is actually using the GPU. Quantization can add some computation overhead.
Model download is slow or fails	Check your network connection and Hugging Face authentication. You can also configure the Hugging Face cache directory using HF_HOME.
bfloat16 is unsupported	Use torch.float16 if it is supported by your GPU.

Project structure

.
├── README.md
├── quantizing_7b_llm.ipynb   # Main notebook
└── requirements.txt          # Optional dependency list

An optional requirements.txt could contain:

torch>=2.0.0
transformers>=4.35.0
accelerate>=0.24.0
bitsandbytes>=0.46.1

These are minimum-version examples, not a fully tested compatibility lockfile. For reproducible results, record the versions that work in your environment.

Benefits and trade-offs

Benefits

• Substantially reduces model-weight memory use compared with FP16.
• Makes it possible to experiment with a 7B model on some consumer GPUs.
• Often preserves useful model quality for inference, although the effect varies by model and task.

Trade-offs

• Output quality may decrease slightly compared with higher-precision weights.
• Generation speed can be slower in some environments because of quantization and dequantization overhead.
• Memory use still depends on input length, generated tokens, and other runtime settings.
• bitsandbytes support depends on the operating system, hardware, and software environment.

License

The code in this project is intended to be released under the MIT License. Include an MIT license file in the repository if you choose that license.

The mistralai/Mistral-7B-v0.1 model has its own license terms. Review the model card on Hugging Face before using or redistributing the model.

References

• Mistral-7B-v0.1 model card
• bitsandbytes documentation
• QLoRA paper
• Transformers bitsandbytes quantization guide