""" Training module for 63.8M SLM on BookCorpusOpen Can be called directly inside a notebook cell: from config import TrainConfig from train import train model, tok = train(train_cfg=TrainConfig()) """ import os import time import math import json import numpy as np import matplotlib.pyplot as plt import torch from torch.utils.data import DataLoader from config import GPTConfig, TrainConfig from model import GPT from dataset import PackedDataset, build_and_cache_dataset, get_tokenizer from checkpoint import save_checkpoint, load_checkpoint, mirror_checkpoint def get_lr(step: int, warmup_steps: int, max_steps: int, peak_lr: float, min_lr: float) -> float: if step < warmup_steps: return peak_lr * (0.01 + 0.99 * (step / max(1, warmup_steps))) if step > max_steps: return min_lr decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps) coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return min_lr + coeff * (peak_lr - min_lr) @torch.no_grad() def evaluate_loss(model, val_loader, eval_iters: int, device: str): model.eval() losses = [] for i, (x, y) in enumerate(val_loader): if i >= eval_iters: break x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True) with torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.float16): _, loss = model(x, targets=y) losses.append(loss.item()) model.train() return float(np.mean(losses)) if losses else float("nan") def save_training_plots(history: dict, output_dir: str): os.makedirs(output_dir, exist_ok=True) fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.5)) ax1.plot(history["steps"], history["train_loss"], label="Train Loss", color="#2563eb", alpha=0.8) if history["eval_steps"]: ax1.plot(history["eval_steps"], history["val_loss"], label="Val Loss", color="#dc2626", marker="o", markersize=4) ax1.set_title("Training Loss Progression") ax1.set_xlabel("Steps") ax1.set_ylabel("Loss") ax1.grid(True, alpha=0.3) ax1.legend() ax2.plot(history["steps"], history["lr"], label="Learning Rate", color="#16a34a") ax2.set_title("Cosine LR Schedule") ax2.set_xlabel("Steps") ax2.set_ylabel("LR") ax2.grid(True, alpha=0.3) ax2.legend() plt.tight_layout() plt.savefig(os.path.join(output_dir, "training_curves.png"), dpi=150) plt.close() def train(train_cfg: TrainConfig = None, model_cfg: GPTConfig = None, resume_from: str = None): """ Main training function callable from Jupyter / Kaggle / Colab notebooks. """ if train_cfg is None: train_cfg = TrainConfig() if model_cfg is None: model_cfg = GPTConfig() if resume_from is not None: train_cfg.resume_from = resume_from device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") if device == "cuda": print(f"GPU Name: {torch.cuda.get_device_name(0)}") os.makedirs(train_cfg.checkpoint_dir, exist_ok=True) if train_cfg.drive_checkpoint_dir: os.makedirs(train_cfg.drive_checkpoint_dir, exist_ok=True) # 1. Dataset setup train_path, val_path = build_and_cache_dataset( cache_dir=train_cfg.data_dir, max_books=train_cfg.max_books, dataset_name=train_cfg.dataset_name, corpus_dir=train_cfg.corpus_dir, target_tokens=train_cfg.target_tokens, ) train_dataset = PackedDataset(train_path, context_length=model_cfg.context_length) val_dataset = PackedDataset(val_path, context_length=model_cfg.context_length) train_loader = DataLoader(train_dataset, batch_size=train_cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=train_cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True) # 2. Model setup model = GPT(model_cfg).to(device) print(f"[*] Total Model Parameters: {model.get_num_params():,} (~{model.get_num_params()/1e6:.2f}M)") # 3. Optimizer setup with weight decay splitting optimizer = model.configure_optimizers( weight_decay=train_cfg.weight_decay, learning_rate=train_cfg.peak_lr, betas=(train_cfg.beta1, train_cfg.beta2), device_type=device ) scaler = torch.amp.GradScaler("cuda", enabled=(device == "cuda" and train_cfg.use_fp16)) tokenizer = get_tokenizer() start_step = 0 best_val_loss = float("inf") # 4. Resume from Checkpoint if specified if train_cfg.resume_from is not None: start_step, best_val_loss = load_checkpoint( train_cfg.resume_from, model, optimizer=optimizer, scaler=scaler, map_location=device ) if best_val_loss is None: best_val_loss = float("inf") # Save model config with open(os.path.join(train_cfg.checkpoint_dir, "model_config.json"), "w") as f: json.dump(model_cfg.__dict__, f, indent=2) history = {"steps": [], "train_loss": [], "eval_steps": [], "val_loss": [], "lr": []} step = start_step start_time = time.time() train_iter = iter(train_loader) print(f"[*] Training from step {start_step + 1} to {train_cfg.max_steps}...") model.train() while step < train_cfg.max_steps: optimizer.zero_grad(set_to_none=True) accum_loss = 0.0 lr = get_lr(step, train_cfg.warmup_steps, train_cfg.max_steps, train_cfg.peak_lr, train_cfg.min_lr) for param_group in optimizer.param_groups: param_group["lr"] = lr for _ in range(train_cfg.grad_accum): try: x, y = next(train_iter) except StopIteration: train_iter = iter(train_loader) x, y = next(train_iter) x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True) device_type = "cuda" if "cuda" in device else "cpu" with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=train_cfg.use_fp16): _, loss = model(x, targets=y) loss = loss / train_cfg.grad_accum accum_loss += loss.item() if train_cfg.use_fp16 and device == "cuda": scaler.scale(loss).backward() else: loss.backward() if train_cfg.use_fp16 and device == "cuda": scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip) scaler.step(optimizer) scaler.update() else: torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip) optimizer.step() step += 1 # Logging if step % 100 == 0: elapsed = time.time() - start_time print(f"step {step:6d} | train_loss {accum_loss:.4f} | lr {lr:.2e} | {elapsed:.1f}s") history["steps"].append(step) history["train_loss"].append(accum_loss) history["lr"].append(lr) # Evaluation & Checkpointing if step % train_cfg.eval_interval == 0 or step == train_cfg.max_steps: val_loss = evaluate_loss(model, val_loader, train_cfg.eval_iters, device) history["eval_steps"].append(step) history["val_loss"].append(val_loss) print(f"step {step:6d} | val_loss {val_loss:.4f}") # Save periodic checkpoint ckpt_path = os.path.join(train_cfg.checkpoint_dir, f"ckpt_step{step}.pt") save_checkpoint(ckpt_path, model, optimizer, scaler, step=step, val_loss=val_loss, config=model_cfg) print(f"Saved checkpoint at step {step} -> {ckpt_path}") mirrored_ckpt = mirror_checkpoint(ckpt_path, train_cfg.drive_checkpoint_dir) if mirrored_ckpt: print(f"Mirrored checkpoint -> {mirrored_ckpt}") # Save best model if val_loss < best_val_loss: best_val_loss = val_loss best_path = os.path.join(train_cfg.checkpoint_dir, "best_model.pt") save_checkpoint(best_path, model, optimizer, scaler, step=step, val_loss=val_loss, config=model_cfg) mirrored_best = mirror_checkpoint(best_path, train_cfg.drive_checkpoint_dir) if mirrored_best: print(f"Mirrored best checkpoint -> {mirrored_best}") weights_path = os.path.join(train_cfg.checkpoint_dir, "slm_weights.pt") torch.save(model.state_dict(), weights_path) mirrored_weights = mirror_checkpoint(weights_path, train_cfg.drive_checkpoint_dir) if mirrored_weights: print(f"Mirrored weights -> {mirrored_weights}") save_training_plots(history, train_cfg.checkpoint_dir) # Sample generation if step % train_cfg.sample_interval == 0: model.eval() prompt_tokens = tokenizer.encode(train_cfg.sample_prompt) idx = torch.tensor(prompt_tokens, dtype=torch.long, device=device).unsqueeze(0) sample_idx = model.generate(idx, max_new_tokens=60, temperature=0.8, top_k=40) print("-" * 60) print(f"[Sample @ step {step}]: " + tokenizer.decode(sample_idx[0].tolist())) print("-" * 60) model.train() print("Training complete.") return model, tokenizer if __name__ == "__main__": train()