From 38f13fa299e4943d3aed42d9e2a9dcc18e7380c2 Mon Sep 17 00:00:00 2001 From: DragonDuck24 Date: Sat, 3 Oct 2026 13:19:47 -0400 Subject: [PATCH] Virtual loss and batched mcts --- Cargo.lock | 28 ++++++++++ engine/Cargo.toml | 5 +- engine/src/main.rs | 4 +- engine/src/mcts.rs | 126 +++++++++++++++++++++++++++---------------- engine/src/train.rs | 25 ++++----- engine/src/worker.rs | 116 +++++++++++++++++++-------------------- 6 files changed, 178 insertions(+), 126 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6e49b24..7d9cc5d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2169,6 +2169,7 @@ dependencies = [ "burn", "chess", "ctrlc", + "futures", "moka", "rand 0.10.1", "rand_distr", @@ -2478,6 +2479,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.34" @@ -2494,6 +2510,17 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.34" @@ -2548,6 +2575,7 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", diff --git a/engine/Cargo.toml b/engine/Cargo.toml index b84420a..d876bf8 100644 --- a/engine/Cargo.toml +++ b/engine/Cargo.toml @@ -11,5 +11,8 @@ rand = "0.10.1" ctrlc = "3.5.2" serde_json = "1.0.150" serde = { version = "1.0.228", features = ["derive"] } -tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread"] } +tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread", "time"] } + moka = { version = "0.12.16", features = ["sync"] } +futures = "0.3" + diff --git a/engine/src/main.rs b/engine/src/main.rs index 1163813..18ac96f 100644 --- a/engine/src/main.rs +++ b/engine/src/main.rs @@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig}; fn main() { let device = Device::cuda(DeviceIndex::default()); - let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128); + let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 16); let adam_config = AdamConfig::new(); @@ -24,7 +24,7 @@ fn main() { max_time_s: None, num_iters: None, max_depth: 100, // unused - model_name: String::from("threadedModel"), + model_name: String::from("96"), load_model: true, hidden_channels: 48, num_blocks: 4, diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index 4da924a..37165e8 100644 --- a/engine/src/mcts.rs +++ b/engine/src/mcts.rs @@ -164,6 +164,8 @@ fn softmax(logits: &[f32]) -> Vec { probs } +use std::future::Future; + impl Mcts { pub async fn search( &mut self, @@ -180,7 +182,8 @@ impl Mcts { let root = 0; nodes.push(Node::new(0.0, board_state.clone(), None)); - self.expand(root, &mut nodes, &model_inference).await; + let (root_policy, root_value) = model_inference(&board_state).await; + self.expand_with_eval(root, &mut nodes, &root_policy, root_value); if training { // APPLY DIRICHLET NOISE HERE @@ -188,48 +191,81 @@ impl Mcts { } let num_sims = self.config.num_simulations; + let mut sim_count = 0; - for _ in 0..num_sims { - let mut path = vec![root]; - let mut current = root; + while sim_count < num_sims { + let current_batch_size = std::cmp::min(self.config.batch_size, num_sims - sim_count); + sim_count += current_batch_size; - // for batch in 0..self.config.num_batches { - // let mut pending = Vec::new(); - // - // // Select several leaves - // for _ in 0..self.config.batch_size { - // let path = select_leaf_with_virtual_loss(...); - // - // let board = nodes[path.last()].board_state.clone(); - // - // pending.push(PendingEval { - // path, - // board, - // }); - // } - // - // // Send all boards to NN server - // let evaluations = model_inference_batch(pending.boards()).await; - // - // // Expand and backpropagate - // for (pending, (policy, value)) in - // pending.into_iter().zip(evaluations) - // { - // expand_with_policy(...); - // remove_virtual_loss(...); - // backpropagate(...); - // } - // } + let mut pending = Vec::with_capacity(current_batch_size); - while !nodes[current].children.is_empty() { - current = nodes[current].select_child(&nodes, &self.config.c_puct); - path.push(current); + for _ in 0..current_batch_size { + let mut path = vec![root]; + let mut current = root; + + while !nodes[current].children.is_empty() { + current = nodes[current].select_child(&nodes, &self.config.c_puct); + path.push(current); + } + + // Apply virtual loss to path + for &idx in &path { + nodes[idx].visit_count += 1; + nodes[idx].value_sum += 1.0; + } + + let status = nodes[current].board_state.status; + let stm = nodes[current].board_state.board.side_to_move(); + let term_val = terminal_value(status, stm); + + pending.push((path, current, term_val)); } - let value: f32 = self.expand(current, &mut nodes, &model_inference).await; + // Gather unique non-terminal leaves needing evaluation + let mut eval_indices = Vec::new(); + for (_, leaf_idx, term_val) in &pending { + if term_val.is_none() && !eval_indices.contains(leaf_idx) { + eval_indices.push(*leaf_idx); + } + } - let color = nodes[current].board_state.board.side_to_move(); - self.backpropagate(&mut nodes, &path, value, color); + // Execute inference for all non-terminal leaves concurrently + let mut eval_map = HashMap::new(); + if !eval_indices.is_empty() { + let futures: Vec<_> = eval_indices + .iter() + .map(|&idx| model_inference(&nodes[idx].board_state)) + .collect(); + + let results = futures::future::join_all(futures).await; + + for (&idx, res) in eval_indices.iter().zip(results.into_iter()) { + eval_map.insert(idx, res); + } + } + + // Revert virtual loss, expand leaves, and backpropagate + for (path, leaf_idx, term_val) in pending { + // Revert virtual loss + for &idx in &path { + nodes[idx].visit_count -= 1; + nodes[idx].value_sum -= 1.0; + } + + let (value, color) = if let Some(tv) = term_val { + let color = nodes[leaf_idx].board_state.board.side_to_move(); + (tv, color) + } else { + let (policy, val) = &eval_map[&leaf_idx]; + if nodes[leaf_idx].children.is_empty() { + self.expand_with_eval(leaf_idx, &mut nodes, policy, *val); + } + let color = nodes[leaf_idx].board_state.board.side_to_move(); + (*val, color) + }; + + self.backpropagate(&mut nodes, &path, value, color); + } } // Build compact move distribution @@ -256,27 +292,22 @@ impl Mcts { MctsResults::new(board_state.clone(), move_dist, nodes[root].value()) } - async fn expand( + fn expand_with_eval( &mut self, node_idx: usize, arena: &mut Vec, - model_inference: &F, - ) -> f32 - where - F: Fn(&BoardState) -> Fut, - Fut: Future, f32)>, - { + policy_logits: &[f32], + value_head: f32, + ) -> f32 { let status = arena[node_idx].board_state.status; let stm = arena[node_idx].board_state.board.side_to_move(); if let Some(term_val) = terminal_value(status, stm) { term_val } else { - let (policy_logits, value_head) = model_inference(&arena[node_idx].board_state).await; - let legal_moves: Vec = MoveGen::new_legal(&arena[node_idx].board_state.board).collect(); - let probs = softmax(&policy_logits); + let probs = softmax(policy_logits); let mut move_priors = Vec::with_capacity(legal_moves.len()); let mut prior_sum = 0.0f32; @@ -308,6 +339,7 @@ impl Mcts { } } + fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) { for &idx in path { let node = &mut nodes[idx]; diff --git a/engine/src/train.rs b/engine/src/train.rs index a947d52..b6da12d 100644 --- a/engine/src/train.rs +++ b/engine/src/train.rs @@ -3,7 +3,7 @@ use crate::net::encoding::decode_move; use crate::net::model::{ ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample, }; -use crate::worker::{NNRequest, NNResponse, NNServer, NNServerCmd, Worker, WorkerCmd, WorkerData}; +use crate::worker::{NNRequest, NNServer, NNServerCmd, Worker, WorkerCmd, WorkerData}; use burn::data::dataloader::batcher::Batcher; use burn::module::AutodiffModule; use burn::optim::{AdamConfig, GradientsParams}; @@ -17,9 +17,10 @@ use std::cmp::min; use std::collections::{HashMap, VecDeque}; use std::fs::File; use std::io::BufReader; + use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::mpsc::{Receiver, Sender}; use std::sync::{mpsc, Arc}; + use std::thread; use std::time::Instant; use std::time::{SystemTime, UNIX_EPOCH}; @@ -124,18 +125,11 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let mut episodes_dispatched: usize = 0; let mut episodes_recieved: usize = 0; - let mut worker_map: HashMap> = HashMap::new(); + let mut worker_map: HashMap> = HashMap::new(); let (server_cmd_tx, server_cmd_rx) = mpsc::channel::(); let (data_tx, data_rx) = mpsc::channel::(); let (server_request_tx, server_request_rx) = mpsc::channel::(); - let mut nn_channels: VecDeque> = VecDeque::new(); - - for i in 0..training_config.num_workers { - let (tx, rx) = mpsc::channel::(); - nn_channels.push_back(rx); - } - let server_device = device.clone(); let server_handle = thread::spawn(move || { let server = NNServer::new( @@ -155,7 +149,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { for i in 0..training_config.num_workers { // println!("initializing worker {}", i); if training_config.num_episodes > episodes_dispatched { - let (cmd_tx, cmd_rx) = mpsc::channel::(); + let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::(100); worker_map.insert(i, cmd_tx); let episodes_to_dispatch = min( training_config.episodes_per_worker, @@ -163,7 +157,6 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { ); let req_tx = server_request_tx.clone(); - // let res_rx = nn_channels.pop_front().unwrap(); let data_tx_thread = data_tx.clone(); let mcts_config_thread = training_config.mcts_config.clone(); let device_thread = device.clone(); @@ -179,9 +172,12 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { mcts_config_thread, ); - tokio::runtime::Runtime::new() + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() .unwrap() .block_on(worker.run_worker(cmd_rx)); + })); episodes_dispatched += episodes_to_dispatch; workers_dispatched += 1; @@ -209,11 +205,12 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { worker_map .get(&data.worker_id) .unwrap() - .send(WorkerCmd::Dispatch) + .blocking_send(WorkerCmd::Dispatch) .expect("Expected dispatch to be sent"); episodes_dispatched += 1; println!("Dispatched episode #{}", episodes_dispatched); } + // else { // worker_map // .get(&data.worker_id) diff --git a/engine/src/worker.rs b/engine/src/worker.rs index 303dc65..c519c9f 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -1,10 +1,9 @@ use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults}; -use crate::net::encoding::encode_board_state_perspective; use crate::net::model::{ChessModel, TrainingSample}; use crate::train::{apply_temperature, sample_move}; use burn::tensor::Device; -use burn::Tensor; use chess::Color; + use moka::sync::Cache; use rand::prelude::SmallRng; use rand::SeedableRng; @@ -169,14 +168,13 @@ impl Worker { // (p, v) // } - pub async fn run_worker(self, cmd_rx: Receiver) { + pub async fn run_worker(self, mut cmd_rx: tokio::sync::mpsc::Receiver) { let mut tasks = Vec::new(); // println!("Worker started"); for _ in 0..self.num_games { let mcts_config = self.mcts_config.clone(); let data_tx = self.data_tx.clone(); let request_tx = self.request_tx.clone(); - let device = self.device.clone(); tasks.push(tokio::spawn(async move { Self::run_game( @@ -184,9 +182,7 @@ impl Worker { mcts_config, |board_state| { let request_tx = request_tx.clone(); - let device = device.clone(); - - model_inference_static(board_state.clone(), request_tx, device) + model_inference_static(board_state.clone(), request_tx) }, data_tx, ) @@ -203,66 +199,56 @@ impl Worker { break; } - match cmd_rx.recv_timeout(Duration::from_millis(10)) { - Ok(cmd) => match cmd { - WorkerCmd::Shutdown => break, + match tokio::time::timeout(Duration::from_millis(10), cmd_rx.recv()).await { + Ok(Some(WorkerCmd::Shutdown)) => break, - WorkerCmd::Dispatch => { - let mcts_config = self.mcts_config.clone(); - let request_tx = self.request_tx.clone(); - let device = self.device.clone(); - let data_tx = self.data_tx.clone(); + Ok(Some(WorkerCmd::Dispatch)) => { + let mcts_config = self.mcts_config.clone(); + let request_tx = self.request_tx.clone(); + let data_tx = self.data_tx.clone(); - tasks.push(tokio::spawn(async move { - Self::run_game( - self.id, - mcts_config, - |board_state| { - let request_tx = request_tx.clone(); - let device = device.clone(); - - model_inference_static(board_state.clone(), request_tx, device) - }, - data_tx, - ) - .await; - })); - } - }, - - Err(RecvTimeoutError::Timeout) => { - continue; + tasks.push(tokio::spawn(async move { + Self::run_game( + self.id, + mcts_config, + |board_state| { + let request_tx = request_tx.clone(); + model_inference_static(board_state.clone(), request_tx) + }, + data_tx, + ) + .await; + })); } - Err(RecvTimeoutError::Disconnected) => { - println!("Request channel disconnected"); - break; + Ok(None) => break, + + Err(_) => { + tokio::task::yield_now().await; } } } for task in tasks { - task.await.expect("Worker game task panicked"); + let _ = task.await; } println!("Worker {} shut down", self.id); } + } async fn model_inference_static( board_state: BoardState, request_tx: Sender, - device: Device, ) -> (Vec, f32) { let cache_key = board_state.nn_cache_key(); - let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>(); - let (tx, rx) = tokio::sync::oneshot::channel::(); request_tx .send(NNRequest { cache_key, - request_tensor: state, + board_state, response_tx: tx, }) .unwrap(); @@ -317,7 +303,7 @@ struct NNJobData { } pub struct NNRequest { cache_key: NNCacheKey, - request_tensor: Tensor<4>, + board_state: BoardState, response_tx: tokio::sync::oneshot::Sender, } @@ -387,7 +373,7 @@ impl NNServer { let mut jobs_to_send: Vec = Vec::new(); - for i in 0..self.batch_size { + for _ in 0..self.batch_size { if let Some(job) = queue.pop_front() { jobs_to_send.push(job); } else { @@ -395,36 +381,42 @@ impl NNServer { } } - let batched_tensor = Tensor::cat( - jobs_to_send - .iter() - .map(|j| j.request_tensor.clone()) - .collect(), - 0, - ); + if jobs_to_send.is_empty() { + continue; + } - // println!("Sending {} jobs to gpu", jobs_to_send.len()); + let b_len = jobs_to_send.len(); + 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]); + } + + let batched_tensor = burn::tensor::Tensor::<4>::from_floats( + burn::tensor::TensorData::new(buffer, [b_len, 19, 8, 8]), + &self.device, + ); let (p, v) = self.model.forward(batched_tensor); batches_ran += 1; - let p_split = p.chunk(jobs_to_send.len(), 0); - let v_split = v.chunk(jobs_to_send.len(), 0); + + let p_vec = p.into_data().to_vec::().unwrap(); + let v_vec = v.into_data().to_vec::().unwrap(); + + const POLICY_DIM: usize = 73 * 8 * 8; // 4672 for (i, job) in jobs_to_send.into_iter().enumerate() { - let p = p_split[i].clone().into_data().to_vec::().unwrap(); - - let v = v_split[i].clone().into_data().to_vec::().unwrap()[0]; + let p_slice = p_vec[i * POLICY_DIM..(i + 1) * POLICY_DIM].to_vec(); + let v_val = v_vec[i]; let res = NNResponse { - policy: p, - value: v, + policy: p_slice, + value: v_val, }; cache.insert(job.cache_key.clone(), res.clone()); - job.response_tx - .send(res) - .expect("worker dropped inference request"); + let _ = job.response_tx.send(res); } } }