Generative AI Google Brain

open-models-fine-tuning-gpt-oss-20b-finetuning-with-unsloth

3 min read

Fine-tuning GPT-OSS 20B with Unsloth on Vertex AI Colab Enterprise and Nvidia A100 40GB GPU

Source notebook

Repo path: open-models/fine-tuning/gpt_oss_20B_finetuning_with_unsloth.ipynb · Open on GitHub · intermediate

Fine-tunes GPT-OSS 20B with Unsloth LoRA on Vertex AI Colab Enterprise using an A100 GPU.

Summary

This notebook teaches how to fine-tune the GPT-OSS 20B open model with Unsloth in a Vertex AI Colab Enterprise GPU runtime. The workflow installs pinned training libraries, initializes Google GenAI for Vertex AI, loads unsloth/gpt-oss-20b in 4-bit mode, adds LoRA adapters, formats the HuggingFaceH4/Multilingual-Thinking dataset, trains with TRL SFTTrainer, reports GPU memory stats, and runs inference with the fine-tuned model.

Key code patterns

Vertex AI GenAI client

from google import genai
 
client = genai.Client(
    vertexai=True,
    project=PROJECT_ID,
    location=LOCATION,
)

Configures the notebook to use Vertex AI with the Google GenAI SDK.

Load GPT-OSS 20B in 4-bit

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gpt-oss-20b",
    dtype=None,
    max_seq_length=1024,
    load_in_4bit=True,
    full_finetuning=False,
)

Loads the model with 4-bit quantization to reduce GPU memory use.

Add LoRA adapters

model = FastLanguageModel.get_peft_model(
    model,
    r=8,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    use_gradient_checkpointing="unsloth",
)

Enables parameter-efficient fine-tuning of selected projection modules.

Format chat dataset

dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")
dataset = standardize_sharegpt(dataset)
dataset = dataset.map(formatting_prompts_func, batched=True)

Converts multilingual reasoning conversations into tokenizer chat-template text.

Train with SFTTrainer

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=SFTConfig(max_steps=30, learning_rate=2e-4, optim="adamw_8bit"),
)
trainer_stats = trainer.train()

Runs supervised fine-tuning with TRL using a short 30-step demo configuration.

Models & APIs used

  • Models: unsloth/gpt-oss-20b
  • APIs / services: Vertex AI
  • SDKs / libraries: google-genai, unsloth, torch, datasets, trl, transformers, peft, accelerate, bitsandbytes, wandb

When to use this

Use this pattern to prototype parameter-efficient fine-tuning of GPT-OSS 20B on Vertex AI Colab Enterprise with an A100 GPU.

Gotchas & caveats

  • Requires a Google Cloud project with the Vertex AI API enabled.
  • Requires a Vertex AI Colab Enterprise runtime with a GPU such as NVIDIA A100.
  • Notebook states it is educational, not suitable for production, and use at your own risk.
  • Installs pinned git revisions and package versions, including numpy<2.0 and protobuf⇐3.20.3, to avoid environment conflicts.
  • The notebook restarts the kernel after installation.
  • The demo trains only 30 steps; a full run requires changing training settings such as num_train_epochs.

Best practices

  • Use 4-bit loading to reduce memory usage for the 20B model.
  • Use LoRA adapters for parameter-efficient fine-tuning instead of full fine-tuning.
  • Use Unsloth gradient checkpointing to reduce VRAM and support longer context.
  • Pin package versions and git revisions for compatibility in the managed notebook environment.
  • Track GPU memory before and after training.

Graph View

Backlinks