colab/bnb_4bit_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/greengerong/awesome-llm/blob/main/colab/bnb_4bit_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# `transformers` meets `bitsandbytes` for democratzing Large Language Models (LLMs) through 4bit quantization\n",
        "\n",
        "<center>\n",
        "<img src=\"https://github.com/huggingface/blog/blob/main/assets/96_hf_bitsandbytes_integration/Thumbnail_blue.png?raw=true\" alt=\"drawing\" width=\"700\" class=\"center\"/>\n",
        "</center>\n",
        "\n",
        "Welcome to this notebook that goes through the recent `bitsandbytes` integration that includes the work from XXX that introduces no performance degradation 4bit quantization techniques, for democratizing LLMs inference and training.\n",
        "\n",
        "In this notebook, we will learn together how to load a large model in 4bit (`gpt-neo-x-20b`) and train it using Google Colab and PEFT library from Hugging Face 🤗.\n",
        "\n",
        "[In the general usage notebook](https://colab.research.google.com/drive/1ge2F1QSK8Q7h0hn3YKuBCOAS0bK8E0wf?usp=sharing), you can learn how to propely load a model in 4bit with all its variants.\n",
        "\n",
        "If you liked the previous work for integrating [*LLM.int8*](https://arxiv.org/abs/2208.07339), you can have a look at the [introduction blogpost](https://huggingface.co/blog/hf-bitsandbytes-integration) to lean more about that quantization method.\n"
      ],
      "metadata": {
        "id": "XIyP_0r6zuVc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FuXIFTFapAMI"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U bitsandbytes\n",
        "!pip install -q -U git+https://github.com/huggingface/transformers.git\n",
        "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
        "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
        "!pip install -q datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First let's load the model we are going to use - GPT-neo-x-20B! Note that the model itself is around 40GB in half precision"
      ],
      "metadata": {
        "id": "MJ-5idQwzvg-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
        "\n",
        "model_id = \"EleutherAI/gpt-neox-20b\"\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={\"\":0})"
      ],
      "metadata": {
        "id": "E0Nl5mWL0k2T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we have to apply some preprocessing to the model to prepare it for training. For that use the `prepare_model_for_kbit_training` method from PEFT."
      ],
      "metadata": {
        "id": "Mp2gMi1ZzGET"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from peft import prepare_model_for_kbit_training\n",
        "\n",
        "model.gradient_checkpointing_enable()\n",
        "model = prepare_model_for_kbit_training(model)"
      ],
      "metadata": {
        "id": "a9EUEDAl0ss3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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} || trainable%: {100 * trainable_params / all_param}\"\n",
        "    )"
      ],
      "metadata": {
        "id": "gkIcwsSU01EB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from peft import LoraConfig, get_peft_model\n",
        "\n",
        "config = LoraConfig(\n",
        "    r=8,\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"query_key_value\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\"\n",
        ")\n",
        "\n",
        "model = get_peft_model(model, config)\n",
        "print_trainable_parameters(model)"
      ],
      "metadata": {
        "id": "Ybeyl20n3dYH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's load a common dataset, english quotes, to fine tune our model on famous quotes."
      ],
      "metadata": {
        "id": "FCc64bfnmd3j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "data = load_dataset(\"Abirate/english_quotes\")\n",
        "data = data.map(lambda samples: tokenizer(samples[\"quote\"]), batched=True)"
      ],
      "metadata": {
        "id": "s6f4z8EYmcJ6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to run the training! For the sake of the demo, we just ran it for few steps just to showcase how to use this integration with existing tools on the HF ecosystem."
      ],
      "metadata": {
        "id": "_0MOtwf3zdZp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import transformers\n",
        "\n",
        "# needed for gpt-neo-x tokenizer\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "trainer = transformers.Trainer(\n",
        "    model=model,\n",
        "    train_dataset=data[\"train\"],\n",
        "    args=transformers.TrainingArguments(\n",
        "        per_device_train_batch_size=1,\n",
        "        gradient_accumulation_steps=4,\n",
        "        warmup_steps=2,\n",
        "        max_steps=10,\n",
        "        learning_rate=2e-4,\n",
        "        fp16=True,\n",
        "        logging_steps=1,\n",
        "        output_dir=\"outputs\",\n",
        "        optim=\"paged_adamw_8bit\"\n",
        "    ),\n",
        "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
        ")\n",
        "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "jq0nX33BmfaC"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}