rapidfireai 0.9.10__py3-none-any.whl → 0.9.11__py3-none-any.whl

tutorial_notebooks/rf-tutorial-dpo-alignment-lite.ipynb

@@ -0,0 +1,412 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "668ecb93",
+   "metadata": {},
+   "source": [
+    "### RapidFire AI Tutorial Use Case: DPO for Alignment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb9cb2b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from rapidfireai import Experiment\n",
+    "from rapidfireai.automl import List, RFGridSearch, RFModelConfig, RFLoraConfig, RFDPOConfig"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9bb751f2",
+   "metadata": {},
+   "source": [
+    "### Load Dataset and Sample Train Subset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e4c3655f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "# Select a subset of the dataset for demo purposes\n",
+    "train_dataset = load_dataset(\n",
+    "    \"trl-lib/ultrafeedback_binarized\", \n",
+    "    split=\"train\").select(range(100))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2485d399",
+   "metadata": {},
+   "source": [
+    "### Initialize Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cf47412f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Every experiment instance must be uniquely named\n",
+    "experiment = Experiment(experiment_name=\"exp1-sft-dpo-lite\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3378be4c",
+   "metadata": {},
+   "source": [
+    "#### Direct Preference Optimization (DPO) Training \n",
+    "\n",
+    "The first step is to train a **Supervised Fine-Tuning (SFT)** model to ensure the data we train on is in-distribution for the DPO algorithm.\n",
+    "\n",
+    "Fine-tuning a language model via DPO consists of **two steps** and is significantly easier than PPO:\n",
+    "\n",
+    "1. Data Collection: Gather a preference dataset containing:\n",
+    "    - **Positive examples**: High-quality generations given a prompt\n",
+    "    - **Negative examples**: Lower-quality generations for the same prompt\n",
+    "    - **Paired comparisons**: Each prompt has both chosen and rejected responses\n",
+    "\n",
+    "2. Optimization: Maximize the log-likelihood of the DPO loss directly.\n",
+    "\n",
+    "Since our goal is to demonstrate DPO, we will assume you already have **pretrained model or adapters from SFT training** and use them as our starting point."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "51df13d9",
+   "metadata": {},
+   "source": [
+    "### Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abbbd1f3",
+   "metadata": {},
+   "source": [
+    "##### Quantization and LoRA Config"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c46f0565",
+   "metadata": {},
+   "source": [
+    "We supervised fine-tuned (QLoRA) the `mistralai/Mistral-7B-Instruct-v0.3` base model by using preferred completions for supervision from the `trl-lib/ultrafeedback_binarized`. This ensures that the output paths sampled during the alignment training already exist in the model's distribution. We do not want the model to have to learn a new task from ground up during the alignment training. The goal of DPO training is to achieve alignment with human preferences by contrastively learning which outputs are considered preferred and which ones are dispreferred.\n",
+    "\n",
+    "The quantization config for `rapidfire-ai-inc/mistral-7b-sft-bnb-4bit` is below:\n",
+    "```python\n",
+    "bnb_config = BitsAndBytesConfig(\n",
+    "    load_in_4bit=True,\n",
+    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
+    "    bnb_4bit_use_double_quant=True,\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
+    ")\n",
+    "```\n",
+    "We will use QLoRA to further fine-tune the SFT model for alignment with DPO"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0ac53010",
+   "metadata": {},
+   "source": [
+    "### Define Multi-Config Knobs for Model, LoRA, and DPO Trainer using RapidFire AI Wrapper APIs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "11960e7a",
+   "metadata": {},
+   "source": [
+    "We create a common base config, since most of the knobs will remain same across the configs in our experiment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af5aab6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import TaskType\n",
+    "\n",
+    "MODEL_NAME_OR_PATH = \"rapidfire-ai-inc/mistral-7b-sft-bnb-4bit\"\n",
+    "\n",
+    "base_lora_config_lite = RFLoraConfig(\n",
+    "    task_type=TaskType.CAUSAL_LM,\n",
+    "    r=64, \n",
+    "    lora_alpha=64, \n",
+    "    lora_dropout=0.1, \n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    bias=\"none\", \n",
+    ")\n",
+    "\n",
+    "base_dpo_config_lite = RFDPOConfig(\n",
+    "    gradient_checkpointing=True,\n",
+    "    gradient_checkpointing_kwargs={\"use_reentrant\": True},\n",
+    "    model_adapter_name=\"default\", # LoRA Adapter from original SFT trained model, if any. This adapter will be trained using DPO.\n",
+    "    ref_adapter_name=\"reference\", # LoRA Adapter from original SFT trained model, if any. This adapter will remain unchanged and will be used for reference model.\n",
+    "    force_use_ref_model=False, \n",
+    "    loss_type=\"sigmoid\", # Uses Bradley-Terry model to calculate the loss.\n",
+    "    beta=0.1, \n",
+    "    max_prompt_length=512,\n",
+    "    max_completion_length=512,\n",
+    "    max_length=1024, # Prompt + completion\n",
+    "    per_device_train_batch_size=2,\n",
+    "    gradient_accumulation_steps=2,\n",
+    "    learning_rate=5e-4,\n",
+    "    warmup_ratio=0.1,\n",
+    "    weight_decay=0,\n",
+    "    lr_scheduler_type=\"linear\",\n",
+    "    num_train_epochs=1, \n",
+    "    optim=\"adamw_8bit\",\n",
+    "    bf16=True,\n",
+    "    save_strategy=\"epoch\",\n",
+    "    logging_strategy=\"steps\",\n",
+    "    logging_steps=1\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d8d056c",
+   "metadata": {},
+   "source": [
+    "Now we create 3 different configs for experimentation by cloning the base config and only modifying the knobs we need."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3f7c2957",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Config 1 - Basic Bradley-Terry model with medium LoRA and large beta\n",
+    "lora_config_1 = base_lora_config_lite.copy()\n",
+    "lora_config_1.r = 16\n",
+    "lora_config_1.lora_alpha = 16\n",
+    "\n",
+    "dpo_config_1 = base_dpo_config_lite.copy()\n",
+    "dpo_config_1.loss_type = \"sigmoid\"\n",
+    "dpo_config_1.beta = 0.1\n",
+    "\n",
+    "# Config 2 - Assumes noisy preference data, use label smoothing to handle it with Robust loss\n",
+    "lora_config_2 = base_lora_config_lite.copy()\n",
+    "lora_config_2.r = 16\n",
+    "lora_config_2.lora_alpha = 16\n",
+    "\n",
+    "dpo_config_2 = base_dpo_config_lite.copy()\n",
+    "dpo_config_2.loss_type = \"robust\"\n",
+    "dpo_config_2.beta = 0.1\n",
+    "dpo_config_2.label_smoothing = 0.5\n",
+    "\n",
+    "# Config 3 - Use a combined loss function with weighted sum\n",
+    "lora_config_3 = base_lora_config_lite.copy()\n",
+    "lora_config_3.r = 16\n",
+    "lora_config_3.lora_alpha = 16\n",
+    "\n",
+    "dpo_config_3 = base_dpo_config_lite.copy()\n",
+    "dpo_config_3.loss_type = [\"sigmoid\", \"bco_pair\", \"sft\"]\n",
+    "dpo_config_3.loss_weights = [0.8, 0.2, 1.0]\n",
+    "dpo_config_3.beta = 0.1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5992c776",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "# List of 3 separate configs\n",
+    "config_set = List([\n",
+    "    RFModelConfig(\n",
+    "        model_name=MODEL_NAME_OR_PATH,\n",
+    "        ref_model_name=None,\n",
+    "        peft_config=lora_config_1,\n",
+    "        training_args=dpo_config_1,\n",
+    "        model_kwargs={\"device_map\": \"auto\", \"torch_dtype\": torch.bfloat16},\n",
+    "        tokenizer_kwargs={\"model_max_length\": 1024, \"padding_side\": \"left\", \"truncation\": True}\n",
+    "    ),\n",
+    "    RFModelConfig(\n",
+    "        model_name=MODEL_NAME_OR_PATH,\n",
+    "        ref_model_name=None,\n",
+    "        peft_config=lora_config_2,\n",
+    "        training_args=dpo_config_2,\n",
+    "        model_kwargs={ \"device_map\": \"auto\", \"torch_dtype\": torch.bfloat16},\n",
+    "        tokenizer_kwargs={\"model_max_length\": 1024, \"padding_side\": \"left\", \"truncation\": True}\n",
+    "    ),\n",
+    "    RFModelConfig(\n",
+    "        model_name=MODEL_NAME_OR_PATH,\n",
+    "        ref_model_name=None,\n",
+    "        peft_config=lora_config_3,\n",
+    "        training_args=dpo_config_3,\n",
+    "        model_kwargs={ \"device_map\": \"auto\", \"torch_dtype\": torch.bfloat16},\n",
+    "        tokenizer_kwargs={\"model_max_length\": 1024, \"padding_side\": \"left\", \"truncation\": True}\n",
+    "    )\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b02f3152",
+   "metadata": {},
+   "source": [
+    "#### Define Model Creation Function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1854ddf0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def sample_create_model(model_config): \n",
+    "   \"\"\"Function to create model object for any given config; must return tuple of (model, tokenizer)\"\"\"\n",
+    "   from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "\n",
+    "   model = AutoModelForCausalLM.from_pretrained(\n",
+    "      model_config[\"model_name\"], \n",
+    "      **model_config[\"model_kwargs\"]\n",
+    "   )\n",
+    "   model.config.use_cache = False # Disable caching to save memory and stability during training\n",
+    "\n",
+    "   tokenizer = AutoTokenizer.from_pretrained(\n",
+    "      model_config[\"model_name\"], \n",
+    "      **model_config[\"tokenizer_kwargs\"]\n",
+    "   )\n",
+    "   if tokenizer.pad_token is None:\n",
+    "      tokenizer.pad_token = tokenizer.eos_token\n",
+    "\n",
+    "   return (model, tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abf1feb4",
+   "metadata": {},
+   "source": [
+    "#### Generate Config Group"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7c2402aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simple grid search across all sets of config knob values = 3 combinations in total\n",
+    "config_group = RFGridSearch(\n",
+    "    configs=config_set,\n",
+    "    trainer_type=\"DPO\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d0f3d1ab",
+   "metadata": {},
+   "source": [
+    "### Run Multi-Config Training"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b2cc2c0d",
+   "metadata": {},
+   "source": [
+    "##### DPO Training Metrics\n",
+    "\n",
+    "During training and evaluation, we track the following reward metrics:\n",
+    "\n",
+    "| Metric | Description |\n",
+    "|--------|-------------|\n",
+    "| `rewards/chosen` | Mean difference between policy and reference model log probabilities for chosen responses (scaled by β) |\n",
+    "| `rewards/rejected` | Mean difference between policy and reference model log probabilities for rejected responses (scaled by β) |\n",
+    "| `rewards/accuracies` | Percentage of cases where chosen rewards > rejected rewards |\n",
+    "| `rewards/margins` | Mean difference between chosen and rejected rewards |\n",
+    "\n",
+    "##### Mathematical Representation\n",
+    "\n",
+    "- **Chosen Rewards**: $\\beta * (\\log \\pi_\\theta(y_w|x) - \\log \\pi_{ref}(y_w|x))$\n",
+    "- **Rejected Rewards**: $\\beta * (\\log \\pi_\\theta(y_l|x) - \\log \\pi_{ref}(y_l|x))$\n",
+    "- **Accuracy**: $\\mathbb{E}[\\mathbf{1}[r_{chosen} > r_{rejected}]]$\n",
+    "- **Margin**: $\\mathbb{E}[r_{chosen} - r_{rejected}]$\n",
+    "\n",
+    "Where:\n",
+    "- $\\pi_\\theta$ = policy model\n",
+    "- $\\pi_{ref}$ = reference model  \n",
+    "- $y_w$ = chosen (winning) response\n",
+    "- $y_l$ = rejected (losing) response\n",
+    "- $\\beta$ = temperature parameter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8c2dcad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Launch training of all configs in the config_group with swap granularity of 3 chunks\n",
+    "experiment.run_fit(config_group, sample_create_model, train_dataset, eval_dataset=None, num_chunks=3, seed=42)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b0a81b9",
+   "metadata": {},
+   "source": [
+    "### End Current Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e85f16ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "experiment.end()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "lite",
+   "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}