Fine-tuning an LLM with plain PyTorch¶

import os
import math
import random
from dataclasses import dataclass

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    get_linear_schedule_with_warmup,
)
from datasets import load_dataset

import pandas as pd

torch.__version__

@dataclass
class Config:
    
    # Data
    max_length: int = 256        # max tokens per example

    # Optimization
    batch_size: int = 1          # reduced for memory efficiency
    num_epochs: int = 1
    learning_rate: float = 5e-6
    weight_decay: float = 0.01

    warmup_ratio: float = 0.1
    gradient_accumulation_steps: int = 16  # increased to compensate for smaller batch

    seed: int = 42
    device: str = "cuda" if torch.cuda.is_available() else "cpu"

cfg = Config()
cfg

def set_seed(seed: int):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(cfg.seed)

dataset = load_dataset("timdettmers/openassistant-guanaco")

train_data = dataset["train"].to_list()
val_data = dataset["test"].to_list()

train_data = train_data[0:500]
val_data = val_data[0:100]

train_df = pd.DataFrame(train_data)
val_df   = pd.DataFrame(val_data)

print("Columns:", train_df.columns.tolist())
print("Train size:", len(train_df))
print("Validation size:", len(val_df))

train_df.head()

len(train_data)

HF_LLM_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(HF_LLM_MODEL)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("Pad token:", tokenizer.pad_token, "ID:", tokenizer.pad_token_id)

def build_prompt_for_inference(user_instruction: str) -> str:
    """
    Build a Guanaco-style prompt for a NEW instruction at inference time.
    The dataset format looks like:
    "### Human: ...### Assistant: ..."
    """
    return f"### Human: {user_instruction}### Assistant:"

class SupervisedTextDataset(Dataset):
    """
    Each row in train_df/val_df has a 'text' field like:

        "### Human: ...### Assistant: ..."

    For supervised fine-tuning of a causal LM, we feed in the full text and
    ask the model to learn to predict the next token at every position.
    """
    def __init__(self, dataframe, tokenizer, max_length: int = 256):
        self.df = dataframe.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        text = row["text"]

        enc = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )

        input_ids = enc["input_ids"].squeeze(0)
        attention_mask = enc["attention_mask"].squeeze(0)

        # For causal LM SFT: labels = input_ids (shift is handled internally)
        labels = input_ids.clone()
        labels[attention_mask == 0] = -100  # ignore padding in loss

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }

train_dataset = SupervisedTextDataset(train_df, tokenizer, max_length=cfg.max_length)
val_dataset   = SupervisedTextDataset(val_df, tokenizer, max_length=cfg.max_length)

train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True)
val_loader   = DataLoader(val_dataset, batch_size=cfg.batch_size)

len(train_loader), len(val_loader)

batch = next(iter(train_loader))
{k: v.shape for k, v in batch.items()}

model = AutoModelForCausalLM.from_pretrained(
    HF_LLM_MODEL,
    dtype=torch.bfloat16,  # Use bfloat16 for memory efficiency
)

model.to(cfg.device)

# Enable gradient checkpointing for memory efficiency
model.gradient_checkpointing_enable()

n_params = sum(p.numel() for p in model.parameters())
print(f"Number of parameters: {n_params / 1e6:.1f}M")
print(f"Model dtype: {next(model.parameters()).dtype}")

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=cfg.learning_rate,
    weight_decay=cfg.weight_decay,
)

steps_per_epoch = math.ceil(len(train_loader))
total_steps = (steps_per_epoch * cfg.num_epochs) // cfg.gradient_accumulation_steps

warmup_steps = int(cfg.warmup_ratio * total_steps)

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps,
)

total_steps, warmup_steps

model.train()

single_batch = next(iter(train_loader))
single_batch = {k: v.to(cfg.device) for k, v in single_batch.items()}

# Forward pass
out = model(
    input_ids=single_batch["input_ids"],
    attention_mask=single_batch["attention_mask"],
    labels=single_batch["labels"],
)

loss = out.loss
print("Single batch loss:", loss.item())

# Backward pass
loss.backward()

# Parameter update
optimizer.step()
scheduler.step()
optimizer.zero_grad()

def evaluate(model, dataloader):
    model.eval()
    losses = []

    with torch.no_grad():
        for batch in dataloader:
            batch = {k: v.to(cfg.device) for k, v in batch.items()}
            out = model(**batch)
            losses.append(out.loss.item())

    model.train()
    return sum(losses) / len(losses)

