| """ |
| Derived from Andrej Karpathy's nanochat project. |
| |
| MIT License |
| |
| Copyright (c) 2025 Andrej Karpathy |
| |
| Permission is hereby granted, free of charge, to any person obtaining a copy |
| of this software and associated documentation files (the "Software"), to deal |
| in the Software without restriction, including without limitation the rights |
| to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| copies of the Software, and to permit persons to whom the Software is |
| furnished to do so, subject to the following conditions: |
| |
| The above copyright notice and this permission notice shall be included in all |
| copies or substantial portions of the Software. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import random |
| import statistics |
| import time |
| from typing import Callable |
|
|
| import numpy as np |
| import torch |
|
|
| from dropout_decay.experiments.artifacts import write_jsonl_row |
| from dropout_decay.models import DropoutGPT, GPTConfig |
| from dropout_decay.specs import DropoutCondition, ModelSpec |
|
|
|
|
| def set_seed(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| def batch_seed( |
| seed: int, model: ModelSpec, dropout_code: int, stage: int | None |
| ) -> int: |
| return ( |
| seed * 1_000_003 |
| + model.n_layer * 100_003 |
| + model.n_head * 10_007 |
| + model.n_embd * 101 |
| + dropout_code * 37 |
| + (stage or 0) * 997 |
| ) |
|
|
|
|
| def make_batch( |
| tokens: np.ndarray, |
| token_limit: int, |
| batch_size: int, |
| block_size: int, |
| rng: np.random.Generator, |
| device: torch.device, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| limit = min(token_limit, len(tokens)) |
| max_start = limit - block_size - 1 |
| if max_start <= 0: |
| raise ValueError("token_limit is too small for the requested block_size") |
| starts = rng.integers(0, max_start, size=batch_size) |
| x_np = np.stack([tokens[start : start + block_size] for start in starts]).astype( |
| np.int64 |
| ) |
| y_np = np.stack( |
| [tokens[start + 1 : start + 1 + block_size] for start in starts] |
| ).astype(np.int64) |
| return ( |
| torch.tensor(x_np, dtype=torch.long, device=device), |
| torch.tensor(y_np, dtype=torch.long, device=device), |
| ) |
|
|
|
|
| @torch.no_grad() |
| def estimate_loss( |
| model: DropoutGPT, |
| tokens: np.ndarray, |
| token_limit: int, |
| batches: int, |
| args: argparse.Namespace, |
| device: torch.device, |
| rng_seed: int, |
| ) -> float: |
| if batches <= 0: |
| return float("nan") |
| model.eval() |
| rng = np.random.default_rng(rng_seed) |
| losses: list[float] = [] |
| for _ in range(batches): |
| x, y = make_batch( |
| tokens, token_limit, args.batch_size, args.block_size, rng, device |
| ) |
| _, loss = model(x, y) |
| losses.append(float(loss.item())) |
| model.train() |
| return float(statistics.fmean(losses)) |
|
|
|
|
| def train_segment( |
| *, |
| run_mode: str, |
| condition: DropoutCondition, |
| model_spec: ModelSpec, |
| config: GPTConfig, |
| train_tokens: np.ndarray, |
| val_tokens: np.ndarray, |
| token_limit: int, |
| steps: int, |
| seed: int, |
| args: argparse.Namespace, |
| device: torch.device, |
| dropout_fn: Callable[[int], float], |
| metrics_file, |
| trace_file, |
| stage: int | None = None, |
| model: DropoutGPT | None = None, |
| optimizer: torch.optim.Optimizer | None = None, |
| tokens_seen_start: int = 0, |
| ) -> tuple[DropoutGPT, torch.optim.Optimizer, int, dict]: |
| if model is None: |
| set_seed(seed) |
| model = DropoutGPT(config).to(device) |
| optimizer = torch.optim.AdamW( |
| model.parameters(), |
| lr=args.lr, |
| betas=(0.9, 0.95), |
| weight_decay=args.weight_decay, |
| ) |
| else: |
| torch.manual_seed(seed + 10_000 + (stage or 0)) |
| if optimizer is None: |
| raise ValueError("optimizer is required when reusing a model") |
|
|
| model.train() |
| dropout_code = int(round(condition.initial * 10_000)) |
| rng = np.random.default_rng(batch_seed(seed, model_spec, dropout_code, stage)) |
| tokens_seen = tokens_seen_start |
| last_loss = float("nan") |
| active_dropout = condition.initial |
| t0 = time.time() |
|
|
| for step in range(1, steps + 1): |
| active_dropout = dropout_fn(tokens_seen) |
| model.set_dropout(active_dropout) |
| x, y = make_batch( |
| train_tokens, token_limit, args.batch_size, args.block_size, rng, device |
| ) |
| _, loss = model(x, y) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| if args.grad_clip > 0: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| optimizer.step() |
|
|
| tokens_seen += args.batch_size * args.block_size |
| last_loss = float(loss.item()) |
| if args.log_every > 0 and step % args.log_every == 0: |
| write_jsonl_row( |
| trace_file, |
| { |
| "event": "train_step", |
| "run_mode": run_mode, |
| "condition": condition.name, |
| "model_name": model_spec.name, |
| "seed": seed, |
| "stage": stage, |
| "step": step, |
| "steps": steps, |
| "token_limit": int(token_limit), |
| "tokens_seen": int(tokens_seen), |
| "dropout": float(active_dropout), |
| "train_batch_loss": last_loss, |
| }, |
| ) |
| if args.eval_every > 0 and step % args.eval_every == 0: |
| train_eval = estimate_loss( |
| model, |
| train_tokens, |
| token_limit, |
| args.trace_eval_batches, |
| args, |
| device, |
| rng_seed=seed + 20_000 + step, |
| ) |
| val_eval = estimate_loss( |
| model, |
| val_tokens, |
| len(val_tokens), |
| args.trace_eval_batches, |
| args, |
| device, |
| rng_seed=seed + 30_000 + step, |
| ) |
| write_jsonl_row( |
| trace_file, |
| { |
| "event": "eval_step", |
| "run_mode": run_mode, |
| "condition": condition.name, |
| "model_name": model_spec.name, |
| "seed": seed, |
| "stage": stage, |
| "step": step, |
| "steps": steps, |
| "token_limit": int(token_limit), |
| "tokens_seen": int(tokens_seen), |
| "dropout": float(active_dropout), |
| "train_eval_loss": train_eval, |
| "val_eval_loss": val_eval, |
| "generalization_gap": val_eval - train_eval, |
| }, |
| ) |
|
|
| train_eval = estimate_loss( |
| model, |
| train_tokens, |
| token_limit, |
| args.train_eval_batches, |
| args, |
| device, |
| rng_seed=seed + 40_000 + (stage or 0), |
| ) |
| val_eval = estimate_loss( |
| model, |
| val_tokens, |
| len(val_tokens), |
| args.eval_batches, |
| args, |
| device, |
| rng_seed=seed + 50_000 + (stage or 0), |
| ) |
| row = { |
| "run_mode": run_mode, |
| "condition": condition.name, |
| "condition_kind": condition.kind, |
| "seed": seed, |
| "stage": stage, |
| "token_limit": int(token_limit), |
| "steps": int(steps), |
| "tokens_seen": int(tokens_seen), |
| "dropout_initial": float(condition.initial), |
| "dropout_final": float(condition.final), |
| "dropout_schedule": condition.schedule, |
| "dropout_active_final": float(active_dropout), |
| "train_loss_last": last_loss, |
| "train_eval_loss": train_eval, |
| "val_eval_loss": val_eval, |
| "eval_loss": val_eval, |
| "generalization_gap": val_eval - train_eval, |
| "elapsed_sec": time.time() - t0, |
| "parameters": model.num_parameters(), |
| "model_config": config.to_dict(), |
| **model_spec.to_dict(), |
| } |
| write_jsonl_row(metrics_file, row) |
| return model, optimizer, tokens_seen, row |
|
|
