{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ovlmqboji0c",
   "metadata": {
    "id": "ovlmqboji0c"
   },
   "source": "# Fine-Tune NVIDIA Nemotron 3 with SFT and LoRA using TRL\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/notebooks/sft_nemotron_3.ipynb)"
  },
  {
   "cell_type": "markdown",
   "id": "kiezlcdkw2k",
   "metadata": {
    "id": "kiezlcdkw2k"
   },
   "source": [
    "![trl banner](https://huggingface.co/datasets/trl-lib/documentation-images/resolve/main/trl_banner_dark.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ezgoxs28857",
   "metadata": {
    "id": "ezgoxs28857"
   },
   "source": "Fine-tune [**NVIDIA Nemotron 3**](https://huggingface.co/collections/nvidia/nvidia-nemotron-v3) models using **LoRA** and **SFT** with the [**TRL**](https://github.com/huggingface/trl) library.\n\nThe Nemotron 3 family includes hybrid **Mamba2/Transformer** models in different sizes (Nano, Super, etc.), natively supported in `transformers` (no `trust_remote_code` needed). See the [full collection](https://huggingface.co/collections/nvidia/nvidia-nemotron-v3) for all available checkpoints.\n\n> **Note:** This notebook requires a GPU with sufficient VRAM (e.g., A100 80GB). It is not designed for free Colab instances.\n\n- [TRL GitHub Repository](https://github.com/huggingface/trl)\n- [Official TRL Examples](https://huggingface.co/docs/trl/example_overview)"
  },
  {
   "cell_type": "markdown",
   "id": "5t9o6vg3jus",
   "metadata": {
    "id": "5t9o6vg3jus"
   },
   "source": "## Install dependencies\n\nWe install **TRL** with the **PEFT** and **quantization** extras, along with **trackio** for experiment tracking. Nemotron 3 models use **Mamba2** layers, which require the `mamba_ssm` and `causal_conv1d` CUDA kernels."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wh29y02x1ni",
   "metadata": {
    "id": "wh29y02x1ni"
   },
   "outputs": [],
   "source": "!pip install -Uq \"trl[peft,quantization]\" trackio\n\n# Nemotron 3 requires transformers>=5.3.0, which may not be installed by default with TRL\n!pip install -Uq \"transformers>=5.3.0\"\n\n# Mamba2 CUDA kernels (--no-build-isolation needed so they find your PyTorch/CUDA installation)\n!pip install --no-build-isolation mamba_ssm==2.2.5  # installation takes around 30 mins\n!pip install --no-build-isolation causal_conv1d==1.5.2"
  },
  {
   "cell_type": "markdown",
   "id": "yic7cr26xoi",
   "metadata": {
    "id": "yic7cr26xoi"
   },
   "source": [
    "### Log in to Hugging Face\n",
    "\n",
    "Log in to your **Hugging Face** account to save your fine-tuned model and track experiments on the Hub. You can find your **access token** on your [account settings page](https://huggingface.co/settings/tokens)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wwocbqwne49",
   "metadata": {
    "id": "wwocbqwne49",
    "outputId": "f8a3f993-63a5-4a11-ea41-5a0093cec19c",
    "colab": {
     "referenced_widgets": [
      "20976278f8564bc080426e851d52f46f"
     ]
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:86: UserWarning: \n",
      "Access to the secret `HF_TOKEN` has not been granted on this notebook.\n",
      "You will not be requested again.\n",
      "Please restart the session if you want to be prompted again.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20976278f8564bc080426e851d52f46f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60ieuf2bc6t",
   "metadata": {
    "id": "60ieuf2bc6t"
   },
   "source": [
    "## Load Dataset\n",
    "\n",
    "We load the [**HuggingFaceH4/Multilingual-Thinking**](https://huggingface.co/datasets/HuggingFaceH4/Multilingual-Thinking) dataset from the Hugging Face Hub. This dataset focuses on **multilingual reasoning**, where the *chain of thought* has been translated into several languages such as French, Spanish, and German. By fine-tuning a reasoning-capable model on this dataset, it learns to **generate reasoning steps in multiple languages**.\n",
    "\n",
    "For efficiency, we'll load only the **training split**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "u6rvgaaf3aj",
   "metadata": {
    "id": "u6rvgaaf3aj",
    "outputId": "79cd32ae-28bd-4ee3-aad0-c4aa98a70558",
    "colab": {
     "referenced_widgets": [
      "c25fda6dab5a4d68b7328e6e45326aa3",
      "a0c98c15f0684abcbce29b6a20dd6a3f",
      "cc1f038d587349929a80a0552ba94e0d"
     ]
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c25fda6dab5a4d68b7328e6e45326aa3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0c98c15f0684abcbce29b6a20dd6a3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data/train-00000-of-00001.parquet:   0%|          | 0.00/5.29M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc1f038d587349929a80a0552ba94e0d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset_name = \"HuggingFaceH4/Multilingual-Thinking\"\n",
    "train_dataset = load_dataset(dataset_name, split=\"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gi8gwuwbjcf",
   "metadata": {
    "id": "gi8gwuwbjcf"
   },
   "source": [
    "This dataset contains different columns. We'll only need `messages` as it contains the conversation and is the one used by the SFT trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aoi32054k7",
   "metadata": {
    "id": "aoi32054k7",
    "outputId": "b6645b9e-ca4e-4d7c-cd19-2f9d07562b07"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['reasoning_language', 'developer', 'user', 'analysis', 'final', 'messages'],\n",
       "    num_rows: 1000\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xo12wzbckok",
   "metadata": {
    "id": "xo12wzbckok"
   },
   "source": [
    "Let's look at a full example to understand the internal structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53pn2jj18m9",
   "metadata": {
    "id": "53pn2jj18m9",
    "outputId": "0701abe5-0988-4572-9154-684aa87d18ca"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'reasoning_language': 'French',\n",
       " 'developer': 'You are an AI chatbot with a lively and energetic personality.',\n",
       " 'user': 'Can you show me the latest trends on Twitter right now?',\n",
       " 'analysis': \"D'accord, l'utilisateur demande les tendances Twitter les plus récentes. Tout d'abord, je dois vérifier si j'ai accès à des données en temps réel. Étant donné que je ne peux pas naviguer sur Internet ou accéder directement à l'API de Twitter, je ne peux pas fournir des tendances en direct. Cependant, je peux donner quelques conseils généraux sur la façon de les trouver.\\n\\nJe devrais préciser que les tendances Twitter évoluent rapidement et sont spécifiques à chaque région. Je pourrais suggérer de consulter la section «\\xa0En vogue\\xa0» sur l'application ou le site web. Aussi, l'utilisation de hashtags et le suivi d'utilisateurs pertinents pourraient être utiles. Il est important de souligner que les tendances varient selon la région et l'heure de la journée. Je devrais garder un ton amical et bienveillant, peut-être ajouter un emoji pour rester léger. Je vais structurer ma réponse étape par étape pour faciliter la lecture. Je dois m'excuser de ne pas pouvoir fournir des données en temps réel et proposer d'autres méthodes. Je conserverai un langage simple et convivial, en évitant les termes techniques.\",\n",
       " 'final': 'Hey there!  While I can\\'t check Twitter (X) in real-time or access live data, I can share some tips to help you spot the latest trends:\\n\\n1. **Open the \"Trending\" tab** on the Twitter app or website – it updates constantly!  \\n2. **Search for hashtags** like #Trending or #Viral to see what’s blowing up.  \\n3. **Follow accounts** that curate trends (e.g., @TrendingNow, @ViralThreads).  \\n4. **Check regional trends** – they often differ by location!  \\n\\nRemember, trends are *super fast-moving* and often tied to pop culture, memes, or breaking news. For example, recent trends have included viral challenges (like the \"Distracted Boyfriend\" meme revival), celebrity drama, or unexpected events (hello, weather disasters!).  \\n\\nWant me to brainstorm *what* might trend next? I’ve got ideas!',\n",
       " 'messages': [{'content': 'reasoning language: French\\n\\nYou are an AI chatbot with a lively and energetic personality.',\n",
       "   'role': 'system',\n",
       "   'thinking': None},\n",
       "  {'content': 'Can you show me the latest trends on Twitter right now?',\n",
       "   'role': 'user',\n",
       "   'thinking': None},\n",
       "  {'content': 'Hey there!  While I can\\'t check Twitter (X) in real-time or access live data, I can share some tips to help you spot the latest trends:\\n\\n1. **Open the \"Trending\" tab** on the Twitter app or website – it updates constantly!  \\n2. **Search for hashtags** like #Trending or #Viral to see what’s blowing up.  \\n3. **Follow accounts** that curate trends (e.g., @TrendingNow, @ViralThreads).  \\n4. **Check regional trends** – they often differ by location!  \\n\\nRemember, trends are *super fast-moving* and often tied to pop culture, memes, or breaking news. For example, recent trends have included viral challenges (like the \"Distracted Boyfriend\" meme revival), celebrity drama, or unexpected events (hello, weather disasters!).  \\n\\nWant me to brainstorm *what* might trend next? I’ve got ideas!',\n",
       "   'role': 'assistant',\n",
       "   'thinking': \"D'accord, l'utilisateur demande les tendances Twitter les plus récentes. Tout d'abord, je dois vérifier si j'ai accès à des données en temps réel. Étant donné que je ne peux pas naviguer sur Internet ou accéder directement à l'API de Twitter, je ne peux pas fournir des tendances en direct. Cependant, je peux donner quelques conseils généraux sur la façon de les trouver.\\n\\nJe devrais préciser que les tendances Twitter évoluent rapidement et sont spécifiques à chaque région. Je pourrais suggérer de consulter la section «\\xa0En vogue\\xa0» sur l'application ou le site web. Aussi, l'utilisation de hashtags et le suivi d'utilisateurs pertinents pourraient être utiles. Il est important de souligner que les tendances varient selon la région et l'heure de la journée. Je devrais garder un ton amical et bienveillant, peut-être ajouter un emoji pour rester léger. Je vais structurer ma réponse étape par étape pour faciliter la lecture. Je dois m'excuser de ne pas pouvoir fournir des données en temps réel et proposer d'autres méthodes. Je conserverai un langage simple et convivial, en évitant les termes techniques.\"}]}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "kytx863dy3j",
   "metadata": {
    "id": "kytx863dy3j"
   },
   "source": [
    "The `messages` column contains an extra `thinking` field that we need to merge into the message content using `<think>...</think>` tags. We also remove the extra columns that are not needed for training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oxzelzs2kf",
   "metadata": {
    "id": "oxzelzs2kf"
   },
   "outputs": [],
   "source": [
    "def merge_thinking_and_remove_key(example):\n",
    "    new_messages = []\n",
    "    for msg in example[\"messages\"]:\n",
    "        content = msg[\"content\"]\n",
    "        thinking = msg.get(\"thinking\")\n",
    "        if thinking and isinstance(thinking, str) and thinking.strip():\n",
    "            content = f\"<think>\\n{thinking}\\n</think>\\n{content}\"\n",
    "        new_messages.append({\"role\": msg[\"role\"], \"content\": content})\n",
    "    example[\"messages\"] = new_messages\n",
    "    return example\n",
    "\n",
    "train_dataset = train_dataset.remove_columns([\"reasoning_language\", \"developer\", \"user\", \"analysis\", \"final\"])\n",
    "train_dataset = train_dataset.map(merge_thinking_and_remove_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gyh82o2b7i7",
   "metadata": {
    "id": "gyh82o2b7i7"
   },
   "source": "## Load model and configure LoRA\n\nLoad an **NVIDIA Nemotron 3** model. These models are natively supported in `transformers` so no `trust_remote_code` is needed. We use `attn_implementation=\"eager\"` as required by the hybrid Mamba2/Transformer architecture."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "asrptlebp6a",
   "metadata": {
    "id": "asrptlebp6a"
   },
   "outputs": [],
   "source": "import torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n# Select a Nemotron 3 checkpoint below\nmodel_id = \"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16\"\n# model_id = \"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16\"\n\noutput_dir = \"nemotron-3-sft\"\n\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_id,\n    attn_implementation=\"eager\",\n    dtype=torch.bfloat16,\n)\ntokenizer = AutoTokenizer.from_pretrained(model_id)"
  },
  {
   "cell_type": "markdown",
   "id": "90a8pzer8rc",
   "metadata": {
    "id": "90a8pzer8rc"
   },
   "source": [
    "Configure the **LoRA adapter**. Instead of modifying the original weights, we fine-tune a lightweight LoRA adapter for efficient and memory-friendly training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pyukji9o3g",
   "metadata": {
    "id": "pyukji9o3g"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=16,\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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rkxovj1wywh",
   "metadata": {
    "id": "rkxovj1wywh"
   },
   "source": "## Train model\n\nConfigure **SFT** using `SFTConfig`. Note that `gradient_checkpointing` is set to `False` because Nemotron 3 (`NemotronHForCausalLM`) does not support it. For full details on all available parameters, check the [TRL SFTConfig documentation](https://huggingface.co/docs/trl/sft_trainer#trl.SFTConfig)."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tdmoy72oow",
   "metadata": {
    "id": "tdmoy72oow"
   },
   "outputs": [],
   "source": [
    "from trl import SFTConfig\n",
    "\n",
    "training_args = SFTConfig(\n",
    "    # Training schedule / optimization\n",
    "    per_device_train_batch_size=1,                          # Batch size per GPU\n",
    "    gradient_accumulation_steps=4,                          # Effective batch size = per_device_train_batch_size * gradient_accumulation_steps\n",
    "    num_train_epochs=1,                                     # Number of full dataset passes\n",
    "    learning_rate=2e-4,                                     # Learning rate for the optimizer\n",
    "    optim=\"paged_adamw_8bit\",                               # Memory-efficient 8-bit optimizer\n",
    "\n",
    "    # Logging / reporting\n",
    "    logging_steps=10,                                       # Log training metrics every N steps\n",
    "    report_to=\"trackio\",                                    # Experiment tracking tool\n",
    "    trackio_space_id=output_dir,                            # HF Space where the experiment tracking will be saved\n",
    "    output_dir=output_dir,                                  # Where to save model checkpoints and logs\n",
    "\n",
    "    max_length=128,                                         # Kept short due to VRAM constraints; increase for better results\n",
    "    gradient_checkpointing=False,                           # NemotronH does not support gradient checkpointing\n",
    "\n",
    "    # Hub integration\n",
    "    push_to_hub=True,                                       # Push the trained model to the Hugging Face Hub\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ciesqgt2ch",
   "metadata": {
    "id": "4ciesqgt2ch"
   },
   "source": [
    "Configure the SFT Trainer. We pass the previously configured `training_args` and the `peft_config` for LoRA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "h0i5c9sd1ip",
   "metadata": {
    "id": "h0i5c9sd1ip"
   },
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    peft_config=peft_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "q1t7v0lba6l",
   "metadata": {
    "id": "q1t7v0lba6l"
   },
   "source": [
    "Show memory stats before training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e0fseivvu",
   "metadata": {
    "id": "b4e0fseivvu",
    "outputId": "94e1ffbb-639c-4dba-d2b7-2d1de0f493bf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU = NVIDIA A100-SXM4-80GB. Max memory = 79.251 GB.\n",
      "58.939 GB of memory reserved.\n"
     ]
    }
   ],
   "source": [
    "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",
    "\n",
    "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n",
    "print(f\"{start_gpu_memory} GB of memory reserved.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1x3fdnogsvh",
   "metadata": {
    "id": "1x3fdnogsvh"
   },
   "source": [
    "And train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ev9j7op2mi",
   "metadata": {
    "id": "3ev9j7op2mi",
    "outputId": "d6db6cdb-8d43-4cc6-e024-e77c266090f7"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 11, 'pad_token_id': None}.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Trackio project initialized: huggingface\n",
      "* Trackio metrics will be synced to Hugging Face Dataset: sergiopaniego/nemotron-3-sft-dataset\n",
      "* Creating new space: https://huggingface.co/spaces/sergiopaniego/nemotron-3-sft\n",
      "* View dashboard by going to: https://sergiopaniego-nemotron-3-sft.hf.space/\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://sergiopaniego-nemotron-3-sft.hf.space/\" width=\"100%\" height=\"1000px\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* NVIDIA GPU detected, enabling automatic GPU metrics logging\n",
      "* Created new run: sergiopaniego-1773149018\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.12/dist-packages/transformers/models/nemotron_h/modeling_nemotron_h.py:1215: FutureWarning: `input_embeds` is deprecated and will be removed in version 5.6.0 for `create_causal_mask`. Use `inputs_embeds` instead.\n",
      "  causal_mask = create_causal_mask(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='250' max='250' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [250/250 08:53, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>2.379619</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>1.849189</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>1.466606</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>1.207237</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>1.064882</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>1.067329</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>1.085463</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>0.964451</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>0.993209</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.061509</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>110</td>\n",
       "      <td>0.976885</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>0.939225</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>130</td>\n",
       "      <td>1.011990</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>0.989627</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>0.946731</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>0.939989</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>170</td>\n",
       "      <td>0.920945</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>1.003975</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>190</td>\n",
       "      <td>0.945233</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.958358</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>210</td>\n",
       "      <td>1.022524</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>220</td>\n",
       "      <td>0.953711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>230</td>\n",
       "      <td>1.033258</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>240</td>\n",
       "      <td>0.955309</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>250</td>\n",
       "      <td>0.959105</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Run finished. Uploading logs to Trackio Space (please wait...)\n"
     ]
    }
   ],
   "source": [
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ngrmymxpbp",
   "metadata": {
    "id": "ngrmymxpbp"
   },
   "source": [
    "Show memory stats after training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08gpt482fn1k",
   "metadata": {
    "id": "08gpt482fn1k",
    "outputId": "e4808857-5756-4f5d-d34c-29c81107a618"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "542.949 seconds used for training.\n",
      "9.05 minutes used for training.\n",
      "Peak reserved memory = 61.922 GB.\n",
      "Peak reserved memory for training = 2.983 GB.\n",
      "Peak reserved memory % of max memory = 78.134 %.\n",
      "Peak reserved memory for training % of max memory = 3.764 %.\n"
     ]
    }
   ],
   "source": [
    "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",
    "\n",
    "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n",
    "print(f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\")\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",
   "id": "pkp4wfpwpxh",
   "metadata": {
    "id": "pkp4wfpwpxh"
   },
   "source": [
    "## Save fine-tuned model\n",
    "\n",
    "Save the fine-tuned model both **locally** and to the **Hugging Face Hub**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "g2j7qp1ycxl",
   "metadata": {
    "id": "g2j7qp1ycxl",
    "outputId": "5791fc37-589b-4bb9-ca66-85434b303787",
    "colab": {
     "referenced_widgets": [
      "20def6e82e1d4b69ae1b677a39a54bfc",
      "53fe5574751540b68f2d6ef047cccc3c",
      "3661c026040f44f3870921e8a2db7bd2",
      "61b4f98716b649dc88332367852ec89c",
      "a61820f54e5445ceaf27efa956212271",
      "d761f8d2a0ff4de2a8728a1998e5f919",
      "3fea49f4e519416f81af219023d45c34",
      "240c5e45b5f245f984cdaaea9ff3a380",
      "7ba75fdca7da467cacc407da59382376",
      "2b3e22abd6704ef6b5cb2455c6fa98ab"
     ]
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20def6e82e1d4b69ae1b677a39a54bfc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processing Files (0 / 0)      : |          |  0.00B /  0.00B            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53fe5574751540b68f2d6ef047cccc3c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "New Data Upload               : |          |  0.00B /  0.00B            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3661c026040f44f3870921e8a2db7bd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...n-3-sft/training_args.bin: 100%|##########| 5.58kB / 5.58kB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61b4f98716b649dc88332367852ec89c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...adapter_model.safetensors: 100%|##########| 13.2MB / 13.2MB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a61820f54e5445ceaf27efa956212271",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...tron-3-sft/tokenizer.json: 100%|##########| 17.1MB / 17.1MB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d761f8d2a0ff4de2a8728a1998e5f919",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processing Files (0 / 0)      : |          |  0.00B /  0.00B            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fea49f4e519416f81af219023d45c34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "New Data Upload               : |          |  0.00B /  0.00B            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "240c5e45b5f245f984cdaaea9ff3a380",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...n-3-sft/training_args.bin: 100%|##########| 5.58kB / 5.58kB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ba75fdca7da467cacc407da59382376",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...tron-3-sft/tokenizer.json: 100%|##########| 17.1MB / 17.1MB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2b3e22abd6704ef6b5cb2455c6fa98ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  ...adapter_model.safetensors: 100%|##########| 13.2MB / 13.2MB            "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/sergiopaniego/nemotron-3-sft/commit/dd05a083358e10113019b06da76c3e9ab4c4c3e4', commit_message='End of training', commit_description='', oid='dd05a083358e10113019b06da76c3e9ab4c4c3e4', pr_url=None, repo_url=RepoUrl('https://huggingface.co/sergiopaniego/nemotron-3-sft', endpoint='https://huggingface.co', repo_type='model', repo_id='sergiopaniego/nemotron-3-sft'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.save_model(output_dir)\n",
    "trainer.push_to_hub(dataset_name=dataset_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "jp3u13oqlye",
   "metadata": {
    "id": "jp3u13oqlye"
   },
   "source": [
    "## Inference\n",
    "\n",
    "Let's run the fine-tuned model using standard `transformers` generation (`model.generate`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1tb9jai824gi",
   "metadata": {
    "id": "1tb9jai824gi"
   },
   "outputs": [],
   "source": [
    "model = trainer.model\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3zp0i0zw4je",
   "metadata": {
    "id": "3zp0i0zw4je"
   },
   "outputs": [],
   "source": [
    "messages = [{\"role\": \"user\", \"content\": \"Continue the sequence: 1, 1, 2, 3, 5, 8,\"}]\n",
    "text = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    ")\n",
    "\n",
    "from transformers import TextStreamer\n",
    "\n",
    "_ = model.generate(\n",
    "    **tokenizer(text, return_tensors=\"pt\").to(\"cuda\"),\n",
    "    max_new_tokens=128,\n",
    "    temperature=0.7,\n",
    "    top_p=0.8,\n",
    "    top_k=20,\n",
    "    streamer=TextStreamer(tokenizer, skip_prompt=True),\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": [],
   "machine_shape": "hm",
   "gpuType": "A100"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}