From 03f11e2ed806d35367e31336afd636fb5ebed67f Mon Sep 17 00:00:00 2001 From: DragonDuck24 Date: Sat, 17 Oct 2026 15:34:35 -0400 Subject: [PATCH] improvements --- Cargo.lock | 36 ++++++----------- engine/Cargo.toml | 2 +- engine/src/main.rs | 35 +++++++++++------ engine/src/mcts.rs | 20 ++++++++-- engine/src/net/model.rs | 85 ++++++++++++++++++++++++----------------- engine/src/worker.rs | 37 +++++++++++------- 6 files changed, 126 insertions(+), 89 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7d9cc5d..44910c2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2170,7 +2170,7 @@ dependencies = [ "chess", "ctrlc", "futures", - "moka", + "lru 0.12.5", "rand 0.10.1", "rand_distr", "serde", @@ -3531,6 +3531,15 @@ dependencies = [ "imgref", ] +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru" version = "0.16.4" @@ -3674,23 +3683,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b" -[[package]] -name = "moka" -version = "0.12.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" -dependencies = [ - "crossbeam-channel", - "crossbeam-epoch", - "crossbeam-utils", - "equivalent", - "parking_lot", - "portable-atomic", - "smallvec", - "tagptr", - "uuid", -] - [[package]] name = "moxcms" version = "0.8.1" @@ -4726,7 +4718,7 @@ dependencies = [ "indoc", "itertools", "kasuari", - "lru", + "lru 0.16.4", "strum 0.27.2", "thiserror 2.0.18", "unicode-segmentation", @@ -5643,12 +5635,6 @@ dependencies = [ "windows", ] -[[package]] -name = "tagptr" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" - [[package]] name = "tar" version = "0.4.46" diff --git a/engine/Cargo.toml b/engine/Cargo.toml index d876bf8..9d826aa 100644 --- a/engine/Cargo.toml +++ b/engine/Cargo.toml @@ -13,6 +13,6 @@ serde_json = "1.0.150" serde = { version = "1.0.228", features = ["derive"] } tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread", "time"] } -moka = { version = "0.12.16", features = ["sync"] } +lru = "0.12" futures = "0.3" diff --git a/engine/src/main.rs b/engine/src/main.rs index 18ac96f..9476621 100644 --- a/engine/src/main.rs +++ b/engine/src/main.rs @@ -3,7 +3,7 @@ use burn::optim::AdamConfig; use burn::tensor::{Device, DeviceIndex}; use engine::mcts::MctsConfig; -use engine::train::{train, TrainingConfig}; +use engine::train::{TrainingConfig, train}; // fn main() { // type MyBackend = Wgpu; // @@ -14,9 +14,20 @@ use engine::train::{train, TrainingConfig}; // } fn main() { + // Ensure spawned background threads (e.g. CubeCL device servers like DSD-0-0) have sufficient stack space + unsafe { + if std::env::var("RUST_MIN_STACK").is_err() { + std::env::set_var("RUST_MIN_STACK", "33554432"); // 32MB + } + } + let device = Device::cuda(DeviceIndex::default()); - let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 16); + // Fast-validation config: 400 sims + 48 games ≈ 22 min/iteration on a 3080. + // Early training with a near-random model benefits more from data diversity + // than from deep search. Once the model has real shape (~150 iterations), + // bump back to MctsConfig::new(800, ..., 32) and num_episodes: 96. + let mcts_config = MctsConfig::new(400, 1.0, 0.3, 0.25, 16); let adam_config = AdamConfig::new(); @@ -24,23 +35,23 @@ fn main() { max_time_s: None, num_iters: None, max_depth: 100, // unused - model_name: String::from("96"), - load_model: true, - hidden_channels: 48, - num_blocks: 4, - batch_size: 1024, // num positions sampled to a batch (2048) + model_name: String::from("128x8"), + load_model: true, // architecture changed — must start with fresh weights + hidden_channels: 128, + num_blocks: 8, + batch_size: 1024, minibatch_size: 256, - num_train_steps: 10, - num_episodes: 96, // num games generated per iteration (5000) - buffer_max_size: 200_000, // max number of samples in the buffer + num_train_steps: 40, + num_episodes: 48, + buffer_max_size: 200_000, mcts_config, optimizer: adam_config, lr: 2e-4, seed: None, num_workers: 8, nn_batch_size: 64, - episodes_per_worker: 12, - nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB. + episodes_per_worker: 6, + nn_cache_size: 1000_000, }; train(training_config, device); diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index 37165e8..0279452 100644 --- a/engine/src/mcts.rs +++ b/engine/src/mcts.rs @@ -185,6 +185,13 @@ impl Mcts { let (root_policy, root_value) = model_inference(&board_state).await; self.expand_with_eval(root, &mut nodes, &root_policy, root_value); + // Seed the root node with the network's own value estimate. + // Without this, the root starts at (visit_count=0, value_sum=0) and the + // first ~N simulations have no reliable value signal — the root value is + // otherwise silently discarded. + let root_color = board_state.board.side_to_move(); + self.backpropagate(&mut nodes, &[root], root_value, root_color); + if training { // APPLY DIRICHLET NOISE HERE self.add_dirichlet_noise(root, &mut nodes); @@ -244,14 +251,19 @@ impl Mcts { } } - // Revert virtual loss, expand leaves, and backpropagate - for (path, leaf_idx, term_val) in pending { - // Revert virtual loss - for &idx in &path { + // Phase 1: Revert ALL virtual losses before touching the tree with real values. + // Doing this in one pass ensures that when we backpropagate simulation N, + // no other simulation's artificial visit counts or value sums are polluting + // the nodes on its path. + for (path, _, _) in &pending { + for &idx in path { nodes[idx].visit_count -= 1; nodes[idx].value_sum -= 1.0; } + } + // Phase 2: Expand leaves and backpropagate with a clean tree. + for (path, leaf_idx, term_val) in pending { let (value, color) = if let Some(tv) = term_val { let color = nodes[leaf_idx].board_state.board.side_to_move(); (tv, color) diff --git a/engine/src/net/model.rs b/engine/src/net/model.rs index 79f9efe..3233330 100644 --- a/engine/src/net/model.rs +++ b/engine/src/net/model.rs @@ -1,5 +1,4 @@ use crate::mcts::{BoardState, MctsResults}; -use crate::net::encoding::encode_board_state_perspective; use burn::data::dataloader::batcher::Batcher; use burn::nn::conv::Conv2dConfig; use burn::nn::loss::{MseLoss, Reduction}; @@ -304,44 +303,38 @@ pub struct ChessBatch { impl Batcher for ChessBatcher { fn batch(&self, items: Vec, device: &Device) -> ChessBatch { - let state_tensors = items - .iter() - .map(|item| { - encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 8, 8]) - }) - .collect::>(); + let b = items.len(); - let policy_target_tensors = items - .iter() - .cloned() - .map(|item| { - let mut policy = vec![0.0f32; 4672]; - for (idx, prob) in item.policy_target.iter() { - policy[*idx] = *prob; + let mut states_buf = vec![0.0f32; b * 1216]; + let mut policy_buf = vec![0.0f32; b * 4672]; + let mut value_buf = vec![0.0f32; b]; + + for (i, item) in items.iter().enumerate() { + crate::net::encoding::encode_board_state_to_slice( + &item.board_state, + &mut states_buf[i * 1216..(i + 1) * 1216], + ); + + let policy_slice = &mut policy_buf[i * 4672..(i + 1) * 4672]; + for &(idx, prob) in &item.policy_target { + if idx < 4672 { + policy_slice[idx] = prob; } - - // Normalize - let sum: f32 = policy.iter().sum(); - if sum > 0.0 { - for p in &mut policy { - *p /= sum; - } + } + let sum: f32 = policy_slice.iter().sum(); + if sum > 0.0 { + for p in policy_slice.iter_mut() { + *p /= sum; } - Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device) - }) - .collect::>(); + } - let value_target_tensors = items - .iter() - .map(|item| { - Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device) - .reshape([1, 1]) - }) - .collect::>(); + value_buf[i] = item.value_target; + } - let states = Tensor::cat(state_tensors, 0); // [B, 19, 8, 8] - let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672] - let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1] + let states = Tensor::<4>::from_floats(TensorData::new(states_buf, [b, 19, 8, 8]), device); + let policy_targets = + Tensor::<2>::from_floats(TensorData::new(policy_buf, [b, 4672]), device); + let value_targets = Tensor::<2>::from_floats(TensorData::new(value_buf, [b, 1]), device); ChessBatch { states, @@ -350,3 +343,27 @@ impl Batcher for ChessBatcher { } } } + +#[cfg(test)] +mod tests { + use super::*; + use burn::optim::{AdamConfig, GradientsParams}; + use burn::tensor::{Device, DeviceIndex}; + + #[test] + fn test_train_step() { + let device = Device::cuda(DeviceIndex::default()); + let autodiff_device = device.clone().autodiff(); + let model = ChessModelConfig::init(8, 128, &autodiff_device); + let mut optim = AdamConfig::new().init(); + let batcher = ChessBatcher {}; + let samples: Vec = (0..256) + .map(|_| TrainingSample::new(BoardState::default(), vec![(0, 1.0)], 0.0)) + .collect(); + let batch = batcher.batch(samples, &autodiff_device); + let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets); + let grads = output.loss.backward(); + let grads = GradientsParams::from_grads(grads, &model); + let _model = optim.step(2e-4, model, grads); + } +} diff --git a/engine/src/worker.rs b/engine/src/worker.rs index c519c9f..20b051b 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -4,9 +4,9 @@ use crate::train::{apply_temperature, sample_move}; use burn::tensor::Device; use chess::Color; -use moka::sync::Cache; -use rand::prelude::SmallRng; +use lru::LruCache; use rand::SeedableRng; +use rand::prelude::SmallRng; use std::collections::VecDeque; use std::hash::Hash; use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender}; @@ -83,8 +83,12 @@ impl Worker { // println!("{}", before.elapsed().as_millis()); episode_buffer.push(results); - let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3); - // let temp = if ply_count < 30 { 1.0 } else { 0.2 }; + // AlphaZero temperature schedule: explore freely for the first 30 plies, + // then play deterministically (argmax) for the rest of the game. + // The previous schedule never dropped below 0.3, keeping games noisy + // throughout and contaminating value targets with random late-game play. + let temp = if ply_count < 30 { 1.0 } else { 0.0 }; + // let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3); let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp); @@ -125,7 +129,8 @@ impl Worker { _ => 0.0, }; - let blended_target = 0.75 * outcome + 0.25 * result.value; + let blended_target = outcome; + // let blended_target = 0.75 * outcome + 0.25 * result.value; samples.push(TrainingSample::from_mcts_with_outcome( result.clone(), @@ -235,7 +240,6 @@ impl Worker { println!("Worker {} shut down", self.id); } - } async fn model_inference_static( @@ -307,6 +311,7 @@ pub struct NNRequest { response_tx: tokio::sync::oneshot::Sender, } +// ~18.7KB #[derive(Debug, Clone)] pub struct NNResponse { policy: Vec, @@ -315,7 +320,9 @@ pub struct NNResponse { impl NNServer { pub fn run_server(&self) { - let mut cache: Cache = Cache::new(self.cache_size); + let cap = + std::num::NonZeroUsize::new(self.cache_size as usize).expect("cache_size must be > 0"); + let mut cache: LruCache = LruCache::new(cap); let mut queue: VecDeque = VecDeque::new(); let mut request_counter: usize = 0; let mut cache_hits: usize = 0; @@ -327,7 +334,7 @@ impl NNServer { match cmd { NNServerCmd::Shutdown => { println!( - "NN server shutting down. Served {} requests, with {} cache hits ({:2}%). Ran {} batches (avg {}requests/batch).", + "NN server shutting down. Served {} requests, with {} cache hits ({:.2}%). Ran {} batches (avg {:.2} requests/batch).", request_counter, cache_hits, cache_hits as f64 / request_counter as f64 * 100.0, @@ -346,6 +353,7 @@ impl NNServer { match self.request_rx.recv_timeout(Duration::from_millis(10)) { Ok(first) => { request_counter += 1; + // queue.push_back(first); if try_use_cache(first, &mut cache, &mut queue) { cache_hits += 1; continue; @@ -365,9 +373,9 @@ impl NNServer { while let Ok(req) = self.request_rx.try_recv() { request_counter += 1; + // queue.push_back(req); if try_use_cache(req, &mut cache, &mut queue) { cache_hits += 1; - continue; } } @@ -389,7 +397,10 @@ impl NNServer { let mut buffer = vec![0.0f32; b_len * 1216]; for (i, job) in jobs_to_send.iter().enumerate() { let start = i * 1216; - crate::net::encoding::encode_board_state_to_slice(&job.board_state, &mut buffer[start..start + 1216]); + crate::net::encoding::encode_board_state_to_slice( + &job.board_state, + &mut buffer[start..start + 1216], + ); } let batched_tensor = burn::tensor::Tensor::<4>::from_floats( @@ -414,7 +425,7 @@ impl NNServer { value: v_val, }; - cache.insert(job.cache_key.clone(), res.clone()); + cache.put(job.cache_key, res.clone()); let _ = job.response_tx.send(res); } @@ -424,12 +435,12 @@ impl NNServer { fn try_use_cache( req: NNRequest, - cache: &mut Cache, + cache: &mut LruCache, queue: &mut VecDeque, ) -> bool { if let Some(cached) = cache.get(&req.cache_key) { // println!("Cache hit"); - req.response_tx.send(cached).unwrap(); + req.response_tx.send(cached.clone()).unwrap(); true } else { queue.push_back(req);