Fine-Tuning LLM - A Step-by-Step Guide of Empowering Customization with Mistral-7B

Fine-tuning a large language model is a powerful way to tailor AI models to specific use cases, whether it's for medical, legal, or other domains. In this tutorial, we will explore the fine-tuning process using the Mistral-7B model and a step-by-step case study of fine-tuning for a specific task, generating prompts. Fine-tuning is a process that leverages pre-trained models to make them more domain-specific, and it can significantly improve the model's performance for your specific application.

Introduction

Fine-tuning involves taking a base model that has been pre-trained on a large corpus of text data and training it further on a smaller dataset related to your specific task. This process helps the model to adapt to the particular language and behavior required for your use case. In this tutorial, we'll discuss how to fine-tune the Mistral-7B model.

mistralai/Mistral-7B-v0.1 is a 7-billion-parameter language model engineered for superior performance and efficiency. Mistral 7B outperforms Llama 2 13B across all evaluated benchmarks, and Llama 1 34B in reasoning, mathematics, and code generation.

mistralai/Mistral-7B-v0.1 · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

If you have more powerful resource, the model upstage/SOLAR-0-70b-16bit is definitely worth trying.

upstage/SOLAR-0-70b-16bit · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

In this article, we will use Google Colab to fine-tune the Mistral-7B model with a dataset of Stable Diffusion preferred prompts.

Step 1: Choose the Base Large Language Model

Before diving into fine-tuning, you need to select the base model that you'll build upon. Hugging Face provides a leaderboard of open-source pre-trained models, and you can choose the one that best suits your needs. For this tutorial, we'll use the Mistral-7B model, which is a highly powerful and efficient model for conversation tasks.

Step 2: Prepare Training Data

One of the critical steps in fine-tuning is to gather and prepare your training dataset. The quality of your dataset directly impacts the performance of your fine-tuned model. There are two types of datasets you can use: public datasets available on platforms like Kaggle or data you collect privately. Your dataset doesn't need to be enormous; even 100 to 200 rows of data can be sufficient for fine-tuning. Demo training dataset for this article

You can also use the large language model itself to generate training data. For example, you can prompt it to generate user instructions and then collect the generated responses as training data. This can be done efficiently on platforms like Relevance AI, where you can run prompts in bulk.

Step 3: Set Up Google Colab

For fine-tuning the Mistral-7B model, we'll use Google Colab as our platform. Ensure that you have a compatible GPU for efficient training. For the free plan you can switch to T4 GPU, to speed up the process.

• Essentials

!pip install -Uqqq pip
!pip install -qqq bitsandbytes==0.39.0
!pip install -qqq torch==2.1.0
!pip install -qqq -U git+https://github.com/huggingface/transformers.git
!pip install -qqq -U git+https://github.com/huggingface/peft.git
!pip install -qqq -U git+https://github.com/huggingface/accelerate.git
!pip install -qqq datasets==2.12.0
!pip install -qqq loralib==0.1.1
!pip install -qqq einops==0.6.1

• Import libraries

# Import necessary Python modules and libraries

import json
import os
from pprint import pprint
import bitsandbytes as bnb
import torch
import torch.nn as nn
import transformers
from datasets import load_dataset
from huggingface_hub import notebook_login
from peft import (
    LoraConfig,
    PeftConfig,
    PeftModel,
    get_peft_model,
    prepare_model_for_kbit_training
)
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)

# Set the environment variable to specify which GPU to use
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

• Log in HuggingFace

# Log in to the Hugging Face Hub using notebook_login()

notebook_login()

Note: If you don’t have any HuggingFace token, create one

Step 4.1: Load the Base Mistral-7B Model

Next, you'll load the Mistral-7B model, which serves as the base for fine-tuning. This model can be further fine-tuned for your specific task. We'll be using a method called Low-Rank Adapters, which is an efficient way to fine-tune large language models.

Before fine-tuning, it's essential to test your base model to see how well it performs on your task. In our case, generating Stable Diffusion prompts. However, if the base model doesn't perform well, fine-tuning can help improve its performance.

Note: If you have low RAM issue, you can also try the sharded version. It shards the original model to 2GB maximum parts in order to reduce the RAM required when loading.

alexsherstinsky/Mistral-7B-v0.1-sharded · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

• Load Mistral-7B Model & Tokenizer

# Define the model name to be used
MODEL_NAME = "alexsherstinsky/Mistral-7B-v0.1-sharded"

# Create a configuration for the 'bitsandbytes' library
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load a pre-trained model using the specified model name, 'bitsandbytes' quantization config, and other options
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    trust_remote_code=True,
    quantization_config=bnb_config
)

