Jamba_Finetuning_Colab-Pro.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sOfeXQXBtLVO"
      },
      "source": [
        "# Fine-tune Jamba-v0.1 on A100 GPU - 40GB VRAM using QLoRA in Google Colab Pro"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "panZSjSVq6Cc"
      },
      "outputs": [],
      "source": [
        "!pip install ninja packaging\n",
        "!pip install flash-attn --no-build-isolation\n",
        "!pip install -U \"transformers>=4.39.0\"\n",
        "!pip install mamba-ssm \"causal-conv1d>=1.2.0\"\n",
        "!pip install peft trl bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zhmfYJMopcB3"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "from trl import SFTTrainer\n",
        "from peft import LoraConfig\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model_id = \"ai21labs/Jamba-v0.1\"\n",
        "\n",
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    llm_int4_skip_modules=[\"mamba\"]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "📌 The `llm_int8_skip_modules` parameter allows you to specify a list of module names that should be skipped or excluded from the quantization process. These modules will remain in their original data type (typically 32-bit floating-point) without being converted to 8-bit integers.\n",
        "\n",
        "📌 This parameter is particularly useful for models with multiple heads or output layers in different parts of the model architecture, not necessarily at the last position. For example, in the case of Causal Language Models (CausalLM), the last layer (often called `lm_head`) is responsible for generating the final output logits or probabilities. Keeping this layer in its original data type can help preserve accuracy and prevent potential degradation caused by quantizing this crucial output layer.\n",
        "\n",
        "📌 The strategic positioning of these heads is crucial for the model's effectiveness. Placing a head at a specific layer allows it to leverage the representations formed up to that point, which might contain the most relevant information for the head's task. \n",
        "\n",
        "📌 In the context of models like Jukebox, which may have multiple output heads positioned at various locations within the model architecture, it's crucial to maintain the precision of certain modules for maintaining overall model performance. For example, if a model utilizes various heads to generate different aspects of output, converting all modules to int8 might degrade the quality of the output or the model's effectiveness due to reduced numerical precision."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "ktiiRqlgp8bI"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    trust_remote_code=True,\n",
        "    device_map='auto',\n",
        "    attn_implementation=\"flash_attention_2\",\n",
        "    quantization_config=quantization_config,\n",
        "    use_mamba_kernels=False #Disabling the mamba kernels since I have a recurrent error.\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "WNZbOoUKrCBt"
      },
      "outputs": [],
      "source": [
        "model.save_pretrained(\"/content/drive/MyDrive/jamba_ft\")\n",
        "tokenizer.save_pretrained(\"/content/drive/MyDrive/jamba_ft\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "I2TWny5tqU8I"
      },
      "outputs": [],
      "source": [
        "dataset = load_dataset(\"Abirate/english_quotes\", split=\"train\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "ee1IMXszqWyU"
      },
      "outputs": [],
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\",\n",
        "    num_train_epochs=3,\n",
        "    per_device_train_batch_size=4,\n",
        "    logging_dir='./logs',\n",
        "    logging_steps=10,\n",
        "    learning_rate=2e-3 # 2.5e-5\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "iyipI7FyqZHI"
      },
      "outputs": [],
      "source": [
        "lora_config = LoraConfig(\n",
        "    r=8,\n",
        "    target_modules=[\"embed_tokens\", \"x_proj\", \"in_proj\", \"out_proj\"],\n",
        "    task_type=\"CAUSAL_LM\",\n",
        "    bias=\"none\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Note the structure of our training dataset\n",
        "\n",
        "`english_quotes` is a dataset of all the quotes retrieved from goodreads quotes. This dataset can be used for multi-label text classification and text generation. \n",
        "\n",
        "The dataset can be used to train a model to generate quotes by fine-tuning an existing pretrained model or to train a model for text-classification, which consists of classifying quotes by author as well as by topic (using tags). \n",
        "\n",
        "![](assets/2024-03-31-19-30-22.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "6gMx9J6Eqas3"
      },
      "outputs": [],
      "source": [
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    args=training_args,\n",
        "    peft_config=lora_config,\n",
        "    train_dataset=dataset,\n",
        "    dataset_text_field=\"quote\",\n",
        "    max_seq_length=256\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`dataset_text_field` (Optional[str]) — The name of the text field of the dataset, in case this is passed by a user, the trainer will automatically create a ConstantLengthDataset based on the dataset_text_field argument.\n",
        "\n",
        "\n",
        "So `dataset_text_field` will be used for training only if `formatting_func` is `None`.\n",
        "\n",
        "You should be careful because if you do this:\n",
        "\n",
        "`dataset_text_field='instruction'`\n",
        "\n",
        "SFTTrainer will only read the text saved in `train_dataset['instruction']`.\n",
        "\n",
        "So that the model being trained will only learn to predict the instructions without the answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9gGKIIyQrZTi"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}