{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "w2TnJ6ta-2zj" }, "source": [ "# GRPO Fine-tuning for verifiable tasks\n", "\n", "Fine-tuning requires a GPU. If you don't have one locally, you can run this notebook for free on [Google Colab](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/grpo_for_verifiable_tasks.ipynb) using a free NVIDIA T4 GPU instance.\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/grpo_for_verifiable_tasks.ipynb)\n", "\n", "### What's in this notebook?\n", "\n", "In this notebook you will learn how to fine-tune a Small Text-to-text Language Model for verifiable tasks using Group Relative Policy Optimization (GRPO).\n", "\n", "GRPO is a Reinforcement Learning algorithm widely used by AI labs and practitioners to fine-tune models for easily verifiables tasks like\n", "\n", "- Mathematical problem solving with numeric verification\n", "- Code generation with unit test validation\n", "- Structured output tasks (JSON, SQL) with schema validation\n", "- Question answering with ground truth answers\n", "\n", "\n", "We will cover\n", "- Environment setup\n", "- Data preparation\n", "- Model training\n", "- Local inference with your new model\n", "- Model saving and exporting it into the format you need for **deployment**.\n", "\n", "### Deployment options\n", "\n", "LFM2.5 models are small and efficient, enabling deployment across a wide range of platforms:\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Deployment TargetUse Case
📱 AndroidMobile apps on Android devices
📱 iOSMobile apps on iPhone/iPad
🍎 Apple Silicon MacLocal inference on Mac with MLX
🦙 llama.cppLocal deployments on any hardware
🦙 OllamaLocal inference with easy setup
🖥️ LM StudioDesktop app for local inference
⚡ vLLMCloud deployments with high throughput
☁️ ModalServerless cloud deployment
🏗️ BasetenProduction ML infrastructure
🚀 FalFast inference API
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Need help building with our models and tools?\n", "Join the Liquid AI Discord Community and ask.\n", "\n", "\"Join\n", "\n", "And now, let the fine tune begin!" ] }, { "cell_type": "markdown", "metadata": { "id": "9RFq6Op7rjc3" }, "source": [ "## 📦 Installation & Setup\n", "\n", "First, let's install all the required packages:\n", "\n", "We'll install **TRL** with the **PEFT** extra, which ensures all main dependencies such as **Transformers** and **PEFT** (a package for parameter-efficient fine-tuning, e.g., LoRA/QLoRA) are included.\n", "\n", "Additionally, we'll install\n", "- **trackio** to log and monitor our experiments\n", "- **bitsandbytes** to enable quantization of LLMs, reducing memory consumption for both inference and training, and\n", "- **liger-kernel** for more efficient training." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "c2jy45nfWbdo", "outputId": "ec580e3b-30c5-4192-bf3d-e50587d801f2" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.1/59.1 MB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m41.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m209.1/209.1 kB\u001b[0m \u001b[31m16.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m230.6/230.6 kB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.3/24.3 MB\u001b[0m \u001b[31m35.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.7/55.7 kB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.9/9.9 MB\u001b[0m \u001b[31m45.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m532.9/532.9 kB\u001b[0m \u001b[31m11.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hCollecting flash-attn\n", " Downloading flash_attn-2.8.3.tar.gz (8.4 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.4/8.4 MB\u001b[0m \u001b[31m93.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25h Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n", "Requirement already satisfied: torch in /usr/local/lib/python3.12/dist-packages (from flash-attn) (2.9.0+cu126)\n", "Requirement already satisfied: einops in /usr/local/lib/python3.12/dist-packages (from flash-attn) (0.8.1)\n", "Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (3.20.3)\n", "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (4.15.0)\n", "Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (75.2.0)\n", "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (1.14.0)\n", "Requirement already satisfied: networkx>=2.5.1 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (3.6.1)\n", "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (3.1.6)\n", "Requirement already satisfied: fsspec>=0.8.5 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (2025.3.0)\n", "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.77)\n", "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.77)\n", "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.80)\n", "Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (9.10.2.21)\n", "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.4.1)\n", "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (11.3.0.4)\n", "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (10.3.7.77)\n", "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (11.7.1.2)\n", "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.5.4.2)\n", "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (0.7.1)\n", "Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (2.27.5)\n", "Requirement already satisfied: nvidia-nvshmem-cu12==3.3.20 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (3.3.20)\n", "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.77)\n", "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (12.6.85)\n", "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (1.11.1.6)\n", "Requirement already satisfied: triton==3.5.0 in /usr/local/lib/python3.12/dist-packages (from torch->flash-attn) (3.5.0)\n", "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch->flash-attn) (1.3.0)\n", "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch->flash-attn) (3.0.3)\n", "Building wheels for collected packages: flash-attn\n", " Building wheel for flash-attn (setup.py) ... \u001b[?25l\u001b[?25hdone\n", " Created wheel for flash-attn: filename=flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl size=253780426 sha256=4e2f9e39313266b1544b68138b15b91ee6221eccf14f7902b7c6620351340810\n", " Stored in directory: /root/.cache/pip/wheels/3d/59/46/f282c12c73dd4bb3c2e3fe199f1a0d0f8cec06df0cccfeee27\n", "Successfully built flash-attn\n", "Installing collected packages: flash-attn\n", "Successfully installed flash-attn-2.8.3\n" ] } ], "source": [ "!pip install -Uq \"trl[peft]\" bitsandbytes trackio math_verify liger-kernel\n", "!pip install flash-attn --no-build-isolation\n", "\n", "# Optional: for experiment tracking\n", "# !pip install trackio" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import transformers\n", "import trl\n", "import os\n", "os.environ[\"WANDB_DISABLED\"] = \"true\"\n", "\n", "print(f\"📦 PyTorch version: {torch.__version__}\")\n", "print(f\"🤗 Transformers version: {transformers.__version__}\")\n", "print(f\"📊 TRL version: {trl.__version__}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "cTEw4xlFrhnQ" }, "source": [ "## Load the dataset\n", "\n", "In this step, we load the [**AI-MO/NuminaMath-TIR**](https://huggingface.co/datasets/AI-MO/NuminaMath-TIR) dataset from the Hugging Face Hub using the `datasets` library.\n", "This dataset focuses on **mathematical reasoning**, featuring problems that require step-by-step logical solutions.\n", "By fine-tuning a model that does not yet exhibit strong reasoning capabilities, it can learn to **generate structured reasoning steps**, enhancing both the model's **accuracy** and **interpretability** on math-related tasks.\n", "\n", "For efficiency, we'll load only a **small portion of the training split**:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 177, "referenced_widgets": [ "a0c1df08396a476bbf695d1c39f9f7fb", "29fa5d7eab434f3b82bd23a2c66360f7", "b33292fcc71e4ab58ea47275c212d697", "cc64a073b0d84a908d0b7cee3180f7d5", "df0e51acaf144f15a19680da5e1d2026", "28c67efc530c413ab67f4ff72b2243fa", "f01a5690de164e2cacb0a09b10890c83", "61e5aea021bc412982a9d5b46b29f150", "7aade0edb6e142ec9e094deeaac9ae3c", "90bf0081c15240f899eefa0cdcaf20a7", "de14c76e12a84b0989d984992dcb2d35", "57387e8d9d084e9ba407dd334c61155b", "274ce29e906b44a7aa950b44fce15e8b", "86123462d43d48759dc54587eb54cc0c", "55139debbd5e4affbebe69dd77cb389b", "59cb8f743dca4a25868b5e26d4dccfcf", "db62e68b6ee34fbbadbf96f5d909abb1", "e30967f2674f4826b8507ecc9903c041", "30f23982feca4d37a7fac7f3145ccf4e", "775a1a16d01f408a95e8966fd4fb9488", "a98f85f5a995425398d47d567e8b211e", "861153263478420782f0794cffe7108b", "86a0209365774b0b90d2769580a6c011", "475e0d86fab542c3b7d3b8d6f8fc34bf", "4cd13573d878419cb4ebfdcf87998fed", "b928511ea22f4a8d819f4a22f1000049", "051953dc8ccb4921bc92add66417635e", "0991796f27c24d27ab3b2732058b6b88", "cfa780a55f3944879213ee46867000ed", "be6159e6766d4335a49c76cedc3e686e", "4899851f1cfb4f61b6b9656202298393", "c77b01bd946f4e3dbcfdb5d8f1df7b94", "be80d2d444cf4fbbb7a5ee28bdc9008e", "23fadf2b87f1429c9a744245db961d68", "4c1cb5f9517348f3a68f91ca28ff5d56", "91993ad8459c4880a6bda5652d48cd87", "e66cee7c058f47bda6f7a05f91198b26", "ace2d4e384364ac8bf9898cbd2fbb97b", "fe136fcdce5e449fb2f7103bc72ead35", "9653aad1deb44d7995d086c82daa07d6", "f8f4fefd5a704f8788f0dcb02c957a8c", "abbe01f79d7043da9f4c720ec0a4be18", "e0df14a39b07429e90e5884b72abb4f1", "fa2e15a2e3d8415b981a5ac30984e997", "eb86a9e0e1964a66b71856ef72041e95", "15e9a659c62b495886e8864751bd9922", "a27c9ae11ccc458a80a0ab8b9f8eaefb", "f74fc66788394347844947149a27c4a7", "a4030c4c302d454d8b5455eb42c02b48", "c2495842a4764d66a2a4f2be4a82631a", "7dd51a09e81e43f5bae2fef5ff598ae5", "95086541873147d9881b22d1018f1785", "6e194a220cb4465aae47a51977c9298c", "df98e95bf79949358b72226ba090ee08", "d2d49e0727094bebb347c5425cedb6f3" ] }, "id": "zU5icx67Wbdp", "outputId": "0971f226-56bd-4c99-e92c-e4aecea9c591" }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "a0c1df08396a476bbf695d1c39f9f7fb", "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": "57387e8d9d084e9ba407dd334c61155b", "version_major": 2, "version_minor": 0 }, "text/plain": [ "data/train-00000-of-00001.parquet: 0%| | 0.00/147M [00:00 and tags. \"\n", " \"After closing , the assistant MUST provide the final answer in plain text.\"\n", ")\n", "\n", "\n", "def make_conversation(example):\n", " return {\n", " \"prompt\": [\n", " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n", " {\"role\": \"user\", \"content\": example[\"problem\"]},\n", " ],\n", " }\n", "\n", "train_dataset = train_dataset.map(make_conversation)\n", "\n", "# remove unused columns\n", "train_dataset = train_dataset.remove_columns(['messages', 'problem'])\n", "\n", "# Check the structure of the dataset\n", "print(train_dataset)\n", "print(train_dataset[0]['prompt'])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load the model" ] }, { "cell_type": "markdown", "metadata": { "id": "aw__94OWDnER" }, "source": [ "This notebook can be used with two fine-tuning methods. By default, it is set up for **QLoRA**, which includes quantization using `BitsAndBytesConfig`. If you prefer to use standard **LoRA** without quantization, simply comment out the `BitsAndBytesConfig` configuration (training without quantization consumes more memory)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "o86TnTchWbdp" }, "outputs": [], "source": [ "import torch\n", "from transformers import AutoModelForCausalLM, BitsAndBytesConfig\n", "\n", "model_id, output_dir = \"LiquidAI/LFM2.5-1.2B-Instruct\", \"LFM2.5-1.2B-Instruct-GRPO\"\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " model_id,\n", " # attn_implementation=\"flash_attention_2\", # Change to Flash Attention if GPU has support\n", " dtype=\"bfloat16\", # Change to bfloat16 if GPU has support\n", " # quantization_config=BitsAndBytesConfig(\n", " # load_in_4bit=True, # Load the model in 4-bit precision to save memory\n", " # bnb_4bit_compute_dtype=torch.float16, # Data type used for internal computations in quantization\n", " # bnb_4bit_use_double_quant=True, # Use double quantization to improve accuracy\n", " # bnb_4bit_quant_type=\"nf4\" # Type of quantization. \"nf4\" is recommended for recent LLMs\n", " # )\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Define LoRA adapters" ] }, { "cell_type": "markdown", "metadata": { "id": "AM-G0_QmDyZC" }, "source": [ "The following cell defines LoRA (or QLoRA if needed). When training with LoRA/QLoRA, we use a **base model** (the one selected above) and, instead of modifying its original weights, we fine-tune a **LoRA adapter**, a lightweight layer that enables efficient and memory-friendly training. The **`target_modules`** specify which parts of the model (e.g., attention or projection layers) will be adapted by LoRA during fine-tuning." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "WIz2pmX6Wbdp" }, "outputs": [], "source": [ "from peft import LoraConfig\n", "peft_config = LoraConfig(\n", " r=32,\n", " lora_alpha=32,\n", " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"out_proj\", \"in_proj\", \"w1\", \"w2\", \"w3\"],\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load reward functions or define your own ones\n", "\n", "GRPO requires **reward functions** to guide the learning process. For convenience, we can directly load pre-defined rewards from `trl.rewards`, which already includes a [collection of ready-to-use rewards](https://huggingface.co/docs/trl/rewards).\n", "\n", "If you want to create your own custom reward functions to teach the model, a reward function is simply a Python function that takes the generated completions and returns a list of floats. For example, the following function, which we use in this notebook, rewards completions that correctly follow the `` format:\n", "\n", "```python\n", "def think_format_reward(completions: list[list[dict[str, str]]], **kwargs) -> list[float]:\n", " pattern = r\"^(?!.*)(.*?).*$\"\n", " completion_contents = [completion[0][\"content\"] for completion in completions]\n", " matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE) for content in completion_contents]\n", " return [1.0 if match else 0.0 for match in matches]\n", "```\n", "\n", "In this notebook, we will use both `think_format_reward`, which rewards completions that correctly follow the `` format, and `reasoning_accuracy_reward`, which evaluates the correctness of the model's solution to the mathematical problem. Together, these rewards guide the model to generate **structured reasoning** while producing **accurate answers**." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from trl.rewards import think_format_reward, reasoning_accuracy_reward" ] }, { "cell_type": "markdown", "metadata": { "id": "prKnAp-Esyiq" }, "source": [ "## Train model" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "lj42Qs5vWbdp" }, "outputs": [], "source": [ "from trl import GRPOConfig\n", "\n", "# Configure training arguments using GRPOConfig\n", "training_args = GRPOConfig(\n", " # Training schedule / optimization\n", " learning_rate=2e-5, # Learning rate for the optimizer\n", " #num_train_epochs=1,\n", " max_steps=30, # Number of dataset passes. For full trainings, use `num_train_epochs` instead\n", "\n", " # Parameters that control GRPO training (you can adapt them)\n", " per_device_train_batch_size = 8,\n", " max_completion_length=256, # default: 256 # Max completion length produced during training\n", " num_generations=8, # default: 8 # Number of generations produced during trainig for comparison\n", "\n", " # Optimizations\n", " optim = \"paged_adamw_8bit\", # Optimizer\n", " use_liger_kernel=True, # Enable Liger kernel optimizations for faster training\n", " gradient_checkpointing=True, # Save memory by re-computing activations during backpropagation\n", "\n", " # Parameters related to reporting and saving\n", " output_dir=output_dir, # Where to save model checkpoints and logs\n", " logging_steps=10, # Log training metrics every N steps\n", " report_to=[],\n", " # report_to=\"trackio\", # Experiment tracking tool\n", " # trackio_space_id=output_dir, # HF Space where the experiment tracking will be saved\n", " log_completions=False, # Return model completions during training\n", "\n", " # Hub integration\n", " # push_to_hub=True, # Automatically push the trained model to the Hugging Face Hub\n", " # The model will be saved under your Hub account in the repository named `output_dir`\n", " # vLLM params\n", " # Enable with `use_vllm=True` and customize with the remaining params faster training\n", " #use_vllm=False,\n", " #vllm_mode='colocate',\n", " #vllm_gpu_memory_utilization=0.1,\n", " #vllm_enable_sleep_mode=True\n", ")\n", "\n", "from trl import GRPOTrainer\n", "trainer = GRPOTrainer(\n", " model=model,\n", " reward_funcs=[think_format_reward, reasoning_accuracy_reward],\n", " args=training_args,\n", " train_dataset=train_dataset,\n", " peft_config=peft_config,\n", ")\n", "\n", "trainer_stats = trainer.train()" ] }, { "cell_type": "markdown", "metadata": { "id": "ibO4f7tuLboQ" }, "source": [ "## (Optional) Save fine tuned model\n", "\n", "In this step, we save the fine-tuned model **locally**." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "background_save": true }, "id": "itpVDjy0Wbdt" }, "outputs": [], "source": [ "trainer.save_model(output_dir)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "machine_shape": "hm", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }