Training and Fine-Tuning

Training and Fine-Tuning

Overview

Fine-tune pre-trained models on custom datasets using the Trainer API. The Trainer handles training loops, gradient accumulation, mixed precision, logging, and checkpointing.

Basic Fine-Tuning Workflow

Step 1: Load and Preprocess Data

from datasets import load_dataset

# Load dataset
dataset = load_dataset("yelp_review_full")
train_dataset = dataset["train"]
eval_dataset = dataset["test"]

# Tokenize
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=512
    )

train_dataset = train_dataset.map(tokenize_function, batched=True)
eval_dataset = eval_dataset.map(tokenize_function, batched=True)

Step 2: Load Model

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=5  # Number of classes
)

Step 3: Define Metrics

import evaluate
import numpy as np

metric = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

Step 4: Configure Training

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
)

Step 5: Create Trainer and Train

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)

# Start training
trainer.train()

# Evaluate
results = trainer.evaluate()
print(results)

Step 6: Save Model

trainer.save_model("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")

# Or push to Hub
trainer.push_to_hub("username/my-finetuned-model")

TrainingArguments Parameters

Essential Parameters

output_dir: Directory for checkpoints and logs

output_dir="./results"

num_train_epochs: Number of training epochs

num_train_epochs=3

per_device_train_batch_size: Batch size per GPU/CPU

per_device_train_batch_size=8

learning_rate: Optimizer learning rate

learning_rate=2e-5  # Common for BERT-style models
learning_rate=5e-5  # Common for smaller models

weight_decay: L2 regularization

weight_decay=0.01

Evaluation and Saving

eval_strategy: When to evaluate ("no", "steps", "epoch")

eval_strategy="epoch"  # Evaluate after each epoch
eval_strategy="steps"  # Evaluate every eval_steps

save_strategy: When to save checkpoints

save_strategy="epoch"
save_strategy="steps"
save_steps=500

load_best_model_at_end: Load best checkpoint after training

load_best_model_at_end=True
metric_for_best_model="accuracy"  # Metric to compare

Optimization

gradient_accumulation_steps: Accumulate gradients over multiple steps

gradient_accumulation_steps=4  # Effective batch size = batch_size * 4

fp16: Enable mixed precision (NVIDIA GPUs)

fp16=True

bf16: Enable bfloat16 (newer GPUs)

bf16=True

gradient_checkpointing: Trade compute for memory

gradient_checkpointing=True  # Slower but uses less memory

optim: Optimizer choice

optim="adamw_torch"  # Default
optim="adamw_8bit"    # 8-bit Adam (requires bitsandbytes)
optim="adafactor"     # Memory-efficient alternative

Learning Rate Scheduling

lr_scheduler_type: Learning rate schedule

lr_scheduler_type="linear"       # Linear decay
lr_scheduler_type="cosine"       # Cosine annealing
lr_scheduler_type="constant"     # No decay
lr_scheduler_type="constant_with_warmup"

warmup_steps or warmup_ratio: Warmup period

warmup_steps=500
# Or
warmup_ratio=0.1  # 10% of total steps

Logging

logging_dir: TensorBoard logs directory

logging_dir="./logs"

logging_steps: Log every N steps

logging_steps=10

report_to: Logging integrations

report_to=["tensorboard"]
report_to=["wandb"]
report_to=["tensorboard", "wandb"]

Distributed Training

ddp_backend: Distributed backend

ddp_backend="nccl"  # For multi-GPU

deepspeed: DeepSpeed config file

deepspeed="ds_config.json"

Data Collators

Handle dynamic padding and special preprocessing:

DataCollatorWithPadding

Pad sequences to longest in batch:

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=data_collator,
)

DataCollatorForLanguageModeling

For masked language modeling:

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15
)

DataCollatorForSeq2Seq

For sequence-to-sequence tasks:

from transformers import DataCollatorForSeq2Seq

data_collator = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=model,
    padding=True
)

Custom Training

Custom Trainer

Override methods for custom behavior:

from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits

        # Custom loss computation
        loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights)
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))

        return (loss, outputs) if return_outputs else loss

Custom Callbacks

Monitor and control training:

from transformers import TrainerCallback

class CustomCallback(TrainerCallback):
    def on_epoch_end(self, args, state, control, **kwargs):
        print(f"Epoch {state.epoch} completed")
        # Custom logic here
        return control

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    callbacks=[CustomCallback],
)

Advanced Training Techniques

Parameter-Efficient Fine-Tuning (PEFT)

Use LoRA for efficient fine-tuning:

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["query", "value"],
    lora_dropout=0.05,
    bias="none",
    task_type="SEQ_CLS"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # Shows reduced parameter count

# Train normally with Trainer
trainer = Trainer(model=model, args=training_args, ...)
trainer.train()

Gradient Checkpointing

Reduce memory at cost of speed:

model.gradient_checkpointing_enable()

training_args = TrainingArguments(
    gradient_checkpointing=True,
    ...
)

Mixed Precision Training

training_args = TrainingArguments(
    fp16=True,  # For NVIDIA GPUs with Tensor Cores
    # or
    bf16=True,  # For newer GPUs (A100, H100)
    ...
)

DeepSpeed Integration

For very large models:

# ds_config.json
{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2
  }
}

training_args = TrainingArguments(
    deepspeed="ds_config.json",
    ...
)

Training Tips

Hyperparameter Tuning

Common starting points:

• Learning rate: 2e-5 to 5e-5 for BERT-like models, 1e-4 to 1e-3 for smaller models
• Batch size: 8-32 depending on GPU memory
• Epochs: 2-4 for fine-tuning, more for domain adaptation
• Warmup: 10% of total steps

Use Optuna for hyperparameter search:

def model_init():
    return AutoModelForSequenceClassification.from_pretrained(
        "bert-base-uncased",
        num_labels=5
    )

def optuna_hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 5e-5, log=True),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [8, 16, 32]),
        "num_train_epochs": trial.suggest_int("num_train_epochs", 2, 5),
    }

trainer = Trainer(model_init=model_init, args=training_args, ...)
best_trial = trainer.hyperparameter_search(
    direction="maximize",
    backend="optuna",
    hp_space=optuna_hp_space,
    n_trials=10,
)

Monitoring Training

Use TensorBoard:

tensorboard --logdir ./logs

Or Weights & Biases:

import wandb
wandb.init(project="my-project")

training_args = TrainingArguments(
    report_to=["wandb"],
    ...
)

Resume Training

Resume from checkpoint:

trainer.train(resume_from_checkpoint="./results/checkpoint-1000")

Common Issues

CUDA out of memory:

• Reduce batch size
• Enable gradient checkpointing
• Use gradient accumulation
• Use 8-bit optimizers

Overfitting:

• Increase weight_decay
• Add dropout
• Use early stopping
• Reduce model size or training epochs

Slow training:

• Increase batch size
• Enable mixed precision (fp16/bf16)
• Use multiple GPUs
• Optimize data loading

Best Practices

1. Start small: Test on small dataset subset first
2. Use evaluation: Monitor validation metrics
3. Save checkpoints: Enable save_strategy
4. Log extensively: Use TensorBoard or W&B
5. Try different learning rates: Start with 2e-5
6. Use warmup: Helps training stability
7. Enable mixed precision: Faster training
8. Consider PEFT: For large models with limited resources