Hands-on Fine-Tuning with Unsloth & TRL

+-------------------------------------------------------------------------------+
|                            WHY USE UNSLOTH?                                   |
+-------------------------------------------------------------------------------+
|  Unsloth rewrites PyTorch cross-entropy and attention kernels in custom       |
|  Triton code:                                                                 |
|  - 2x to 5x faster training speed                                             |
|  - 70% to 80% less VRAM consumption                                           |
|  - 0% loss in accuracy (mathematically identical gradients)                  |
|  - Native support for Llama 3, Mistral, Qwen 2.5, and Gemma 2                 |
+-------------------------------------------------------------------------------+

Complete Production Fine-Tuning Script

Below is a complete, runnable script using unsloth and Hugging Face's trl (Transformer Reinforcement Learning) SFTTrainer:

import torch
from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments

# 1. Configuration
max_seq_length = 2048
model_name = "unsloth/meta-llama-3.1-8b-Instruct-bnb-4bit"
dtype = None # Auto-detect (Float16 for Tesla T4, Bfloat16 for Ampere/Hopper)
load_in_4bit = True

# 2. Load Base Model and Tokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_name,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)

# 3. Configure LoRA Parameters
model = FastLanguageModel.get_peft_model(
    model,
    r=16,                         # Rank: 8, 16, 32, 64
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0.0,             # Unsloth supports optimized 0 dropout
    bias="none",
    use_gradient_checkpointing="unsloth", # 30% VRAM reduction
    random_state=42,
)

# 4. Load & Format Dataset
dataset = load_dataset("json", data_files="pyspark_dataset.jsonl", split="train")

def formatting_prompts_func(examples):
    convos = examples["messages"]
    texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
    return {"text": texts}

dataset = dataset.map(formatting_prompts_func, batched=True)

# 5. Initialize SFTTrainer
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, # Set True for 5x faster training on short sequences
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=100, # or num_train_epochs=3
        learning_rate=2e-4,
        fp16=not torch.cuda.is_bf16_supported(),
        bf16=torch.cuda.is_bf16_supported(),
        logging_steps=10,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=42,
        output_dir="outputs",
    ),
)

# 6. Execute Training
trainer_stats = trainer.train()

# 7. Save LoRA Adapters
model.save_pretrained("my_custom_lora_adapters")
tokenizer.save_pretrained("my_custom_lora_adapters")
print("Training complete! Adapters saved to ./my_custom_lora_adapters")