From 478e7568f5eda9780d1409bb8dccc86221126281 Mon Sep 17 00:00:00 2001 From: Dylan Date: Mon, 9 Nov 2026 16:38:13 +1300 Subject: feat: Multi-threaded training; play matches every 10 generations Gradients and losses are computed on rayon's pool in fixed-size chunks summed in order, so the weights are the same on any number of threads. Self-play uses the same pool, and the output shows the training time. The matches against the earlier network and ClaudePlayer now run only every --match-every generations (default 10), against the generation that last played them, read back from its saved file. Co-Authored-By: Claude Opus 5.5 --- src/main.rs | 6 +++- src/network.rs | 88 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++-- src/train.rs | 64 ++++++++++++++++++++++-------------------- 3 files changed, 124 insertions(+), 34 deletions(-) (limited to 'src') diff --git a/src/main.rs b/src/main.rs index 5a3c798..5237100 100644 --- a/src/main.rs +++ b/src/main.rs @@ -145,9 +145,12 @@ struct TrainArgs /// Adam's learning rate #[arg(long, default_value_t = train::Config::default().learning_rate)] learning_rate: f32, - /// Games against ClaudePlayer, and against the previous generation, after each generation + /// Games against ClaudePlayer, and against the generation that last played them #[arg(long, default_value_t = train::Config::default().match_games)] match_games: usize, + /// Play those matches only every this many generations + #[arg(long, default_value_t = train::Config::default().match_every)] + match_every: usize, } impl From for train::Config @@ -171,6 +174,7 @@ impl From for train::Config batch_size: args.batch_size, learning_rate: args.learning_rate, match_games: args.match_games, + match_every: args.match_every, } } } diff --git a/src/network.rs b/src/network.rs index 5836798..0463ccf 100644 --- a/src/network.rs +++ b/src/network.rs @@ -1,5 +1,7 @@ use std::{fs, io, path::Path}; +use rayon::prelude::*; + use crate::board::{GameBoard, Turn}; /// Inputs: a bit per point for tigers (0..25) and for sheep (25..50), the sheep in hand over 20 @@ -24,6 +26,12 @@ const W3: usize = B2 + HIDDEN2; const B3: usize = W3 + HIDDEN2; pub const PARAMS: usize = B3 + 1; +/// Samples per parallel task in `gradient`. The chunks are a fixed size and summed in order, so +/// the result is the same whatever the number of threads. +const GRAD_CHUNK: usize = 32; +/// Samples per parallel task in `loss`, where each one is only a forward pass. +const LOSS_CHUNK: usize = 1024; + const MAGIC: &[u8; 4] = b"TSNN"; const VERSION: u32 = 1; @@ -224,11 +232,30 @@ impl Network err * err / 2.0 } + /// The gradient of the error summed over `batch`, and that summed error, worked out in + /// `GRAD_CHUNK`-sized pieces across all cores. + fn gradient(&self, batch: &[Sample]) -> (Vec, f32) + { + let parts: Vec<(Vec, f32)> = batch.par_chunks(GRAD_CHUNK).map(|chunk| + { + let mut grad = vec![0.0; PARAMS]; + let loss = chunk.iter().map(|s| self.add_gradient(s, &mut grad)).sum(); + (grad, loss) + }).collect(); + let mut grad = vec![0.0; PARAMS]; + let mut loss = 0.0; + for (g, l) in parts + { + grad.iter_mut().zip(&g).for_each(|(a, b)| *a += b); + loss += l; + } + (grad, loss) + } + /// One optimiser step on the mean error over `batch`; returns that mean error. pub fn train_batch(&mut self, batch: &[Sample], adam: &mut Adam) -> f32 { - let mut grad = vec![0.0; PARAMS]; - let loss: f32 = batch.iter().map(|s| self.add_gradient(s, &mut grad)).sum(); + let (mut grad, loss) = self.gradient(batch); let n = batch.len().max(1) as f32; for g in &mut grad { @@ -241,7 +268,11 @@ impl Network /// The mean error over `samples`, without training. pub fn loss(&self, samples: &[Sample]) -> f32 { - let total: f32 = samples.iter().map(|s| { let e = self.forward(&s.board, s.turn) - s.target; e * e / 2.0 }).sum(); + let parts: Vec = samples.par_chunks(LOSS_CHUNK) + .map(|chunk| chunk.iter().map(|s| { let e = self.forward(&s.board, s.turn) - s.target; e * e / 2.0 }).sum()) + .collect(); + // A fold from 0, since an empty f32 `sum` is -0 + let total = parts.into_iter().fold(0.0, |a, b| a + b); total / samples.len().max(1) as f32 } @@ -379,6 +410,57 @@ mod tests assert!(checked > 3000); } + #[test] + fn parallel_gradient_matches_serial() + { + let net = Network::random(5); + // Not a multiple of GRAD_CHUNK, so the last chunk is short + let batch = samples(GRAD_CHUNK * 4 + 7, 9); + let (grad, loss) = net.gradient(&batch); + + // The same chunks, one after another, give exactly the same sums + let mut serial = vec![0.0; PARAMS]; + let mut serial_loss = 0.0; + for chunk in batch.chunks(GRAD_CHUNK) + { + let mut g = vec![0.0; PARAMS]; + serial_loss += chunk.iter().map(|s| net.add_gradient(s, &mut g)).sum::(); + serial.iter_mut().zip(&g).for_each(|(a, b)| *a += b); + } + assert_eq!(grad, serial); + assert_eq!(loss, serial_loss); + + // And one running sum gives nearly the same + let mut plain = vec![0.0; PARAMS]; + let plain_loss: f32 = batch.iter().map(|s| net.add_gradient(s, &mut plain)).sum(); + assert!((loss - plain_loss).abs() < 1e-4 * plain_loss.abs().max(1.0)); + for (a, b) in grad.iter().zip(&plain) + { + assert!((a - b).abs() < 1e-4 * b.abs().max(1.0), "{a} vs {b}"); + } + } + + #[test] + fn training_is_the_same_on_any_number_of_threads() + { + let train = |threads: usize| + { + let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap(); + pool.install(|| + { + let mut net = Network::random(2); + let data = samples(300, 4); + let mut adam = Adam::new(1e-3); + for batch in data.chunks(100) + { + net.train_batch(batch, &mut adam); + } + (net, Network::random(3).loss(&samples(3000, 6))) + }) + }; + assert_eq!(train(1), train(4)); + } + #[test] fn training_can_fit_a_few_samples() { diff --git a/src/train.rs b/src/train.rs index 1e68080..35f8fdd 100644 --- a/src/train.rs +++ b/src/train.rs @@ -1,4 +1,6 @@ -use std::{collections::VecDeque, fs, path::{Path, PathBuf}, sync::atomic::{AtomicUsize, Ordering}, time::{Duration, Instant}}; +use std::{collections::VecDeque, fs, path::{Path, PathBuf}, time::{Duration, Instant}}; + +use rayon::prelude::*; use crate::{board::{repetitions, GameBoard, Move, Outcome, Turn, REPETITION_DRAW, SYMMETRIES}, claude_player::ClaudePlayer, network::{Adam, Network, Rng, Sample}, neural_player::NeuralPlayer, player::{MinimaxPlayer, Player}}; @@ -29,8 +31,11 @@ pub struct Config pub epochs: usize, pub batch_size: usize, pub learning_rate: f32, - /// Games against `ClaudePlayer`, and against the previous generation, after each generation (half as each side). + /// Games against `ClaudePlayer`, and against the generation that last played these matches (half as each side). pub match_games: usize, + /// Play those matches only on generations whose number is a multiple of this, so generation + /// N plays generation N - `match_every`. + pub match_every: usize, } impl Default for Config @@ -54,6 +59,7 @@ impl Default for Config batch_size: 256, learning_rate: 1e-3, match_games: 100, + match_every: 10, } } } @@ -128,27 +134,7 @@ fn td_targets(plies: &[Ply], outcome: Outcome, lambda: f32) -> Vec /// Runs `f(0..n)` across all cores, returning the results in order. fn parallel(n: usize, f: impl Fn(usize) -> T + Sync) -> Vec { - let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(n.max(1)); - let next = AtomicUsize::new(0); - let mut results: Vec<(usize, T)> = std::thread::scope(|s| - { - let handles: Vec<_> = (0..threads).map(|_| s.spawn(|| - { - let mut done = Vec::new(); - loop - { - let i = next.fetch_add(1, Ordering::Relaxed); - if i >= n - { - return done; - } - done.push((i, f(i))); - } - })).collect(); - handles.into_iter().flat_map(|h| h.join().expect("worker thread panicked")).collect() - }); - results.sort_by_key(|(i, _)| *i); - results.into_iter().map(|(_, r)| r).collect() + (0..n).into_par_iter().map(&f).collect() } /// A player that searches exactly `depth` plies. @@ -270,13 +256,15 @@ pub fn run(cfg: &Config) -> std::io::Result<()> } None => Network::random(0x5EED), }; + // The opponent for the first matches, if the generation they should be against isn't saved + let initial = net.clone(); let mut adam = Adam::new(cfg.learning_rate); // For each generation, its positions from tiger-win, sheep-win and drawn games let mut buffer: VecDeque<[Vec; 3]> = VecDeque::new(); let mut rng = Rng::new(first as u64); println!("Self-play results, then the mean error on positions from tiger-win (T) and sheep-win (S) games, before and after training"); - println!("{:>4} {:>7} {:>6} {:>6} {:>6} {:>8} {:>8} {:>8} {:>8} {:>7} {:>9} {:>7}", - "gen", "trained", "tigers", "sheep", "draws", "T before", "S before", "T after", "S after", "vs prev", "vs Claude", "time s"); + println!("{:>4} {:>7} {:>6} {:>6} {:>6} {:>8} {:>8} {:>8} {:>8} {:>7} {:>9} {:>7} {:>7}", + "gen", "trained", "tigers", "sheep", "draws", "T before", "S before", "T after", "S after", "vs prev", "vs Claude", "train s", "time s"); for generation in first..first + cfg.generations { let start = Instant::now(); @@ -290,10 +278,11 @@ pub fn run(cfg: &Config) -> std::io::Result<()> buffer.pop_front(); } + // Training: the losses before and after, and every epoch + let train_start = Instant::now(); let pools: [Vec; 3] = [0, 1, 2].map(|k| buffer.iter().flat_map(|g| &g[k]).copied().collect()); let losses = |net: &Network| (net.loss(&pools[0]), net.loss(&pools[1])); let before = losses(&net); - let previous = net.clone(); let mut trained = 0; for _ in 0..cfg.epochs { @@ -314,13 +303,28 @@ pub fn run(cfg: &Config) -> std::io::Result<()> } } let after = losses(&net); + let train_time = train_start.elapsed(); - let vs_prev = match_score(&net, || neural(&previous, cfg.depth), cfg, seed_base ^ 0xA); - let vs_claude = match_score(&net, || claude(cfg.depth), cfg, seed_base ^ 0xB); + // The matches are slow and only measure progress, so they are only played now and then + let every = cfg.match_every.max(1); + let (vs_prev, vs_claude) = if generation.is_multiple_of(every) + { + // Against the generation that last played them, read back from its file, so a rerun + // that carries on from the out dir still finds it + let last = cfg.out_dir.join(format!("gen-{:03}.bin", generation - every)); + let previous = Network::load(&last).unwrap_or_else(|_| initial.clone()); + let percent = |score: f32| format!("{:.0}%", score * 100.0); + (percent(match_score(&net, || neural(&previous, cfg.depth), cfg, seed_base ^ 0xA)), + percent(match_score(&net, || claude(cfg.depth), cfg, seed_base ^ 0xB))) + } + else + { + ("-".to_string(), "-".to_string()) + }; net.save(&cfg.out_dir.join(format!("gen-{generation:03}.bin")))?; net.save(&cfg.out_dir.join("latest.bin"))?; - println!("{generation:>4} {trained:>7} {tigers:>6} {sheep:>6} {draws:>6} {:>8.4} {:>8.4} {:>8.4} {:>8.4} {:>6.0}% {:>8.0}% {:>7.1}", - before.0, before.1, after.0, after.1, vs_prev * 100.0, vs_claude * 100.0, start.elapsed().as_secs_f64()); + println!("{generation:>4} {trained:>7} {tigers:>6} {sheep:>6} {draws:>6} {:>8.4} {:>8.4} {:>8.4} {:>8.4} {vs_prev:>7} {vs_claude:>9} {:>7.2} {:>7.1}", + before.0, before.1, after.0, after.1, train_time.as_secs_f64(), start.elapsed().as_secs_f64()); } Ok(()) } -- cgit v1.3.1