Agent Skills: ML Fine-Tuning

ML Fine-Tuning

Overview

Use this skill for adapting pretrained models to new tasks, domains, styles, instructions, modalities, or constraints. Fine-tuning is not always the right first step: compare prompting, retrieval-augmented generation, feature extraction, classical heads, and smaller task-specific models before training a large foundation model.

Choose the Adaptation Method

| Method | Use when | Trade-offs | |---|---|---| | Prompting/system instructions | Behavior change is simple and context fits | No training cost; limited persistence and control | | RAG | Need factual/domain knowledge that changes or must be cited | Requires retrieval quality, chunking, and grounding evaluation | | Feature extraction + head | Small labeled dataset and strong pretrained embeddings | Efficient; limited deep adaptation | | Partial fine-tuning | Need domain adaptation with limited compute | Must choose layers carefully | | Full fine-tuning | Large dataset, high task specificity, enough compute | Highest cost and forgetting risk | | LoRA/adapters/PEFT | Need efficient adaptation and many variants | Slight capacity limits; target modules matter | | QLoRA | Fine-tune large LLMs on constrained GPUs | Quantization and optimizer choices affect stability | | Preference tuning/RLHF/DPO | Need behavior alignment to preferences | Reward/preference data quality dominates | | Distillation | Need smaller/faster deployable model | Requires teacher quality and representative data |

Dataset Preparation

Fine-tuning quality is usually data quality. Deduplicate, remove corrupted samples, filter unsafe or irrelevant content, balance tasks, standardize labels, and preserve evaluation holdouts. For instruction tuning, use consistent schemas and chat templates matching the base model. For classification, check label definitions and inter-annotator agreement. For retrieval/embedding fine-tuning, mine hard negatives and avoid false negatives.

Split data by entity, time, source, or document when examples are correlated. Never let near-duplicates, prompt variants, chunks from the same document, or generated paraphrases cross train/validation/test boundaries. For generative tasks, keep a human-reviewable validation set with representative hard cases.

QLoRA and PEFT (Hugging Face Blueprint)

QLoRA keeps the base model frozen in 4-bit while training small low-rank adapters, saving massive amounts of GPU VRAM.

1. Initialize 4-Bit Base Model (BitsAndBytesConfig)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

# Configure 4-bit quantization details
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4", # Normalized Float 4 (optimal for zero-mean normal weights)
    bnb_4bit_compute_dtype=torch.bfloat16 # Compute in BF16 for training stability
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

2. Configure LoRA Parameters & Prepare Model

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Prepare frozen model for quantized adapter training
model = prepare_model_for_kbit_training(model)

# Define Low-Rank Adaptation (LoRA) Config
peft_config = LoraConfig(
    r=16, # Rank dimension (typically 8, 16, 32, or 64)
    lora_alpha=32, # Scaling factor (usually 2 * rank)
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Wrap base model with LoRA adapters
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

TRL (Transformer Reinforcement Learning) Training Setup

TRL provides wrappers such as SFTTrainer for Supervised Fine-Tuning and DPOTrainer for Direct Preference Optimization.

1. Supervised Fine-Tuning (SFTTrainer)

from trl import SFTTrainer
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=10,
    max_steps=500,
    fp16=False,
    bf16=True, # Optimal on Ampere/Hopper architecture GPUs
    optim="paged_adamw_8bit" # Page state out to host memory to prevent OOMs
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_args
)

# trainer.train()

2. Direct Preference Optimization (DPOTrainer)

DPO directly trains a model on pairs of preferred/rejected outputs without needing a separate reward model.

from trl import DPOTrainer

dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None, # TRL will implicitly handle ref_model under PEFT by disabling adapters
    args=training_args,
    beta=0.1, # Temperature parameter for DPO loss (scales difference from reference policy)
    train_dataset=preference_dataset,
    tokenizer=tokenizer,
    max_length=512,
    max_prompt_length=256
)
# dpo_trainer.train()

SentenceTransformers Custom Embedding Model Fine-Tuning

Embedding fine-tuning creates custom vector spaces for retrieval-heavy (RAG) applications.

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

model = SentenceTransformer("all-MiniLM-L6-v2")

# Create positive pairs (e.g., query and matching document)
train_examples = [
    InputExample(texts=["How to initialize FSDP in PyTorch?", "Use FullyShardedDataParallel wrapper..."]),
    InputExample(texts=["What is QLoRA?", "QLoRA trains adapters over a frozen 4-bit base model..."])
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# MultipleNegativesRankingLoss treats other pairs in the batch as implicit negative examples
train_loss = losses.MultipleNegativesRankingLoss(model=model)

# model.fit(
#     train_objectives=[(train_dataloader, train_loss)], 
#     epochs=3, 
#     warmup_steps=100
# )

Training Stability

Use conservative learning rates for fine-tuning. Warmup often helps. Monitor train and validation loss, task metrics, generation quality, and overfitting. For LLMs, watch for format drift, verbosity changes, refusal regressions, hallucination, tool-use regressions, and safety issues. Use gradient clipping for instability and validate that labels are masked correctly for instruction tuning.

Catastrophic forgetting appears when the model improves on new data but loses general ability. Mitigate with smaller LR, fewer epochs, mixed-domain data, replay examples, adapters instead of full fine-tuning, regularization, or multi-task balancing.

Sources

• Hugging Face Transformers documentation: https://huggingface.co/docs/transformers
• Hugging Face PEFT documentation: https://huggingface.co/docs/peft
• Hugging Face TRL documentation: https://huggingface.co/docs/trl
• Hugging Face Diffusers training docs: https://huggingface.co/docs/diffusers/training/overview
• Sentence Transformers documentation: https://www.sbert.net/
• Ray and Accelerate distributed training docs: https://docs.ray.io/ and https://huggingface.co/docs/accelerate