| """
|
| 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
|
|
|
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
| from model import ModelConfig, Retriever500M
|
| from tokenizers import Tokenizer
|
|
|
|
|
| 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")
|
|
|
|
|
|
|
| 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")
|
|
|
|
|
| chunk_size = 1_000_000
|
| 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.
|
| """
|
|
|
| max_start = len(tokens) - seq_len - 1
|
| indices = np.random.randint(0, max_start, size=batch_size)
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
| def setup_optimizer(model: Retriever500M, lr: float, use_8bit: bool = True):
|
| """Set up optimizer — 8-bit AdamW if available, else standard AdamW."""
|
|
|
| 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):
|
|
|
| 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)
|
|
|
|
|
| print("Loading tokenizer...")
|
| tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
|
| vocab_size = tokenizer.get_vocab_size()
|
| print(f"Vocab size: {vocab_size}")
|
|
|
|
|
| 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)
|
|
|
|
|
| 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 = setup_optimizer(model, args.lr, use_8bit=args.use_8bit_adam)
|
|
|
|
|
| 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
|
|
|
| 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
|
|
|
|
|
| effective_batch = args.batch_size * args.grad_accum
|
| max_steps_total = start_step + args.steps
|
| 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()
|
|
|
|
|
| log = {
|
| "config": asdict(config),
|
| "train_args": vars(args),
|
| "total_params": total_params,
|
| "steps": list(prev_log_steps),
|
| }
|
|
|
| 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:
|
|
|
| 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)
|
|
|
|
|
| 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()
|
|
|
|
|
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
|
|
|
|
| optimizer.step()
|
|
|
| avg_loss = total_loss
|
| accum_loss = accum_loss * 0.95 + avg_loss * 0.05
|
|
|
|
|
| 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",
|
| })
|
|
|
|
|
| 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}")
|
|
|
|
|
| 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)
|
|
|
|
|
| if step % 50 == 0 and device.type == "cuda":
|
| torch.cuda.reset_peak_memory_stats()
|
|
|
|
|
| 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()
|
|
|
