Fine-Tuning OpenAI gpt-oss-20B with Unsloth: 4-bit MoE QLoRA and Harmony Chat Templates

Step 1: Model Loading in 4-bit & MoE-Aware LoRA Configuration

from __future__ import annotations

import torch
from unsloth import FastLanguageModel

max_seq_length = 1024

# Load 4-bit quantized gpt-oss-20B
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gpt-oss-20b-bnb-4bit",
    max_seq_length=max_seq_length,
    dtype=None,             # Auto-selects FP32 on T4, BF16 on Ampere+/Hopper
    load_in_4bit=True,
)

# Configure LoRA strictly for attention and shared MLP projections (excluding router)
lora_rank = 8
model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,
    lora_alpha=lora_rank * 2,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

model.print_trainable_parameters()
# Trainable parameters: 1,990,656 / 20,916,747,840 (0.01% trained)

Step 2: Dataset Formatting with Developer-Steered Reasoning

We map HuggingFaceH4/Multilingual-Thinking into Harmony format to steer the model’s internal reasoning channel to French:

from datasets import load_dataset

REASONING_LANGUAGE = "French"
raw_dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")


def format_harmony_stream(examples: dict) -> dict:
    conversations = examples["messages"]
    formatted_texts = []

    for convo in conversations:
        # Inject Harmony developer steering directive prior to user turn
        steered_convo = []
        for msg in convo:
            if msg["role"] == "user":
                steered_convo.append({
                    "role": "developer",
                    "content": f"reasoning language: {REASONING_LANGUAGE}\n\nYou are a helpful assistant.",
                })
            steered_convo.append(msg)

        text = tokenizer.apply_chat_template(
            steered_convo,
            tokenize=False,
            add_generation_prompt=False,
        )
        formatted_texts.append(text)

    return {"text": formatted_texts}


dataset = raw_dataset.map(format_harmony_stream, batched=True)

Step 3: Supervised Fine-Tuning Execution

from transformers import TrainingArguments
from trl import SFTTrainer
from model_init import model, tokenizer, max_seq_length
from dataset_prep import dataset

training_args = TrainingArguments(
    output_dir="outputs/gpt_oss_20b_sft",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,      # Effective batch size = 8
    warmup_steps=5,
    max_steps=60,
    learning_rate=2e-4,
    fp16=not torch.cuda.is_bf16_supported(),
    bf16=torch.cuda.is_bf16_supported(),
    logging_steps=1,
    optim="adamw_8bit",
    weight_decay=0.01,
    lr_scheduler_type="linear",
    seed=3407,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    packing=False,
    args=training_args,
)

trainer.train()

Step 4: Multi-Channel Inference & Reasoning Effort Modulation

from transformers import TextStreamer
from unsloth import FastLanguageModel

# Fuse dequantization kernels for fast generation
FastLanguageModel.for_inference(model)

messages = [
    {"role": "developer", "content": "reasoning language: French\n\nYou are a helpful math assistant."},
    {"role": "user", "content": "Solve x^5 + 3x^4 - 10 = 3."},
]

inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
    reasoning_effort="medium",  # Configures depth: 'low', 'medium', or 'high'
).to("cuda")

streamer = TextStreamer(tokenizer)
_ = model.generate(**inputs, max_new_tokens=1024, streamer=streamer)

Step 5: Exporting Checkpoints (LoRA, 16-bit, and MXFP4)

# 1. Save standalone LoRA adapter (~50 MB)
model.save_pretrained("outputs/gpt_oss_lora")
tokenizer.save_pretrained("outputs/gpt_oss_lora")

# 2. Merge back into native MXFP4 layout (~12 GB, compatible with official OpenAI runner)
model.save_pretrained_merged(
    "outputs/gpt_oss_mxfp4",
    tokenizer,
    save_method="mxfp4",
)

# 3. Merge into full 16-bit precision for vLLM / TGI deployment (~40 GB)
model.save_pretrained_merged(
    "outputs/gpt_oss_16bit",
    tokenizer,
    save_method="merged_16bit",
)

Empirical Benchmark Evaluation

We evaluated gpt-oss-20B across hardware tiers using the 60-step supervised fine-tuning protocol:

Troubleshooting Common Synthesis Faults

1. Silent NaN Divergence on Float16 Hardware

• Symptom: Training loss collapses to NaN after step 100 on Turing or Volta cards.
• Remedy: Ensure dtype=None during model loading. Unsloth will automatically promote adapter gradients to FP32 while retaining the 4-bit base.

2. Gating Instability from LoRA Target Inclusion

• Symptom: Memory consumption exceeds 24 GB and training latency slows by $>3\times$.
• Remedy: Verify that "router" or "moe_gate" is excluded from target_modules. LoRA adapters should attach strictly to shared projections and attention blocks.

3. Missing Developer Reasoning Channel

• Symptom: Model responds in English despite steering instructions.
• Remedy: Format the steering directive using the Harmony "developer" role rather than standard "system" tags.