topos 0.12.0

//! Trains a small character-level transformer — the attention act of
//! the makemore series, and the consumer that closes the transformer
//! rung: every piece is composition over the existing op set.
//!
//! The batch packs its samples into one `[BATCH_LEN * CONTEXT_LEN]`
//! token row, so each head's attention is a single rank-2 matmul pair
//! over the packed axis; a block-diagonal causal mask (an additive
//! `0 / -inf` leaf) keeps samples independent and time causal, which
//! is the sequence-packing idiom and the reason no batched matmul is
//! needed. Masked softmax records as the mask added before the axis
//! softmax, heads record as a loop of rank-2 attentions joined by
//! `concat`, and the per-sample prediction rows come back through a
//! one-hot `gather`. The block is pre-norm: attention and feed-forward
//! each read an `RmsNorm` of their input and add onto the residual
//! stream, and mask-fed dropout (keep probability `KEEP`) guards both
//! residual writes — fed per training step from a seeded factory,
//! while the sampling twin runs unfed on the all-ones identity
//! default. Loss is measured at the last context position only, so
//! the number is comparable to the MLP acts (uniform cost:
//! ln 27 ~ 3.30).
//!
//! Run with: `cargo run --release --example makemore_transformer`

mod chart;
mod corpus;

use std::time::Instant;

use topos::{Dropout, Module, RmsNorm, Shape, Tape, Tensor, Value, concat, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model attends over before
/// predicting the next one.
const CONTEXT_LEN: usize = 8;

/// How many dimensions the residual stream has.
const EMBED_DIM: usize = 32;

/// How many attention heads split the stream.
const HEAD_COUNT: usize = 2;

/// How many dimensions each head reads and writes.
const HEAD_DIM: usize = EMBED_DIM / HEAD_COUNT;

/// How many neurons the feed-forward hidden layer has.
const HIDDEN_LEN: usize = 4 * EMBED_DIM;

/// How many samples each training step packs into the attention row.
const BATCH_LEN: usize = 32;

/// How many token rows the packed training batch holds.
const PACKED_LEN: usize = BATCH_LEN * CONTEXT_LEN;

/// The dropout keep probability: masks carry `1 / KEEP` or zero, and
/// only the training expression is ever fed — the sampling twin runs
/// on the all-ones identity default.
const KEEP: f64 = 0.9;

/// One attention head's projections, each `[EMBED_DIM, HEAD_DIM]`.
struct Head<'tape> {
    query: Value<'tape, f32>,
    key: Value<'tape, f32>,
    value: Value<'tape, f32>,
}

/// The model's parameters as recorded proxies: embeddings, one
/// pre-norm attention block, and the language-model head.
struct Model<'tape> {
    embeddings: Value<'tape, f32>,
    positions: Value<'tape, f32>,
    heads: Vec<Head<'tape>>,
    projection: Value<'tape, f32>,
    attention_norm: RmsNorm<f32>,
    hidden_weights: Value<'tape, f32>,
    output_weights: Value<'tape, f32>,
    hidden_norm: RmsNorm<f32>,
    final_norm: RmsNorm<f32>,
    logit_weights: Value<'tape, f32>,
    logit_bias: Value<'tape, f32>,
    scale: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: embedding tables for
    /// characters and positions, the heads' projections, the block's
    /// two norms and feed-forward, the final norm, and the affine
    /// logit head. The attention scale `1 / sqrt(HEAD_DIM)` rides
    /// along as a single-value leaf.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        let ones = Tensor::filled([EMBED_DIM], 1.0);
        let epsilon = Tensor::filled([], 1e-5);
        Self {
            embeddings: tape.parameter(init::normal(8, 1.0)(&Shape::new([
                VOCABULARY_LEN,
                EMBED_DIM,
            ]))),
            positions: tape.parameter(init::normal(9, 1.0)(&Shape::new([CONTEXT_LEN, EMBED_DIM]))),
            heads: (0..HEAD_COUNT)
                .map(|_| Head {
                    query: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                    key: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                    value: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                })
                .collect(),
            projection: tape.parameter(weights(&Shape::new([EMBED_DIM, EMBED_DIM]))),
            attention_norm: RmsNorm::new(tape, ones.clone(), epsilon.clone()),
            hidden_weights: tape.parameter(weights(&Shape::new([EMBED_DIM, HIDDEN_LEN]))),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, EMBED_DIM]))),
            hidden_norm: RmsNorm::new(tape, ones.clone(), epsilon.clone()),
            final_norm: RmsNorm::new(tape, ones, epsilon),
            logit_weights: tape.parameter(weights(&Shape::new([EMBED_DIM, VOCABULARY_LEN]))),
            logit_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
            scale: tape.leaf(Tensor::filled([], 1.0 / (HEAD_DIM as f32).sqrt())),
        }
    }

    /// Records the block over one packed token row and returns the
    /// normalized residual stream: embeddings in, attention and
    /// feed-forward added on, final norm out.
    ///
    /// `tokens` and `positions` are one-hot selections over the packed
    /// row; `mask` is the additive `0 / -inf` block-causal leaf shaped
    /// `[rows, rows]`.
    fn states(
        &self,
        _tape: &'tape Tape<f32>,
        tokens: Value<'tape, f32>,
        positions: Value<'tape, f32>,
        mask: Value<'tape, f32>,
        dropouts: &[Dropout<f32>; 2],
    ) -> Value<'tape, f32> {
        let stream = self.embeddings.gather(tokens) + self.positions.gather(positions);

        // Pre-norm attention: every head attends over the same
        // normalized stream, and `concat` joins the head outputs along
        // the feature axis.
        let normalized = self.attention_norm.express(stream);
        let heads: Vec<Value<'tape, f32>> = self
            .heads
            .iter()
            .map(|head| {
                let query = normalized.matmul(head.query);
                let key = normalized.matmul(head.key);
                let scores = query.matmul(key.transpose());
                let scaled = scores * self.scale.broadcast_like(scores);
                let weights = (scaled + mask).softmax(1);
                weights.matmul(normalized.matmul(head.value))
            })
            .collect();
        let stream = stream + dropouts[0].express(concat(&heads, 1).matmul(self.projection));

        // Pre-norm feed-forward onto the residual stream.
        let normalized = self.hidden_norm.express(stream);
        let stream = stream
            + dropouts[1].express(
                normalized
                    .matmul(self.hidden_weights)
                    .relu()
                    .matmul(self.output_weights),
            );

        self.final_norm.express(stream)
    }

    /// Records the logits of the rows picked by the one-hot
    /// `extraction` — the last context position of each packed sample.
    fn logits(
        &self,
        states: Value<'tape, f32>,
        extraction: Value<'tape, f32>,
    ) -> Value<'tape, f32> {
        let product = states.gather(extraction).matmul(self.logit_weights);
        product + self.logit_bias.broadcast_along_like(0, product)
    }
}