# Create a tokenizer for the specified model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# Define a function to print the number of trainable parameters in the model
def print_trainable_parameters(model):
  trainable_params = 0
  all_param = 0
  for _, param in model.named_parameters():
    all_param += param.numel()
    if param.requires_grad:
      trainable_params += param.numel()
  print(
      f"trainable params: {trainable_params} || all params: {all_param} || trainables%: {100 * trainable_params / all_param}"
  )

# Enable gradient checkpointing in the model and prepare it for kbit training
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

# Create a configuration for the 'peft' library
config = LoraConfig(
    r=16,  # dimension of the updated matrices
    lora_alpha=64,  # parameter for scaling
    target_modules=[
    "q_proj",
    "up_proj",
    "o_proj",
    "k_proj",
    "down_proj",
    "gate_proj",
    "v_proj"],
    lora_dropout=0.1,  # dropout probability for layers
    bias="none",
    task_type="CAUSAL_LM",
)

# Update the model with 'peft' configurations
model = get_peft_model(model, config)

# Print the number of trainable parameters in the updated model
print_trainable_parameters(model)

Output can be:

trainable params: 41943040 || all params: 3794014208 || trainables%: 1.1055056122762943

(Optional) Step 4.2: Test the Base Mistral-7B Model

• Test prompt

# Define a prompt for text generation
prompt = """
Stable Diffusion prompt for a warrior wearing a futuristic armor
""".strip()

# Configure generation settings
generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config.temperature = 0.7
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id

• Get the result of the base model

# Measure the time taken for model inference
%%time
device = "cuda:0"

# Encode the prompt using the tokenizer and generate text
encoding = tokenizer(prompt, return_tensors="pt").to(device)
with torch.inference_mode():
  outputs = model.generate(
      input_ids = encoding.input_ids,
      attention_mask = encoding.attention_mask,
      generation_config = generation_config
  )

# Print the generated text by decoding the model output
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

• Demo output

Stable Diffusion prompt for a warrior wearing a futuristic armor.

I’m a big fan of the Stable Diffusion AI art generator. I’ve been using it for a while now and I’m always amazed by the results.

I’ve been using it to create art for my blog posts and I’m always impressed by the quality of the images it produces.

I’ve also been using it to create art for my own personal use and I’m always amazed by the results.
CPU times: user 1min 9s, sys: 17 s, total: 1min 26s
Wall time: 1min 32s

As you can see above, the result for “Stable Diffusion prompt for a warrior wearing a futuristic armor” by the original Mistral-7B makes no sense.

Step 5: Load Training Data

Load your prepared training data into Google Colab. The dataset should include pairs of user instructions and corresponding Stable Diffusion prompts. The data will be tokenized for further processing.

• Load training dataset

# Load a dataset from a CSV file
data = load_dataset("csv", data_files="SD_traning_prompt.csv")

# Access the loaded dataset
data

Output should be like this:

DatasetDict({
    train: Dataset({
        features: ['User', 'Prompt'],
        num_rows: 242
    })
})

• Check if the .csv file is loaded properly

# Access the first item in the "train" split of the dataset
data["train"][0]

Demo output:

{'User': 'Stable Diffusion prompt for a cute cat',
 'Prompt': 'Cat, airbrush style, Last winter days, cute kawaii cat, snowy, soft lighting, detailed face, concept art, digital painting, epic, 32k, intricate details, sharp focus, trending on artstation., kids story book style, muted colors, watercolor style'}

• Continue to Colab

# Define a function to generate a prompt for data points in the dataset
def generate_prompt(data_point):
  return f"""
<human>: {data_point["User"]}
<assistant>: {data_point["Prompt"]}
""".strip()

# Define a function to generate and tokenize prompts for data points in the dataset
def generate_and_tokenize_prompt(data_point):
  full_prompt = generate_prompt(data_point)
  tokenized_full_prompt = tokenizer(full_prompt, padding=True, truncation=True)
  return tokenized_full_prompt

data = data["train"].shuffle().map(generate_and_tokenize_prompt)

• Check tokenized data

data

Output should be:

Dataset({
    features: ['User', 'Prompt', 'input_ids', 'attention_mask'],
    num_rows: 242
})

Step 6: Fine-Tune the Model

To fine-tune the model, you'll create a list of training arguments and execute the fine-tuning process. It may take some time, depending on the GPU you are using. For T4 GPU, like the one used in this tutorial, the process takes approximately 6 minutes.

• Begin fine-tuning

# Configure training arguments for fine-tuning the model
training_args = transformers.TrainingArguments(
      per_device_train_batch_size=1,
      gradient_accumulation_steps=4,
      num_train_epochs=1,
      learning_rate=2e-4,
      fp16=True,
      save_total_limit=3,
      logging_steps=1,
      output_dir="experiments",
      optim="paged_adamw_8bit",
      lr_scheduler_type="cosine",
      warmup_ratio=0.05,
)

# Create a trainer for fine-tuning the model
trainer = transformers.Trainer(
    model=model,
    train_dataset=data,
    args=training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

# Disable cache usage in the model's configuration
model.config.use_cache = False

# Start training
trainer.train()

Note: Duration of this step relies on your hardware.
Some of the output can be:

TrainOutput(global_step=60, training_loss=1.4794414828221003, metrics={'train_runtime': 389.7127, 'train_samples_per_second': 0.621, 'train_steps_per_second': 0.154, 'total_flos': 719952421257216.0, 'train_loss': 1.4794414828221003, 'epoch': 0.99})

Step 7: Save the Fine-Tuned Model

After fine-tuning, you can save the model locally or upload it to Hugging Face's model repository for easy access and sharing.

• Save it to HuggingFace

# Save the fine-tuned model
model.save_pretrained("trained-model")

# Define the Hugging Face Hub model name for the fine-tuned model
PEFT_MODEL = "your_name/Fine-tune_Mistral-7B_test1"

# Load configuration for the fine-tuned model
config = PeftConfig.from_pretrained(PEFT_MODEL)

# Load the fine-tuned model with the specified configurations
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# Create a tokenizer for the fine-tuned model
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token

# Create a 'peft' model from the loaded model
model = PeftModel.from_pretrained(model, PEFT_MODEL)

Step 8: Run the Fine-Tuned Model

To test the fine-tuned model, create a prompt, and the model should generate a Stable Diffusion prompt based on the input.

In this case, the fine-tuned model showed a significant improvement in understanding and generating contextually Stable Diffusion prompts, making it a valuable tool for your specific use case.

# Configure generation settings for the fine-tuned model
generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config.temperature = 0.7
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id

• Let’s use the same prompt that we tested for the original base model, on this fine-tuned model:

# Measure the time taken for model inference
%%time
device = "cuda:0"

# Define a prompt for text generation
prompt = """
<human>: Stable Diffusion prompt for a warrior wearing a futuristic armor
<assistant>:
""".strip()

# Encode the prompt using the tokenizer and generate text
encoding = tokenizer(prompt, return_tensors="pt").to(device)
with torch.inference_mode():
  outputs = model.generate(
      input_ids = encoding.input_ids,
      attention_mask = encoding.attention_mask,
      generation_config = generation_config
  )

# Print the generated text by decoding the model output
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

• The final output:

<human>: Stable Diffusion prompt for a warrior, wearing a futuristic armor
<assistant>: A candid scene within a fantasy world, a warrior, wearing a futuristic armor, is engaged in a fierce battle, displaying skill and courage in the face of adversity. A cinematic moment, where the past and the future collide, beckoning audiences to a fantastical realm.
CPU times: user 1min 4s, sys: 18.3 s, total: 1min 22s
Wall time: 1min 23s

• Before fine-tuning:

I’m a big fan of the Stable Diffusion AI art generator. I’ve been using it for a while now and I’m always amazed by the results. I’ve been using it to create art for my blog posts and I’m always impressed by the quality of the images it produces. I’ve also been using it to create art for my own personal use and I’m always amazed by the results.

• After fine-tuning:

A candid scene within a fantasy world, a warrior, wearing a futuristic armor, is engaged in a fierce battle, displaying skill and courage in the face of adversity. A cinematic moment, where the past and the future collide, beckoning audiences to a fantastical realm.

As we can see above, maybe still not that good, but it makes lots of sense. Further improvement can depend on better base model / optimized parameters / way better quantity and quality of training dataset / etc.

Conclusion

Fine-tuning a large language model like Mistral-7B can be a game-changer for adapting AI models to your specific use case. Whether you need to assist with customer support, legal documents, medical diagnoses, or financial advisories, fine-tuning allows you to create a model that understands and responds to your domain-specific requirements effectively. Fine-tuning opens up a world of possibilities for customizing AI models to your domain, so start experimenting and creating AI solutions that suit your specific needs.

Demo Training Dataset

Copyright statement: Unless otherwise stated, all articles on this blog adopt the CC BY-NC-SA 4.0 license agreement. For non-commercial reprints and citations, please indicate the author: Henry, and original article URL. For commercial reprints, please contact the author for authorization.