Training a Transformer

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup
from datasets import load_dataset
from tqdm.auto import tqdm
import numpy as np

# Setup device - GPU if available, CPU otherwise
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(f"Using device: {device}")

if device.type == "cuda":
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

ALPACA_TEMPLATE = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Response:
{response}"""

class SFTDataset(Dataset):
    """Dataset for supervised fine-tuning.
    
    This class does three critical things:
    1. Formats examples with the Alpaca template
    2. Tokenizes text into numbers the model understands
    3. Creates labels that mask the prompt (we only train on responses!)
    """
    
    def __init__(self, data, tokenizer, max_length=512):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        """How many examples do we have?"""
        return len(self.data)
    
    def __getitem__(self, idx):
        """Get one training example, fully prepped and ready."""
        item = self.data[idx]
        
        # Step 1: Format with our template
        formatted = ALPACA_TEMPLATE.format(
            instruction=item['instruction'],
            response=item['output']
        )
        
        # Step 2: Figure out where the response starts
        # We need this because we DON'T want to train on the instruction part
        # Only the response should contribute to the loss
        prompt = ALPACA_TEMPLATE.format(
            instruction=item['instruction'],
            response=''  # Empty response to find where it would start
        )
        
        # Step 3: Tokenize the full text
        full_tokens = self.tokenizer(
            formatted,
            max_length=self.max_length,
            truncation=True,  # Cut off if too long
            padding='max_length',  # Pad to consistent length
            return_tensors='pt'
        )
        
        # Step 4: Tokenize just the prompt to find response boundary
        prompt_tokens = self.tokenizer(
            prompt,
            max_length=self.max_length,
            truncation=True,
            return_tensors='pt'
        )
        
        response_start = prompt_tokens['input_ids'].shape[1]
        
        # Step 5: Create labels with masking
        # Here's the key insight: we copy the input_ids but mask the prompt
        # The -100 value tells PyTorch "don't calculate loss for these tokens"
        labels = full_tokens['input_ids'].clone().squeeze(0)
        labels[:response_start] = -100  # Mask the instruction
        
        # Also mask padding tokens (we don't want to train on padding!)
        labels[labels == self.tokenizer.pad_token_id] = -100
        
        # Return everything the training loop needs
        return {
            'input_ids': full_tokens['input_ids'].squeeze(0),
            'attention_mask': full_tokens['attention_mask'].squeeze(0),
            'labels': labels
        }

# Let's test it with a sample
print("Dataset class defined successfully!")
print("\nKey insight: We mask the instruction with -100 so the model only learns")
print("to predict the response. This is supervised fine-tuning in action.")

from dataclasses import dataclass

@dataclass
class SFTConfig:
    """Configuration for SFT training.
    
    Let's break down what each setting does:
    """
    
    # Which model to start from
    model_name: str = "gpt2"
    
    # Maximum sequence length (longer = more memory)
    max_length: int = 512
    
    # Batch size: how many examples to process at once
    # Bigger = faster but needs more memory
    batch_size: int = 4
    
    # Learning rate: how big of steps to take when updating weights
    # This is the single most important hyperparameter
    learning_rate: float = 2e-4  # 0.0002 - small steps, safe and steady
    
    # Epochs: how many times to loop through the entire dataset
    # One epoch = seeing every training example once
    num_epochs: int = 3
    
    # Warmup steps: gradually increase learning rate at the start
    # Prevents the model from making crazy updates early on
    warmup_steps: int = 100
    
    # Gradient accumulation: simulate bigger batches without more memory
    # Process this many batches before updating weights
    # Effective batch size = batch_size * gradient_accumulation_steps
    gradient_accumulation_steps: int = 4
    
    # Gradient clipping: prevent exploding gradients
    # If gradients get bigger than this, scale them down
    max_grad_norm: float = 1.0
    
    # How often to log progress
    logging_steps: int = 10
    
    # How often to check validation loss
    eval_steps: int = 100
    
    # How often to save checkpoints
    save_steps: int = 500
    
    # Where to save the model
    output_dir: str = "./sft_output"

config = SFTConfig()

print("Training configuration:")
print("=" * 60)
for k, v in vars(config).items():
    print(f"  {k:.<35} {v}")
print("=" * 60)
print(f"\nEffective batch size: {config.batch_size * config.gradient_accumulation_steps}")
print("(That's how many examples we process before each weight update)")

def train_sft(model, tokenizer, train_dataset, eval_dataset, config):
    """Complete SFT training loop.
    
    This function orchestrates the entire training process.
    """
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=0
    )
    
    eval_loader = DataLoader(
        eval_dataset,
        batch_size=config.batch_size,
        shuffle=False
    )
    
    # Setup optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        weight_decay=0.01
    )
    
    # Learning rate scheduler
    total_steps = len(train_loader) * config.num_epochs // config.gradient_accumulation_steps
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=config.warmup_steps,
        num_training_steps=total_steps
    )
    
    # Training state
    model.train()
    global_step = 0
    best_eval_loss = float('inf')
    
    for epoch in range(config.num_epochs):
        epoch_loss = 0
        
        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config.num_epochs}")
        
        for step, batch in enumerate(progress_bar):
            batch = {k: v.to(device) for k, v in batch.items()}
            
            # Forward pass
            outputs = model(
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['labels']
            )
            
            loss = outputs.loss / config.gradient_accumulation_steps
            
            # Backward pass
            loss.backward()
            
            epoch_loss += loss.item() * config.gradient_accumulation_steps
            
            # Update weights
            if (step + 1) % config.gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    config.max_grad_norm
                )
                
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                
                global_step += 1
                
                if global_step % config.logging_steps == 0:
                    avg_loss = epoch_loss / (step + 1)
                    progress_bar.set_postfix({
                        'loss': f'{avg_loss:.4f}',
                        'lr': f'{scheduler.get_last_lr()[0]:.2e}'
                    })
        
        # End of epoch - evaluate
        eval_loss = evaluate(model, eval_loader, device)
        print(f"\nEpoch {epoch+1} - Train Loss: {epoch_loss/len(train_loader):.4f}, Eval Loss: {eval_loss:.4f}")
        
        # Save best model
        if eval_loss < best_eval_loss:
            best_eval_loss = eval_loss
            model.save_pretrained(f"{config.output_dir}/best")
            tokenizer.save_pretrained(f"{config.output_dir}/best")
            print(f"Saved best model with eval loss: {eval_loss:.4f}")
    
    return model


def evaluate(model, eval_loader, device):
    """Evaluate model on validation set.
    
    This tells us how well the model generalizes to new data.
    If training loss goes down but eval loss goes up, that's overfitting.
    """
    model.eval()
    total_loss = 0
    
    with torch.no_grad():
        for batch in eval_loader:
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['labels']
            )
            total_loss += outputs.loss.item()
    
    model.train()
    return total_loss / len(eval_loader)

print("Training loop defined.")
print("\nKey concepts:")
print("  • Forward pass: Model makes predictions")
print("  • Backward pass: Calculate gradients")
print("  • Optimizer step: Update the weights")
print("  • Gradient accumulation: Simulate bigger batches")
print("  • Validation: Check if we're learning or just memorizing")

# Load model and tokenizer
model_name = "gpt2"
print(f"Loading {model_name}...")

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 doesn't have a pad token by default, so we'll use the EOS token
# (This is a common trick - pad tokens are just for batching anyway)
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id

# Move model to GPU
model.to(device)
print(f"✓ Model loaded and moved to {device}")

# Load dataset
print("\nLoading Alpaca dataset...")
raw_data = load_dataset("yahma/alpaca-cleaned", split="train")

# Use a small subset for this demo
# (Training on 1000 examples instead of 52,000 - your GPU will thank me)
raw_data = raw_data.select(range(1000))
print(f"✓ Loaded {len(raw_data)} examples")

# Split into train/eval (90/10 split)
train_size = int(0.9 * len(raw_data))
train_data = raw_data.select(range(train_size))
eval_data = raw_data.select(range(train_size, len(raw_data)))

# Create datasets using our SFTDataset class
# We use shorter sequences (256) to save memory
train_dataset = SFTDataset(train_data, tokenizer, max_length=256)
eval_dataset = SFTDataset(eval_data, tokenizer, max_length=256)

print(f"✓ Created datasets:")
print(f"  - Training: {len(train_dataset)} examples")
print(f"  - Validation: {len(eval_dataset)} examples")

# Let's peek at one example to make sure everything looks right
sample = train_dataset[0]
print(f"\nSample training example:")
print(f"  - Input IDs shape: {sample['input_ids'].shape}")
print(f"  - Labels shape: {sample['labels'].shape}")
print(f"  - Number of masked tokens: {(sample['labels'] == -100).sum().item()}")
print(f"  - Number of trainable tokens: {(sample['labels'] != -100).sum().item()}")

# Adjust config for a quick demo run
config.num_epochs = 1
config.max_length = 256

print("Starting training...")
print("=" * 60)

model = train_sft(model, tokenizer, train_dataset, eval_dataset, config)

print("=" * 60)
print("Training complete.")
print("\nWhat happened:")
print("  • The model saw 900 training examples")
print("  • It adjusted its weights to minimize prediction errors")
print("  • We validated on 100 held-out examples to check generalization")
print("  • The best model was saved to disk")

def generate_response(model, tokenizer, instruction, max_new_tokens=100):
    """Generate a response for an instruction."""
    prompt = ALPACA_TEMPLATE.format(instruction=instruction, response='')
    
    inputs = tokenizer(prompt, return_tensors='pt').to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response.split("### Response:")[-1].strip()
    
    return response

# Test with a few different instructions
test_instructions = [
    "Explain what machine learning is in simple terms.",
    "Write a haiku about programming.",
    "What are the three branches of the US government?"
]

print("Testing the fine-tuned model...")
print("=" * 60)

for instruction in test_instructions:
    print(f"\nInstruction: {instruction}")
    print("-" * 60)
    response = generate_response(model, tokenizer, instruction)
    print(f"Response: {response}")
    print("=" * 60)

print("\nThe model should:")
print("  ✓ Follow the instruction format")
print("  ✓ Stay on topic")
print("  ✓ Generate coherent responses")
print("\nRemember: We only trained for one epoch on 900 examples.")
print("With more data and training time, it would get much better.")