drake/chess_dragon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Virtual loss and batched mcts

Browse Source
This commit is contained in:
drake
2026-10-03 13:19:47 -04:00
parent 01192b9fb8
commit 38f13fa299
6 changed files with 178 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/Cargo.toml
+4 -1
View File
@@ -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
+2 -2
View File
@@ -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
+79 -47
View File
@@ -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,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<F, Fut>(
fn expand_with_eval(
&mut self,
node_idx: usize,
arena: &mut Vec<Node>,
model_inference: &F,
) -> f32
where
F: Fn(&BoardState) -> Fut,
Fut: Future<Output = (Vec<f32>, 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<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
+11 -14
View File
@@ -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
+54 -62
View File
@@ -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, 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<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(
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::<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 v = v_split[i].clone().into_data().to_vec::<f32>().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);
}
}
}