Training entry point
An entry point. Running python train.py calls train.main at line 85.
3 of 3 functions, classes and methods here are reachable from train.main, and 0 from talk.main.
| Pipeline stage | Step | Explained on |
|---|---|---|
| Text to training windows | Text | data pipeline |
| Text to training windows | Tokenizer | tokenization |
| Text to training windows | Ids | tokenization |
| Text to training windows | Windows | data pipeline |
| Loss and update | Clipping | adamw learning rate clipping |
| Loss and update | AdamW with schedule | adamw learning rate clipping |
config.py, data.py, checkpoint.py, generate.py, ops.py, optim.py, tokenizer.py, transformer.pymath, numpy| Concept | What this file does | Lines | Status |
|---|---|---|---|
| Mini-batch loss and gradient | Computes it, Calls it | 13, 14, 15, 19, 22, 24, 61 | Derived from implementation |
| Cross-entropy | Calls it | 31 | Implemented |
| Validation loss | Computes it, Calls it | 32, 34, 70 | Derived from implementation |
| Random choice | Calls it | 37 | Implemented |
| The split | Calls it | 43 | Implemented |
| Shuffling and batching | Writes it in place | 54 | Derived from implementation |
| Step counts | Computes it, Calls it | 54, 55, 64 | Derived from implementation |
| Gradient clipping | Calls it | 62 | Implemented |
| Learning-rate schedule | Calls it | 64 | Implemented |
| AdamW update | Calls it | 66 | Implemented |
| Perplexity | Calls it | 71 | Implemented |
def batch_loss_and_grads(params, xb, yb, cfg: ModelConfig)Lines 12-24GitHub: train.batch_loss_and_gradsLoops over a batch and averages loss and gradients
train.maintrain.maintransformer.loss_and_gradsMean cross-entropy over all validation windows
train.maintrain.mainops.cross_entropy_forward, transformer.forwardFull training pipeline; entry point at lines 84-85
train.maintrain.<module>checkpoint.save_checkpoint, data.encode_corpus, data.iterate_batches, data.load_corpus, data.make_windows, data.split_ids, generate.generate, ops.perplexity, optim.Adam, optim.Adam.__init__, optim.Adam.step, optim.clip_grad_norm, optim.lr_schedule, tokenizer.Tokenizer.decode, tokenizer.build_tokenizer, train.batch_loss_and_grads, train.validation_loss, transformer.ModelConfig, transformer.init_paramsimport mathimport numpy as npimport configfrom data import load_corpus, encode_corpus, split_ids, make_windows, iterate_batchesfrom checkpoint import save_checkpointfrom generate import generatefrom ops import cross_entropy_forward, perplexityfrom optim import Adam, clip_grad_norm, lr_schedulefrom tokenizer import build_tokenizerfrom transformer import ModelConfig, init_params, forward, loss_and_gradsdef batch_loss_and_grads(params, xb, yb, cfg: ModelConfig): B = len(xb) grads = {k: np.zeros_like(v) for k, v in params.items()} total_loss = 0.0 for x, y in zip(xb, yb): loss, g = loss_and_grads(params, x, y, cfg) total_loss += loss / B for k in grads: grads[k] += g[k] / B return total_loss, gradsdef validation_loss(params, cfg, inputs, targets): losses = [] for x, y in zip(inputs.to_numpy(), targets.to_numpy()): logits, _ = forward(params, x, cfg) loss, _ = cross_entropy_forward(logits, y) losses.append(loss) return float(np.mean(losses))def main(prefix="checkpoints/model"): rng = np.random.default_rng(config.seed) # 1. text -> ids -> train/val windows (split BEFORE windowing) text = load_corpus(config.corpus_path, config.corpus_column, config.max_rows) tok = build_tokenizer(text, config.num_merges) ids = encode_corpus(text, tok) train_ids, val_ids = split_ids(ids, config.val_fraction) tr_x, tr_y = make_windows(train_ids, config.seq_len, config.stride) va_x, va_y = make_windows(val_ids, config.seq_len, config.stride) # 2. model + optimiser cfg = ModelConfig(tok.vocab_size, config.context_len, config.emb_dim, config.num_heads, config.ff_dim, config.num_layers) params = init_params(cfg, rng) opt = Adam(params, config.beta1, config.beta2, config.adam_eps, config.weight_decay) print(f"V={cfg.vocab_size} ln V={math.log(cfg.vocab_size):.3f}") steps_per_epoch = math.ceil(len(tr_x) / config.batch_size) total_steps = config.epochs * steps_per_epoch # 3. the loop: gradient -> clip -> learning rate -> Adam step = 0 for epoch in range(config.epochs): for xb, yb in iterate_batches(tr_x, tr_y, config.batch_size, rng): loss, grads = batch_loss_and_grads(params, xb, yb, cfg) norm = clip_grad_norm(grads, config.grad_clip) step += 1 lr = lr_schedule(step, config.warmup_steps, total_steps, config.lr_max, config.lr_min) opt.step(params, grads, lr) if step == 1 or step % 50 == 0: print(f"step {step:5d} loss {loss:.3f} |g| {norm:.2f} lr {lr:.5f}") v = validation_loss(params, cfg, va_x, va_y) print(f"epoch {epoch + 1}: val loss {v:.3f} ppl {perplexity(v):.1f}") seed_ids = [int(i) for i in va_x.iloc[0][:3]] out = generate(params, cfg, seed_ids, 30, config.seq_len, config.temperature, config.top_k, config.top_p, rng) print(" sample:", repr(tok.decode(out))) # 4. keep the trained model so talk.py can load it later save_checkpoint(prefix, params, cfg, tok) print(f"saved checkpoint to {prefix}.*") return params, cfg, tokif __name__ == "__main__": main()