{
  "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",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Ffine-tuning%2Fgpt_oss_20B_finetuning_with_unsloth.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "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
}
