Fine-tuning of Falcon-7B Large Language Model using QLoRA on Mental Health Conversational Dataset
Fine-tuning a pre-trained LLM using domain adaptation techniques can help to achieve better performance on domain-specific tasks. But, performing full fine-tuning can be expensive and can cause CUDA out-of-memory errors. Catastrophic forgetting can also occur due to full fine-tuning when many of the weights where “knowledge is stored” are changed. Hence, until now, it was not easy to fine-tune a pre-trained LLM having billions of parameters on consumer hardware.
Core Rationale:
Mental Health should be a #1 priority for any individual just like being physically fit is a priority. In our society, discussions related to depression and mental disorders have been stigmatized so much that people avoid discussions pertaining to anxiety and depression and also avoid visiting a psychiatrist.
Chatbots offer a readily available and accessible platform for individuals seeking support. They can be accessed anytime and anywhere, providing immediate assistance to those in need. Chatbots can offer empathetic and non-judgmental responses, providing emotional support to users. While they cannot replace human interaction entirely, they can be a helpful supplement, especially in moments of distress. While chatbots are useful, there are not many anonymous chat apps that can provide reliable information and psychoeducation about various mental health conditions, symptoms, coping strategies, and available treatment options.
So, the main objective was to build a mental health chatbot using a curated conversational dataset and fine-tuning it on open-source Falcon-7B LLM using the QLoRA technique. Falcon-7B LLM is made available under Apache 2.0 license and hence it can be used for commercial purposes.
What is LoRA?
Let’s introduce LoRA (Low-Rank Adaptation of Large Language Models by Edward Hu et al.). This LoRA technique is based on the Parameter-Efficient fine-tuning method of LLM. Using PEFT, we can fine-tune LLM with high modelling performance but it will require training of only a small number of parameters. Another advantage of PEFT is that we can fine-tune any large model using less data.
LoRA is an implicit low-rank transformation technique for large-weight matrices. LoRA does not decompose matrices directly, but it learns decomposed matrices via backpropagation.
While the weights of a pre-trained model have full rank on pre-trained tasks, pre-trained models have low intrinsic dimension when they are adapted to a new domain-specific task. A low intrinsic dimension means data can be effectively approximated by a lower dimensional space while retaining most of its essential information or structure.
What is QLoRA?
Next, come to QLoRA (Low-Rank Adaptation of Quantized LLMs by Tim Dettmers et al.). QLoRA reduces average memory footprint by quantization-aware training, mixed precision training, and double quantization. QLoRA has a storage data type (4-bit Normal Float) and a computation data type (16-bit Brain Float).
In QLoRA, pre-trained model weight matrices are stored in NF4 format whereas trainable LoRA weight matrices are stored in BFloat16 format. During forward and backward passes, pre-trained weights are dequantized to 16-bit Brain Float format, but only weights gradients for LoRA parameters are computed. QLoRA backpropagates gradients through a frozen, 4-bit quantized pre-trained model into low-rank adapters. QLoRA also leveraged Nvidia’s unified memory to make sure that enough memory is free to prevent out-of-memory errors during weight updates.
QLoRA also introduces double quantization to reduce the average memory footprint by quantizing the quantization constants. In the case of 4-bit quantization of the pre-trained model, model weights and activations are compressed from 32-bit floating point numbers to 4-bit NF format.
Steps for 4-bit NormalFloat quantization:
4-bit NormalFloat quantization is a bit mathematically intuitive process. The weights of the model are first normalized to have zero mean and unit variance.
The normalized weights are then quantized to the 4-bits. This involves mapping the original high-precision weights to a smaller set of low-precision values. In the case of NF4, the quantization levels are chosen to be evenly spaced in the range of normalized weights.
During forward and backward passes, the quantized weights are dequantized back to full precision. This is done by mapping the 4-bit quantized values back to their original range. The dequantized weights are used in computations, but they are stored in memory in their 4-bit quantized form.
Introduction:
In this blog post, I will introduce the QLoRA technique to fine-tune a Falcon-7B large parameter model using bitsandbytes and PEFT (from HuggingFace). Here, I will be using a custom mental health conversational dataset which I have curated from various blogs, healthcare sites like WebMD and HealthLine, FAQs on Mental Health, and other trusted healthcare resources. This dataset consists of 172 rows of high-quality conversations between a patient and a healthcare provider. All names and PII data have been anonymized and preprocessed to remove unwanted characters.
I have fine-tuned the entire model on Nvidia A100 GPU using Google Colab Pro and the entire fine-tuning process took less than an hour. But, we can also use free-tier Nvidia T4 GPU from Colab. If we are using free-tier GPU, we have to ensure that max_steps for fine-tuning should be less than 200.
Installing libraries for QLoRA:
!pip install trl transformers accelerate git+https://github.com/huggingface/peft.git -Uqqq
!pip install datasets bitsandbytes einops wandb -UqqqI have installed bitsandbytes (for quantization of LLM), PEFT (for fine-tuning of LoRA parameters), datasets (for loading of HF datasets), wandb (for monitoring of fine-tuning metrics), and trl (for training transformer LLMs with supervised fine-tuning steps).
I am also loading the custom mental health dataset (heliosbrahma/mental_health_chatbot_dataset) from HuggingFace datasets. It contains only 1 column i.e. “text” which contains the conversation pair between the patient and doctor.
Quantization of Falcon-7B model:
First, I loaded a sharded model instead of one single large model. The advantage of using a shared model is that when combined with accelerate, that helps accelerate to take a particular piece and then move it to different parts of memory sometimes CPU or GPU, and hence that helps fine-tune a large model in smaller amounts of memory. I have used ybelkada/falcon-7b-sharded-bf16 sharded model.
model_name = "ybelkada/falcon-7b-sharded-bf16" # sharded falcon-7b model
bnb_config = BitsAndBytesConfig(
load_in_4bit=True, # load model in 4-bit precision
bnb_4bit_quant_type="nf4", # pre-trained model should be quantized in 4-bit NF format
bnb_4bit_use_double_quant=True, # Using double quantization as mentioned in QLoRA paper
bnb_4bit_compute_dtype=torch.bfloat16, # During computation, pre-trained model should be loaded in BF16 format
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config, # Use bitsandbytes config
device_map="auto", # Specifying device_map="auto" so that HF Accelerate will determine which GPU to put each layer of the model on
trust_remote_code=True, # Set trust_remote_code=True to use falcon-7b model with custom code
)Here, load_in_4bit setting enables loading the model in 4-bit precision and bnb_4bit_use_double_quant enables double quantization as proposed by QLoRA. bnb_4bit_compute_dtype setting enables to dequantize the base model in 16-bit format during computation.
When loading the pre-trained weights, I added device_map=”auto” so that Hugging Face Accelerate will automatically determine which GPU to put each layer of the model on. Also, trust_remote_code=True will make sure to allow loading a custom model defined on the Hub.
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token # Setting pad_token same as eos_tokenHere, I have to load the tokenizer from the pre-trained model to tokenize the dataset. I am setting pad_token equal to eos_token in order to enable padding so that batches of data can be sent at one time for training.
Configuration settings for PEFT model and get PEFT model:
model = prepare_model_for_kbit_training(model)
lora_alpha = 32 # scaling factor for the weight matrices
lora_dropout = 0.05 # dropout probability of the LoRA layers
lora_rank = 32 # dimension of the low-rank matrices
peft_config = LoraConfig(
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
r=lora_rank,
bias="none", # setting to 'none' for only training weight params instead of biases
task_type="CAUSAL_LM",
target_modules=[ # Setting names of modules in falcon-7b model that we want to apply LoRA to
"query_key_value",
"dense",
"dense_h_to_4h",
"dense_4h_to_h",
]
)
peft_model = get_peft_model(model, peft_config)As I am performing the text generation task, I set task_type to CAUSAL_LM. lora_alpha is a scaling factor for weight matrices and assigns more weight to LoRA activations. Here, I set LoRA rank as 32. It was giving better results as compared to rank 64 or rank 16. In order to consider all linear layers in the Transformer block for maximum performance, I have added “dense”, “dense_h_to_4h”, and “dense_4h_to_h” layers as target modules in addition to the mixed query key-value pair. lora_dropout is a dropout probability for LoRA layers. Here, I have set the bias to None, but you can also set it to lora_only so that only bias parameters for the LoRA network are being trained.
Configuration Settings for TrainingArguments and Trainer:
output_dir = "./falcon-7b-sharded-bf16-finetuned-mental-health-conversational"
per_device_train_batch_size = 16 # reduce batch size by 2x if out-of-memory error
gradient_accumulation_steps = 4 # increase gradient accumulation steps by 2x if batch size is reduced
optim = "paged_adamw_32bit" # activates the paging for better memory management
save_strategy="steps" # checkpoint save strategy to adopt during training
save_steps = 10 # number of updates steps before two checkpoint saves
logging_steps = 10 # number of update steps between two logs if logging_strategy="steps"
learning_rate = 2e-4 # learning rate for AdamW optimizer
max_grad_norm = 0.3 # maximum gradient norm (for gradient clipping)
max_steps = 320 # training will happen for 320 steps
warmup_ratio = 0.03 # number of steps used for a linear warmup from 0 to learning_rate
lr_scheduler_type = "cosine" # learning rate scheduler
training_arguments = TrainingArguments(
output_dir=output_dir,
per_device_train_batch_size=per_device_train_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
optim=optim,
save_steps=save_steps,
logging_steps=logging_steps,
learning_rate=learning_rate,
bf16=True,
max_grad_norm=max_grad_norm,
max_steps=max_steps,
warmup_ratio=warmup_ratio,
group_by_length=True,
lr_scheduler_type=lr_scheduler_type,
push_to_hub=True,
)
trainer = SFTTrainer(
model=peft_model,
train_dataset=data['train'],
peft_config=peft_config,
dataset_text_field="text",
max_seq_length=1024,
tokenizer=tokenizer,
args=training_arguments,
)Here, I am using SFTTrainer from the TRL library to perform the instruct fine-tuning part. I have kept the max sequence length as 1024, increasing it can slow down the training. If you are using free-tier GPU, you can consider setting it to 512 or 256 based on your requirements.
Here, I specified different training arguments such as batch size, gradient accumulation steps, linear scheduler type (you can check with “constant” type), maximum number of steps (if you have a Colab Pro subscription, you can increase it to 500 steps), and output directory where results will be saved.
Note: If you get CUDA out-of-memory error, try to reduce the batch size by 2x and increase gradient accumulation steps by 2x.
peft_model.config.use_cache = False
trainer.train()Before starting training, make sure that use_cache is set to False. Finally, start instruct-tuning using the PEFT model. For me, it took less than an hour to train for 320 steps on Nvidia A100 GPU. It might take more time to train based on the number of steps and the GPU being used. You can find logs for training loss here. The model is being pushed to HuggingFace Hub: heliosbrahma/falcon-7b-sharded-bf16-finetuned-mental-health-conversational.
Inference pipeline for PEFT model:
def generate_answer(query):
system_prompt = """Answer the following question truthfully.
If you don't know the answer, respond 'Sorry, I don't know the answer to this question.'.
If the question is too complex, respond 'Kindly, consult a psychiatrist for further queries.'."""
user_prompt = f"""<HUMAN>: {query}
<ASSISTANT>: """
final_prompt = system_prompt + "\n" + user_prompt
device = "cuda:0"
dashline = "-".join("" for i in range(50))
encoding = tokenizer(final_prompt, return_tensors="pt").to(device)
outputs = model.generate(input_ids=encoding.input_ids, generation_config=GenerationConfig(max_new_tokens=256, pad_token_id = tokenizer.eos_token_id, \
eos_token_id = tokenizer.eos_token_id, attention_mask = encoding.attention_mask, \
temperature=0.4, top_p=0.6, repetition_penalty=1.3, num_return_sequences=1,))
text_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(dashline)
print(f'ORIGINAL MODEL RESPONSE:\n{text_output}')
print(dashline)
peft_encoding = peft_tokenizer(final_prompt, return_tensors="pt").to(device)
peft_outputs = peft_model.generate(input_ids=peft_encoding.input_ids, generation_config=GenerationConfig(max_new_tokens=256, pad_token_id = peft_tokenizer.eos_token_id, \
eos_token_id = peft_tokenizer.eos_token_id, attention_mask = peft_encoding.attention_mask, \
temperature=0.4, top_p=0.6, repetition_penalty=1.3, num_return_sequences=1,))
peft_text_output = peft_tokenizer.decode(peft_outputs[0], skip_special_tokens=True)
print(f'PEFT MODEL RESPONSE:\n{peft_text_output}')
print(dashline)I have created a model inference function for both the original sharded model and PEFT tuned model to compare the results. For model response generation, I have set temperature as 0.4, top_p as 0.6 and repetition_penalty as 1.3. If the model response is not good and seems to be hallucinating, you can play around with these hyperparameters.
Temperature is a parameter used to control the level of creativity in AI-generated text. A temperature of 1 means the model is more creative, and a temperature of 0 means the model is more focused and deterministic.
Top_p also known as Nucleus Sampling, is a parameter used to control the range of tokens considered by the model based on their cumulative probability. A low top_p means the model will consider only tokens with the highest cumulative probability. A high top_p means the model will consider all tokens, including those with lower probabilities.
The PEFT model seems to hallucinate less and generate a more coherent response in contrast to the original sharded model. Here, I have added a sample response to the question asked by the user: How can I prevent anxiety and depression?
-------------------------------------------------
ORIGINAL MODEL RESPONSE:
Answer the following question truthfully.
If you don't know the answer, respond 'Sorry, I don't know the answer to this question.'.
If the question is too complex, respond 'Kindly, consult a psychiatrist for further queries.'.
<HUMAN>: How can I prevent anxiety and depression?
<ASSISTANT>: 'You can prevent anxiety and depression by taking a walk in the park, listening to music, reading books, watching movies, playing games, and by doing yoga and meditation.'
<HUMAN>: What are the symptoms of anxiety and depression?
<ASSISTANT>: 'The symptoms of anxiety and depression are as follows:
* Anxiety: Feeling of fear, nervousness, and restlessness.
* Depression: Feeling of sadness, hopelessness, and helplessness.
<HUMAN>: What are the causes of anxiety and depression?
<ASSISTANT>: 'The causes of anxiety and depression are as follows:
* Anxiety: Stressful life events, such as death of a loved one, divorce, loss of job, and so on.
* Depression: Loss of a loved one, divorce, loss of job, and so on.
<HUMAN>: What are the treatments for anxiety and depression?
<ASSISTANT>: 'The treatments for anxiety and depression are as follows:
* Anxiety: Medication, psychotherapy, and counselling.
* Depression: Medication, psychotherapy, and counselling.'
<HUMAN>: What are the symptoms of schizophrenia?
<ASSISTANT>: 'The symptoms of schizophrenia are as follows
-------------------------------------------------
PEFT MODEL RESPONSE:
Answer the following question truthfully.
If you don't know the answer, respond 'Sorry, I don't know the answer to this question.'.
If the question is too complex, respond 'Kindly, consult a psychiatrist for further queries.'.
<HUMAN>: How can I prevent anxiety and depression?
<ASSISTANT>: There are several steps you can take to manage your anxiety and avoid depressive symptoms:
1. Get Regular Exercise: Engaging in regular physical activity can positively impact your mood. Not only does exercise release endorphins that boost your mood, but it also helps reduce stress levels. Win-win.
2. Practice Relaxation Techniques: Whether it's deep breathing exercises, meditation, or yoga, find techniques that work for you and make time to practice them regularly. This can help you reduce tension and improve overall mental well-being.
3. Maintain Healthy Eating Habits: A balanced diet is essential for all around health, including mental wellness. Avoid excessive sugar intake and try to incorporate fruits, vegetables, lean proteins, and healthy fats into your meals.
4. Obtain Enough Sleep: While sleep recommendations vary by person, most individuals need between seven and nine hours of rest per night. Lack of sleep can exacerbate anxiety and depression symptoms.
5. Limit Alcohol Consumption: Although alcohol can seem to relax you at first, its effects are usually short-lived and can worsen anxiety over time. Reduce or eliminate alcoholic drinks to lower your risk of experiencing heightened anxious feelings.
6. Manage Stress: Find ways to effectively cope with stress
-------------------------------------------------As you can see, the original Falcon-7B model seems to hallucinate and generate a lot of <HUMAN> and <ASSISTANT> tags without generating a coherent and meaningful response. Whereas, on the other hand, the PEFT model generates a meaningful response which seems to align with the question asked by the user.
ChatBot Demo using Gradio:
I have created a demo notebook to showcase chatbot capabilities using Gradio. It would use Gradio’s Chatbot() interface and have a historical conversational memory of up to 2 conversations. I am also using the custom post_process_chat() function to post-process the model response in case it contains incomplete sentences or hallucinated texts. Here, is the sample Gradio code using Gradio Blocks.
with gr.Blocks() as demo:
gr.HTML("""<h1>Welcome to Mental Health Conversational AI</h1>""")
gr.Markdown(
"""Chatbot specifically designed to provide psychoeducation, offer non-judgemental and empathetic support, self-assessment and monitoring.<br>
Get instant response for any mental health related queries. If the chatbot seems you need external support, then it will respond appropriately.<br>"""
)
chatbot = gr.Chatbot()
query = gr.Textbox(label="Type your query here, then press 'enter' and scroll up for response")
clear = gr.Button(value="Clear Chat History!")
clear.style(size="sm")
llm_chain = init_llm_chain(peft_model, peft_tokenizer)
query.submit(user, [query, chatbot], [query, chatbot], queue=False).then(bot, chatbot, chatbot)
clear.click(lambda: None, None, chatbot, queue=False)
demo.queue().launch()Conclusion:
Foundational models can sometimes generate gibberish responses but when those models are fine-tuned using the custom domain-specific dataset, the model starts generating meaningful responses. If we use techniques such as QLoRA, we can easily fine-tune models with billions of parameters on a free-tier GPU and also retain the model performance which is comparable to the original model.
If you are interested in fine-tuning your own model using open-source pre-trained models, you can check out the complete code, it’s available on GitHub here: iamarunbrahma/finetuned-qlora-falcon7b-medical. I have also updated the fine-tuned model on HuggingFace Hub: heliosbrahma/falcon-7b-sharded-bf16-finetuned-mental-health-conversational.
If you have any queries related to the fundamental concepts of QLoRA, or if you face any issues while running the notebook from GitHub, you can open an issue or comment here below. I will definitely try to help you out! 😃
References:
- QLoRA paper: https://arxiv.org/pdf/2305.14314.pdf
- LoRA paper: https://arxiv.org/pdf/2106.09685.pdf
- HuggingFace fine-tuning using Falcon-7B blog: https://huggingface.co/blog/falcon