/// Returns the additive attention mask for `samples` packed windows of
/// `CONTEXT_LEN` tokens: zero where the key's row is in the same
/// sample and not after the query's row, negative infinity elsewhere.
fn block_causal_mask(samples: usize) -> Tensor<f32> {
    let rows = samples * CONTEXT_LEN;
    let mut elements = Vec::with_capacity(rows * rows);
    for query in 0..rows {
        for key in 0..rows {
            let same_sample = query / CONTEXT_LEN == key / CONTEXT_LEN;
            let causal = key % CONTEXT_LEN <= query % CONTEXT_LEN;
            elements.push(if same_sample && causal {
                0.0
            } else {
                f32::NEG_INFINITY
            });
        }
    }
    Tensor::new([rows, rows], elements)
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape = Tape::new();
    let model = Model::new(&tape);

    // The training expression over the packed batch: per-sample
    // prediction rows come back through the fixed one-hot extraction.
    let tokens = tape.input(Tensor::selection(vec![0; PACKED_LEN], VOCABULARY_LEN, 1.0));
    let positions = tape.leaf(Tensor::selection(
        (0..PACKED_LEN)
            .map(|row| row % CONTEXT_LEN)
            .collect::<Vec<_>>(),
        CONTEXT_LEN,
        1.0,
    ));
    let mask = tape.leaf(block_causal_mask(BATCH_LEN));
    let extraction = tape.leaf(Tensor::selection(
        (0..BATCH_LEN)
            .map(|sample| sample * CONTEXT_LEN + CONTEXT_LEN - 1)
            .collect::<Vec<_>>(),
        PACKED_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    // Dropout on both residual writes; the mask inputs default to
    // ones, so only the fed training steps ever drop anything.
    let dropouts = [
        Dropout::new(&tape, [PACKED_LEN, EMBED_DIM]),
        Dropout::new(&tape, [PACKED_LEN, EMBED_DIM]),
    ];
    let states = model.states(&tape, tokens, positions, mask, &dropouts);
    let loss = cross_entropy(model.logits(states, extraction), targets);

    // The sampling twin is the same expression over one window.
    let sample_tokens = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_positions = tape.leaf(Tensor::selection(
        (0..CONTEXT_LEN).collect::<Vec<_>>(),
        CONTEXT_LEN,
        1.0,
    ));
    let sample_mask = tape.leaf(block_causal_mask(1));
    let sample_extraction = tape.leaf(Tensor::selection(vec![CONTEXT_LEN - 1], CONTEXT_LEN, 1.0));
    // The twin's dropouts are never fed: the identity default is the
    // inference mode, with no flag and no second formula.
    let sample_dropouts = [
        Dropout::new(&tape, [CONTEXT_LEN, EMBED_DIM]),
        Dropout::new(&tape, [CONTEXT_LEN, EMBED_DIM]),
    ];
    let sample_states = model.states(
        &tape,
        sample_tokens,
        sample_positions,
        sample_mask,
        &sample_dropouts,
    );
    let sample_probabilities = model.logits(sample_states, sample_extraction).softmax(1);

    let (tokens, targets, loss, sample_tokens, sample_probabilities) = (
        tokens.symbol(),
        targets.symbol(),
        loss.symbol(),
        sample_tokens.symbol(),
        sample_probabilities.symbol(),
    );
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // Entry once: training keeps only the loss, sampling is
    // forward-only.
    let training_plan = network.entry([loss]).backward().lower();
    let sampling_plan = network.entry([sample_probabilities]).lower();

    // The seeded mask factory: two draws per step, one per residual
    // write, deterministic in `(seed, step)` so runs replay bitwise.
    let mut dropout_masks = init::dropout::<f32>(13, KEEP);
    let mask_shape = Shape::new([PACKED_LEN, EMBED_DIM]);

    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_tokens: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        let run = training_plan.forward(
            &parameters,
            [
                (tokens, Tensor::selection(batch_tokens, VOCABULARY_LEN, 1.0)),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
                (dropouts[0].mask(), dropout_masks(&mask_shape)),
                (dropouts[1].mask(), dropout_masks(&mask_shape)),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }

        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("transformer training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = sampling_plan.forward(
                &parameters,
                [(
                    sample_tokens,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}