Fine-tuning Google Gemma 2 on a Single Consumer GPU

Fine-Tuning Gemma 2

I recently tested fine-tuning Google's Gemma 2 (9B) on a standard consumer GPU. Traditionally, a 9B model would OOM on a 24GB card during training, but using Unsloth makes it incredibly efficient. I managed to get LoRA fine-tuning running smoothly on an RTX 3090, and I'll walk you through my exact script.

Why Unsloth?

I found that standard Hugging Face PEFT scripts use way too much VRAM and take forever. Unsloth rewrites the backpropagation kernels, speeding up training by 2x while cutting memory usage in half. For Gemma 2, it's virtually mandatory if you're working locally.

The Training Script

First, install the required packages. I highly recommend a fresh Python environment.

pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes

Here is the Python script I used to set up the model and start training. I used 4-bit quantization (bitsandbytes) to squeeze it into my VRAM.

from unsloth import FastLanguageModel
import torch

max_seq_length = 2048
dtype = None
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma-2-9b",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
)

Formatting the Dataset

For my dataset, I mapped my custom instructions to the standard ChatML format. Making sure the tokenizer handles the EOS token properly is critical, otherwise the model will generate endless text during inference.

Starting the Trainer

I used the standard SFTTrainer from TRL, plugged in the Unsloth model, and kicked off the run. I monitored the VRAM usage and it peaked at around 16GB, well within the limits of my RTX 3090.