from tqdm.auto import tqdm

global_step = 0
best_val_loss = float("inf")
save_dir = "ft_model"

os.makedirs(save_dir, exist_ok=True)

for epoch in range(cfg.num_epochs):
    model.train()
    running_loss = 0.0

    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{cfg.num_epochs}")

    for step, batch in enumerate(progress_bar):
        batch = {k: v.to(cfg.device) for k, v in batch.items()}

        out = model(**batch)
        loss = out.loss / cfg.gradient_accumulation_steps
        loss.backward()

        running_loss += loss.item()

        if (step + 1) % cfg.gradient_accumulation_steps == 0:
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            global_step += 1

            avg_loss = running_loss / cfg.gradient_accumulation_steps
            running_loss = 0.0

            progress_bar.set_postfix(
                train_loss=f"{avg_loss:.4f}",
                lr=f"{scheduler.get_last_lr()[0]:.2e}",
            )

    # Validation at the end of the epoch
    val_loss = evaluate(model, val_loader)
    print(f"\nValidation loss after epoch {epoch+1}: {val_loss:.4f}")

    # Simple checkpointing: keep the best model
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        model.save_pretrained(save_dir)
        tokenizer.save_pretrained(save_dir)
        print(f"→ New best model saved to {save_dir}\n")
    else:
        print("No improvement, keeping previous best model.\n")

# Fine-tuned model – load from checkpoint
ft_model = AutoModelForCausalLM.from_pretrained(
    save_dir,
    dtype=torch.bfloat16,
).to(cfg.device)
ft_model.eval()

def generate_response(model, instruction: str, max_new_tokens: int = 128):
    """
    Generate a reply from the model given a human instruction.
    We create a Guanaco-style prompt:
        "### Human: ...### Assistant:"
    and let the model continue.
    """
    prompt_text = build_prompt_for_inference(instruction)

    inputs = tokenizer(prompt_text, return_tensors="pt").to(cfg.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,  # greedy for clarity
            pad_token_id=tokenizer.pad_token_id,
        )

    text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return text

example_text = val_df.iloc[11]["text"]
example_text

base_model = AutoModelForCausalLM.from_pretrained(
    HF_LLM_MODEL,
    dtype=torch.bfloat16,
).to(cfg.device)

# crude split to get the human message
if "### Human:" in example_text and "### Assistant:" in example_text:
    human_part = example_text.split("### Human:")[1].split("### Assistant:")[0].strip()
    assistant_part = example_text.split("### Assistant:")[1].strip()
else:
    human_part = example_text
    assistant_part = ""

print("### HUMAN (PROMPT) ###")
print(human_part)

print("\n### GROUND TRUTH ASSISTANT ###")
print(assistant_part)

print("\n### BASE MODEL ###")
print(generate_response(base_model, human_part))

print("\n### FINE-TUNED MODEL ###")
print(generate_response(ft_model, human_part))

my_instruction = "WRITE YOUR INSTRUCTION HERE"

print("\nFINE-TUNED MODEL:\n")
print(generate_response(ft_model, my_instruction))

use_amp = torch.cuda.is_available()  # only useful on GPU
print("Using AMP (mixed precision)?", use_amp)

scaler = torch.amp.GradScaler('cuda', enabled=use_amp)

def train_one_epoch_amp(model, dataloader, optimizer, scheduler, scaler, epoch_idx: int):
    model.train()
    running_loss = 0.0
    progress_bar = tqdm(dataloader, desc=f"[AMP] Epoch {epoch_idx+1}/{cfg.num_epochs}")

    for step, batch in enumerate(progress_bar):
        batch = {k: v.to(cfg.device) for k, v in batch.items()}

        with torch.amp.autocast('cuda', enabled=use_amp):
            out = model(**batch)
            loss = out.loss / cfg.gradient_accumulation_steps

        scaler.scale(loss).backward()
        running_loss += loss.item()

        if (step + 1) % cfg.gradient_accumulation_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            scheduler.step()

            avg_loss = running_loss / cfg.gradient_accumulation_steps
            running_loss = 0.0

            progress_bar.set_postfix(
                train_loss=f"{avg_loss:.4f}",
                lr=f"{scheduler.get_last_lr()[0]:.2e}",
            )