"""Manual Tiny GPT training: bounded batches, token-weighted loss and checkpoints.""" import argparse import json import math import random import signal import time import tomllib from contextlib import nullcontext from pathlib import Path import numpy as np import torch from torch.utils.tensorboard import SummaryWriter from vimeml.training.data import SentenceWindowDataset, file_sha, make_loader, prepare_indexes, write_json from vimeml.training.model import GPTConfig, TinyGPT ROOT = Path(__file__).resolve().parents[3] def learning_rate(step, settings): """step is one-based optimizer update; linear warmup followed by cosine.""" warmup, maximum = settings["warmup_steps"], settings["max_steps"] high, low = settings["learning_rate"], settings["min_learning_rate"] if warmup and step <= warmup: return high * step / warmup fraction = min(1.0, max(0.0, (step - warmup) / max(1, maximum - warmup))) return low + (high - low) * (1 + math.cos(math.pi * fraction)) / 2 def worker_init(worker_id): # Ctrl+C goes to the parent, which checkpoints after its current update. signal.signal(signal.SIGINT, signal.SIG_IGN) torch.set_num_threads(1) def amp_context(device, precision): if precision == "fp32": return nullcontext() return torch.autocast(device_type=device.type, dtype=torch.bfloat16 if precision == "bf16" else torch.float16) def resolve_precision(device, requested): if requested == "auto": return "bf16" if device.type == "cuda" and torch.cuda.is_bf16_supported() else ( "fp16" if device.type == "cuda" else "fp32") if requested not in {"fp32", "bf16", "fp16"}: raise ValueError("precision must be auto/fp32/bf16/fp16.") if device.type != "cuda" and requested != "fp32": raise ValueError("CPU checks use fp32; choose precision=fp32 or auto.") if requested == "bf16" and not torch.cuda.is_bf16_supported(): raise ValueError("GPU does not support bf16; choose precision=auto or fp16.") return requested def optimizer_for(model, settings, device): decay, no_decay = [], [] for parameter in model.parameters(): (decay if parameter.ndim >= 2 else no_decay).append(parameter) return torch.optim.AdamW( [{"params": decay, "weight_decay": settings["weight_decay"]}, {"params": no_decay, "weight_decay": 0.0}], lr=settings["learning_rate"], betas=(settings["beta1"], settings["beta2"]), fused=device.type == "cuda") def atomic_checkpoint(path, state): temporary = Path(path).with_suffix(".pt.tmp") torch.save(state, temporary) temporary.replace(path) def rng_state(): return {"python": random.getstate(), "torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []} def restore_rng(state): random.setstate(state["python"]) torch.set_rng_state(state["torch"].cpu()) if state["cuda"]: torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda"]]) @torch.inference_mode() def evaluate(model, loader, device, precision): model.eval() loss_sum = total_tokens = 0 try: for batch in loader: inputs = batch["input_ids"].to(device, non_blocking=True) labels = batch["labels"].to(device, non_blocking=True) with amp_context(device, precision): result = model(inputs, labels) loss_sum += float(result["loss_sum"]) total_tokens += int(batch["lengths"].sum()) if not total_tokens or not math.isfinite(loss_sum): raise ValueError("Empty or non-finite validation result.") return {"loss": loss_sum / total_tokens, "prediction_pairs": total_tokens, "perplexity": math.exp(min(50, loss_sum / total_tokens))} finally: model.train() def validate_config(config): model = GPTConfig(**config["model"]) settings = config["training"] positive = ("cpu_threads", "batch_size", "bucket_multiplier", "gradient_accumulation", "max_steps", "log_every", "eval_every", "eval_batches", "checkpoint_every") if any(settings[name] < 1 for name in positive) or settings["num_workers"] < 0: raise ValueError("Invalid training batch/worker/interval configuration.") if model.context_length % 8 or not 0 <= settings["warmup_steps"] <= settings["max_steps"]: raise ValueError("Context must be divisible by 8; warmup must fit training steps.") if (not 0 < settings["min_learning_rate"] <= settings["learning_rate"] or settings["weight_decay"] < 0 or settings["grad_clip"] <= 0 or not 0 <= settings["beta1"] < 1 or not 0 <= settings["beta2"] < 1 or not 0 <= settings["benchmark_warmup_steps"] < settings["max_steps"]): raise ValueError("Invalid optimizer or benchmark configuration.") return model, settings def validate_epoch_schedule(settings, windows): """An explicit epoch promise must match a complete, loss-preserving pass.""" batches = math.ceil(windows / settings["batch_size"]) epochs = settings.get("epochs") if epochs is not None: if not isinstance(epochs, int) or isinstance(epochs, bool) or epochs < 1: raise ValueError("epochs must be a positive integer.") # Current epoch mode ends on a committed update without borrowing data # from the next epoch to fill a partial accumulation group. if settings["gradient_accumulation"] != 1 or settings["max_steps"] != epochs * batches: raise ValueError("Epoch mode requires gradient_accumulation=1 and max_steps=epochs*batches_per_epoch.") return batches def run(config, resume=False, dry_run=False): model_config, settings = validate_config(config) token_dir, index_dir = ROOT / config["token_dir"], ROOT / config["index_dir"] output, log_dir = ROOT / config["output_dir"], ROOT / config["log_dir"] token_manifest = json.loads((token_dir / "manifest.json").read_text(encoding="utf-8")) index_manifest = json.loads((index_dir / "manifest.json").read_text(encoding="utf-8")) if (token_manifest["vocab_size"] != model_config.vocab_size or index_manifest["context_length"] != model_config.context_length): raise ValueError("Model vocabulary/context must match token data/window index.") signatures = {"tokens": file_sha(token_dir / "manifest.json"), "windows": file_sha(index_dir / "manifest.json"), "training_code": {name: file_sha(Path(__file__).with_name(name)) for name in ("train.py", "model.py", "data.py")}} if token_manifest.get("status") != "complete" or index_manifest.get("status") != "complete": raise ValueError("Complete token store and window index required.") if signatures["tokens"] != index_manifest["token_manifest_sha256"]: raise ValueError("Window index belongs to a different token dataset.") batches_per_epoch = validate_epoch_schedule(settings, index_manifest["splits"]["train"]["windows"]) device = torch.device(settings["device"]) if device.type not in {"cpu", "cuda"}: raise ValueError("Initial trainer supports CPU or CUDA.") if device.type == "cuda" and not torch.cuda.is_available(): raise ValueError("CUDA unavailable; verify installation before training.") precision = resolve_precision(device, settings["precision"]) torch.set_num_threads(settings["cpu_threads"]) random.seed(settings["seed"]) np.random.seed(settings["seed"]) torch.manual_seed(settings["seed"]) if device.type == "cuda": torch.cuda.manual_seed_all(settings["seed"]) model = TinyGPT(model_config) plan = {"model": model.configuration(), "parameters": model.parameter_count(), "device": str(device), "precision": precision, "max_steps": settings["max_steps"], "training_prediction_pairs_per_epoch": token_manifest["splits"]["train"]["prediction_pairs"], "training_batches_per_epoch": batches_per_epoch, "planned_epochs": settings.get("epochs"), "output": str(output), "tensorboard": str(log_dir), "validation": "Fixed random subset, token-weighted loss; test unused.", "corpus_quality_mode": token_manifest.get("corpus_quality_mode"), "full_validation_at_end": settings.get("validate_full_at_end", False), "purpose": "First full-corpus baseline." if settings.get("epochs") else "Short pipeline/throughput run; not a deployment-ready model."} if dry_run: print(json.dumps(plan, ensure_ascii=False, indent=2)) return plan last_path = output / "last.pt" if resume: if not last_path.exists(): raise ValueError("No last.pt checkpoint to resume.") elif output.exists() and any(output.iterdir()): raise ValueError("Output directory is nonempty; use --resume or a new output directory.") elif log_dir.exists() and any(log_dir.iterdir()): raise ValueError("TensorBoard directory is nonempty; use a new log directory.") prepare_indexes(token_dir, index_dir, model_config.context_length) model.to(device) optimizer = optimizer_for(model, settings, device) scaler = torch.amp.GradScaler("cuda", enabled=precision == "fp16") step = epoch = batch_cursor = total_tokens = total_windows = 0 best_val = math.inf benchmark_tokens = benchmark_steps = 0 benchmark_seconds = 0.0 restored_rng = None if resume: saved = torch.load(last_path, map_location="cpu", weights_only=True) if saved["config"] != config or saved["signatures"] != signatures or saved["precision"] != precision: raise ValueError("Resume requires the same configuration, precision and dataset signatures.") model.load_state_dict(saved["model"]) optimizer.load_state_dict(saved["optimizer"]) scaler.load_state_dict(saved["scaler"]) step, epoch, batch_cursor = saved["step"], saved["epoch"], saved["batch_cursor"] total_tokens, best_val = saved["total_tokens"], saved["best_val"] total_windows = saved["total_windows"] benchmark_tokens, benchmark_steps, benchmark_seconds = ( saved["benchmark_tokens"], saved["benchmark_steps"], saved["benchmark_seconds"]) restored_rng = saved["rng"] if step >= settings["max_steps"]: summary_path = output / "summary.json" final_eval_complete = ((output / "full-validation.json").exists() and summary_path.exists() and json.loads(summary_path.read_text(encoding="utf-8")).get("status") == "complete") if not settings.get("validate_full_at_end", False) or final_eval_complete: print(f"Already complete at step {step}; no training started.") return plan print("Training updates complete; resuming unfinished full validation.", flush=True) output.mkdir(parents=True, exist_ok=True) write_json(output / "config.json", config) write_json(output / "plan.json", plan) environment = {"torch_version": str(torch.__version__), "cuda_runtime": torch.version.cuda, "device": str(device), "precision": precision} if device.type == "cuda": environment["gpu"] = torch.cuda.get_device_name(device) write_json(output / "environment.json", environment) train_data = SentenceWindowDataset(token_dir, index_dir, "train") val_data = SentenceWindowDataset(token_dir, index_dir, "validation") train_loader = make_loader(train_data, settings["batch_size"], settings["num_workers"], settings["seed"], pin_memory=device.type == "cuda", bucket_multiplier=settings["bucket_multiplier"], start_batch=batch_cursor, worker_init_fn=worker_init) train_loader.batch_sampler.set_epoch(epoch, batch_cursor) val_size = min(len(val_data), settings["batch_size"] * settings["eval_batches"]) val_indices = random.Random(settings["seed"] + 1_000_019).sample(range(len(val_data)), val_size) val_loader = make_loader(val_data, settings["batch_size"], settings["num_workers"], settings["seed"], shuffle=False, pin_memory=device.type == "cuda", bucket_multiplier=settings["bucket_multiplier"], indices=val_indices, worker_init_fn=worker_init) writer = SummaryWriter(str(log_dir), purge_step=step + 1 if resume else None, flush_secs=5) writer.add_text("run/config", json.dumps(config, ensure_ascii=False, indent=2), step) stop_requested = False previous_signal = signal.getsignal(signal.SIGINT) def stop_handler(signum, frame): nonlocal stop_requested if stop_requested: raise KeyboardInterrupt stop_requested = True print("Stop requested; finishing the current update, then saving checkpoint.", flush=True) signal.signal(signal.SIGINT, stop_handler) iterator = full_loader = None start_time = time.perf_counter() latest_validation = None def save_checkpoint(path): atomic_checkpoint(path, {"format": "vimeml_tiny_gpt_v1", "model": model.state_dict(), "model_config": model.configuration(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(), "config": config, "signatures": signatures, "precision": precision, "step": step, "epoch": epoch, "batch_cursor": batch_cursor, "total_tokens": total_tokens, "total_windows": total_windows, "best_val": best_val, "rng": rng_state(), "benchmark_tokens": benchmark_tokens, "benchmark_steps": benchmark_steps, "benchmark_seconds": benchmark_seconds}) def validation(): result = evaluate(model, val_loader, device, precision) writer.add_scalar("loss/validation", result["loss"], step) writer.add_scalar("perplexity/validation", result["perplexity"], step) writer.flush() print(f"[validation] step={step} loss={result['loss']:.4f} tokens={result['prediction_pairs']:,}", flush=True) return result interval_loss = interval_tokens = interval_positions = 0 interval_seconds = interval_data_seconds = 0.0 try: if not resume: latest_validation = validation() best_val = latest_validation["loss"] save_checkpoint(last_path) # Initial state can resume even before step 1. save_checkpoint(output / "best.pt") if restored_rng is not None: restore_rng(restored_rng) if device.type == "cuda": torch.cuda.reset_peak_memory_stats(device) model.train() print(f"Training {plan['parameters']:,} parameters; {precision}; max_steps={settings['max_steps']}", flush=True) while step < settings["max_steps"] and not stop_requested: update_start = time.perf_counter() microbatches = [] for _ in range(settings["gradient_accumulation"]): if iterator is None: iterator = iter(train_loader) try: batch = next(iterator) except StopIteration: epoch += 1 batch_cursor = 0 train_loader.batch_sampler.set_epoch(epoch) iterator = iter(train_loader) batch = next(iterator) microbatches.append(batch) batch_cursor += 1 data_seconds = time.perf_counter() - update_start token_count = sum(int(batch["lengths"].sum()) for batch in microbatches) positions = sum(batch["input_ids"].numel() for batch in microbatches) rate = learning_rate(step + 1, settings) for group in optimizer.param_groups: group["lr"] = rate optimizer.zero_grad(set_to_none=True) loss_sums = [] for batch in microbatches: inputs = batch["input_ids"].to(device, non_blocking=True) labels = batch["labels"].to(device, non_blocking=True) with amp_context(device, precision): result = model(inputs, labels) loss = result["loss_sum"] / token_count loss_sums.append(result["loss_sum"].detach()) scaler.scale(loss).backward() scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), settings["grad_clip"], error_if_nonfinite=True) old_scale = scaler.get_scale() scaler.step(optimizer) scaler.update() if scaler.get_scale() < old_scale: raise FloatingPointError("FP16 update was skipped; resume last committed checkpoint or choose bf16.") if device.type == "cuda": torch.cuda.synchronize(device) loss_sum = float(torch.stack(loss_sums).sum()) if not math.isfinite(loss_sum): raise FloatingPointError("Non-finite training loss.") seconds = time.perf_counter() - update_start step += 1 total_tokens += token_count total_windows += sum(batch["input_ids"].shape[0] for batch in microbatches) interval_loss += loss_sum interval_tokens += token_count interval_positions += positions interval_seconds += seconds interval_data_seconds += data_seconds if step > settings["benchmark_warmup_steps"]: benchmark_tokens += token_count benchmark_steps += 1 benchmark_seconds += seconds if step == 1 or step % settings["log_every"] == 0 or step == settings["max_steps"] or stop_requested: metrics = {"loss/train": interval_loss / interval_tokens, "optimizer/learning_rate": rate, "optimizer/grad_norm": float(grad_norm), "throughput/effective_tokens_per_second": interval_tokens / interval_seconds, "data/padding_fraction": 1 - interval_tokens / interval_positions, "data/wait_fraction": interval_data_seconds / interval_seconds, "progress/trained_tokens": total_tokens, "progress/percent": 100 * step / settings["max_steps"]} if device.type == "cuda": metrics["memory/peak_allocated_mib"] = torch.cuda.max_memory_allocated(device) / 1024**2 metrics["memory/peak_reserved_mib"] = torch.cuda.max_memory_reserved(device) / 1024**2 for tag, value in metrics.items(): writer.add_scalar(tag, value, step) elapsed = time.perf_counter() - start_time progress = {"status": "training", "step": step, "max_steps": settings["max_steps"], "epoch": epoch, "batch_cursor": batch_cursor, "total_tokens": total_tokens, "elapsed_seconds_this_session": elapsed, "metrics": metrics} write_json(output / "progress.json", progress) with (output / "metrics.jsonl").open("a", encoding="utf-8") as stream: stream.write(json.dumps(progress, ensure_ascii=False) + "\n") print(f"step {step}/{settings['max_steps']} loss={metrics['loss/train']:.4f} " f"tokens/s={metrics['throughput/effective_tokens_per_second']:,.0f} " f"padding={metrics['data/padding_fraction']:.1%} elapsed={elapsed:.1f}s", flush=True) interval_loss = interval_tokens = interval_positions = 0 interval_seconds = interval_data_seconds = 0.0 if step % settings["eval_every"] == 0 or step == settings["max_steps"]: latest_validation = validation() if latest_validation["loss"] < best_val: best_val = latest_validation["loss"] save_checkpoint(output / "best.pt") if step % settings["checkpoint_every"] == 0 or step == settings["max_steps"]: save_checkpoint(last_path) save_checkpoint(last_path) status = "complete" if step >= settings["max_steps"] else "interrupted" if status == "complete" and settings.get("epochs"): expected_tokens = settings["epochs"] * token_manifest["splits"]["train"]["prediction_pairs"] expected_windows = settings["epochs"] * index_manifest["splits"]["train"]["windows"] if total_tokens != expected_tokens or total_windows != expected_windows: raise ValueError("Training did not cover the exact promised windows/prediction pairs.") full_validation = None if status == "complete" and settings.get("validate_full_at_end", False): print("[validation_full] Evaluating all validation windows for last/best checkpoints...", flush=True) write_json(output / "progress.json", {"status": "evaluating_full_validation", "step": step, "max_steps": settings["max_steps"], "total_tokens": total_tokens}) # Release sampled validation workers before creating the full loader. val_loader = None full_loader = make_loader(val_data, settings["batch_size"], settings["num_workers"], settings["seed"], shuffle=False, pin_memory=device.type == "cuda", bucket_multiplier=settings["bucket_multiplier"], worker_init_fn=worker_init) last_full = evaluate(model, full_loader, device, precision) writer.add_scalar("loss/validation_full_last", last_full["loss"], step) best_saved = torch.load(output / "best.pt", map_location="cpu", weights_only=True) best_step = best_saved["step"] if best_step == step: best_full = dict(last_full) else: model.load_state_dict(best_saved["model"]) best_full = evaluate(model, full_loader, device, precision) del best_saved expected_validation = token_manifest["splits"]["validation"]["prediction_pairs"] if any(result["prediction_pairs"] != expected_validation for result in (last_full, best_full)): raise ValueError("Full validation did not cover every validation prediction pair.") writer.add_scalar("loss/validation_full_best", best_full["loss"], step) writer.flush() full_validation = {"last": last_full, "best": best_full, "best_checkpoint_step": best_step, "best_selection": "Lowest fixed-subset validation loss during training."} write_json(output / "full-validation.json", full_validation) print(f"[validation_full] last={last_full['loss']:.4f} best={best_full['loss']:.4f} " f"tokens={expected_validation:,}", flush=True) speed = benchmark_tokens / benchmark_seconds if benchmark_seconds else None summary = {"status": status, "step": step, "parameters": model.parameter_count(), "total_trained_tokens": total_tokens, "best_validation_loss": best_val, "total_trained_windows": total_windows, "completed_data_passes": total_tokens / token_manifest["splits"]["train"]["prediction_pairs"], "last_validation": latest_validation, "benchmark_steps": benchmark_steps, "full_validation": full_validation, "benchmark_effective_tokens_per_second": speed, "benchmark_seconds_excluding_validation_checkpoints_and_first_warmup_steps": benchmark_seconds, "estimated_one_epoch_training_hours_excluding_validation_checkpoints": token_manifest["splits"]["train"]["prediction_pairs"] / speed / 3600 if speed else None, "estimate_note": "Short-run measurement; length mix, GPU power/temperature and validation/checkpoint overhead affect full-run time.", "elapsed_seconds_this_session": time.perf_counter() - start_time, "checkpoint": str(last_path), "tensorboard": str(log_dir)} if device.type == "cuda": summary["peak_allocated_mib"] = torch.cuda.max_memory_allocated(device) / 1024**2 summary["peak_reserved_mib"] = torch.cuda.max_memory_reserved(device) / 1024**2 write_json(output / "summary.json", summary) write_json(output / "progress.json", summary) write_json(output / "manifest.json", {"status": status, "format": "vimeml_tiny_gpt_run_v1", "plan": plan, "signatures": signatures, "step": step}) print(json.dumps(summary, ensure_ascii=False, indent=2)) print(f"Summary: {output / 'summary.json'}") return summary except BaseException as error: write_json(output / "progress.json", {"status": "failed", "step": step, "error": str(error), "resume": "Resume the last committed last.pt checkpoint."}) raise finally: signal.signal(signal.SIGINT, previous_signal) writer.close() # Persistent iterators are owned by their loaders; delete both owners. del iterator, train_loader, val_loader, full_loader train_data.close() val_data.close() def main(argv=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=ROOT / "configs/train-smoke.toml") parser.add_argument("--dry-run", action="store_true", help="Print plan only; no training or output files.") parser.add_argument("--resume", action="store_true", help="Resume last.pt with identical configuration.") args = parser.parse_args(argv) config = tomllib.loads(args.config.read_text(encoding="utf-8")) run(config, resume=args.resume, dry_run=args.dry_run) if __name__ == "__main__": main()