Finetune_codellama-34B-with-QLoRA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetune codellama-34B with QLoRA\n",
    "\n",
    "### Checkout my [Twitter(@rohanpaul_ai)](https://twitter.com/rohanpaul_ai) for daily LLM bits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig\n",
    "from datasets import load_dataset\n",
    "from trl import SFTTrainer\n",
    "from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "import bitsandbytes as bnb\n",
    "\n",
    "def find_all_linear_names(model):\n",
    "    cls = bnb.nn.Linear4bit\n",
    "    lora_module_names = set()\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, cls):\n",
    "            names = name.split('.')\n",
    "            lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "\n",
    "    return list(lora_module_names)\n",
    "\n",
    "\n",
    "def print_trainable_parameters(model):\n",
    "  \"\"\"\n",
    "  Prints the number of trainable parameters in the model.\n",
    "  \"\"\"\n",
    "  trainable_params = 0\n",
    "  all_param = 0\n",
    "  for _, param in model.named_parameters():\n",
    "    all_param += param.numel()\n",
    "    if param.requires_grad:\n",
    "      trainable_params += param.numel()\n",
    "  print(\n",
    "      f\"trainable params: {trainable_params} || all params: {all_param} || trainables%: {100 * trainable_params / all_param}\"\n",
    "  )\n",
    "\n",
    "def setup_environment():\n",
    "    \"\"\" Sets up necessary imports and configurations. \"\"\"\n",
    "    output_dir = \"./results\"\n",
    "    model_name = \"codellama/CodeLlama-34b-hf\"\n",
    "    return output_dir, model_name\n",
    "\n",
    "def load_and_prepare_dataset():\n",
    "    \"\"\" Loads the dataset and prepares it for training. \"\"\"\n",
    "    return load_dataset('timdettmers/openassistant-guanaco', split=\"train\")\n",
    "\n",
    "def initialize_tokenizer(model_name):\n",
    "    \"\"\" Initializes and configures the tokenizer. \"\"\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "    tokenizer.padding_side = \"right\"\n",
    "    return tokenizer\n",
    "\n",
    "def create_base_model(model_name):\n",
    "    \"\"\" Creates and configures the base model with low-bit quantization. \"\"\"\n",
    "    bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "    )\n",
    "    base_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=bnb_config)\n",
    "    base_model.config.use_cache = False\n",
    "    return prepare_model_for_kbit_training(base_model)\n",
    "\n",
    "def apply_peft_to_model(base_model):\n",
    "    \"\"\" Applies prompt engineering for fine-tuning (PeFT) using LoRA to the base model. \"\"\"\n",
    "    peft_config = LoraConfig(\n",
    "        r=32,\n",
    "        lora_alpha=16,\n",
    "        target_modules=find_all_linear_names(base_model),\n",
    "        lora_dropout=0.05,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "    )\n",
    "    return get_peft_model(base_model, peft_config)\n",
    "\n",
    "def setup_training(base_model, dataset, tokenizer):\n",
    "    \"\"\" Configures training arguments and initializes the trainer. \"\"\"\n",
    "    training_args = TrainingArguments(\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=1,\n",
    "        gradient_checkpointing=True,\n",
    "        max_grad_norm=0.3,\n",
    "        num_train_epochs=3,\n",
    "        learning_rate=1e-4,\n",
    "        bf16=True,\n",
    "        save_total_limit=3,\n",
    "        logging_steps=300,\n",
    "        output_dir=output_dir,\n",
    "        optim=\"paged_adamw_32bit\",\n",
    "        lr_scheduler_type=\"constant\",\n",
    "        warmup_ratio=0.05,\n",
    "    )\n",
    "    return SFTTrainer(\n",
    "        base_model,\n",
    "        train_dataset=dataset,\n",
    "        dataset_text_field=\"text\",\n",
    "        tokenizer=tokenizer,\n",
    "        max_seq_length=512,\n",
    "        args=training_args\n",
    "    )\n",
    "\n",
    "def train_and_save_model(trainer, output_dir):\n",
    "    \"\"\" Handles the training process and saves the model. \"\"\"\n",
    "    trainer.train()\n",
    "    trainer.save_model(output_dir)\n",
    "    final_output_dir = os.path.join(output_dir, \"final_checkpoint\")\n",
    "    trainer.model.save_pretrained(final_output_dir)\n",
    "    tokenizer.save_pretrained(final_output_dir)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir, model_name = setup_environment()\n",
    "\n",
    "dataset = load_and_prepare_dataset()\n",
    "\n",
    "tokenizer = initialize_tokenizer(model_name)\n",
    "\n",
    "base_model = create_base_model(model_name)\n",
    "\n",
    "base_model = apply_peft_to_model(base_model)\n",
    "\n",
    "print_trainable_parameters(base_model)\n",
    "\n",
    "trainer = setup_training(base_model, dataset, tokenizer)\n",
    "\n",
    "train_and_save_model(trainer, output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------\n",
    "\n",
    "### Explanations of the Key terms of the BitsAndBytesConfig\n",
    "\n",
    "\n",
    "```py\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "    )\n",
    "\n",
    "```\n",
    "\n",
    "👉 **`load_in_4bit` parameter** is for loading the model in 4 bits precision\n",
    "\n",
    "This means that the weights and activations of the model are represented using 4 bits instead of the usual 32 bits. This can significantly reduce the memory footprint of the model. 4-bit precision models can use up to 16x less memory than full precision models and can be up to 2x faster than full precision models.\n",
    "\n",
    "However, if you need the highest possible accuracy, then you may want to use full precision models.\n",
    "\n",
    "--------------\n",
    "\n",
    "👉 `bnb_4bit_use_double_quant=True` : This parameter enables double quantization or also called nested quantization, which applies a second quantization after the initial one. It saves an additional 0.4 bits per parameter.\n",
    "\n",
    "--------------\n",
    "\n",
    "👉 `use_nested_quant`: A flag will be applied to determine if nested (or double) quantization.\n",
    "\n",
    "--------------\n",
    "\n",
    "👉 `bnb_4bit_quant_type=\"nf4\"` : This parameter specifies the type of 4-bit quantization to be used. In this case, \"nf4\" refers to normalized float 4, which is the default quantization type.\n",
    "\n",
    "--------------\n",
    "\n",
    "👉 `bnb_4bit_compute_dtype=torch.bfloat16` : This parameter determines the compute data type used during the computation. It specifies the use of the bfloat16 data type for faster training. The compute data type can be chosen from options like float16, bfloat16, float32, etc.\n",
    "\n",
    "This configuration is needed because, while 4-bit bitsandbytes stores weights in 4-bits, the computation still happens in 16 or 32-bit and here any combination can be chosen (float16, bfloat16, float32 etc).\n",
    "\n",
    "The matrix multiplication and training will be faster if one uses a 16-bit compute dtype (and actually the default value for this parameter is torch.float32).\n",
    "\n",
    "--------------\n",
    "\n",
    "Does Floating Point 4-bit precision quantization have any hardware requirements?\n",
    "\n",
    "Note that this method is only compatible with GPUs, hence it is not possible to quantize models in 4bit on a CPU. Among GPUs, there should not be any hardware requirement about this method, therefore any GPU could be used to run the 4bit quantization as long as you have CUDA>=11.2 installed. Keep also in mind that the computation is not done in 4bit, the weights and activations are compressed to that format and the computation is still kept in the desired or native dtype.\n",
    "\n",
    "=====================\n",
    "\n",
    "FP8 and FP4 stand for Floating Point 8-bit and 4-bit precision, respectively. They are part of the minifloats family of floating point values (among other precisions, the minifloats family also includes bfloat16 and float16).\n",
    "\n",
    "----------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Further possibilities for improving / re-organizing the code\n",
    "\n",
    "### Performance and Memory Optimization\n",
    "\n",
    "1. **Batch Size and Gradient Accumulation**: we're using a batch size of 1 with gradient accumulation. If hardware permits, increasing the batch size can improve training efficiency. Balancing between batch size and gradient accumulation steps is key for optimal GPU utilization.\n",
    "\n",
    "3. **Model Parallelism**: If we're working with very large models and have access to multiple GPUs, implementing model parallelism can be beneficial. This involves splitting the model across different GPUs.\n",
    "\n",
    "4. **Data Loading Optimization**: Optimizing data loading can have a significant impact on training speed. Consider using techniques like prefetching, multi-threaded data loading, and ensuring your dataset is stored in a fast-access storage medium.\n",
    "\n",
    "### Hyperparameter Tuning\n",
    "\n",
    "1. **Learning Rate Scheduler**: we're using a constant learning rate. Experimenting with different learning rate schedules like linear decay or cyclical learning rates might yield better results.\n",
    "\n",
    "2. **Optimizer Tweaks**: While we are using `paged_adamw_32bit`, exploring other optimizers like `AdamW` or `SGD` with momentum could offer different performance characteristics.\n",
    "\n",
    "### Advanced Techniques\n",
    "\n",
    "1. **Regularization Techniques**: Implementing regularization methods like dropout, weight decay, or more advanced techniques like data augmentation (if applicable to your task) can prevent overfitting.\n",
    "\n",
    "3. **Evaluation Strategy**: Ensure a robust evaluation strategy is in place, including validation during training and possibly more nuanced evaluation metrics tailored to your specific application.\n",
    "\n",
    "4. **Experiment Tracking**: If not already in place, integrating an experiment tracking system like Weights & Biases or TensorBoard can be very helpful for monitoring training progress and comparing different training runs."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}