| """Tiny GPT training: bounded batches, token-weighted loss and checkpoints.""" |
| import argparse |
| import json |
| import math |
| import random |
| import signal |
| import subprocess |
| import sys |
| 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_factory import checkpoint_format, configuration_for, create_model |
| from vimeml.training.data_v2 import PrefixCropWindowDataset |
|
|
| ROOT = Path(__file__).resolve().parents[3] |
|
|
|
|
| def learning_rate(step, settings): |
| """One-based update; optional stable phase before the cosine decay.""" |
| 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 |
| decay_start = (settings.get("stable_steps", 0) |
| if settings.get("learning_rate_schedule", "cosine") == "stable_decay" |
| else warmup) |
| fraction = min(1.0, max(0.0, (step - decay_start) / max(1, maximum - decay_start))) |
| return low + (high - low) * (1 + math.cos(math.pi * fraction)) / 2 |
|
|
|
|
| def worker_init(worker_id): |
| |
| 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, characters=None): |
| 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.") |
| metrics = {"loss": loss_sum / total_tokens, "nll_sum": loss_sum, |
| "prediction_pairs": total_tokens, |
| "perplexity": math.exp(min(50, loss_sum / total_tokens))} |
| if characters is not None: |
| if characters < 1: |
| raise ValueError("Validation characters must be positive.") |
| metrics.update(characters=characters, bpc=loss_sum / (characters * math.log(2)), |
| bpc_policy="All prediction targets including EOS; original Unicode characters excluding BOS/EOS.") |
| return metrics |
| finally: |
| model.train() |
|
|
|
|
| def validate_config(config): |
| model = configuration_for(config.get("architecture", "tiny_gpt_v1"), config["model"]) |
| settings = dict(config["training"]) |
| crop_probability = settings.get("prefix_crop_probability", 0.0) |
| crop_minimum = settings.get("prefix_crop_min_remaining_tokens", 8) |
| if (not math.isfinite(crop_probability) or not 0 <= crop_probability <= 1 or |
| type(crop_minimum) is not int or crop_minimum < 1): |
| raise ValueError("Invalid prefix crop probability or minimum token count.") |
| if crop_probability and config.get("architecture", "tiny_gpt_v1") != "tiny_gpt_v2": |
| raise ValueError("Prefix crops are only enabled for tiny_gpt_v2.") |
| if settings.get("early_stopping_patience", 0) < 0 or settings.get("early_stopping_min_delta", 0) < 0: |
| raise ValueError("Invalid epoch degradation stopping configuration.") |
| if type(settings.get('data_epoch_offset', 0)) is not int or settings.get('data_epoch_offset', 0) < 0: |
| raise ValueError('data_epoch_offset must be a nonnegative integer.') |
| if config.get('runtime', {}).get('compile_backbone', False) and config.get('architecture') != 'tiny_gpt_v2': |
| raise ValueError('Backbone compilation is only supported for V2.') |
| automatic_steps = settings.get("max_steps") == 0 and settings.get("epochs", 0) > 0 |
| schedule = settings.get("learning_rate_schedule", "cosine") |
| stable_steps = settings.get("stable_steps", 0) |
| if schedule not in {"cosine", "stable_decay"} or type(stable_steps) is not int: |
| raise ValueError("Invalid learning rate schedule.") |
| if (stable_steps < 0 or (schedule == "stable_decay" and |
| (stable_steps < settings["warmup_steps"] or |
| (not automatic_steps and stable_steps >= settings["max_steps"])))): |
| raise ValueError("Stable phase must leave a nonempty decay interval after warmup.") |
| baseline_bpc = config.get("initialization", {}).get("baseline_validation_bpc") |
| if baseline_bpc is not None and (not math.isfinite(baseline_bpc) or baseline_bpc <= 0): |
| raise ValueError("Initialization baseline BPC must be finite and positive.") |
| if config.get("initialization", {}).get("restore_optimizer", False) and not config["initialization"].get("checkpoint"): |
| raise ValueError("Optimizer continuation requires an initialization checkpoint.") |
| positive = ("cpu_threads", "batch_size", "bucket_multiplier", "gradient_accumulation", |
| "log_every", "eval_every", "eval_batches", "checkpoint_every") |
| if any(settings[name] < 1 for name in positive) or settings["num_workers"] < 0 or (not automatic_steps and settings["max_steps"] < 1): |
| raise ValueError("Invalid training batch/worker/interval configuration.") |
| if model.context_length % 8 or settings["warmup_steps"] < 0 or (not automatic_steps and 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 |
| settings["benchmark_warmup_steps"] < 0 or |
| (not automatic_steps and 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.") |
| |
| |
| 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) |
| architecture = config.get("architecture", "tiny_gpt_v1") |
| 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.") |
| code_names = ["train.py", "model.py", "data.py", "model_factory.py"] |
| if architecture == "tiny_gpt_v2": |
| code_names.extend(("model_v2.py", "data_v2.py")) |
| if settings.get("validate_full_each_epoch", False): |
| code_names.append("evaluation_v2.py") |
| if config.get('runtime', {}).get('compile_backbone', False) or config.get('initialization', {}).get('checkpoint'): |
| code_names.append('runtime_v2.py') |
| 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 code_names}} |
| 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.") |
| windows_per_epoch = index_manifest["splits"]["train"]["windows"] |
| if settings["max_steps"] == 0: |
| settings["max_steps"] = settings["epochs"] * math.ceil(windows_per_epoch / settings["batch_size"]) |
| _, settings = validate_config({**config, "training": settings}) |
| batches_per_epoch = validate_epoch_schedule(settings, windows_per_epoch) |
| 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 = create_model(architecture, model_config) |
| plan = {"architecture": architecture, "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"), |
| "initialization": config.get('initialization'), |
| "runtime": config.get('runtime'), |
| "data_epoch_offset": settings.get('data_epoch_offset', 0), |
| "prefix_crop": {"probability": settings.get("prefix_crop_probability", 0.0), |
| "min_remaining_tokens": settings.get("prefix_crop_min_remaining_tokens", 8), |
| "determinism": "seed, zero-based epoch and sample index; epoch-tagged worker inputs"}, |
| "output": str(output), "tensorboard": str(log_dir), |
| "validation": "Fixed random sentence-complete subset plus full validation each epoch; BPC includes EOS; test unused." |
| if settings.get('validate_full_each_epoch', False) else "Fixed random subset, token-weighted loss; test unused.", |
| "full_validation_each_epoch": settings.get('validate_full_each_epoch', False), |
| "corpus_quality_mode": token_manifest.get("corpus_quality_mode"), |
| "full_validation_at_end": settings.get("validate_full_at_end", False), |
| "purpose": "Continued pretraining with compatible AdamW state." if config.get('initialization', {}).get('restore_optimizer') else |
| "Continued pretraining from frozen weights with a fresh optimizer." if config.get('initialization') else |
| "Full-corpus experiment." 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, |
| verification=config.get("verification", {}).get("mode", "sha256")) |
| model.to(device) |
| optimizer = optimizer_for(model, settings, device) |
| initialization = None |
| if config.get('initialization', {}).get('checkpoint') and not resume: |
| from vimeml.training.runtime_v2 import initialize_weights |
| initialization = initialize_weights(model, ROOT / config['initialization']['checkpoint'], signatures, |
| optimizer=optimizer if config['initialization'].get('restore_optimizer', False) else None, |
| precision=precision) |
| print(f"Initializing from step {initialization['source_step']}; {initialization['optimizer']}.", flush=True) |
| scaler = torch.amp.GradScaler("cuda", enabled=precision == "fp16") |
| step = epoch = batch_cursor = total_tokens = total_windows = 0 |
| total_dropped_tokens = total_cropped_windows = 0 |
| best_val = math.inf |
| best_epoch_bpc = config.get('initialization', {}).get('baseline_validation_bpc', math.inf) |
| degrading_epochs = 0 |
| evaluated_epochs = [] |
| early_stopped = False |
| 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"] |
| total_dropped_tokens = saved.get("total_dropped_tokens", 0) |
| total_cropped_windows = saved.get("total_cropped_windows", 0) |
| best_epoch_bpc = saved.get("best_epoch_bpc", math.inf) |
| degrading_epochs = saved.get("degrading_epochs", 0) |
| evaluated_epochs = saved.get("evaluated_epochs", []) |
| early_stopped = saved.get("early_stopped", False) |
| initialization = saved.get('initialization') |
| 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} |
| environment['initialization'] = initialization |
| if config.get('runtime', {}).get('compile_backbone', False): |
| from vimeml.training.runtime_v2 import compile_backbone |
| environment['runtime_optimization'] = compile_backbone(model) |
| source_provenance = ROOT / config.get('source_provenance', 'source-provenance.json') |
| if source_provenance.exists(): |
| environment['source_provenance'] = json.loads(source_provenance.read_text(encoding='utf-8')) |
| if device.type == "cuda": |
| environment["gpu"] = torch.cuda.get_device_name(device) |
| write_json(output / "environment.json", environment) |
| if settings.get("prefix_crop_probability", 0.0): |
| train_data = PrefixCropWindowDataset(token_dir, index_dir, "train", seed=settings["seed"], |
| probability=settings["prefix_crop_probability"], |
| min_remaining_tokens=settings.get("prefix_crop_min_remaining_tokens", 8)) |
| else: |
| 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 + settings.get('data_epoch_offset', 0), 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_characters = None |
| if settings.get("validate_full_each_epoch", False): |
| from vimeml.training.evaluation_v2 import validation_subset |
| val_indices, val_characters = validation_subset(val_data, val_indices, ROOT / config['evaluation']['tokenizer_dir']) |
| 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": checkpoint_format(architecture), "architecture": architecture, |
| "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, |
| "source_provenance": environment.get('source_provenance'), |
| "initialization": initialization, |
| "total_tokens": total_tokens, "total_windows": total_windows, |
| "total_dropped_tokens": total_dropped_tokens, "total_cropped_windows": total_cropped_windows, |
| "best_val": best_val, "rng": rng_state(), |
| "best_epoch_bpc": best_epoch_bpc, "degrading_epochs": degrading_epochs, |
| "evaluated_epochs": evaluated_epochs, "early_stopped": early_stopped, |
| "benchmark_tokens": benchmark_tokens, "benchmark_steps": benchmark_steps, |
| "benchmark_seconds": benchmark_seconds, |
| "scheduler": {"name": settings.get("learning_rate_schedule", "cosine"), "step": step, "max_steps": settings["max_steps"], |
| "stable_steps": settings.get("stable_steps", 0), |
| "warmup_steps": settings["warmup_steps"], "learning_rate": settings["learning_rate"], |
| "min_learning_rate": settings["min_learning_rate"]}}) |
|
|
| def validation(): |
| result = evaluate(model, val_loader, device, precision, val_characters) |
| writer.add_scalar("loss/validation", result["loss"], step) |
| writer.add_scalar("perplexity/validation", result["perplexity"], step) |
| if 'bpc' in result: |
| writer.add_scalar("bpc/validation_subset", result['bpc'], step) |
| writer.flush() |
| print(f"[validation] step={step} loss={result['loss']:.4f} tokens={result['prediction_pairs']:,}", flush=True) |
| return result |
|
|
| def evaluate_epoch(number): |
| nonlocal full_loader, best_epoch_bpc, degrading_epochs, early_stopped |
| from vimeml.training.evaluation_v2 import epoch_ime |
| print(f"[epoch {number}] Full validation and frozen IME evaluations...", flush=True) |
| save_checkpoint(last_path) |
| if full_loader is 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) |
| result = evaluate(model, full_loader, device, precision, token_manifest['splits']['validation']['characters']) |
| if result['prediction_pairs'] != token_manifest['splits']['validation']['prediction_pairs']: |
| raise ValueError('Incomplete epoch validation coverage.') |
| improved = result['bpc'] < best_epoch_bpc |
| degrading_epochs = (degrading_epochs + 1 if result['bpc'] > best_epoch_bpc + |
| settings.get('early_stopping_min_delta', .01) else 0) |
| if improved: |
| best_epoch_bpc = result['bpc'] |
| ime = epoch_ime(model, ROOT / config['evaluation']['tokenizer_dir'], |
| config['evaluation']['benchmarks'], ROOT, output, number) |
| for name, value in result.items(): |
| if isinstance(value, (int, float)): |
| writer.add_scalar(f'epoch_validation/{name}', value, step) |
| for benchmark, values in ime.items(): |
| for name, value in values.items(): |
| if isinstance(value, (int, float)): |
| writer.add_scalar(f'ime/{benchmark}/{name}', value, step) |
| writer.add_scalar('progress/completed_epochs', number, step) |
| writer.flush() |
| patience = settings.get('early_stopping_patience', 0) |
| early_stopped = bool(patience and degrading_epochs >= patience) |
| report = {'epoch': number, 'step': step, 'validation': result, 'ime': ime, |
| 'best_epoch_bpc': best_epoch_bpc, 'degrading_epochs': degrading_epochs, |
| 'early_stopped': early_stopped} |
| epoch_output = output / 'epoch-evaluation' / f'epoch-{number}' |
| epoch_output.mkdir(parents=True, exist_ok=True) |
| write_json(epoch_output / 'summary.json', report) |
| evaluated_epochs.append(number) |
| save_checkpoint(output / f'epoch-{number}.pt') |
| if improved: |
| save_checkpoint(output / 'best-epoch.pt') |
| save_checkpoint(last_path) |
| print(f"[epoch {number}] BPC={result['bpc']:.6f} IME={ime}", flush=True) |
| return result |
|
|
| interval_loss = interval_tokens = interval_positions = 0 |
| interval_windows = interval_dropped_tokens = interval_cropped_windows = 0 |
| interval_seconds = interval_data_seconds = 0.0 |
| try: |
| if not resume: |
| latest_validation = validation() |
| best_val = latest_validation["loss"] |
| save_checkpoint(last_path) |
| save_checkpoint(output / "best.pt") |
| if math.isfinite(best_epoch_bpc): |
| save_checkpoint(output / "best-epoch.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) |
| if settings.get('validate_full_each_epoch', False) and batch_cursor == batches_per_epoch: |
| if epoch + 1 not in evaluated_epochs: |
| evaluate_epoch(epoch + 1) |
| while step < settings["max_steps"] and not stop_requested and not early_stopped: |
| 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 + settings.get('data_epoch_offset', 0)) |
| 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) |
| dropped_tokens = sum(int(batch["crop_offset"].sum()) for batch in microbatches if "crop_offset" in batch) |
| cropped_windows = sum(int((batch["crop_offset"] > 0).sum()) for batch in microbatches if "crop_offset" in batch) |
| window_count = sum(batch["input_ids"].shape[0] 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 += window_count |
| total_dropped_tokens += dropped_tokens |
| total_cropped_windows += cropped_windows |
| interval_loss += loss_sum |
| interval_tokens += token_count |
| interval_positions += positions |
| interval_windows += window_count |
| interval_dropped_tokens += dropped_tokens |
| interval_cropped_windows += cropped_windows |
| 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, |
| "throughput/samples_per_second": interval_windows / interval_seconds, |
| "data/non_padding_tokens_per_update": token_count, |
| "data/prefix_crop_fraction": interval_cropped_windows / interval_windows, |
| "data/prefix_crop_dropped_tokens": interval_dropped_tokens, |
| "data/padding_fraction": 1 - interval_tokens / interval_positions, |
| "data/wait_fraction": interval_data_seconds / interval_seconds, |
| "timing/update_seconds": seconds, |
| "progress/epoch_fraction": total_windows / windows_per_epoch, |
| "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, |
| "total_dropped_tokens": total_dropped_tokens, "total_cropped_windows": total_cropped_windows, |
| "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_windows = interval_dropped_tokens = interval_cropped_windows = 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) |
| if settings.get('validate_full_each_epoch', False) and batch_cursor == batches_per_epoch: |
| completed_epoch = epoch + 1 |
| if completed_epoch not in evaluated_epochs: |
| evaluate_epoch(completed_epoch) |
| save_checkpoint(last_path) |
| status = "early_stopped" if early_stopped else ("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 + total_dropped_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 in {'complete', 'early_stopped'} 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}) |
| |
| val_loader = None |
| if full_loader is 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) |
| val_char_count = token_manifest['splits']['validation'].get('characters') |
| epoch_report = output / 'epoch-evaluation' / f'epoch-{epoch + 1}' / 'summary.json' |
| if settings.get('validate_full_each_epoch', False) and epoch_report.exists(): |
| cached_epoch = json.loads(epoch_report.read_text(encoding='utf-8')) |
| last_full = cached_epoch['validation'] if cached_epoch['step'] == step else None |
| else: |
| last_full = None |
| if last_full is None: |
| last_full = evaluate(model, full_loader, device, precision, val_char_count) |
| 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, val_char_count) |
| 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_windows / windows_per_epoch, |
| "effective_token_passes": total_tokens / token_manifest["splits"]["train"]["prediction_pairs"], |
| "prefix_crop_dropped_tokens": total_dropped_tokens, |
| "prefix_cropped_windows": total_cropped_windows, |
| "uncropped_prediction_pairs_seen": total_tokens + total_dropped_tokens, |
| "last_validation": latest_validation, "benchmark_steps": benchmark_steps, |
| "full_validation": full_validation, |
| "evaluated_epochs": evaluated_epochs, "best_epoch_bpc": best_epoch_bpc if math.isfinite(best_epoch_bpc) else None, |
| "early_stopped": early_stopped, |
| "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_v2" if architecture == "tiny_gpt_v2" else "vimeml_tiny_gpt_run_v1", |
| "architecture": architecture, |
| "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() |
| |
| iterator = train_loader = val_loader = full_loader = None |
| 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")) |
| tracking = config.get("tracking", {}) |
| if tracking.get("enabled", False): |
| command = [sys.executable, "-X", "utf8", "-u", |
| str(ROOT / "scripts/training/train_wandb.py"), "--config", str(args.config.resolve())] |
| if args.resume: |
| command.append("--resume") |
| if args.dry_run: |
| command.append("--dry-run") |
| if tracking.get("mode", "online") == "offline": |
| command.append("--offline") |
| subprocess.run(command, check=True, cwd=ROOT) |
| return |
| run(config, resume=args.resume, dry_run=args.dry_run) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
