14. Training Loop

Chapter 14 of 24 · 20 min

KEY INSIGHT

The training loop is the execution engine that translates hyperparameter choices into learned weights. A reliable implementation handles data loading, gradient computation, optimizer updates, and logging with clean separation of concerns. The core sequence: load a batch â†’ compute loss â†’ backpropagate â†’ update weights â†’ repeat. Each iteration must preserve numerical stability and avoid gradient accumulation errors that produce silent failures. ```python from torch.utils.data import DataLoader from transformers import get_linear_schedule_with_warmup from torch.optim import AdamW import torch.nn.functional as F def train_epoch(model, dataloader, optimizer, scheduler, device, gradient_accumulation_steps=4): model.train() total_loss = 0 optimizer.zero_grad() for step, batch in enumerate(dataloader): input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) # Scale loss for gradient accumulation loss = outputs.loss / gradient_accumulation_steps loss.backward() if (step + 1) % gradient_accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss += outputs.loss.detach().cpu().item() return total_loss / len(dataloader) ``` Gradient clipping prevents exploding gradientsâ€”a common failure mode with LoRA adapters on certain data distributions. The threshold of 1.0 works for most transformer architectures, but verify with your specific use case. **Optimizer selection**: AdamW with weight decay handles the regularization correctly. For LoRA, bias and layer norm parameters typically get zero weight decay. QLoRA introduces the gradient scaling factor that must be accounted for in loss calculations. ```python # Correct optimizer setup for LoRA optimizer = AdamW( [{"params": model.lora_parameters, "weight_decay": 0.01}, {"params": model.other_parameters, "weight_decay": 0.1}], lr=2e-4 ) ```

EXERCISE

: Build a Complete Training Loop

Extend the training loop to include checkpoint saving and early stopping:

def train_with_validation(
    model, train_loader, val_loader, optimizer, scheduler,
    device, epochs=3, checkpoint_dir="./checkpoints"
):
    best_val_loss = float("inf")
    patience = 2
    patience_counter = 0
    
    for epoch in range(epochs):
        train_loss = train_epoch(model, train_loader, optimizer, scheduler, device)
        val_loss = evaluate(model, val_loader, device)
        
        print(f"Epoch {epoch}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}")
        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            save_checkpoint(model, optimizer, epoch, checkpoint_dir)
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print("Early stopping triggered")
                break
    
    return best_val_loss

Implement the evaluate() and save_checkpoint() functions to complete this pattern.