""" Base pretraining script for Retriever500M. Memory optimizations for 8GB VRAM (RTX 4070 Laptop): - bf16 mixed precision (autocast) - Gradient checkpointing (recompute activations during backward) - 8-bit AdamW optimizer (bitsandbytes) — halves optimizer state memory - Gradient accumulation (effective batch size > micro batch size) - Short sequence length (512 tokens) for base training - Tied embeddings (shared input/output weight) - Flash Attention via torch SDPA Usage: python src/train.py [--steps N] [--seq_len N] [--batch_size N] [--grad_accum N] """ import argparse import json import os import sys import time from dataclasses import asdict import numpy as np import torch import torch.nn.functional as F from tqdm import tqdm # Add src to path sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from model import ModelConfig, Retriever500M from tokenizers import Tokenizer # ─── Paths ─────────────────────────────────────────────────────────────────── PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) DATA_DIR = os.path.join(PROJECT_DIR, "data") TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer") CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints") LOGS_DIR = os.path.join(PROJECT_DIR, "logs") CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt") CURATED_CORPUS_PATH = os.path.join(DATA_DIR, "corpus_curated.txt") TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json") # ─── Data Loading ──────────────────────────────────────────────────────────── def load_and_tokenize(corpus_path: str, tokenizer: Tokenizer) -> np.ndarray: """Load corpus, tokenize everything, return a flat numpy array of token IDs.""" print(f"Loading corpus from {corpus_path}...") with open(corpus_path, "r", encoding="utf-8") as f: text = f.read() print(f"Corpus size: {len(text) / 1e6:.1f} MB") # Tokenize in chunks to avoid memory issues chunk_size = 1_000_000 # 1MB chunks all_tokens = [] print("Tokenizing corpus...") for i in tqdm(range(0, len(text), chunk_size)): chunk = text[i : i + chunk_size] encoded = tokenizer.encode(chunk) all_tokens.extend(encoded.ids) tokens = np.array(all_tokens, dtype=np.int32) print(f"Total tokens: {len(tokens):,}") return tokens def get_batch( tokens: np.ndarray, batch_size: int, seq_len: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: """Sample a random batch of sequences from the token array. Returns (input_ids, targets) where targets are shifted by 1. """ # Random starting indices max_start = len(tokens) - seq_len - 1 indices = np.random.randint(0, max_start, size=batch_size) # Gather sequences input_ids = np.stack([tokens[i : i + seq_len] for i in indices]) targets = np.stack([tokens[i + 1 : i + seq_len + 1] for i in indices]) input_ids = torch.from_numpy(input_ids).long().to(device) targets = torch.from_numpy(targets).long().to(device) return input_ids, targets # ─── Training ──────────────────────────────────────────────────────────────── def setup_optimizer(model: Retriever500M, lr: float, use_8bit: bool = True): """Set up optimizer — 8-bit AdamW if available, else standard AdamW.""" # Separate embedding params (no weight decay) from rest decay_params = [] no_decay_params = [] for name, param in model.named_parameters(): if not param.requires_grad: continue if "embedding" in name or "norm" in name or "weight" in name and ".weight" not in name: no_decay_params.append(param) else: decay_params.append(param) param_groups = [ {"params": decay_params, "weight_decay": 0.1}, {"params": no_decay_params, "weight_decay": 0.0}, ] if use_8bit: try: import bitsandbytes as bnb optimizer = bnb.optim.AdamW8bit( param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, ) print("Using 8-bit AdamW (bitsandbytes)") return optimizer except Exception as e: print(f"8-bit optimizer unavailable ({e}), falling back to AdamW") optimizer = torch.optim.AdamW( param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, ) print("Using standard AdamW") return optimizer def resume_from_checkpoint( model: Retriever500M, optimizer, resume_path: str, device: torch.device, ) -> tuple[int, float]: """Load model + optimizer state from a checkpoint. Returns (start_step, best_loss) so the training loop can continue. """ print(f"Resuming from {resume_path}") ckpt = torch.load(resume_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model_state_dict"]) print(f" Loaded model weights (step {ckpt.get('step', '?')})") if "optimizer_state_dict" in ckpt: try: optimizer.load_state_dict(ckpt["optimizer_state_dict"]) print(" Loaded optimizer state") except Exception as e: print(f" Could not load optimizer state ({e}); starting fresh optimizer") start_step = int(ckpt.get("step", 0)) best_loss = float(ckpt.get("loss", float("inf"))) print(f" Resuming at step {start_step} (best_loss={best_loss:.4f})") return start_step, best_loss def get_lr(step: int, warmup_steps: int, max_steps: int, max_lr: float, min_lr: float) -> float: """Cosine learning rate schedule with linear warmup.""" if step < warmup_steps: return max_lr * (step + 1) / warmup_steps if step > max_steps: return min_lr decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps) coeff = 0.5 * (1.0 + np.cos(np.pi * decay_ratio)) return min_lr + coeff * (max_lr - min_lr) def train(args): # ─── Setup ─────────────────────────────────────────────────────────────── device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") if device.type == "cuda": print(f"GPU: {torch.cuda.get_device_name(0)}") print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB") os.makedirs(CHECKPOINT_DIR, exist_ok=True) os.makedirs(LOGS_DIR, exist_ok=True) # ─── Tokenizer ─────────────────────────────────────────────────────────── print("Loading tokenizer...") tokenizer = Tokenizer.from_file(TOKENIZER_PATH) vocab_size = tokenizer.get_vocab_size() print(f"Vocab size: {vocab_size}") # ─── Data ──────────────────────────────────────────────────────────────── if args.corpus == "curated": corpus_path = CURATED_CORPUS_PATH if not os.path.exists(corpus_path): raise FileNotFoundError(f"Curated corpus not found: {corpus_path}. Run src/curate.py first.") print(f"Using CURATED corpus: {corpus_path}") elif args.corpus == "default": corpus_path = CORPUS_PATH else: corpus_path = args.corpus tokens = load_and_tokenize(corpus_path, tokenizer) # ─── Model ─────────────────────────────────────────────────────────────── config = ModelConfig( vocab_size=vocab_size, d_model=1_280, n_layers=23, n_heads=20, d_ff=3_456, max_seq_len=args.seq_len, dropout=0.0, tie_embeddings=True, ) model = Retriever500M(config).to(device) total_params = model.count_parameters() print(f"Model parameters: {total_params:,} ({total_params / 1e6:.1f}M)") # ─── Optimizer ─────────────────────────────────────────────────────────── optimizer = setup_optimizer(model, args.lr, use_8bit=args.use_8bit_adam) # ─── Resume from checkpoint ────────────────────────────────────────────── start_step = 0 best_loss = float("inf") prev_log_steps = [] if args.resume: resume_path = args.resume_path or os.path.join(CHECKPOINT_DIR, "latest.pt") if not os.path.exists(resume_path): raise FileNotFoundError(f"Cannot resume: {resume_path} does not exist") start_step, best_loss = resume_from_checkpoint(model, optimizer, resume_path, device) accum_loss = best_loss # continue EMA from saved loss # Load previous log entries so we append rather than overwrite history prev_log_path = os.path.join(LOGS_DIR, "training_log.json") if os.path.exists(prev_log_path): try: with open(prev_log_path, "r") as f: prev_log = json.load(f) prev_log_steps = prev_log.get("steps", []) print(f" Loaded {len(prev_log_steps)} previous log entries") except Exception as e: print(f" Could not load previous log ({e})") else: accum_loss = 0.0 # ─── Training loop ─────────────────────────────────────────────────────── effective_batch = args.batch_size * args.grad_accum max_steps_total = start_step + args.steps # for LR schedule continuity print(f"\nTraining configuration:") print(f" Micro batch size: {args.batch_size}") print(f" Gradient accum: {args.grad_accum}") print(f" Effective batch: {effective_batch}") print(f" Sequence length: {args.seq_len}") print(f" Learning rate: {args.lr}") print(f" Steps this run: {args.steps}") print(f" Start step: {start_step}") print(f" Target step: {max_steps_total}") print(f" Warmup steps: {args.warmup}") print(f" Grad checkpointing: {args.grad_checkpoint}") print() # Training log log = { "config": asdict(config), "train_args": vars(args), "total_params": total_params, "steps": list(prev_log_steps), # carry over previous entries } model.train() step = start_step start_time = time.time() pbar = tqdm(range(start_step, max_steps_total), desc="Training", initial=start_step, total=max_steps_total) for step in pbar: # Learning rate schedule (uses absolute step for continuity) lr = get_lr(step, args.warmup, max_steps_total, args.lr, args.lr * 0.1) for pg in optimizer.param_groups: pg["lr"] = lr optimizer.zero_grad(set_to_none=True) # Gradient accumulation total_loss = 0.0 for micro_step in range(args.grad_accum): input_ids, targets = get_batch(tokens, args.batch_size, args.seq_len, device) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): out = model(input_ids, targets=targets, use_checkpoint=args.grad_checkpoint) loss = out["loss"] / args.grad_accum loss.backward() total_loss += loss.item() # Gradient clipping torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Optimizer step optimizer.step() avg_loss = total_loss # already divided by grad_accum accum_loss = accum_loss * 0.95 + avg_loss * 0.05 # EMA # Logging if step % args.log_every == 0 or step == max_steps_total - 1: elapsed = time.time() - start_time steps_this_run = step - start_step + 1 steps_per_sec = steps_this_run / elapsed vram_used = torch.cuda.max_memory_allocated() / 1e9 if device.type == "cuda" else 0 log_entry = { "step": step, "loss": avg_loss, "ema_loss": accum_loss, "lr": lr, "elapsed_s": elapsed, "steps_per_sec": steps_per_sec, "vram_gb": vram_used, } log["steps"].append(log_entry) pbar.set_postfix({ "loss": f"{avg_loss:.4f}", "ema": f"{accum_loss:.4f}", "lr": f"{lr:.2e}", "vram": f"{vram_used:.1f}G", }) # Save checkpoint if (step + 1) % args.save_every == 0 or step == max_steps_total - 1: ckpt_path = os.path.join(CHECKPOINT_DIR, f"model_step_{step + 1}.pt") torch.save({ "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "config": asdict(config), "step": step + 1, "loss": accum_loss, }, ckpt_path) print(f"\n Saved checkpoint: {ckpt_path}") # Also save latest latest_path = os.path.join(CHECKPOINT_DIR, "latest.pt") torch.save({ "model_state_dict": model.state_dict(), "config": asdict(config), "step": step + 1, "loss": accum_loss, }, latest_path) if accum_loss < best_loss: best_loss = accum_loss best_path = os.path.join(CHECKPOINT_DIR, "best.pt") torch.save({ "model_state_dict": model.state_dict(), "config": asdict(config), "step": step + 1, "loss": accum_loss, }, best_path) # Reset peak memory stats periodically if step % 50 == 0 and device.type == "cuda": torch.cuda.reset_peak_memory_stats() # ─── Save training log ─────────────────────────────────────────────────── log_path = os.path.join(LOGS_DIR, "training_log.json") with open(log_path, "w") as f: json.dump(log, f, indent=2) print(f"\nTraining log saved to {log_path}") total_time = time.time() - start_time print(f"\nTraining complete!") print(f" Total time: {total_time:.1f}s ({total_time/60:.1f} min)") print(f" Final EMA loss: {accum_loss:.4f}") print(f" Best loss: {best_loss:.4f}") print(f" Steps/sec: {args.steps / total_time:.2f}") return model, log def main(): parser = argparse.ArgumentParser(description="Train Retriever500M base model") parser.add_argument("--steps", type=int, default=2000, help="Total training steps") parser.add_argument("--batch_size", type=int, default=4, help="Micro batch size") parser.add_argument("--grad_accum", type=int, default=8, help="Gradient accumulation steps") parser.add_argument("--seq_len", type=int, default=512, help="Sequence length") parser.add_argument("--lr", type=float, default=3e-4, help="Peak learning rate") parser.add_argument("--warmup", type=int, default=100, help="Warmup steps") parser.add_argument("--save_every", type=int, default=500, help="Save checkpoint every N steps") parser.add_argument("--log_every", type=int, default=10, help="Log every N steps") parser.add_argument("--grad_checkpoint", action="store_true", default=True, help="Use gradient checkpointing") parser.add_argument("--no_grad_checkpoint", dest="grad_checkpoint", action="store_false") parser.add_argument("--use_8bit_adam", action="store_true", default=True, help="Use 8-bit AdamW") parser.add_argument("--no_8bit_adam", dest="use_8bit_adam", action="store_false") parser.add_argument("--resume", action="store_true", help="Resume training from latest checkpoint") parser.add_argument("--resume_path", type=str, default=None, help="Specific checkpoint to resume from (default: checkpoints/latest.pt)") parser.add_argument("--corpus", type=str, default="default", help="Corpus to use: 'default', 'curated', or a path") args = parser.parse_args() train(args) if __name__ == "__main__": main()