{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a3PTFH-H9Ozk"
   },
   "source": [
    "# 💧 LFM2 - DPO with TRL\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/dpo_with_trl.ipynb)\n",
    "\n",
    "\n",
    "This tutorial demonstrates how to fine-tune our LFM2 models, e.g. [`LiquidAI/LFM2.5-1.2B-Instruct`](https://huggingface.co/LiquidAI/LFM2.5-1.2B-Instruct), using the TRL library.\n",
    "\n",
    "Follow along if it's your first time using trl, or take single code snippets for your own workflow\n",
    "\n",
    "## 🎯 What You'll Find:\n",
    "- **DPO** (Direct Preference Optimization) - Align with human preferences  \n",
    "\n",
    "## 📋 Prerequisites:\n",
    "- **GPU Runtime**: Select GPU in `Runtime` → `Change runtime type`\n",
    "- **Hugging Face Account**: For accessing models and datasets\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x0RPLu2h9ome"
   },
   "source": [
    "# 📦 Installation & Setup\n",
    "\n",
    "First, let's install all the required packages:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3FIcp_wo9nsR"
   },
   "outputs": [],
   "source": [
    "!pip install \"transformers>=5.11.0\" \"trl>=1.5.1\" \"peft>=0.15.2\" \"torchao>=0.16.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "41UEf1uxCd6m"
   },
   "source": [
    "Let's now verify the packages are installed correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 17861,
     "status": "ok",
     "timestamp": 1781338300945,
     "user": {
      "displayName": "Maxime Labonne",
      "userId": "06977289147474683598"
     },
     "user_tz": -60
    },
    "id": "bSJgYtHT_Os4",
    "outputId": "cde2735f-a451-4eca-e33d-b70b6f72b3c4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📦 PyTorch version: 2.11.0+cu128\n",
      "🤗 Transformers version: 5.12.0\n",
      "📊 TRL version: 1.6.0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import transformers\n",
    "import trl\n",
    "import os\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
    "\n",
    "print(f\"📦 PyTorch version: {torch.__version__}\")\n",
    "print(f\"🤗 Transformers version: {transformers.__version__}\")\n",
    "print(f\"📊 TRL version: {trl.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v_uXLzxQ_rnK"
   },
   "source": [
    "# Loading the model from Transformers 🤗\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iA3erKM4-HhS"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "\n",
    "# Select a model to fine-tune from the list\n",
    "lfm_models = [\n",
    "    \"LiquidAI/LFM2.5-230M\",\n",
    "    \"LiquidAI/LFM2.5-350M\",\n",
    "    \"LiquidAI/LFM2.5-1.2B-Thinking\",\n",
    "    \"LiquidAI/LFM2.5-1.2B-Instruct\",\n",
    "    \"LiquidAI/LFM2.5-1.2B-JP-202606\",\n",
    "    \"LiquidAI/LFM2.5-2.6B\",\n",
    "    \"LiquidAI/LFM2-700M\",\n",
    "]\n",
    "\n",
    "# Model to fine-tune\n",
    "model_id = \"LiquidAI/LFM2.5-350M\"\n",
    "\n",
    "print(\"📚 Loading tokenizer...\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "\n",
    "print(\"🧠 Loading model...\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=\"auto\",\n",
    ")\n",
    "\n",
    "print(\"✅ Local model loaded successfully!\")\n",
    "print(f\"🔢 Parameters: {model.num_parameters():,}\")\n",
    "print(f\"📖 Vocab size: {len(tokenizer)}\")\n",
    "print(f\"💾 Model size: ~{model.num_parameters() * 2 / 1e9:.1f} GB (bfloat16)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o83NqFPNA_nk"
   },
   "source": [
    "# 🎯 Direct Preference Optimization (DPO + LoRA)\n",
    "\n",
    "DPO aligns the model with human preferences by learning from preference pairs (chosen vs rejected responses). This typically follows SFT training.\n",
    "\n",
    "DPO might be too compute heavy if you're running on one of the free-tier colab GPUs. Hence we use LoRA (Low-Rank Adaptation) to finetune the model by only training a small number of additional parameters. Perfect for limited compute resources!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0w-40n_XJL9H"
   },
   "source": [
    "## Load a DPO Dataset\n",
    "\n",
    "We will use [mlabonne/orpo-dpo-mix-40k](https://huggingface.co/datasets/mlabonne/orpo-dpo-mix-40k), limiting ourselves to the first 2k samples for brevity. Feel free to change the limit by changing the slicing index in the parameter `split`. The size of the validation data can be adjusted by changing `test_size`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ajq6EABUBAv_"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "print(\"📥 Loading DPO dataset...\")\n",
    "\n",
    "dataset_dpo = load_dataset(\"mlabonne/orpo-dpo-mix-40k\", split=\"train[:2000]\")\n",
    "dataset_dpo = dataset_dpo.train_test_split(test_size=0.1, seed=42)\n",
    "train_dataset_dpo, eval_dataset_dpo = dataset_dpo['train'], dataset_dpo['test']\n",
    "\n",
    "print(\"✅ DPO Dataset loaded:\")\n",
    "print(f\"   📚 Train samples: {len(train_dataset_dpo)}\")\n",
    "print(f\"   🧪 Eval samples: {len(eval_dataset_dpo)}\")\n",
    "\n",
    "sample = train_dataset_dpo[0]\n",
    "print(\"\\n📝 Single Sample:\")\n",
    "print(f\"   Prompt: {sample['prompt'][:100]}...\")\n",
    "print(f\"   ✅ Chosen: {sample['chosen'][:100]}...\")\n",
    "print(f\"   ❌ Rejected: {sample['rejected'][:100]}...\")\n",
    "\n",
    "def to_conversational_dpo(example):\n",
    "    return {\n",
    "        \"prompt\": example[\"chosen\"][:-1],       # shared conversation up to the final turn\n",
    "        \"chosen\": [example[\"chosen\"][-1]],      # preferred final assistant turn\n",
    "        \"rejected\": [example[\"rejected\"][-1]],  # dispreferred final assistant turn\n",
    "    }\n",
    "\n",
    "train_dataset_dpo = train_dataset_dpo.map(\n",
    "    to_conversational_dpo, remove_columns=train_dataset_dpo.column_names\n",
    ")\n",
    "eval_dataset_dpo = eval_dataset_dpo.map(\n",
    "    to_conversational_dpo, remove_columns=eval_dataset_dpo.column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v8or_m-UdKGc"
   },
   "source": [
    "## Wrap the model with PEFT\n",
    "\n",
    "We specify target modules that will be finetuned while the rest of the models weights remains frozen. Feel free to modify the `r` (rank) value:\n",
    "- higher -> better approximation of full-finetuning\n",
    "- lower -> needs even less compute resources\n",
    "\n",
    "You can skip this part if you have a premium GPU and want to go for a full finetune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 595,
     "status": "ok",
     "timestamp": 1781338353442,
     "user": {
      "displayName": "Maxime Labonne",
      "userId": "06977289147474683598"
     },
     "user_tz": -60
    },
    "id": "z7x5RhC5eLmc",
    "outputId": "a6307272-1103-4caa-8fb5-13c87f3b0f90"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 2,998,272 || all params: 357,482,240 || trainable%: 0.8387\n",
      "✅ LoRA configuration applied!\n",
      "🎛️  LoRA rank: 8\n",
      "📊 LoRA alpha: 16\n",
      "🎯 Target modules: {'out_proj', 'w3', 'in_proj', 'w2', 'v_proj', 'k_proj', 'w1', 'q_proj'}\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "GLU_MODULES = [\"w1\", \"w2\", \"w3\"]\n",
    "MHA_MODULES = [\"q_proj\", \"k_proj\", \"v_proj\", \"out_proj\"]\n",
    "CONV_MODULES = [\"in_proj\", \"out_proj\"]\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    inference_mode=False,\n",
    "    r=8,  # <- lower values = fewer parameters\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=GLU_MODULES + MHA_MODULES + CONV_MODULES,\n",
    "    bias=\"none\",\n",
    "    modules_to_save=None,\n",
    ")\n",
    "\n",
    "lora_model = get_peft_model(model, lora_config)\n",
    "lora_model.print_trainable_parameters()\n",
    "\n",
    "print(\"✅ LoRA configuration applied!\")\n",
    "print(f\"🎛️  LoRA rank: {lora_config.r}\")\n",
    "print(f\"📊 LoRA alpha: {lora_config.lora_alpha}\")\n",
    "print(f\"🎯 Target modules: {lora_config.target_modules}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yxuxfHu7Jopx"
   },
   "source": [
    "## Launch Training\n",
    "\n",
    "We are now ready to launch a DPO run with `DPOTrainer`, feel free to modify `DPOConfig` to play around with different configurations.\n",
    "\n",
    "Replace `lora_model` with `model` if you have a premium-tier in colab and want to run a full finetune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nGBgkf9LBHp6"
   },
   "outputs": [],
   "source": [
    "from trl import DPOConfig, DPOTrainer\n",
    "\n",
    "# DPO Training configuration\n",
    "dpo_config = DPOConfig(\n",
    "    output_dir=\"./lfm2-dpo\",\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=1,\n",
    "    learning_rate=1e-6,\n",
    "    lr_scheduler_type=\"linear\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"epoch\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),\n",
    ")\n",
    "\n",
    "# Create DPO trainer\n",
    "print(\"🏗️  Creating DPO trainer...\")\n",
    "dpo_trainer = DPOTrainer(\n",
    "    model=lora_model,\n",
    "    args=dpo_config,\n",
    "    train_dataset=train_dataset_dpo,\n",
    "    eval_dataset=eval_dataset_dpo,\n",
    "    processing_class=tokenizer,\n",
    ")\n",
    "\n",
    "# Start DPO training\n",
    "print(\"\\n🚀 Starting DPO training...\")\n",
    "dpo_trainer.train()\n",
    "\n",
    "print(\"🎉 DPO training completed!\")\n",
    "\n",
    "# Save the DPO model\n",
    "dpo_trainer.save_model()\n",
    "print(f\"💾 DPO model saved to: {dpo_config.output_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WIrCmjUsfNIg"
   },
   "source": [
    "## Save merged model\n",
    "\n",
    "If you have used LoRA. Merge the extra weights learned with LoRA back into the model to obtain a \"normal\" model checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RUJoU3BwfPTg"
   },
   "outputs": [],
   "source": [
    "print(\"\\n🔄 Merging LoRA weights...\")\n",
    "merged_model = lora_model.merge_and_unload()\n",
    "merged_model.save_pretrained(\"./lfm2-lora-merged\")\n",
    "tokenizer.save_pretrained(\"./lfm2-lora-merged\")\n",
    "print(\"💾 Merged model saved to: ./lfm2-lora-merged\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}