File size: 40,586 Bytes
29f25be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 | """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):
# 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, 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.")
# 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)
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) # Initial state can resume even before step 1.
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})
# Release sampled validation workers before creating the full loader.
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()
# Release iterators and loaders before closing the memory maps.
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()
|