topos 0.12.0

//! Trains the same character-level MLP as `makemore_mlp_facade`, data
//! parallel: every step splits its minibatch into equal shards, runs
//! one forward and backward pass per shard concurrently on the shared
//! network, and averages the shard gradients into a single update.
//!
//! This leans on the engine's concurrency contract: the sealed network
//! is immutable, feeds are run state rather than graph state, and each
//! shard's backward projects onto the shared parameter slots
//! (`Field::parameters`), so gradients from concurrent runs over one
//! parameter state combine with the `Parameters` algebra — summing
//! slot-sized tables, never graph-sized buffers. Because `cross_entropy`
//! normalizes each shard by its own mass, the average of equal-sized
//! shard gradients equals the full-batch gradient exactly, and summing
//! the shards in a fixed pairwise tree keeps the run deterministic
//! regardless of thread scheduling.
//!
//! Run with: `cargo run --release --example makemore_mlp_parallel`

mod chart;
mod corpus;

use std::time::Instant;

use rayon::prelude::*;

use topos::{Activation, Mlp, Module, Parameters, Shape, Tape, Tensor, 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 sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many concurrent shards each training step fans out.
///
/// The count is fixed rather than detected — the shard partition
/// decides the arithmetic, so a detected count would make runs differ
/// across machines. Rayon still schedules the fixed shards over every
/// core it has. The count trades per-run fixed cost (favoring fewer,
/// larger shards) against load balancing over heterogeneous cores
/// (favoring more, smaller ones); eight-by-eight measured fastest on an
/// eight-performance-core machine, where sixteen-by-four raised
/// utilization but paid more in run overhead than it recovered.
const SHARD_COUNT: usize = 8;

/// How many samples each shard carries; the effective batch is
/// `SHARD_COUNT * SHARD_LEN`, matching the serial examples' 64.
const SHARD_LEN: usize = 8;

/// Sums shard gradients as a pairwise tree whose shape depends only on
/// the shard count: the reduction runs its pairs concurrently and
/// finishes in logarithmic depth, while the tree — not the scheduler —
/// decides the order of additions, keeping the result deterministic.
fn tree_sum(mut layer: Vec<Parameters<f32>>) -> Parameters<f32> {
    while layer.len() > 1 {
        layer = layer
            .par_chunks(2)
            .map(|pair| match pair {
                [left, right] => left + right,
                [single] => single.clone(),
                _ => unreachable!("chunks of two hold one or two tables"),
            })
            .collect();
    }
    layer.into_iter().next().expect("at least one shard ran")
}

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<f32> = Tape::new();

    // The same model as the serial examples, recorded at shard shape:
    // the batch size is baked into the graph, so the parallel plan is
    // one shard-shaped expression run once per shard, not a wider one.
    let embeddings = tape.parameter(init::normal(8, 1.0)(&Shape::new([
        VOCABULARY_LEN,
        EMBED_DIM,
    ])));
    let mlp = Mlp::new(
        &tape,
        &[CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN, VOCABULARY_LEN],
        Activation::Tanh,
        init::xavier(7),
    );

    let contexts = tape.input(Tensor::selection(
        vec![0; SHARD_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; SHARD_LEN], VOCABULARY_LEN, 1.0));
    let embedded = embeddings
        .gather(contexts)
        .reshape([SHARD_LEN, CONTEXT_LEN * EMBED_DIM]);
    let loss = cross_entropy(mlp.express(embedded), targets);

    // The sampling twin: the same parameters expressed over a single
    // context row, with the composite softmax on top.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_embedded = embeddings
        .gather(sample_context)
        .reshape([1, CONTEXT_LEN * EMBED_DIM]);
    let sample_probabilities = mlp.express(sample_embedded).softmax(1);

    let contexts_symbol = contexts.symbol();
    let targets_symbol = targets.symbol();
    let loss_symbol = loss.symbol();
    let sample_context_symbol = sample_context.symbol();
    let sample_probabilities_symbol = sample_probabilities.symbol();
    let network = tape.into_network();
    let recorded_nodes = network.len();

    let batch_len = SHARD_COUNT * SHARD_LEN;
    let shard_inverse = Tensor::new([], [1.0 / SHARD_COUNT as f32]);
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut parameters = network.parameters();
    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];

        // Fan out: one immutable forward and backward run per shard,
        // all reading the same network and parameter state.
        let shard_results: Vec<(f32, Parameters<f32>)> = (0..SHARD_COUNT)
            .into_par_iter()
            .map(|shard| {
                let rows = &batch[shard * SHARD_LEN..(shard + 1) * SHARD_LEN];
                let shard_contexts: Vec<usize> = rows
                    .iter()
                    .flat_map(|(context, _)| context.iter().copied())
                    .collect();
                let shard_targets: Vec<usize> = rows.iter().map(|&(_, next)| next).collect();

                // Slice the run to the loss: the sampling twin on the
                // same tape is skipped during training.
                let run = network.entry([loss_symbol]).interpret(
                    &parameters,
                    [
                        (
                            contexts_symbol,
                            Tensor::selection(shard_contexts, VOCABULARY_LEN, 1.0),
                        ),
                        (
                            targets_symbol,
                            Tensor::selection(shard_targets, VOCABULARY_LEN, 1.0),
                        ),
                    ],
                );
                let shard_loss = run.of(loss_symbol).scalar();
                (
                    shard_loss,
                    run.backward(loss_symbol).parameters(&parameters),
                )
            })
            .collect();

        let (shard_losses, shard_gradients): (Vec<f32>, Vec<Parameters<f32>>) =
            shard_results.into_iter().unzip();
        let batch_loss = shard_losses.iter().sum::<f32>() / SHARD_COUNT as f32;
        let gradients = tree_sum(shard_gradients)
            .map(|gradient| gradient.clone() * shard_inverse.broadcast_like(gradient));

        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 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 on {SHARD_COUNT} shards 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("mlp (data parallel) 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 = network.entry([sample_probabilities_symbol]).interpret(
                &parameters,
                [(
                    sample_context_symbol,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities_symbol).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}");
    }
}