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 <noreply@anthropic.com>

3 files changed, 124 insertions, 34 deletions

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<TrainArgs> for train::Config
@@ -171,6 +174,7 @@ impl From<TrainArgs> 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>, f32)
+	{
+		let parts: Vec<(Vec<f32>, 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<f32> = 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
 	}
 
@@ -380,6 +411,57 @@ mod tests
 	}
 
 	#[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::<f32>();
+			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()
 	{
 		let mut net = Network::random(1);
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<f32>
 /// Runs `f(0..n)` across all cores, returning the results in order.
 fn parallel<T: Send>(n: usize, f: impl Fn(usize) -> T + Sync) -> Vec<T>
 {
-	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<Sample>; 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<Sample>; 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(())
 }