{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "ur8xi4C7S06n" }, "outputs": [], "source": [ "# Copyright 2025 Google LLC\n", "#\n", "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "JAPoU8Sm5E6e" }, "source": [ "# Fine-tuning GPT-OSS 20B with Unsloth on Vertex AI Colab Enterprise and Nvidia A100 40GB GPU\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Vertex AI Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
\n", "\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", "" ] }, { "cell_type": "markdown", "metadata": { "id": "84f0f73a0f76" }, "source": [ "| Author(s) |\n", "| --- |\n", "| [Fred Molina](https://github.com/mltuto) |" ] }, { "cell_type": "markdown", "metadata": { "id": "tvgnzT1CKxrO" }, "source": [ "## Overview\n", "\n", "This notebook provides a step-by-step guide to fine-tuning the GPT-OSS 20B model using Unsloth.\n", " The process involves:\n", "\n", " 01. Installing the required libraries.\n", " 02. Loading the GPT-OSS 20B model.\n", " 03. Adding LoRA adapters to the model for fine-tuning.\n", " 04. Preparing the dataset for fine-tuning.\n", " 05. Fine-tuning the model on the dataset.\n", "\n", " **DISCLAIMER**\n", " This notebook is intended for educational purposes only.\n", "\n", " - Date: Aug 2025\n", " - Not suitable for production environments.\n", " - Use at your own risk.\n", " - This notebook is an adaptation of the original Unsloth team Notebook that runs on Colab public with T4 GPUS: https://docs.unsloth.ai/get-started/unsloth-notebooks all credits to them!\n", " Some minor changes were done in how to install the required packages as Vertex AI Colab Enterprise manages the environments differently than local or Colab public\n", " \n", "Requirements:\n", "A Vertex AI colab enterprise environment running on a Runtime that have a GPU (e.g., NVIDIA A100)." ] }, { "cell_type": "markdown", "metadata": { "id": "61RBz8LLbxCR" }, "source": [ "## Get started" ] }, { "cell_type": "markdown", "metadata": { "id": "No17Cw5hgx12" }, "source": [ "### Install Google Gen AI SDK and other required packages\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tFy3H3aPgx12" }, "outputs": [], "source": [ "%pip install --upgrade --quiet google-genai" ] }, { "cell_type": "markdown", "metadata": { "id": "dmWOrTJ3gx13" }, "source": [ "### Authenticate your notebook environment (Colab only)\n", "\n", "If you're running this notebook on Google Colab, run the cell below to authenticate your environment." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "NyKGtVQjgx13" }, "outputs": [], "source": [ "import sys\n", "\n", "if \"google.colab\" in sys.modules:\n", " from google.colab import auth\n", "\n", " auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": { "id": "DF4l8DTdWgPY" }, "source": [ "### Set Google Cloud project information\n", "\n", "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", "\n", "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Nqwi-5ufWp_B" }, "outputs": [], "source": [ "# Use the environment variable if the user doesn't provide Project ID.\n", "import os\n", "\n", "PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n", "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n", " PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n", "\n", "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"global\")\n", "\n", "from google import genai\n", "\n", "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)" ] }, { "cell_type": "markdown", "metadata": { "id": "5303c05f7aa6" }, "source": [ "### Import libraries" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6fc324893334" }, "outputs": [], "source": [ "from IPython.display import Markdown, display\n", "# 1. Upgrade uv, the fast package installer\n", "!pip install --upgrade -qqq uv\n", "# 2. Use uv to install all packages in a single, consolidated command.\n", "# THIS VERSION FORCES NUMPY to a version < 2.0 to solve the TensorFlow conflict.\n", "print(\"⏳ Installing all required libraries with NumPy compatibility fix...\")\n", "!uv pip install --system --upgrade \\\n", " \"numpy<2.0\" \\\n", " \"torch>=2.8.0\" \\\n", " \"triton>=3.4.0\" \\\n", " \"torchvision==0.23.0\" \\\n", " \"bitsandbytes==0.46.1\" \\\n", " \"unsloth @ git+https://github.com/unslothai/unsloth.git@79b46f71b249600488842511c9ee40f27a3989f2\" \\\n", " \"unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@26615eb3021b92abbfc8f895da4cd6803322b658\" \\\n", " \"peft @ git+https://github.com/huggingface/peft.git@a90003f0edd6353f489f48bd2c35080d27bb6974\" \\\n", " \"accelerate @ git+https://github.com/huggingface/accelerate.git@23cf4ef8a3b58f016f63eeb158b4aa2c3e79fe6f\" \\\n", " \"transformers @ git+https://github.com/huggingface/transformers.git@f4d57f2f0cdff0f63ee74a1f16f442dfaf525231\" \\\n", " \"protobuf<=3.20.3\" \\\n", " \"setuptools==69.5.1\" \\\n", " \"wandb==0.21.1\"\n", "print(\"✅✅✅ Installation complete!\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "twtmUqlBSSXR" }, "outputs": [], "source": [ "# Restart Notebook Kernel\n", "import os\n", "\n", "os.kill(os.getpid(), 9)" ] }, { "cell_type": "markdown", "metadata": { "id": "EdvJRUWRNGHE" }, "source": [ "##OpenAI GPT-OSS 20B finetuning on Vertex AI Colab Enterprise with Unsloth!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "n_9JN_wFSgAr" }, "outputs": [], "source": [ "import torch\n", "from unsloth import FastLanguageModel\n", "\n", "max_seq_length = 1024\n", "dtype = None\n", "\n", "model, tokenizer = FastLanguageModel.from_pretrained(\n", " model_name=\"unsloth/gpt-oss-20b\",\n", " dtype=dtype, # None for auto detection\n", " max_seq_length=max_seq_length, # Choose any for long context!\n", " load_in_4bit=True, # 4 bit quantization to reduce memory\n", " full_finetuning=False, # [NEW!] We have full finetuning now!\n", " # token = \"hf_...\", # use one if using gated models\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "B_jeQ2mISmdM" }, "source": [ "We now add LoRA adapters for parameter efficient finetuning - this allows us to only efficiently train 1% of all parameters." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "sW7jpAQJSg2v" }, "outputs": [], "source": [ "model = FastLanguageModel.get_peft_model(\n", " model,\n", " r=8, # 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", " ],\n", " lora_alpha=16,\n", " lora_dropout=0, # Supports any, but = 0 is optimized\n", " bias=\"none\", # Supports any, but = \"none\" is optimized\n", " # [NEW] \"unsloth\" uses 30% less VRAM, fits 2x larger batch sizes!\n", " use_gradient_checkpointing=\"unsloth\", # True or \"unsloth\" for very long context\n", " random_state=3407,\n", " use_rslora=False, # We support rank stabilized LoRA\n", " loftq_config=None, # And LoftQ\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "L4XmgbNQYKyU" }, "source": [ "###Data Prep" ] }, { "cell_type": "markdown", "metadata": { "id": "D5AWtlaPYOxy" }, "source": [ "The HuggingFaceH4/Multilingual-Thinking dataset will be utilized as our example. This dataset, available on Hugging Face, contains reasoning chain-of-thought examples derived from user questions that have been translated from English into four other languages. It is also the same dataset referenced in OpenAI's cookbook for fine-tuning. The purpose of using this dataset is to enable the model to learn and develop reasoning capabilities in these four distinct languages." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "E_VJ7N7PYH-o" }, "outputs": [], "source": [ "def formatting_prompts_func(examples):\n", " convos = examples[\"messages\"]\n", " texts = [\n", " tokenizer.apply_chat_template(\n", " convo, tokenize=False, add_generation_prompt=False\n", " )\n", " for convo in convos\n", " ]\n", " return {\n", " \"text\": texts,\n", " }\n", "\n", "\n", "from datasets import load_dataset\n", "\n", "dataset = load_dataset(\"HuggingFaceH4/Multilingual-Thinking\", split=\"train\")\n", "dataset" ] }, { "cell_type": "markdown", "metadata": { "id": "gX2BBFGJYRYN" }, "source": [ "To format our dataset, we will apply our version of the GPT OSS prompt" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "cQ1WXAN7YUaC" }, "outputs": [], "source": [ "from unsloth.chat_templates import standardize_sharegpt\n", "\n", "dataset = standardize_sharegpt(dataset)\n", "dataset = dataset.map(\n", " formatting_prompts_func,\n", " batched=True,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "hGzLfYvXYWwC" }, "source": [ "Let's take a look at the dataset, and check what the 1st example shows" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Po_E_X7OYYCz" }, "outputs": [], "source": [ "print(dataset[0][\"text\"])" ] }, { "cell_type": "markdown", "metadata": { "id": "KonBtw8DYanM" }, "source": [ "What is unique about GPT-OSS is that it uses OpenAI Harmony format which supports conversation structures, reasoning output, and tool calling." ] }, { "cell_type": "markdown", "metadata": { "id": "PnpsupSBYd_D" }, "source": [ "### Train the model\n", "\n", "Now let's use Huggingface TRL's SFTTrainer! More docs here: TRL SFT docs. We do 30 steps to speed things up, but you can set num_train_epochs=1 for a full run, and turn off max_steps=None." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "auIEyRD-YYGa" }, "outputs": [], "source": [ "from trl import SFTConfig, SFTTrainer\n", "\n", "trainer = SFTTrainer(\n", " model=model,\n", " tokenizer=tokenizer,\n", " train_dataset=dataset,\n", " args=SFTConfig(\n", " per_device_train_batch_size=1,\n", " gradient_accumulation_steps=4,\n", " warmup_steps=5,\n", " # num_train_epochs = 1, # Set this for 1 full training run.\n", " max_steps=30,\n", " learning_rate=2e-4,\n", " logging_steps=1,\n", " optim=\"adamw_8bit\",\n", " weight_decay=0.01,\n", " lr_scheduler_type=\"linear\",\n", " seed=3407,\n", " output_dir=\"outputs\",\n", " report_to=\"none\",\n", " ),\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "kC1IdtxjYnlf" }, "source": [ "# Show current memory stats" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "24EDdt5mYYJk" }, "outputs": [], "source": [ "# @title Show current memory stats\n", "gpu_stats = torch.cuda.get_device_properties(0)\n", "start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n", "max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n", "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n", "print(f\"{start_gpu_memory} GB of memory reserved.\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "OCoiRZPzYYMW" }, "outputs": [], "source": [ "trainer_stats = trainer.train()" ] }, { "cell_type": "markdown", "metadata": { "id": "qFI-dVv8Ytx0" }, "source": [ "# Show final memory and time stats" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "SJNr5q_FYYPL" }, "outputs": [], "source": [ "# @title Show final memory and time stats\n", "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n", "used_memory_for_lora = round(used_memory - start_gpu_memory, 3)\n", "used_percentage = round(used_memory / max_memory * 100, 3)\n", "lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)\n", "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n", "print(\n", " f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\"\n", ")\n", "print(f\"Peak reserved memory = {used_memory} GB.\")\n", "print(f\"Peak reserved memory for training = {used_memory_for_lora} GB.\")\n", "print(f\"Peak reserved memory % of max memory = {used_percentage} %.\")\n", "print(f\"Peak reserved memory for training % of max memory = {lora_percentage} %.\")" ] }, { "cell_type": "markdown", "metadata": { "id": "iuIt6vBZYy9u" }, "source": [ "## Inference\n", "Let's run the model! You can change the instruction and input - leave the output blank!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "bx9XjEPdY2TN" }, "outputs": [], "source": [ "messages = [\n", " {\n", " \"role\": \"system\",\n", " \"content\": \"You are a helpful assistant that can solve mathematical problems.\",\n", " },\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", " reasoning_effort=\"medium\",\n", ").to(model.device)\n", "from transformers import TextStreamer\n", "\n", "_ = model.generate(**inputs, max_new_tokens=128, streamer=TextStreamer(tokenizer))" ] }, { "cell_type": "markdown", "metadata": { "id": "mPzuhAFSY4vL" }, "source": [ "We just saw how to Fine Tune GPT - OSS 20B with an A100 40GB on Vertex AI Colab Enterprise using Unsloth. Unsloth has a Discord channel If you like Unsloth optimizations, show your support and ⭐️ Star Unsloth on Github ⭐️" ] } ], "metadata": { "colab": { "name": "gpt_oss_20B_finetuning_with_unsloth.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }