Fine-Tuning Mistral 7B using QLoRA with PyTorch pt. 1: The Model | ML Engineering

    

Hi All

Today we're working with a popular and slightly bigger model than our previous example. Mistral 7B is capable of chat and light coding tasks, for older hardware it's a winner for sure. 

Here's a complete, runnable example of fine-tuning Mistral 7B using QLoRA with the peft, transformers, and bitsandbytes libraries. This example assumes you're working with a single GPU (eg. an A100 or similar).

First install the required packages:

pip install -q bitsandbytes datasets accelerate peft transformers trl

View full script below, also available here:  

Full breakdown of the script above, block-by-block.

1.    Dataset Loading

dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")

*    Loads a preprocessed instruction-following dataset (Guanco, derived from OpenAssistant).

*    split="train" selects the training portion

*    The dataset is in a conversational format (instruction & response pairs).

*    Each entry has a text field combining the prompt and response.

 

2. Model and Tokenizer Setup

model_id = "mistralai/Mistral-7B-v0.1"
bnb_config = BitsAndBytesConfig(...)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

*    Configures the model and tokenizer with 4-bit quantization.

*    Key Components:

            *    bnb_config:

                        *    load_in_4bit=True: Enables 4-bit quantization.

                        *    bnb_4bit_quant_type="nf4": Uses NormalFloat (optimal for LLMs)

                        *     bnb_4bit_compute_dtype=torch.bfloat16: Uses BF16 for computations.

*    device_map="auto": Automatically distributes layers across GPU's (or uses CPU if needed).

*    Tokenizer: Sets the padding token to the EOS token for batching.   

        

3. LoRA Configuration

peft_config = LoraConfig(
    r=16,                # Rank of LoRA matrices
    lora_alpha=32,       # Scaling factor
    lora_dropout=0.05,   # Dropout for regularization
    bias="none",         # No bias adaptation
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Targeted layers
)

3.1    Defines the LoRA (Low Rank Adaptation) parameters.

3.2    Key parameters:

        *    r: Rank of the low-rank matrices (smaller = more efficient)

        *    target_modules: Only adapts attention projection layers (nt MLP layers).

        *    task_type: Specifies causal language modelling (for autoregressive tasks).

 

4. Training Arguments  

training_args = TrainingArguments(
    output_dir="mistral-7b-qlora-finetuned",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    optim="paged_adamw_8bit",  # Memory-efficient AdamW
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    save_strategy="epoch",
    logging_steps=10,
    num_train_epochs=3,
    max_steps=100,  # For demo (remove in practice)
    fp16=True,      # Mixed precision training
)

*    Configures training hyperparameters and behaviour.

*    Key Parameters:

                    *    Batch Size: 4 per device, with gradient accumulation to simulate 16.

                    *    Optimizer: paged_adamw_8bit (8-bit AdamW with memory paging).

                    *    Learning Rate: 2e-4 (typical  for LoRA).

                    *    Scheduler: Cosine annealing for stable convergence.

                    *    Mixed Precision: fp16=True for faster training.

 

5. Trainer Initialization

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",  # Column in dataset containing text
    tokenizer=tokenizer,
    packing=True,  # Packs sequences to max length for efficiency
)

*    Sets up the training loop with PEFT (Parameter-Efficient Fine-Tuning)  

*    Key Features:

            *    SFTTrainer: A subclass of Trainer optimized for instruction tuning.

            *    packing=True: Combines multiple short sequences into one long sequence (reduces padding overhead).

 

6. Training Execution

trainer.train()

*    Starts the fine-tuning process

*    Loads batches from dataset

*    Applies LoRA updates to the frozen base model.

*    Logs metrics (loss, etc.) every logging_steps. 

 

7. Saving the Model

trainer.model.save_pretrained("mistral-7b-qlora-finetuned")

*    Saves the fine-tuned model and LoRA weights.

*    Output: directory with:

        *    adapter_config.json (LoRA config).

        *    adapter_model.bin (LoRA weights)

        *    config.json (base model config). 

Final Notes:

1.    Quantization Impact:

        *    The base model is loaded in 4-bit, but gradients are computed in BF16 (via compute_dtype).

        *    LoRA adapters are trained in full precision (no quantization).

2.    Memory Savings:

        *    QLoRA reduces memory usage by ~90% compared to full fine-tuning (eg. 40GB vs 80GB for Mistral 7B).

3.    Customization:

        *    For your own dataset, replace Dataset_text_field with your column name.

        *    Adjust target_modules if you want to adapt other layers (eg. MLP).

 The full source is available here. Stay tuned for part 2 where we'll dive into some Ops on the Host (ahem Linux). That's all for now folks.

 

Later

 

-    Ed