10 KiB
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
- Start small: Test on small dataset subset first
- Use evaluation: Monitor validation metrics
- Save checkpoints: Enable save_strategy
- Log extensively: Use TensorBoard or W&B
- Try different learning rates: Start with 2e-5
- Use warmup: Helps training stability
- Enable mixed precision: Faster training
- Consider PEFT: For large models with limited resources