""" Step-by-step training script for nano GPT. What this script does: 1. Load the preprocessed data (train / val tokens) 2. Build the GPT model with our config 3. Define a batching function that grabs random chunks of text 4. Set up an AdamW optimizer with cosine learning-rate schedule 5. Train loop: sample batch -> forward -> loss -> backward -> step 6. Periodically evaluate on validation set and print metrics 7. Save the best model checkpoint 8. Generate a sample from the model after training """ import os import math import time import torch # Import our model from model import GPT, GPTConfig # --------------------------------------------------------------------------- # 1. Hyperparameters & Config # --------------------------------------------------------------------------- # Feel free to tweak these! For a tutorial we keep things small and fast. BATCH_SIZE = 64 # how many sequences to process in parallel BLOCK_SIZE = 256 # max context length for each sequence (must match model!) MAX_ITERS = 5000 # total training steps LEARNING_RATE = 1e-3 # starting learning rate WARMUP_ITERS = 200 # linear warmup steps (gradually increase LR) LR_DECAY_ITERS = 5000 # when to reach min LR (usually = MAX_ITERS) MIN_LR = 1e-4 # minimum learning rate at end of cosine schedule EVAL_INTERVAL = 500 # how often to run validation EVAL_ITERS = 200 # how many val batches to average for a stable loss estimate GRAD_CLIP = 1.0 # max gradient norm (prevents exploding gradients) # Device selection device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # --------------------------------------------------------------------------- # 2. Load Data # --------------------------------------------------------------------------- # We load the dictionary saved by prepare.py data_path = os.path.join(os.path.dirname(__file__), "data.pt") data = torch.load(data_path, weights_only=False) train_data = data["train"] val_data = data["val"] vocab_size = data["vocab_size"] chars = data["chars"] stoi = data["stoi"] itos = data["itos"] print(f"Vocab size : {vocab_size}") print(f"Train tokens: {len(train_data):,}") print(f"Val tokens : {len(val_data):,}") # --------------------------------------------------------------------------- # 3. Batch sampling # --------------------------------------------------------------------------- # For language modeling, each training example is a random contiguous chunk # of text. The input is tokens[0:T-1], the target is tokens[1:T]. def get_batch(split: str): """Sample a single batch from train or val data.""" data_split = train_data if split == "train" else val_data ix = torch.randint(len(data_split) - BLOCK_SIZE, (BATCH_SIZE,)) x = torch.stack([data_split[i : i + BLOCK_SIZE] for i in ix]) y = torch.stack([data_split[i + 1 : i + BLOCK_SIZE + 1] for i in ix]) x, y = x.to(device), y.to(device) return x, y # --------------------------------------------------------------------------- # 4. Helper: Learning-rate schedule (cosine with linear warmup) # --------------------------------------------------------------------------- # Warmup is crucial for transformers — it prevents early spikes in loss # caused by large gradients when the model is still random. def get_lr(iteration: int) -> float: if iteration < WARMUP_ITERS: # Linear warmup return LEARNING_RATE * (iteration + 1) / WARMUP_ITERS if iteration > LR_DECAY_ITERS: return MIN_LR # Cosine decay after warmup decay_ratio = (iteration - WARMUP_ITERS) / (LR_DECAY_ITERS - WARMUP_ITERS) coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return MIN_LR + coeff * (LEARNING_RATE - MIN_LR) # --------------------------------------------------------------------------- # 5. Model Setup # --------------------------------------------------------------------------- # We match block_size to our training hyperparameter above. # For tiny Shakespeare, even a 4-layer model can learn structure. config = GPTConfig( block_size=BLOCK_SIZE, vocab_size=vocab_size, n_layer=6, # deeper = more capacity to learn patterns n_head=6, n_embd=384, dropout=0.0, ) model = GPT(config) model.to(device) # Count parameters param_count = sum(p.numel() for p in model.parameters()) print(f"\nModel config: {config}") print(f"Total parameters: {param_count / 1e6:.2f} M") # --------------------------------------------------------------------------- # 6. Optimizer # --------------------------------------------------------------------------- # We separate parameters that should get weight decay (2D weights) # from those that should not (1D biases, LayerNorm scales). # This is standard practice and slightly improves training. decay_params = [] no_decay_params = [] for name, param in model.named_parameters(): if param.dim() >= 2: decay_params.append(param) else: no_decay_params.append(param) optim_groups = [ {"params": decay_params, "weight_decay": 0.1}, {"params": no_decay_params, "weight_decay": 0.0}, ] optimizer = torch.optim.AdamW(optim_groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8) # --------------------------------------------------------------------------- # 7. Evaluation helper # --------------------------------------------------------------------------- # We average the loss over multiple validation batches for a stable estimate. # torch.no_grad() disables gradient computation -> faster and less memory. @torch.no_grad() def estimate_loss(): out = {} model.eval() # set model to evaluation mode for split in ["train", "val"]: losses = torch.zeros(EVAL_ITERS) for k in range(EVAL_ITERS): xb, yb = get_batch(split) _, loss = model(xb, yb) losses[k] = loss.item() out[split] = losses.mean() model.train() # set model back to training mode return out # --------------------------------------------------------------------------- # 8. Training Loop # --------------------------------------------------------------------------- print("\n" + "=" * 60) print("Starting training...") print("=" * 60) best_val_loss = float("inf") start_time = time.time() for iter_num in range(MAX_ITERS): # --- Learning rate scheduling --- lr = get_lr(iter_num) for param_group in optimizer.param_groups: param_group["lr"] = lr # --- Periodic evaluation --- if iter_num % EVAL_INTERVAL == 0 or iter_num == MAX_ITERS - 1: losses = estimate_loss() elapsed = time.time() - start_time print( f"step {iter_num:5d} | " f"train loss {losses['train']:.4f} | " f"val loss {losses['val']:.4f} | " f"lr {lr:.2e} | " f"time {elapsed:.1f}s" ) # Save the best checkpoint if losses["val"] < best_val_loss: best_val_loss = losses["val"] checkpoint_path = os.path.join(os.path.dirname(__file__), "best.pt") torch.save({ "model_state_dict": model.state_dict(), "config": config, "vocab_size": vocab_size, "chars": chars, "stoi": stoi, "itos": itos, }, checkpoint_path) print(f" -> Saved new best model (val_loss={best_val_loss:.4f})") # --- Training step --- xb, yb = get_batch("train") # Forward logits, loss = model(xb, yb) # Backward optimizer.zero_grad(set_to_none=True) loss.backward() # Gradient clipping (prevents exploding gradients) torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) # Optimizer step optimizer.step() # --------------------------------------------------------------------------- # 9. Final evaluation # --------------------------------------------------------------------------- losses = estimate_loss() print(f"\nFinal -> train loss {losses['train']:.4f} | val loss {losses['val']:.4f}") # --------------------------------------------------------------------------- # 10. Generate text from the trained model # --------------------------------------------------------------------------- print("\n" + "=" * 60) print("Generating sample text...") print("=" * 60) model.eval() # Start from a newline character (index of '\n' in our vocab) start_token = stoi["\n"] context = torch.zeros((1, 1), dtype=torch.long, device=device) context[0, 0] = start_token with torch.no_grad(): generated = model.generate(context, max_new_tokens=500, temperature=1.0, top_k=40) # Rebuild decode function from saved mappings decode = lambda l: "".join([itos[i] for i in l]) # Decode to text print("\n--- Generated text ---\n") print(decode(generated[0].tolist())) print("\n--- End ---") print("\nTraining complete! Best checkpoint saved to: best.pt")