camenduru's_peft_finetune_transformers_int8.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/notebooks/blob/main/camenduru's_peft_finetune_transformers_int8.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZsJUxc0aRsAf"
      },
      "outputs": [],
      "source": [
        "# https://github.com/huggingface/peft example modified\n",
        "\n",
        "!pip install -q bitsandbytes datasets accelerate loralib\n",
        "!pip install -q git+https://github.com/huggingface/transformers.git@main \n",
        "!pip install -q git+https://github.com/huggingface/peft.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import bitsandbytes as bnb\n",
        "from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"facebook/opt-6.7b\",\n",
        "    load_in_8bit=True,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-6.7b\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from peft import prepare_model_for_int8_training\n",
        "\n",
        "model = prepare_model_for_int8_training(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from peft import LoraConfig, get_peft_model\n",
        "\n",
        "config = LoraConfig(\n",
        "    r=16, lora_alpha=32, target_modules=[\"q_proj\", \"v_proj\"], lora_dropout=0.05, bias=\"none\", task_type=\"CAUSAL_LM\"\n",
        ")\n",
        "\n",
        "model = get_peft_model(model, config)\n",
        "print_trainable_parameters(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import transformers\n",
        "from datasets import load_dataset\n",
        "\n",
        "data = load_dataset(\"Abirate/english_quotes\")\n",
        "data = data.map(lambda samples: tokenizer(samples[\"quote\"]), batched=True)\n",
        "\n",
        "trainer = transformers.Trainer(\n",
        "    model=model,\n",
        "    train_dataset=data[\"train\"],\n",
        "    args=transformers.TrainingArguments(\n",
        "        per_device_train_batch_size=4,\n",
        "        gradient_accumulation_steps=4,\n",
        "        warmup_steps=100,\n",
        "        max_steps=200,\n",
        "        learning_rate=2e-4,\n",
        "        fp16=True,\n",
        "        logging_steps=1,\n",
        "        output_dir=\"outputs\",\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()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model.push_to_hub(\"ybelkada/opt-6.7b-lora\", use_auth_token=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from peft import PeftModel, PeftConfig\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "peft_model_id = \"ybelkada/opt-6.7b-lora\"\n",
        "config = PeftConfig.from_pretrained(peft_model_id)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map=\"auto\"\n",
        ")\n",
        "tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
        "\n",
        "# Load the Lora model\n",
        "model = PeftModel.from_pretrained(model, peft_model_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "batch = tokenizer(\"Two things are infinite: \", return_tensors=\"pt\")\n",
        "\n",
        "with torch.cuda.amp.autocast():\n",
        "    output_tokens = model.generate(**batch, max_new_tokens=50)\n",
        "\n",
        "print(\"\\n\\n\", tokenizer.decode(output_tokens[0], skip_special_tokens=True))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}