model-registry-201/Logging_Models_HuggingFace.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/edu/blob/main/model-registry-201/Logging_Models_HuggingFace.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "<!--- @wandbcode{modelreg201-hf} -->"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Outline\n",
    "- Fine-tune a lightweight LLM (OPT-125M) with LoRA and 8-bit quantization using Launch\n",
    "- Checkpoint the LoRA adapter weights as artifacts\n",
    "- Link the best checkpoint in Model Registry\n",
    "- Linkage triggers an automation in github actions or Modal to quantize the model for inference\n",
    "  - CI job should also profile the model and generate a W&B report\n",
    "- Add `production` alias to registered model version\n",
    "- Triggers automation to deploy model as Fast API inference server in Modal\n",
    "\n",
    "Stretch: interact with model and log inputs/outputs with W&B prompts\n",
    "<!--- @wandbcode{wandb201-hf} -->\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune large models using 🤗 `peft` adapters, `transformers` & `bitsandbytes`\n",
    "\n",
    "In this tutorial we will cover how we can fine-tune large language models using the very recent `peft` library and `bitsandbytes` for loading large models in 8-bit.\n",
    "The fine-tuning method will rely on a recent method called \"Low Rank Adapters\" (LoRA), instead of fine-tuning the entire model you just have to fine-tune these adapters and load them properly inside the model.\n",
    "After fine-tuning the model you can also share your adapters on the 🤗 Hub and load them very easily. Let's get started!\n",
    "\n",
    "**TODO:** Turn this section of code into a launch job"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install requirements\n",
    "\n",
    "First, run the cells below to install the requirements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q bitsandbytes datasets accelerate loralib\n",
    "!pip install -q git+https://github.com/huggingface/transformers.git@main git+https://github.com/huggingface/peft.git\n",
    "!pip install -q wandb\n",
    "!pip install -q ctranslate2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model loading\n",
    "\n",
    "Here let's load the `opt-6.7b` model, its weights in half-precision (float16) are about 13GB on the Hub! If we load them in 8-bit we would require around 7GB of memory instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\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",
    "project_name = \"model-registry-201\" #@param\n",
    "entity = \"wandb\" #@param\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"facebook/opt-125m\",\n",
    "    load_in_8bit=True,\n",
    "    device_map='auto',\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-125m\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Post-processing on the model\n",
    "\n",
    "Finally, we need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in model.parameters():\n",
    "  param.requires_grad = False  # freeze the model - train adapters later\n",
    "  if param.ndim == 1:\n",
    "    # cast the small parameters (e.g. layernorm) to fp32 for stability\n",
    "    param.data = param.data.to(torch.float32)\n",
    "\n",
    "model.gradient_checkpointing_enable()  # reduce number of stored activations\n",
    "model.enable_input_require_grads()\n",
    "\n",
    "class CastOutputToFloat(nn.Sequential):\n",
    "  def forward(self, x): return super().forward(x).to(torch.float32)\n",
    "model.lm_head = CastOutputToFloat(model.lm_head)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply LoRA\n",
    "\n",
    "Here comes the magic with `peft`! Let's load a `PeftModel` and specify that we are going to use low-rank adapters (LoRA) using `get_peft_model` utility function from `peft`."
   ]
  },
  {
   "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,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Log Checkpoints Automatically with Hugging Face\n",
    "\n",
    "Logging your Hugging Face model to W&B Artifacts can be done by setting a W&B environment variable called `WANDB_LOG_MODEL`\n",
    "- `WANDB_LOG_MODEL='end'` - logs only the final model\n",
    "- `WANDB_LOG_MODEL='checkpoint'` - logs the model checkpoints every `save_steps` in the `TrainingArguments`\n",
    "- Optionally use the wandb artifacts api to implement your own checkpointing logic using HF's Callbacks\n",
    "\n",
    "See more details on our Hugging Face integration [here](https://docs.wandb.ai/guides/integrations/huggingface)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "from datasets import load_dataset\n",
    "import wandb\n",
    "\n",
    "os.environ[\"WANDB_LOG_MODEL\"] = \"checkpoint\"\n",
    "\n",
    "wandb.init(project=project_name,\n",
    "           entity=entity,\n",
    "           job_type=\"training\")\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",
    "        report_to=\"wandb\",\n",
    "        warmup_steps=5,\n",
    "        max_steps=25,\n",
    "        learning_rate=2e-4,\n",
    "        fp16=True,\n",
    "        logging_steps=1,\n",
    "        save_steps=5,\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()\n",
    "last_run_id = wandb.run.id\n",
    "wandb.finish()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Model Weights to Model Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(project=project_name, entity=entity, job_type=\"registering_best_model\")\n",
    "best_model = wandb.use_artifact(f'{entity}/{project_name}/checkpoint-{last_run_id}:latest')\n",
    "registered_model_name = \"Review Summarization\" #@param {type: \"string\"}\n",
    "wandb.run.link_artifact(best_model, f'{entity}/model-registry/{registered_model_name}', aliases=['staging'])\n",
    "wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}