Virtual loss and batched mcts

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",

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"
+

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,

engine/src/mcts.rs

@@ -164,6 +164,8 @@ fn softmax(logits: &[f32]) -> Vec<f32> {
     probs
 }
 
+use std::future::Future;
+
 impl Mcts {
     pub async fn search<F, Fut>(
         &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,49 +191,82 @@ impl Mcts {
         }
 
         let num_sims = self.config.num_simulations;
+        let mut sim_count = 0;
 
-        for _ in 0..num_sims {
+        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;
+
+            let mut pending = Vec::with_capacity(current_batch_size);
+
+            for _ in 0..current_batch_size {
                 let mut path = vec![root];
                 let mut current = root;
 
-            // 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(...);
-            //     }
-            // }
-
                 while !nodes[current].children.is_empty() {
                     current = nodes[current].select_child(&nodes, &self.config.c_puct);
                     path.push(current);
                 }
 
-            let value: f32 = self.expand(current, &mut nodes, &model_inference).await;
+                // 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));
+            }
+
+            // 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);
+                }
+            }
+
+            // 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)
+                };
 
-            let color = nodes[current].board_state.board.side_to_move();
                 self.backpropagate(&mut nodes, &path, value, color);
             }
+        }
 
         // Build compact move distribution
         let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
@@ -256,27 +292,22 @@ impl Mcts {
         MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
     }
 
-    async fn expand<F, Fut>(
+    fn expand_with_eval(
         &mut self,
         node_idx: usize,
         arena: &mut Vec<Node>,
-        model_inference: &F,
+        policy_logits: &[f32],
-    ) -> f32
+        value_head: f32,
-    where
+    ) -> f32 {
-        F: Fn(&BoardState) -> Fut,
-        Fut: Future<Output = (Vec<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<ChessMove> =
                 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];

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<Device>) {
         let mut episodes_dispatched: usize = 0;
         let mut episodes_recieved: usize = 0;
 
-        let mut worker_map: HashMap<usize, Sender<WorkerCmd>> = HashMap::new();
+        let mut worker_map: HashMap<usize, tokio::sync::mpsc::Sender<WorkerCmd>> = HashMap::new();
         let (server_cmd_tx, server_cmd_rx) = mpsc::channel::<NNServerCmd>();
         let (data_tx, data_rx) = mpsc::channel::<WorkerData>();
         let (server_request_tx, server_request_rx) = mpsc::channel::<NNRequest>();
 
-        let mut nn_channels: VecDeque<Receiver<NNResponse>> = VecDeque::new();
-
-        for i in 0..training_config.num_workers {
-            let (tx, rx) = mpsc::channel::<NNResponse>();
-            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<Device>) {
         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::<WorkerCmd>();
+                let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<WorkerCmd>(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<Device>) {
                 );
 
                 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<Device>) {
                         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<Device>) {
                 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)

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<WorkerCmd>) {
+    pub async fn run_worker(self, mut cmd_rx: tokio::sync::mpsc::Receiver<WorkerCmd>) {
         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)
-
-                        model_inference_static(board_state.clone(), request_tx, device)
                     },
                     data_tx,
                 )
@@ -203,14 +199,12 @@ impl Worker {
                 break;
             }
 
-            match cmd_rx.recv_timeout(Duration::from_millis(10)) {
+            match tokio::time::timeout(Duration::from_millis(10), cmd_rx.recv()).await {
-                Ok(cmd) => match cmd {
+                Ok(Some(WorkerCmd::Shutdown)) => break,
-                    WorkerCmd::Shutdown => break,
 
-                    WorkerCmd::Dispatch => {
+                Ok(Some(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();
 
                     tasks.push(tokio::spawn(async move {
@@ -219,50 +213,42 @@ impl Worker {
                             mcts_config,
                             |board_state| {
                                 let request_tx = request_tx.clone();
-                                    let device = device.clone();
+                                model_inference_static(board_state.clone(), request_tx)
-
-                                    model_inference_static(board_state.clone(), request_tx, device)
                             },
                             data_tx,
                         )
                         .await;
                     }));
                 }
-                },
 
-                Err(RecvTimeoutError::Timeout) => {
+                Ok(None) => break,
-                    continue;
-                }
 
-                Err(RecvTimeoutError::Disconnected) => {
+                Err(_) => {
-                    println!("Request channel disconnected");
+                    tokio::task::yield_now().await;
-                    break;
                 }
             }
         }
 
         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<NNRequest>,
-    device: Device,
 ) -> (Vec<f32>, 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::<NNResponse>();
 
     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<NNResponse>,
 }
 
@@ -387,7 +373,7 @@ impl NNServer {
 
             let mut jobs_to_send: Vec<NNRequest> = 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(
+            if jobs_to_send.is_empty() {
-                jobs_to_send
+                continue;
-                    .iter()
+            }
-                    .map(|j| j.request_tensor.clone())
-                    .collect(),
-                0,
-            );
 
-            // 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::<f32>().unwrap();
+            let v_vec = v.into_data().to_vec::<f32>().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::<f32>().unwrap();
+                let p_slice = p_vec[i * POLICY_DIM..(i + 1) * POLICY_DIM].to_vec();
-
+                let v_val = v_vec[i];
-                let v = v_split[i].clone().into_data().to_vec::<f32>().unwrap()[0];
 
                 let res = NNResponse {
-                    policy: p,
+                    policy: p_slice,
-                    value: v,
+                    value: v_val,
                 };
 
                 cache.insert(job.cache_key.clone(), res.clone());
 
-                job.response_tx
+                let _ = job.response_tx.send(res);
-                    .send(res)
-                    .expect("worker dropped inference request");
             }
         }
     }