{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "Oh3eu2ULxWAn"
},
"source": [
"To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n",
"<div class=\"align-center\">\n",
"<a href=\"https://unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
"<a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord button.png\" width=\"145\"></a>\n",
"<a href=\"https://docs.unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a></a> Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐\n",
"</div>\n",
"\n",
"To install Unsloth on your own computer, follow the installation instructions on our Github page [here](https://docs.unsloth.ai/get-started/installing-+-updating).\n",
"\n",
"You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & [how to save it](#Save)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rNxpti3ExWAp"
},
"source": [
"### News"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yo-qTlvNxWAp"
},
"source": [
"**Read our [Gemma 3 blog](https://unsloth.ai/blog/gemma3) for what's new in Unsloth and our [Reasoning blog](https://unsloth.ai/blog/r1-reasoning) on how to train reasoning models.**\n",
"\n",
"Visit our docs for all our [model uploads](https://docs.unsloth.ai/get-started/all-our-models) and [notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks).\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MrkTn5S9xWAq"
},
"source": [
"### Installation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NV0BlxygxWAq"
},
"outputs": [],
"source": [
"%%capture\n",
"import os\n",
"\n",
"if \"COLAB_\" not in \"\".join(os.environ.keys()):\n",
" !pip install unsloth vllm\n",
"else:\n",
" # [NOTE] Do the below ONLY in Colab! Use [[pip install unsloth vllm]]\n",
" !pip install --no-deps unsloth vllm"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8lIsllEIxWAr"
},
"outputs": [],
"source": [
"# @title Colab Extra Install { display-mode: \"form\" }\n",
"%%capture\n",
"import os\n",
"\n",
"if \"COLAB_\" not in \"\".join(os.environ.keys()):\n",
" !pip install unsloth vllm\n",
"else:\n",
" !pip install --no-deps unsloth vllm\n",
" # [NOTE] Do the below ONLY in Colab! Use [[pip install unsloth vllm]]\n",
" # Skip restarting message in Colab\n",
" import sys\n",
" import re\n",
" import requests\n",
"\n",
" modules = list(sys.modules.keys())\n",
" for x in modules:\n",
" sys.modules.pop(x) if \"PIL\" in x or \"google\" in x else None\n",
" !pip install --no-deps bitsandbytes accelerate xformers==0.0.29.post3 peft trl triton cut_cross_entropy unsloth_zoo\n",
" !pip install sentencepiece protobuf datasets huggingface_hub hf_transfer\n",
"\n",
" # vLLM requirements - vLLM breaks Colab due to reinstalling numpy\n",
" f = requests.get(\n",
" \"https://raw.githubusercontent.com/vllm-project/vllm/refs/heads/main/requirements/common.txt\"\n",
" ).content\n",
" with open(\"vllm_requirements.txt\", \"wb\") as file:\n",
" file.write(re.sub(rb\"(transformers|numpy|xformers)[^\\n]{1,}\\n\", b\"\", f))\n",
" !pip install -r vllm_requirements.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "a6uTCUXFxWAr"
},
"source": [
"### Unsloth"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Joje4qPsyxM9"
},
"source": [
"Load up `Qwen 2.5 3B Instruct`, and set parameters"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000,
"referenced_widgets": [
"b46d7118f71043109c2dd00fb0f40557",
"32defd049d8a473c8ad8fc954c0e3d13",
"424078e3ad6c405c925bedfb76d18b3f",
"8d937714a44d422aa9851c02c2345e5c",
"25b26c398e314cb79a2f8bbb7e32e158",
"694494752ab2458ead83b0331af20122",
"1cc89fa4d7c744c69c3be1a3d0b3ba2d",
"15116ef21fae463ba9013c8b9dbd40b5",
"43fdf1fe592147538ec43410a1bf8e09",
"52ba857bdfc742a2aae9e9a949054e5c",
"589a3857411f49d4ae019465f536da45",
"5dafeefcf436462e81a9affe7f6c66d9",
"98d1697fb2054fd88432f027e592ecc0",
"b1adc20d1d8740f99a6ebd589f240688",
"f45a696d73054cd6af181a7ce35a62e4",
"372b687ac774400088d7252467f9cef2",
"a422e26291bb47d08e36b7d94fffcc9b",
"799dc8fe6f7d4e4db68ebd6b1d8fdbd5",
"e61a718f06a045b6b028fbd33382d37a",
"7c712c0011954872834e90cc908b8cf6",
"9929aaffe7614506ae934a14d2b5c3ca",
"1060550e9dec4fa59e2b9b488af61b48",
"ebd36fe7f8dd47b99e69176ace9b985d",
"42a3c5c410cd49a38400b788980ad428",
"e5373b31cbfc4e87af314bf4898d2b72",
"4379f29ff4a74df784bb2fb3989c64cc",
"87b42e656263458cbb33ee3ca29128a8",
"b245d8e4c115438583ef8b433d4d16be",
"df7f9d244f70440aaf401c6a856a6ed0",
"5d39d9e8cb724af59d7af29ec9c41cf4",
"6292f7a4de1948d09bebaf5064fa0d42",
"62cbedb2d67d4a2c8888b46e56e39f65",
"57f2c072b0574f74a7dbe2a9d303cbd4",
"78345adf6d8143e3a649a05c5d476115",
"f89bfff6526c43bcb75c93f987bfba72",
"a7244280fc934b8484b3dfd7eecca161",
"ec7ae2abc92f49f98ab0c52953bf82af",
"092355509a67474e9cc13a6b4718a192",
"c69730dfeec14cedaf8917ee9ad27dcb",
"00177cac24374c7d9144e3856594ddaf",
"1b1e49195c5b48a1b6e4364db01c5c87",
"97af130f6bd54d618e0af88ed9e85819",
"1f88a409f3c64ea291861837e830ce39",
"e9881edb2c7846fea2070c8e724b25cb",
"dfc8cc5ee6704f1fb3b4df0507ea471b",
"16bb453bd4934bd9840a9d365be604cd",
"1ee1a56d04fa4e85a20e4cde46c164e5",
"9872da92026e4debb96b731c892408e2",
"451d181b9ef24b0e9fc3914a4466b9f1",
"6a49aa5a2b3646d4b772e2398eacf72d",
"edeb0f283b0b495eab6e06a15fa38af0",
"0ed98c973a9d4d30b31ab751a7ecf3f6",
"a0b7dab4f65e448f816e00d12dc3b469",
"3b6d4b41992b43acac5daa31f119cd30",
"c5b4960fc3ba4459a6eaff43772050a3",
"71f9e574addd491cbbe65ddbe0eb23ad",
"442e58ceec4e44a3a4a2ebadd145ee73",
"4382289897b54e63b22541d0b45c9bf3",
"ba6998ac52544199b4e828d1a7f99d0d",
"29399d40d44d46c8b3648218e8378e89",
"c6b2d407256549a5bc7543cb6d7475b4",
"545251749b10408d82c9a60b914dce34",
"1bbadef2fca5427fa1a4975153968f56",
"d267ec5dd1954b9a88e566dff4a9af9d",
"4ddc8fa0c9834cd3b36274a9e0e05504",
"ec13219478284bbdbe69cbcd3dff5c34",
"562637f216114f42aa352cabd59cab57",
"e37bcfc1604b4a788f53825c5a7e04b7",
"16e05d45084c44b68697316bc8791b88",
"9b80621f94a0484791b33a13b3c6bade",
"fc87fbccba304554be90701b39429918",
"aaf79572906c4443b5542bef5dee8055",
"bc81dfa869f94f059545a7d59b7eb70a",
"b3c84c45c0804e4cb703f9c512c6f90d",
"6520db76749d4a8aaa6d26005f07c43c",
"29a3f867c8bd4769b8e649b8912776a2",
"2317c9a94dfe4a62a8605038ebc2210a",
"bd27949a1e3c43e5912a01b04b6621a4",
"94ed75e3a9fe4cbdb373d4038af0c32b",
"adfd8a554ada48f9a3dc2fab69467b79",
"0f9dcd9b698b4af699cfafdf2874612c",
"f8700d095ee04a40a15c6a88d7daf1ae",
"1ccb12282661417383a7f96bd5ea4eac",
"32a16b0cbbe548f4869cf665d737b9ba",
"3f546bf0aa124e45b39daa0328a1d23d",
"61f6966f8b6d4070a9ea3604373cb96a",
"f8198301b2db406685189589e59efd9a",
"e813e26451e3476a82a9b75668b638a1",
"8aa00f7a2de64af8b0a5b8cd5fefb14e",
"3404fb0c4c4b4ca0ae6d9eeec6722e02",
"f268eda021e040e6b25889e72ee8dd1e",
"b61f7faf19594e3c81e1429045f341de",
"a4ecf424e6ed44c597f25acee7b43df2",
"c314f7a0c6084790a2620ec231297525",
"950c1b85a8bd456ab928df4cc3710b21",
"806b17b8ce074563941dbba9ca2c9422",
"0b42f72fac844c19adea7b6689e79432",
"37e2f38471d4481eb93031ca52bdb785",
"067aaf24d87a400b9f636762051598f3",
"a03f0a986af84eacbb5c333f49726105",
"abf0e1d60e9f4fadbcda876dd28f6d14",
"81a7650cc4124b58b4e191093a5433e1",
"74f24411d35e4c918c608a1afd799156",
"a8501b8551d4431fbb708d6556e7dd85",
"24757877452d48f0afe92bb66e2c48ac",
"b64be6a0d6df44899ba6fe0a47f676e1",
"c1588f6fe9ca48a5a7512734f26ae9a3",
"374765208bf345089f58e925d2a70984",
"78774135f2d1468d9172fa7588afef20",
"9b3233c52ecc4b259deed2d5d6f332eb",
"2567e718b28f417e834a10c5310f92b2",
"587ea317bd0a4831a9417a3980207a73",
"9e857c0cb85d4b50ad73c670436b48a5",
"f8705a1680534caa801dc6e14c3601e3",
"7e982757835048ae85abf6cba5d27ecb",
"e14a085ca7e6410ca78fe0f27533be6d",
"6d204a683e2d49deb5f504f91ee72044",
"6659b3fef8644c6e96ed9d43adb744d5",
"4c130bb196db422c8b9c309cbe59bac1",
"32f42609c389450bbaad10f0a884c527",
"a9e874a076234974a102b02551ec8666",
"e505a51d928542b883743a6d48c5cbd3",
"32ae8645c975473b81788c17c376367a",
"7cb758be0be64a48930379b2cfa9fb44",
"fb4bf57c071d4ef68046988a4965b312",
"f1af6e23e42c4051a1c7fd4864450eaa",
"5cf88ae3c2fc4419a71eb45f28f67384",
"327fa24e70ee4cae903fe5673bcf3bfc",
"d1d5afc36de043afb8638dec0001288e",
"fffe74f9a1184de5b1297c659c7891b6",
"878702a019e34789b5b2f68eb321f2c8",
"91998be216844c079144c18278f660d7",
"5a1c360e82204ad99238901d366828cf",
"fc98ab1d58d84f69b6c36dbc19f5edd7",
"e2e86147b7d44f66a2cdfe02d63ee141",
"2f5178d399ec48648b435bac95faf188",
"c13ff304d7bc4c67a1f51da5d2303f96",
"0e8c4086734c4296884db9c5b67ddd33",
"c3a97f3130bc42c1be2f8c7dc81ca6cc",
"0043184800574bfdbfb3c82f4842c3ed",
"32719fe47c8543f69bcddc9b9bc0f719",
"784a5c0d201d46459a160e0dc8857fd9",
"d753c1ecc6034237ba4401a007451bfb",
"55fd60765d5e46c0baf1ac075820fc87",
"771af8b28988459ba4f2dbaadd0d5ce1",
"0cc22299405f4361b1b5f6234f35bae3",
"0f22a4136dce46f89e9efbf0581057a2",
"5cda6bd66e3248c4be477f1c1b40198d",
"269b131121554544a21d47596a24aa58",
"73758a3ca9c14d7eaa6504d8b9b6e333",
"5a16728cee0f4b75be08696261d8d8cc",
"da77a7b205b1446a9dacab0550aafebe",
"d6bd8a9af6e84500956bb565f2fe2c4e",
"6643bd32eb984a929159c8d441530aaf",
"8358f98f790d47fdb783eb5ca210c1fc",
"0c94630c2693424d9edd8b9921333aac",
"edf185cd11544b2c8bce478fb2f0b565",
"a7b4b7d3f1a34b089d72eb2a4ab169b0",
"09a607655a054ed5ad7f36716f8ee4c6",
"9258ed2d94924de3b944105e280fbf5f",
"1c68fbbc833d4c3fbbb686f15b9ba939",
"fa996ce2f8f24a5d80914c4ef5b6e788",
"f77e9e785cc04c00822fc8b7ef2959da",
"5d190801b10b49a698c3a69c7cd1f253",
"0112cfe0a27e43fd9509ef8da050e188",
"4c6b2159bbc04b8d811fadff3a22f4a6",
"c705d699568543829441df2258f2e615",
"eb8f921c8d67442b9bc2a6a67930d94c",
"8a89212807f243d4b135a3188f4fb6d0",
"cc018dfe29144d0182058816134e1fd9",
"c404db8298364742b8ce076cddcb7641",
"03aca6d92597461481419eac9bdb6635",
"59b90e05822f480793685902b9c6d055",
"d13cc9053e7a4982b6fc7f00474f25c7",
"c9b781db7ee2422b8dc4eb6aa2575d87",
"5100eb13df1c453980942ec78d0e810b"
]
},
"id": "DkIvEkIIkEyB",
"outputId": "8bf13df5-67ab-4ea2-db01-5719298b2dc9"
},
"outputs": [],
"source": [
"from unsloth import FastLanguageModel, is_bfloat16_supported\n",
"\n",
"max_seq_length = 1024 # Can increase for longer reasoning traces\n",
"lora_rank = 64 # Larger rank = smarter, but slower\n",
"\n",
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
" model_name=\"Qwen/Qwen2.5-3B-Instruct\",\n",
" max_seq_length=max_seq_length,\n",
" load_in_4bit=True, # False for LoRA 16bit\n",
" fast_inference=True, # Enable vLLM fast inference\n",
" max_lora_rank=lora_rank,\n",
" gpu_memory_utilization=0.5, # Reduce if out of memory\n",
")\n",
"\n",
"model = FastLanguageModel.get_peft_model(\n",
" model,\n",
" r=lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128\n",
" target_modules=[\n",
" \"q_proj\",\n",
" \"k_proj\",\n",
" \"v_proj\",\n",
" \"o_proj\",\n",
" \"gate_proj\",\n",
" \"up_proj\",\n",
" \"down_proj\",\n",
" ], # Remove QKVO if out of memory\n",
" lora_alpha=lora_rank,\n",
" use_gradient_checkpointing=\"unsloth\", # Enable long context finetuning\n",
" random_state=3407,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0Y56ln_izS9E"
},
"source": [
"### Data Prep\n",
"<a name=\"Data\"></a>\n",
"\n",
"We directly leverage [@willccbb](https://gist.github.com/willccbb/4676755236bb08cab5f4e54a0475d6fb) for data prep and all reward functions. You are free to create your own!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 209,
"referenced_widgets": [
"6b2cf7a462884bc4a3c2f8670f8982ed",
"8a358b880c194e6baab24f711937f695",
"0ec5dfa2fa6b4be2bfb45692ac20c053",
"1a619a0a53c149c3b9b0cbd53f1d2034",
"fdcbb976f35e4aa383d9ebcd7304849f",
"b131891a909743f5826429a84a74b4f9",
"3745da58146d445abe863a248e86d9be",
"22d29d39ef8149dba932ed9af68def67",
"413b4818b6de4bd995944df2f73bc3d5",
"c626fa7f5c624ec5a414e6f58d71ea96",
"2a8fb4b980ab45be81ae39d5c6b7bad4",
"ec1837ba407c4ecb86335d9c0c239ad9",
"1793da35014840b08d776b4abc9a1212",
"a06bde82466c43a49b58c0d942071be0",
"a1a3ee663a154319ba186954816e6932",
"97a9db845a9c45caa88ea1cd0ab6c40d",
"231be59acb2a4921ba53ba97bdf5888a",
"5aa0e699560c4fe89250b4ccac6401cb",
"7529f4d84b9f46709ec5f0e22b75baf4",
"8b8fd84911e94b5d8d627911fa9a64f7",
"420154f0bbe344a7979769a8ef066df7",
"1ee34b9171604805887aee3baa4ff099",
"3c95b86a26bc49b4a65589df0e7865e2",
"ef6d7a6a7fb344fab5befe656b98c7b7",
"c3d523c23ebf4b769c59eda950a3fc1f",
"bffc24c5bcdb47988f2e309c7fbc21b8",
"11052ab3d1f441248458cd6ab8b810db",
"0c2e39b20d5941969101b793b94fdbfa",
"f0c74de390814c2b9f38f7cc02458426",
"6c9502ae4cbc4248b14590fc6efe832f",
"8972cb9e823342bcb5c493e487dc64b6",
"71898e427b1245d682178a90c6752e8e",
"90295f1308054c66b557da231d25d3d5",
"54f444e6a25349fcb8fb35b9efb33e68",
"2f9a64ad318248bf99fccdedafaddfc0",
"d3ce3f3bee5646d893242758444acd6a",
"5b4f99c20a8a4ecf873b0a8ef8b150cd",
"70e141512e91492bae4713400879d019",
"7c67cf4d83e34cd6a475405be45f345e",
"c184e417ad134e94995a55d4bc7c94b5",
"f1c8d7356b584fb8965d2a62ccec52c8",
"9bcb0000da6e434d9b241a09d9b46244",
"11bd7b89fa794a899db9c83cb2a99df5",
"0e9d3d2244a54eb781311c01965c33e4",
"00d1bd9a327241e8b434ef7b3bde30df",
"cd17f032bbb44575977797caaa92bd54",
"04b74b7851334892b0db68bea1c8fd7e",
"44b6ecc9b9fc436e92701ec2409e0dd1",
"46998bc2b1ee4858a8ee1108ce36b9b7",
"95fa6a54e7e543c7b2dc5d5cf04daf58",
"8ee15a3d02954d31899d0b6a8dbf7988",
"1d5a58d63c51427bbbfb729216ea8468",
"e2790d1e5e714b3e9081b4edcef24535",
"96519710bf5047239f6800100471ee6e",
"e9b020ce2dda442b871b41e417210b52",
"4d18ef3f7c8046ddb0b03314861c3414",
"176ccdb40b994ee3b381f7ea0f379c3c",
"3519fdf3abac4d7690c45710f617e5d0",
"dc320ae9e2f7426588d49848b0871ca4",
"6b5a1287c457454fb367e5dc03d8b157",
"989519f3e61b4132b3f88fbdf29703a9",
"e0f87f6fd7c84261818624defe575a1c",
"e6d2b6aaa9794971b01c778b1a0596db",
"334c4833051e4ccea727dc2cdf3871c3",
"d36027362b41413fa3ee4670dd820a88",
"d94e3315b99747c39190fee9b79f00ac"
]
},
"id": "cXk993X6C2ZZ",
"outputId": "2d545127-0313-405e-d8f9-a43d7dba8d11"
},
"outputs": [],
"source": [
"import re\n",
"from datasets import load_dataset, Dataset\n",
"\n",
"# Load and prep dataset\n",
"SYSTEM_PROMPT = \"\"\"\n",
"Respond in the following format:\n",
"<reasoning>\n",
"...\n",
"</reasoning>\n",
"<answer>\n",
"...\n",
"</answer>\n",
"\"\"\"\n",
"\n",
"XML_COT_FORMAT = \"\"\"\\\n",
"<reasoning>\n",
"{reasoning}\n",
"</reasoning>\n",
"<answer>\n",
"{answer}\n",
"</answer>\n",
"\"\"\"\n",
"\n",
"\n",
"def extract_xml_answer(text: str) -> str:\n",
" answer = text.split(\"<answer>\")[-1]\n",
" answer = answer.split(\"</answer>\")[0]\n",
" return answer.strip()\n",
"\n",
"\n",
"def extract_hash_answer(text: str) -> str | None:\n",
" if \"####\" not in text:\n",
" return None\n",
" return text.split(\"####\")[1].strip()\n",
"\n",
"\n",
"# uncomment middle messages for 1-shot prompting\n",
"def get_gsm8k_questions(split=\"train\") -> Dataset:\n",
" data = load_dataset(\"openai/gsm8k\", \"main\")[split] # type: ignore\n",
" data = data.map(\n",
" lambda x: { # type: ignore\n",
" \"prompt\": [\n",
" {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
" {\"role\": \"user\", \"content\": x[\"question\"]},\n",
" ],\n",
" \"answer\": extract_hash_answer(x[\"answer\"]),\n",
" }\n",
" ) # type: ignore\n",
" return data # type: ignore\n",
"\n",
"\n",
"dataset = get_gsm8k_questions()\n",
"\n",
"\n",
"# Reward functions\n",
"def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:\n",
" responses = [completion[0][\"content\"] for completion in completions]\n",
" q = prompts[0][-1][\"content\"]\n",
" extracted_responses = [extract_xml_answer(r) for r in responses]\n",
" print(\n",
" \"-\" * 20,\n",
" f\"Question:\\n{q}\",\n",
" f\"\\nAnswer:\\n{answer[0]}\",\n",
" f\"\\nResponse:\\n{responses[0]}\",\n",
" f\"\\nExtracted:\\n{extracted_responses[0]}\",\n",
" )\n",
" return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]\n",
"\n",
"\n",
"def int_reward_func(completions, **kwargs) -> list[float]:\n",
" responses = [completion[0][\"content\"] for completion in completions]\n",
" extracted_responses = [extract_xml_answer(r) for r in responses]\n",
" return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]\n",
"\n",
"\n",
"def strict_format_reward_func(completions, **kwargs) -> list[float]:\n",
" \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n",
" pattern = r\"^<reasoning>\\n.*?\\n</reasoning>\\n<answer>\\n.*?\\n</answer>\\n$\"\n",
" responses = [completion[0][\"content\"] for completion in completions]\n",
" matches = [re.match(pattern, r) for r in responses]\n",
" return [0.5 if match else 0.0 for match in matches]\n",
"\n",
"\n",
"def soft_format_reward_func(completions, **kwargs) -> list[float]:\n",
" \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n",
" pattern = r\"<reasoning>.*?</reasoning>\\s*<answer>.*?</answer>\"\n",
" responses = [completion[0][\"content\"] for completion in completions]\n",
" matches = [re.match(pattern, r) for r in responses]\n",
" return [0.5 if match else 0.0 for match in matches]\n",
"\n",
"\n",
"def count_xml(text) -> float:\n",
" count = 0.0\n",
" if text.count(\"<reasoning>\\n\") == 1:\n",
" count += 0.125\n",
" if text.count(\"\\n</reasoning>\\n\") == 1:\n",
" count += 0.125\n",
" if text.count(\"\\n<answer>\\n\") == 1:\n",
" count += 0.125\n",
" count -= len(text.split(\"\\n</answer>\\n\")[-1]) * 0.001\n",
" if text.count(\"\\n</answer>\") == 1:\n",
" count += 0.125\n",
" count -= (len(text.split(\"\\n</answer>\")[-1]) - 1) * 0.001\n",
" return count\n",
"\n",
"\n",
"def xmlcount_reward_func(completions, **kwargs) -> list[float]:\n",
" contents = [completion[0][\"content\"] for completion in completions]\n",
" return [count_xml(c) for c in contents]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bTnL_tJnzh2L"
},
"source": [
"<a name=\"Train\"></a>\n",
"### Train the model\n",
"\n",
"Now set up GRPO Trainer and all configurations!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ptqkXK2D4d6p",
"outputId": "7d36a642-4cc8-48ee-c8cc-1470978d4503"
},
"outputs": [],
"source": [
"from trl import GRPOConfig, GRPOTrainer\n",
"\n",
"training_args = GRPOConfig(\n",
" use_vllm=True, # use vLLM for fast inference!\n",
" learning_rate=5e-6,\n",
" adam_beta1=0.9,\n",
" adam_beta2=0.99,\n",
" weight_decay=0.1,\n",
" warmup_ratio=0.1,\n",
" lr_scheduler_type=\"cosine\",\n",
" optim=\"adamw_8bit\",\n",
" logging_steps=1,\n",
" bf16=is_bfloat16_supported(),\n",
" fp16=not is_bfloat16_supported(),\n",
" per_device_train_batch_size=1,\n",
" gradient_accumulation_steps=1, # Increase to 4 for smoother training\n",
" num_generations=8, # Decrease if out of memory\n",
" max_prompt_length=256,\n",
" max_completion_length=200,\n",
" # num_train_epochs = 1, # Set to 1 for a full training run\n",
" max_steps=250,\n",
" save_steps=250,\n",
" max_grad_norm=0.1,\n",
" report_to=\"none\", # Can use Weights & Biases\n",
" output_dir=\"outputs\",\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "X_71Y0eKz5yE"
},
"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": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "vzOuSVCL_GA9",
"outputId": "c6fdbbe8-7062-45a2-a48c-649f2b878ac4"
},
"outputs": [],
"source": [
"trainer = GRPOTrainer(\n",
" model=model,\n",
" processing_class=tokenizer,\n",
" reward_funcs=[\n",
" xmlcount_reward_func,\n",
" soft_format_reward_func,\n",
" strict_format_reward_func,\n",
" int_reward_func,\n",
" correctness_reward_func,\n",
" ],\n",
" args=training_args,\n",
" train_dataset=dataset,\n",
")\n",
"trainer.train()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yUbluAAhD0Lg"
},
"source": [
"<a name=\"Inference\"></a>\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": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 72
},
"id": "IqzsdZzeDM_m",
"outputId": "a7637565-80fe-4ee9-cdc9-33b9affb65dc"
},
"outputs": [],
"source": [
"text = tokenizer.apply_chat_template(\n",
" [\n",
" {\"role\": \"user\", \"content\": \"How many r's are in strawberry?\"},\n",
" ],\n",
" tokenize=False,\n",
" add_generation_prompt=True,\n",
")\n",
"\n",
"from vllm import SamplingParams\n",
"\n",
"sampling_params = SamplingParams(\n",
" temperature=0.8,\n",
" top_p=0.95,\n",
" max_tokens=1024,\n",
")\n",
"output = (\n",
" model.fast_generate(\n",
" [text],\n",
" sampling_params=sampling_params,\n",
" lora_request=None,\n",
" )[0]\n",
" .outputs[0]\n",
" .text\n",
")\n",
"\n",
"output"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "G4lzJD7REFjs"
},
"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": "YC9BBT0RESln"
},
"outputs": [],
"source": [
"model.save_lora(\"grpo_saved_lora\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LherO2vzEbMt"
},
"source": [
"Now we load the LoRA and test:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 125
},
"id": "SDKIhhvN6lAF",
"outputId": "8bc37894-678f-40d7-874b-365300685836"
},
"outputs": [],
"source": [
"text = tokenizer.apply_chat_template(\n",
" [\n",
" {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
" {\"role\": \"user\", \"content\": \"How many r's are in strawberry?\"},\n",
" ],\n",
" tokenize=False,\n",
" add_generation_prompt=True,\n",
")\n",
"\n",
"from vllm import SamplingParams\n",
"\n",
"sampling_params = SamplingParams(\n",
" temperature=0.8,\n",
" top_p=0.95,\n",
" max_tokens=1024,\n",
")\n",
"output = (\n",
" model.fast_generate(\n",
" text,\n",
" sampling_params=sampling_params,\n",
" lora_request=model.load_lora(\"grpo_saved_lora\"),\n",
" )[0]\n",
" .outputs[0]\n",
" .text\n",
")\n",
"\n",
"output"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_ZBnvg2f9Nlg"
},
"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": "RphEZRSfFhru"
},
"source": [
"<a name=\"Save\"></a>\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": "GwNY9_PrFiXZ"
},
"outputs": [],
"source": [
"# Merge to 16bit\n",
"if False:\n",
" model.save_pretrained_merged(\n",
" \"model\",\n",
" tokenizer,\n",
" save_method=\"merged_16bit\",\n",
" )\n",
"if False:\n",
" model.push_to_hub_merged(\n",
" \"hf/model\", tokenizer, save_method=\"merged_16bit\", token=\"\"\n",
" )\n",
"\n",
"# Merge to 4bit\n",
"if False:\n",
" model.save_pretrained_merged(\n",
" \"model\",\n",
" tokenizer,\n",
" save_method=\"merged_4bit\",\n",
" )\n",
"if False:\n",
" 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_merged(\n",
" \"model\",\n",
" tokenizer,\n",
" save_method=\"lora\",\n",
" )\n",
"if False:\n",
" model.push_to_hub_merged(\"hf/model\", tokenizer, save_method=\"lora\", token=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BDUGPiL3Fkkq"
},
"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": "AGo4dbWvFk4M"
},
"outputs": [],
"source": [
"# Save to 8bit Q8_0\n",
"if False:\n",
" model.save_pretrained_gguf(\n",
" \"model\",\n",
" tokenizer,\n",
" )\n",
"# Remember to go to https://huggingface.co/settings/tokens for a token!\n",
"# And change hf to your username!\n",
"if False:\n",
" model.push_to_hub_gguf(\"hf/model\", tokenizer, token=\"\")\n",
"\n",
"# Save to 16bit GGUF\n",
"if False:\n",
" model.save_pretrained_gguf(\"model\", tokenizer, quantization_method=\"f16\")\n",
"if False:\n",
" model.push_to_hub_gguf(\"hf/model\", tokenizer, quantization_method=\"f16\", token=\"\")\n",
"\n",
"# Save to q4_k_m GGUF\n",
"if False:\n",
" model.save_pretrained_gguf(\"model\", tokenizer, quantization_method=\"q4_k_m\")\n",
"if False:\n",
" model.push_to_hub_gguf(\n",
" \"hf/model\", tokenizer, quantization_method=\"q4_k_m\", token=\"\"\n",
" )\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=[\n",
" \"q4_k_m\",\n",
" \"q8_0\",\n",
" \"q5_k_m\",\n",
" ],\n",
" token=\"\",\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bhYfRM8PxWAu"
},
"source": [
"Now, use the `model-unsloth.gguf` file or `model-unsloth-Q4_K_M.gguf` file in llama.cpp or a UI based system like Jan or Open WebUI. You can install Jan [here](https://github.com/janhq/jan) and Open WebUI [here](https://github.com/open-webui/open-webui)\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",
"<div class=\"align-center\">\n",
" <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
" <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
" <a href=\"https://docs.unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
"\n",
" Join Discord if you need help + ⭐️ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐️\n",
"</div>\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}