{ "cells": [ { "cell_type": "markdown", "id": "cell-0", "metadata": {}, "source": [ "# GRPO Fine-tuning with Unsloth\n", "\n", "Fine-tuning requires a GPU. If you don't have one locally, you can run this notebook for free on [Google Colab](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/grpo_with_unsloth.ipynb) using a free NVIDIA T4 GPU instance.\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/grpo_with_unsloth.ipynb)\n", "\n", "### What's in this notebook?\n", "\n", "In this notebook you will learn how to use Group Relative Policy Optimization (GRPO) for reinforcement learning-based fine-tuning with Unsloth.\n", "We will use the [LFM2.5-1.2B-Instruct](https://docs.liquid.ai/docs/models/lfm25-1.2b-instruct) model and train it on mathematical reasoning tasks using the Open R1 Math dataset. GRPO is ideal for verifiable tasks where you can programmatically evaluate the correctness of model outputs, such as math problems, code generation, and structured data tasks.\n", "\n", "We will cover\n", "- Environment setup\n", "- Data preparation for GRPO\n", "- Reward function design\n", "- Model training with GRPO\n", "- Local inference with your new model\n", "- Model saving and exporting it into the format you need for **deployment**.\n", "\n", "### Deployment options\n", "\n", "LFM2.5 models are small and efficient, enabling deployment across a wide range of platforms:\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Deployment TargetUse Case
📱 AndroidMobile apps on Android devices
📱 iOSMobile apps on iPhone/iPad
🍎 Apple Silicon MacLocal inference on Mac with MLX
🦙 llama.cppLocal deployments on any hardware
🦙 OllamaLocal inference with easy setup
🖥️ LM StudioDesktop app for local inference
⚡ vLLMCloud deployments with high throughput
☁️ ModalServerless cloud deployment
🏗️ BasetenProduction ML infrastructure
🚀 FalFast inference API
" ] }, { "cell_type": "markdown", "id": "cell-0b", "metadata": {}, "source": [ "### Need help building with our models and tools?\n", "Join the Liquid AI Discord Community and ask.\n", "\n", "\"Join\n", "\n", "And now, let the fine tune begin!" ] }, { "cell_type": "markdown", "metadata": { "id": "jPzMt5cOM4ze" }, "source": [ "### Installation" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "AgvsTCcGM4zf" }, "outputs": [], "source": [ "%%capture\n", "import os\n", "os.environ[\"UNSLOTH_VLLM_STANDBY\"] = \"1\" # [NEW] Extra 30% context lengths!\n", "if \"COLAB_\" not in \"\".join(os.environ.keys()):\n", " # If you're not in Colab, just use pip install or uv pip install\n", " !pip install unsloth vllm\n", "else:\n", " pass # For Colab / Kaggle, we need extra instructions hidden below \\/" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "iK3fwlTDM4zg" }, "outputs": [], "source": [ "#@title Colab Extra Install { display-mode: \"form\" }\n", "%%capture\n", "import os\n", "!pip install --upgrade -qqq uv\n", "if \"COLAB_\" not in \"\".join(os.environ.keys()):\n", " # If you're not in Colab, just use pip install!\n", " !pip install unsloth vllm\n", "else:\n", " try: import numpy, PIL; get_numpy = f\"numpy=={numpy.__version__}\"; get_pil = f\"pillow=={PIL.__version__}\"\n", " except: get_numpy = \"numpy\"; get_pil = \"pillow\"\n", " try: import subprocess; is_t4 = \"Tesla T4\" in str(subprocess.check_output([\"nvidia-smi\"]))\n", " except: is_t4 = False\n", " get_vllm, get_triton = (\"vllm==0.9.2\", \"triton==3.2.0\") if is_t4 else (\"vllm==0.11.2\", \"triton\")\n", " !uv pip install -qqq --upgrade \\\n", " unsloth {get_vllm} {get_numpy} {get_pil} torchvision bitsandbytes xformers\n", " !uv pip install -qqq {get_triton}\n", "!uv pip install transformers==4.57.1\n", "!uv pip install --no-deps trl==0.22.2" ] }, { "cell_type": "markdown", "metadata": { "id": "ZkH_y8UC9lvv" }, "source": [ "### Unsloth" ] }, { "cell_type": "markdown", "metadata": { "id": "jN75nmdx9lvw" }, "source": [ "Goal: To convert `LFM2.5-1.2B-Instruct` into a reasoning model via GRPO by using OpenR1's Math dataset.\n", "\n", "We first pre fine-tune the model to make GRPO skip trying to match formatting - this speeds GRPO up." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 383, "referenced_widgets": [ "36088fd6d4e54b95880f8fa8853f482a", "b0eb1dfbf88d4add8489983432000355", "e1d521db18044397bbd7692f142e9d6f", "a684e247871c4bb4819449697f5ef22b", "747fcfdf0d224809a132d9c500c0a430", "605f8f517f8445b59bbba86e523f0b3f", "e6e18ccd377a4fa98ecbf8e9e39bc10e", "f9fcd0069f79442085af7543ed1a0476", "0ae2300faff140459aeb70ef11d0f46d", "f549fcda34ff4c0297bc045043f0e6a6", "b17f3485f5634e6d9a05c71872a63500", "2e4cd3d5b38348ca8d6ecb0930053b66", "5b2fa6d7eda74f06b470331885d7d3b3", "1720694002b24a65b8a873a362f01c8e", "bd24744a0bc3496d91f70b9c679aeb3d", "8986deb8008646eb8082a28f9acab624", "ffec9897780440c08f4f60d3c8a12567", "d56567c6a508494b9c347d921747b717", "2280fcd994674fca9782ddef9e6f1b00", "a402e572908b4475b9a444789735958c", "e9417284b574497686f6485c794f4c88", "af30f0e3fcd34347a9382ba5c7276fd6", "ac4de5a1188844b4839be598a8c8d7d7", "026e6f202b264e5c937188357c97d986", "c1bbde8076ea4c99af9ed56aacfca691", "bef2ece9232342e7b4595f6e84b8fcaf", "953de5996a164b81990e0c76e602f722", "2711e4e096014f2f84e224c7da9a3c3c", "78a964d7fa87456fb50f8807b72fd103", "1a25dfeea87344b9ad22c5721c01205e", "0c796830b7d347149e473a377135a9be", "6046821f64374a40b8f2a8ca9b2ec090", "cd3d8febb1974865a8ded0a952516bea", "31eeb3541dd44ceda5bab3655444b987", "c91f8939836c40bcb98a10c9b7ea81f1", "6053c031745b432fa1c6a8c8118a16a7", "b882eb13f7de4631b48e8f049c22dd0b", "7c7f9da3e166420e9aab092d3395a1db", "f71dd3aa24d940e591a01f0fa11b6181", "554349c2ef374e39a04eac5254684e04", "d80b68cd197c4378897edeebb98f10c2", "848412c9479e44dc85b23ad80057b3b6", "fe8beb196f4f4afd8b42aaa89f4af897", "6c4d87d154f74290a711052aceef6fed", "1b2faed22c7e4184833741e7867e0e8d", "c35c689ae7f34906912ae6b10ec4a540", "c98e4b176d474df6b606338949736b54", "4fc7b609f9b444708fb6ab5138d84f3a", "bf2750788a6e42e2932de13e4fae3ccd", "f7e9ab3ad6884dafb901caf2ce7eb62d", "938f4721443f4473a47410e9320500d5", "f7c5c04b28234671952c0ae34eb45b47", "c47f624d82c14904927aa738aa22a8f2", "8852af8cb55544ee80c4f8ae8da4a1a5", "451f082df6054b62afc07131e4aa7e5c", "4a7008393d764ac8a25f12a9618f902a", "5f3f47bac0084ee994a6d758b9f5263f", "f6a8c52e0c6a486a8be532969071e3a4", "c9d74ca122ea49ffb9d4a94ce32d6365", "b9eae3713f29461887ce77ddd7d3d6d2", "917652e750d64d13a5dbbeaf4204148c", "54833f826bed4ba3aa378632d39b99c7", "07e9b432ee38416b8847b037bdbe334b", "d90520d20d224119b23d42321f075fb2", "b5c9b4d9bfd444859dc78c8b16235ea0", "cd31e44f371e40e2aa0a4c56283a9253" ] }, "id": "DkIvEkIIkEyB", "outputId": "b879d9a8-a5d2-430b-b712-adaa0c2a3bc4" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n", "🦥 Unsloth Zoo will now patch everything to make training faster!\n", "==((====))== Unsloth 2026.1.2: Fast Lfm2 patching. Transformers: 4.57.1.\n", " \\\\ /| NVIDIA L4. Num GPUs = 1. Max memory: 22.161 GB. Platform: Linux.\n", "O^O/ \\_/ \\ Torch: 2.9.1+cu128. CUDA: 8.9. CUDA Toolkit: 12.8. Triton: 3.5.1\n", "\\ / Bfloat16 = TRUE. FA [Xformers = 0.0.33.post2. FA2 = False]\n", " \"-____-\" Free license: http://github.com/unslothai/unsloth\n", "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n", "Unsloth: QLoRA and full finetuning all not selected. Switching to 16bit LoRA.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "36088fd6d4e54b95880f8fa8853f482a", "version_major": 2, "version_minor": 0 }, "text/plain": [ "model.safetensors: 0%| | 0.00/2.34G [00:00 0 ! Suggested 8, 16, 32, 64, 128\n", " target_modules = [\n", " \"q_proj\", \"k_proj\", \"v_proj\", \"out_proj\", \"in_proj\",\n", " \"w1\", \"w2\", \"w3\"\n", " ],\n", " lora_alpha = lora_rank*2, # *2 speeds up training\n", " use_gradient_checkpointing = \"unsloth\", # Reduces memory usage\n", " random_state = 3407,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "6mojUOPYE_1x", "outputId": "8b242b4c-569b-4fbe-fe76-59dddff8adb0" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "<|startoftext|><|im_start|>user\n", "Solve x^5 + 3x^4 - 10 = 3.<|im_end|>\n", "<|im_start|>assistant\n", "We are given the equation:\n", "\n", "$$\n", "x^5 + 3x^4 - 10 = 3\n", "$$\n", "\n", "### Step 1: Move all terms to one side\n", "\n", "$$\n", "x^5 + 3x^4 - 10 - 3 = 0\n", "$$\n", "\n", "$$\n", "x^5 + 3x^4 - 13 = 0\n", "$$\n", "\n", "Now we have the equation:\n", "\n", "$$\n", "x^5 + 3x^4 - 13 = 0\n", "$$\n", "\n", "This is a fifth-degree polynomial equation, which is generally difficult to solve algebraically. Let's try to find rational roots using the Rational Root Theorem.\n", "\n", "### Step 2: Apply Rational Root Theorem\n", "\n", "Possible rational roots are factors of the constant term (-13) divided by factors of the leading coefficient (1):\n", "\n", "Possible rational roots: $ \\pm1, \\pm13 $\n", "\n", "Try $ x = 1 $:\n", "\n", "$$\n", "1^5 + 3(1)^4 - 13 = 1 + 3 - 13 = -9 \\neq 0\n", "$$\n", "\n", "Try $ x = -1 $:\n", "\n", "$$\n", "(-1)^5 + 3(-1)^4 - 13 = -1 + 3 - 13 = -11 \\neq 0\n", "$$\n", "\n", "Try $ x = 13 $: too large, not likely.\n", "\n", "Try $ x = -13 $: also too large.\n", "\n", "So no rational roots. Let's try to estimate the solution numerically.\n", "\n", "### Step " ] } ], "source": [ "messages = [\n", " {\"role\": \"user\", \"content\": \"Solve x^5 + 3x^4 - 10 = 3.\"},\n", "]\n", "inputs = tokenizer.apply_chat_template(\n", " messages,\n", " add_generation_prompt = True,\n", " return_tensors = \"pt\",\n", " return_dict = True,\n", ").to(model.device)\n", "from transformers import TextStreamer\n", "_ = model.generate(**inputs, max_new_tokens = 4096, use_cache = True, streamer = TextStreamer(tokenizer))\n" ] }, { "cell_type": "markdown", "metadata": { "id": "W9DuiVRLhMco" }, "source": [ "### GRPO chat template\n", "Since we're using an instruct model, we need to ask the model to output `` and `` tokens in the system prompt. This will guide the model's responses and make the final answers easier to parse." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6UjowCbT-cFz" }, "outputs": [], "source": [ "reasoning_start = \"\"\n", "reasoning_end = \"\"\n", "\n", "system_prompt = \\\n", "f\"Think step by step between {reasoning_start} and {reasoning_end} and then\"\\\n", "\"output your answer. Output the final answer in \\\\boxed{} format.\"\n", "system_prompt" ] }, { "cell_type": "markdown", "metadata": { "id": "vEcLdymBEHdk" }, "source": [ "Let's see how the chat template behaves on an example:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "BciEDYSSYFNj" }, "outputs": [], "source": [ "tokenizer.apply_chat_template([\n", " {\"role\" : \"system\", \"content\" : system_prompt},\n", " {\"role\" : \"user\", \"content\" : \"What is 1+1?\"},\n", " {\"role\" : \"assistant\", \"content\" : f\"{reasoning_start}I think it's 2.{reasoning_end}2\"},\n", " {\"role\" : \"user\", \"content\" : \"What is 2+2?\"},\n", "], tokenize = False, add_generation_prompt = True)" ] }, { "cell_type": "markdown", "metadata": { "id": "_mdsuGjxHrjT" }, "source": [ "### Pre fine-tuning for formatting\n", "We now use a subset of NVIDIA's [Open Math Reasoning dataset](https://huggingface.co/datasets/nvidia/OpenMathReasoning) which was filtered to only include high quality DeepSeek R1 traces.\n", "\n", "We'll only filter ~59 or so examples to first \"prime\" / pre fine-tune the model to understand our custom GRPO formatting." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "AXxM2lStVIkd" }, "outputs": [], "source": [ "from datasets import load_dataset\n", "import pandas as pd\n", "import numpy as np\n", "\n", "dataset = load_dataset(\"unsloth/OpenMathReasoning-mini\", split = \"cot\")\n", "dataset = dataset.to_pandas()[\n", " [\"expected_answer\", \"problem\", \"generated_solution\"]\n", "]\n", "\n", "# Try converting to number - if not, replace with NaN\n", "is_number = pd.to_numeric(pd.Series(dataset[\"expected_answer\"]), errors = \"coerce\").notnull()\n", "# Select only numbers\n", "dataset = dataset.iloc[np.where(is_number)[0]]\n", "\n", "dataset" ] }, { "cell_type": "markdown", "metadata": { "id": "JVRFqoSdIEVK" }, "source": [ "We have to format the dataset to follow our GRPO style formatting:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Z9ydcV_Abfi6" }, "outputs": [], "source": [ "def format_dataset(x):\n", " problem = x[\"problem\"]\n", " solution = x[\"generated_solution\"]\n", "\n", " return [\n", " {\"role\" : \"system\", \"content\" : system_prompt},\n", " {\"role\" : \"user\", \"content\" : problem},\n", " {\"role\" : \"assistant\", \"content\" : solution},\n", " ]\n", "\n", "dataset[\"Messages\"] = dataset.apply(format_dataset, axis = 1)" ] }, { "cell_type": "markdown", "metadata": { "id": "X5NI47rOIRP2" }, "source": [ "Check to see if it worked:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "LTdXBKcslhRH" }, "outputs": [], "source": [ "tokenizer.apply_chat_template(dataset[\"Messages\"][0], tokenize = False)" ] }, { "cell_type": "markdown", "metadata": { "id": "iHV9BXYiIYaq" }, "source": [ "Let's truncate the pre fine-tuning dataset to `max_seq_length/2` since we don't want too long reasoning traces.\n", "\n", "Note this might take 2 minutes!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "MBHFlRbae9_s" }, "outputs": [], "source": [ "dataset[\"N\"] = dataset[\"Messages\"].apply(lambda x: len(tokenizer.apply_chat_template(x)))\n", "\n", "dataset = dataset.loc[dataset[\"N\"] <= max_seq_length/2].copy()\n", "dataset.shape" ] }, { "cell_type": "markdown", "metadata": { "id": "E6NkUCAGIj8N" }, "source": [ "We then tokenize the messages and convert it to a Hugging Face compatible dataset format:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "3rgdtiV_f5hx" }, "outputs": [], "source": [ "from datasets import Dataset\n", "\n", "dataset[\"text\"] = tokenizer.apply_chat_template(dataset[\"Messages\"].values.tolist(), tokenize = False)\n", "dataset = Dataset.from_pandas(dataset)\n", "dataset" ] }, { "cell_type": "markdown", "metadata": { "id": "bAQJjQrYKzOk" }, "source": [ "Let's now pre fine-tune the model so it follows our custom GRPO formatting!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "woYi0SSygpqp" }, "outputs": [], "source": [ "from trl import SFTTrainer, SFTConfig\n", "trainer = SFTTrainer(\n", " model = model,\n", " tokenizer = tokenizer,\n", " train_dataset = dataset,\n", " args = SFTConfig(\n", " dataset_text_field = \"text\",\n", " per_device_train_batch_size = 8,\n", " gradient_accumulation_steps = 1, # Use GA to mimic batch size!\n", " warmup_steps = 10,\n", " # num_train_epochs = 1, # Set this for 1 full training run.\n", " max_steps=100,\n", " learning_rate = 2e-4, # Reduce to 2e-5 for long training runs\n", " logging_steps = 5,\n", " optim = \"adamw_8bit\",\n", " weight_decay = 0.01,\n", " lr_scheduler_type = \"linear\",\n", " seed = 3407,\n", " report_to = \"none\", # Use this for WandB etc\n", " ),\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "1JHECEwE6XYk" }, "outputs": [], "source": [ "trainer.train()" ] }, { "cell_type": "markdown", "metadata": { "id": "DRMBNUBgLC8T" }, "source": [ "Let's check if the model has learnt to follow the custom format:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "0jdu99rRnUq1" }, "outputs": [], "source": [ "import gc\n", "for _ in range(5):\n", " torch.cuda.empty_cache()\n", " gc.collect()" ] }, { "cell_type": "markdown", "metadata": { "id": "6E4qYyb2znTp" }, "source": [ "Let's use Unsloth normal inference (not vLLM):" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "9HJxrS76h3Ds" }, "outputs": [], "source": [ "text = tokenizer.apply_chat_template(\n", " dataset[0][\"Messages\"][:2],\n", " tokenize = False,\n", " add_generation_prompt = True, # Must add for generation\n", ")\n", "\n", "from transformers import TextStreamer\n", "_ = model.generate(\n", " **tokenizer(text, return_tensors = \"pt\").to(\"cuda\"),\n", " temperature = 1,\n", " max_new_tokens = 4096,\n", " streamer = TextStreamer(tokenizer, skip_prompt = False),\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "AtZ3qGOALF95" }, "source": [ "Yes it did follow the formatting! Great! Let's remove some items before the GRPO step" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "YWSZ0DET7bob" }, "outputs": [], "source": [ "del dataset\n", "torch.cuda.empty_cache()\n", "import gc\n", "gc.collect()" ] }, { "cell_type": "markdown", "metadata": { "id": "7KGgPgk_5S8r" }, "source": [ "### Data Prep\n", "\n", "\n", "We're using Hugging Face's [Open R1 Math dataset](https://huggingface.co/datasets/open-r1/DAPO-Math-17k-Processed). You can also utilize OpenAI's famous [GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "o7-eUrQn-OzE" }, "outputs": [], "source": [ "from datasets import load_dataset\n", "dataset = load_dataset(\"open-r1/DAPO-Math-17k-Processed\", \"en\", split = \"train\")\n", "dataset" ] }, { "cell_type": "markdown", "metadata": { "id": "3b00gUsS-ROW" }, "source": [ "Let's look at the first row:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "siopxjG8-ReF" }, "outputs": [], "source": [ "dataset[0][\"prompt\"]" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "KGupRQqD-Wcf" }, "outputs": [], "source": [ "dataset[0][\"solution\"]" ] }, { "cell_type": "markdown", "metadata": { "id": "CmnXj6hn-Ydi" }, "source": [ "In GSM8K, we notice all answers like about have a ####, so we extract it. But for the Open R1 dataset, we can skip the below." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "8JJGXKdJ-Zl_" }, "outputs": [], "source": [ "def extract_hash_answer(text):\n", " # if \"####\" not in text: return None\n", " # return text.split(\"####\")[1].strip()\n", " return text\n", "extract_hash_answer(dataset[0][\"solution\"])" ] }, { "cell_type": "markdown", "metadata": { "id": "K30CygaU-dir" }, "source": [ "Let's map the dataset! and see the first row:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "qyEVI972-d3n" }, "outputs": [], "source": [ "dataset = dataset.map(lambda x: {\n", " \"prompt\" : [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": x[\"prompt\"]},\n", " ],\n", " \"answer\": extract_hash_answer(x[\"solution\"]),\n", "})\n", "dataset[0]" ] }, { "cell_type": "markdown", "metadata": { "id": "-9m8eR9T-gMh" }, "source": [ "We create a regex format to match the reasoning sections and answers:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "iQwjTjNz-gY_" }, "outputs": [], "source": [ "import re\n", "\n", "# Match the answer after , extracting \\boxed{} content\n", "solution_end_regex = r\"[\\s]*\" + \\\n", " \"(?:\" + re.escape(tokenizer.eos_token) + \")?\"\n", "\n", "match_format = re.compile(\n", " r\"[\\s]*\"\n", " r\".*?\"\n", " r\"\\\\boxed\\{([^{}]*(?:\\{[^{}]*\\}[^{}]*)*)\\}\"\n", " r\"[\\s]*$\",\n", " flags=re.MULTILINE | re.DOTALL\n", ")\n", "match_format" ] }, { "cell_type": "markdown", "metadata": { "id": "OycMneOq-iNC" }, "source": [ "We verify it works:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ndzHnQ_6-jHt" }, "outputs": [], "source": [ "match_format.findall(\n", " \"Let me think!\"\\\n", " \"\\boxed\\{2\\}\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "eRMDAzDk2x6t" }, "outputs": [], "source": [ "match_format.findall(\n", " \"Let me think!\"\\\n", " \"\\boxed\\{2\\} \\n\\n\",\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "weOjmO5l-kl3" }, "source": [ "We now want to create a reward function to match the format exactly - we reward it with 3 points if it succeeds:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "qgFNXORy-lpO" }, "outputs": [], "source": [ "def match_format_exactly(completions, **kwargs):\n", " scores = []\n", " for completion in completions:\n", " score = 0\n", " response = completion[0][\"content\"]\n", " # Match if format is seen exactly!\n", " if match_format.search(response) is not None: score += 3.0\n", " scores.append(score)\n", " return scores" ] }, { "cell_type": "markdown", "metadata": { "id": "Gf69i2WT-m4K" }, "source": [ "If it fails, we want to reward the model if it at least follows the format partially, by counting each symbol:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "cUfHzCVx-nGK" }, "outputs": [], "source": [ "def match_format_approximately(completions, **kwargs):\n", " scores = []\n", " for completion in completions:\n", " score = 0\n", " response = completion[0][\"content\"]\n", " # Count how many keywords are seen - we penalize if too many!\n", " # If we see 1, then plus some points!\n", "\n", " # No need to reward since we always prepend it!\n", " score += 0.5 if response.count(reasoning_start) == 1 else -1.0\n", " score += 0.5 if response.count(reasoning_end) == 1 else -1.0\n", " scores.append(score)\n", " return scores" ] }, { "cell_type": "markdown", "metadata": { "id": "9wAUWwtE-s6n" }, "source": [ "Finally, we want to extract the generated answer, and reward or penalize it! We also reward it based on how close the answer is to the true one via ratios:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "hmtI_8gg-uIE" }, "outputs": [], "source": [ "def check_answer(prompts, completions, answer, **kwargs):\n", " question = prompts[0][-1][\"content\"]\n", " responses = [completion[0][\"content\"] for completion in completions]\n", "\n", " extracted_responses = [\n", " guess.group(1)\n", " if (guess := match_format.search(r)) is not None else None \\\n", " for r in responses\n", " ]\n", "\n", " scores = []\n", " for guess, true_answer in zip(extracted_responses, answer):\n", " score = 0\n", " if guess is None:\n", " scores.append(-2.0)\n", " continue\n", " # Correct answer gets 5 points!\n", " if guess == true_answer:\n", " score += 5.0\n", " # Match if spaces are seen, but less reward\n", " elif guess.strip() == true_answer.strip():\n", " score += 3.5\n", " else:\n", " # We also reward it if the answer is close via ratios!\n", " # Ie if the answer is within some range, reward it!\n", " try:\n", " ratio = float(guess) / float(true_answer)\n", " if ratio >= 0.9 and ratio <= 1.1: score += 2.0\n", " elif ratio >= 0.8 and ratio <= 1.2: score += 1.5\n", " else: score -= 2.5 # Penalize wrong answers\n", " except:\n", " score -= 4.5 # Penalize\n", " scores.append(score)\n", " return scores" ] }, { "cell_type": "markdown", "metadata": { "id": "atMyfhXh-v3R" }, "source": [ "Also sometimes it might not be 1 number as the answer, but like a sentence for example \"The solution is $20\" -> we extract 20.\n", "\n", "We also remove possible commas for example as in 123,456" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "AVW0kL8q-wL5" }, "outputs": [], "source": [ "match_numbers = re.compile(\n", " r\".*?[\\s]{0,}([-]?[\\d\\.\\,]{1,})\",\n", " flags = re.MULTILINE | re.DOTALL\n", ")\n", "print(match_numbers.findall(\"\\boxed{0.34}\"))\n", "print(match_numbers.findall(\"Answer is \\(\\boxed{44}\\).\"))\n", "print(match_numbers.findall(\"The answer is -0.234\"))\n", "print(match_numbers.findall(\"17\"))" ] }, { "cell_type": "markdown", "metadata": { "id": "RbfaaAywNHHh" }, "source": [ "We now prepare our main function which will print out the generated responses and the true answer, along with another reward function which converts text to float via `float` and sees if it's the same." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "GjBFrttr-y1_" }, "outputs": [], "source": [ "global PRINTED_TIMES\n", "PRINTED_TIMES = 0\n", "global PRINT_EVERY_STEPS\n", "PRINT_EVERY_STEPS = 5\n", "\n", "def check_numbers(prompts, completions, answer, **kwargs):\n", " question = prompts[0][-1][\"content\"]\n", " responses = [completion[0][\"content\"] for completion in completions]\n", "\n", " extracted_responses = [\n", " guess.group(1)\n", " if (guess := match_numbers.search(r)) is not None else None \\\n", " for r in responses\n", " ]\n", "\n", " scores = []\n", " # Print only every few steps\n", " global PRINTED_TIMES\n", " global PRINT_EVERY_STEPS\n", " if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:\n", " print(\n", " '*'*20 + f\"Question:\\n{question}\", f\"\\nAnswer:\\n{answer[0]}\", f\"\\nResponse:\\n{responses[0]}\", f\"\\nExtracted:\\n{extracted_responses[0]}\"\n", " )\n", " PRINTED_TIMES += 1\n", "\n", " for guess, true_answer in zip(extracted_responses, answer):\n", " if guess is None:\n", " scores.append(-2.5)\n", " continue\n", " # Convert to numbers\n", " try:\n", " true_answer = float(true_answer.strip())\n", " # Remove commas like in 123,456\n", " guess = float(guess.strip().replace(\",\", \"\"))\n", " scores.append(3.5 if guess == true_answer else -1.5)\n", " except:\n", " scores.append(0)\n", " continue\n", " return scores" ] }, { "cell_type": "markdown", "metadata": { "id": "fgOR3wJ_AyLr" }, "source": [ "Get the top 90% prompt length so we don't accidentally truncate them!\n", "\n", "Ie we'll remove the top 10% long prompts." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6EgAi4Q5fGE-" }, "outputs": [], "source": [ "tokenized = dataset.map(\n", " lambda x: {\"tokens\" : tokenizer.apply_chat_template(x[\"prompt\"], add_generation_prompt = True, tokenize = True)},\n", " batched = True,\n", ")\n", "print(tokenizer.decode(tokenized[0][\"tokens\"]))\n", "tokenized = tokenized.map(lambda x: {\"L\" : len(x[\"tokens\"])})\n", "\n", "import numpy as np\n", "maximum_length = int(np.quantile(tokenized[\"L\"], 0.9))\n", "print(\"Max Length = \", maximum_length)\n", "\n", "# Filter only samples smaller than 90% max length\n", "dataset = dataset.select(np.where(np.array(tokenized[\"L\"]) <= maximum_length)[0])\n", "del tokenized" ] }, { "cell_type": "markdown", "metadata": { "id": "9-IOMhVg-2AM" }, "source": [ "\n", "### Train the model\n", "\n", "Now set up GRPO Trainer and all configurations!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ptqkXK2D4d6p" }, "outputs": [], "source": [ "max_prompt_length = maximum_length + 1 # + 1 just in case!\n", "max_completion_length = max_seq_length - max_prompt_length\n", "\n", "from vllm import SamplingParams\n", "vllm_sampling_params = SamplingParams(\n", " min_p = 0.1,\n", " top_p = 1.0,\n", " top_k = -1,\n", " seed = 3407,\n", " stop = [tokenizer.eos_token],\n", " include_stop_str_in_output = True,\n", ")\n", "\n", "from trl import GRPOConfig, GRPOTrainer\n", "training_args = GRPOConfig(\n", " vllm_sampling_params = vllm_sampling_params,\n", " temperature = 1.0,\n", " learning_rate = 5e-6,\n", " weight_decay = 0.01,\n", " warmup_ratio = 0.1,\n", " lr_scheduler_type = \"linear\",\n", " optim = \"adamw_8bit\",\n", " logging_steps = 1,\n", " per_device_train_batch_size = 4,\n", " gradient_accumulation_steps = 1, # Increase to 4 for smoother training\n", " num_generations = 4, # Decrease if out of memory\n", " max_prompt_length = max_prompt_length,\n", " max_completion_length = max_completion_length,\n", " # num_train_epochs = 1, # Set to 1 for a full training run\n", " max_steps = 10,\n", " save_steps = 100,\n", " report_to = \"none\", # Can use Weights & Biases\n", " output_dir = \"outputs\",\n", "\n", " # For optional training + evaluation\n", " # fp16_full_eval = True,\n", " # per_device_eval_batch_size = 4,\n", " # eval_accumulation_steps = 1,\n", " # eval_strategy = \"steps\",\n", " # eval_steps = 1,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "r9Mv8UZO5hz-" }, "source": [ "And let's run the trainer! If you scroll up, you'll see a table of rewards. The goal is to see the `reward` column increase!\n", "\n", "You might have to wait 150 to 200 steps for any action. You'll probably get 0 reward for the first 100 steps. Please be patient!\n", "\n", "| Step | Training Loss | reward | reward_std | completion_length | kl |\n", "|------|---------------|-----------|------------|-------------------|----------|\n", "| 1 | 0.000000 | 0.125000 | 0.000000 | 200.000000 | 0.000000 |\n", "| 2 | 0.000000 | 0.072375 | 0.248112 | 200.000000 | 0.000000 |\n", "| 3 | 0.000000 | -0.079000 | 0.163776 | 182.500000 | 0.000005 |\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "vzOuSVCL_GA9" }, "outputs": [], "source": [ "# For optional training + evaluation\n", "# new_dataset = dataset.train_test_split(test_size = 0.01)\n", "\n", "trainer = GRPOTrainer(\n", " model = model,\n", " processing_class = tokenizer,\n", " reward_funcs = [\n", " match_format_exactly,\n", " match_format_approximately,\n", " check_answer,\n", " check_numbers,\n", " ],\n", " args = training_args,\n", " train_dataset = dataset,\n", "\n", " # For optional training + evaluation\n", " # train_dataset = new_dataset[\"train\"],\n", " # eval_dataset = new_dataset[\"test\"],\n", ")\n", "trainer.train()" ] }, { "cell_type": "markdown", "metadata": { "id": "tlaUdxC_VHpz" }, "source": [ "\n", "### Inference\n", "Now let's try the model we just trained! First, let's first try the model without any GRPO trained:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "qtcz_lpbVC92" }, "outputs": [], "source": [ "messages = [\n", " {\"role\": \"user\", \"content\": \"What is the sqrt of 101?\"},\n", "]\n", "inputs = tokenizer.apply_chat_template(\n", " messages,\n", " add_generation_prompt = True,\n", " return_tensors = \"pt\",\n", " return_dict = True,\n", ").to(model.device)\n", "_ = model.generate(\n", " **inputs,\n", " temperature = 0.1,\n", " top_k = 50,\n", " top_p = 0.1,\n", " repetition_penalty=1.05,\n", " max_new_tokens = 1024,\n", " use_cache = True,\n", " streamer = TextStreamer(tokenizer)\n", ")\n" ] }, { "cell_type": "markdown", "metadata": { "id": "Colxz9TAVMsi" }, "source": [ "And now with the LoRA we just trained with GRPO - we first save the LoRA first!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "AL-BcuB1VLIv" }, "outputs": [], "source": [ "model.save_lora(\"grpo_saved_lora\")" ] }, { "cell_type": "markdown", "metadata": { "id": "a4LMOBl8boGX" }, "source": [ "Verify LoRA is actually trained!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "4SfdI-ERbpiw" }, "outputs": [], "source": [ "from safetensors import safe_open\n", "\n", "tensors = {}\n", "with safe_open(\"grpo_saved_lora/adapter_model.safetensors\", framework = \"pt\") as f:\n", " # Verify both A and B are non zero\n", " for key in f.keys():\n", " tensor = f.get_tensor(key)\n", " n_zeros = (tensor == 0).sum() / tensor.numel()\n", " assert(n_zeros.item() != tensor.numel())" ] }, { "cell_type": "markdown", "metadata": { "id": "CwpbwnDBVRLg" }, "source": [ "Now we load the LoRA and test:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "zf_OY5WMVOxF" }, "outputs": [], "source": [ "messages = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": \"What is the sqrt of 101?\"},\n", "]\n", "\n", "inputs = tokenizer.apply_chat_template(\n", " messages,\n", " add_generation_prompt = True,\n", " return_tensors = \"pt\",\n", " return_dict = True,\n", ").to(model.device)\n", "\n", "_ = model.generate(\n", " **inputs,\n", " temperature = 0.1,\n", " top_k = 50,\n", " top_p = 0.1,\n", " repetition_penalty=1.05,\n", " max_tokens = 1024,\n", " use_cache = True,\n", " streamer = TextStreamer(tokenizer)\n", ")\n" ] }, { "cell_type": "markdown", "metadata": { "id": "6aDgFfhFYIAS" }, "source": [ "Our reasoning model is much better - it's not always correct, since we only trained it for an hour or so - it'll be better if we extend the sequence length and train for longer!" ] }, { "cell_type": "markdown", "metadata": { "id": "-NUEmHFSYNTp" }, "source": [ "\n", "### Saving to float16 for VLLM\n", "\n", "We also support saving to `float16` directly. Select `merged_16bit` for float16 or `merged_4bit` for int4. We also allow `lora` adapters as a fallback. Use `push_to_hub_merged` to upload to your Hugging Face account! You can go to https://huggingface.co/settings/tokens for your personal tokens." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "NjXGTkp7YNtB" }, "outputs": [], "source": [ "# Merge to 16bit\n", "if False: model.save_pretrained_merged(\"model\", tokenizer, save_method = \"merged_16bit\",)\n", "if False: model.push_to_hub_merged(\"hf/model\", tokenizer, save_method = \"merged_16bit\", token = \"\")\n", "\n", "# Merge to 4bit\n", "if False: model.save_pretrained_merged(\"model\", tokenizer, save_method = \"merged_4bit\",)\n", "if False: model.push_to_hub_merged(\"hf/model\", tokenizer, save_method = \"merged_4bit\", token = \"\")\n", "\n", "# Just LoRA adapters\n", "if False:\n", " model.save_pretrained(\"model\")\n", " tokenizer.save_pretrained(\"model\")\n", "if False:\n", " model.push_to_hub(\"hf/model\", token = \"\")\n", " tokenizer.push_to_hub(\"hf/model\", token = \"\")\n" ] }, { "cell_type": "markdown", "metadata": { "id": "52WMb3k_YPt8" }, "source": [ "### GGUF / llama.cpp Conversion\n", "To save to `GGUF` / `llama.cpp`, we support it natively now! We clone `llama.cpp` and we default save it to `q8_0`. We allow all methods like `q4_k_m`. Use `save_pretrained_gguf` for local saving and `push_to_hub_gguf` for uploading to HF.\n", "\n", "Some supported quant methods (full list on our [Wiki page](https://github.com/unslothai/unsloth/wiki#gguf-quantization-options)):\n", "* `q8_0` - Fast conversion. High resource use, but generally acceptable.\n", "* `q4_k_m` - Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q4_K.\n", "* `q5_k_m` - Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q5_K.\n", "\n", "[**NEW**] To finetune and auto export to Ollama, try our [Ollama notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-Ollama.ipynb)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "QyEjW-WuYQIm" }, "outputs": [], "source": [ "# Save to 8bit Q8_0\n", "if False: model.save_pretrained_gguf(\"model\", tokenizer,)\n", "# Remember to go to https://huggingface.co/settings/tokens for a token!\n", "# And change hf to your username!\n", "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, token = \"\")\n", "\n", "# Save to 16bit GGUF\n", "if False: model.save_pretrained_gguf(\"model\", tokenizer, quantization_method = \"f16\")\n", "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, quantization_method = \"f16\", token = \"\")\n", "\n", "# Save to q4_k_m GGUF\n", "if False: model.save_pretrained_gguf(\"model\", tokenizer, quantization_method = \"q4_k_m\")\n", "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, quantization_method = \"q4_k_m\", token = \"\")\n", "\n", "# Save to multiple GGUF options - much faster if you want multiple!\n", "if False:\n", " model.push_to_hub_gguf(\n", " \"hf/model\", # Change hf to your username!\n", " tokenizer,\n", " quantization_method = [\"q4_k_m\", \"q8_0\", \"q5_k_m\",],\n", " token = \"\",\n", " )" ] }, { "cell_type": "markdown", "metadata": { "id": "V15Yhj1V9lwG" }, "source": [ "Now, use the `model-unsloth.gguf` file or `model-unsloth-Q4_K_M.gguf` file in llama.cpp.\n", "\n", "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", "\n", "Some other links:\n", "1. Train your own reasoning model - Llama GRPO notebook [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb)\n", "2. Saving finetunes to Ollama. [Free notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-Ollama.ipynb)\n", "3. Llama 3.2 Vision finetuning - Radiography use case. [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(11B)-Vision.ipynb)\n", "6. See notebooks for DPO, ORPO, Continued pretraining, conversational finetuning and more on our [documentation](https://docs.unsloth.ai/get-started/unsloth-notebooks)!\n", "\n", "
\n", " \n", " \n", " \n", "\n", " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", "
\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "L4", "machine_shape": "hm", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }