diff --git a/engine/src/lib.rs b/engine/src/lib.rs index 02effbd..9e637b4 100644 --- a/engine/src/lib.rs +++ b/engine/src/lib.rs @@ -1,6 +1,7 @@ pub mod mcts; mod net; pub mod train; +mod worker; use chess::Game; diff --git a/engine/src/main.rs b/engine/src/main.rs index 2b05886..187f2dd 100644 --- a/engine/src/main.rs +++ b/engine/src/main.rs @@ -25,11 +25,11 @@ fn main() { num_iters: None, max_depth: 100, // unused model_name: String::from("Test1"), - load_model: false, + load_model: true, hidden_channels: 48, num_blocks: 4, batch_size: 1024, // num positions sampled to a batch (2048) - num_episodes: 100, // num games generated per iteration (5000) + num_episodes: 96, // num games generated per iteration (5000) buffer_max_size: 25_000, // max number of samples in the buffer mcts_config, optimizer: adam_config, diff --git a/engine/src/train.rs b/engine/src/train.rs index 3b9d86e..7be0f5d 100644 --- a/engine/src/train.rs +++ b/engine/src/train.rs @@ -3,6 +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}; use burn::data::dataloader::batcher::Batcher; use burn::module::AutodiffModule; use burn::optim::{AdamConfig, GradientsParams}; @@ -12,11 +13,14 @@ use chess::{ChessMove, Color}; use rand::rngs::SmallRng; use rand::seq::SliceRandom; use rand::{RngExt, SeedableRng}; -use std::collections::VecDeque; +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::Arc; +use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{mpsc, Arc}; +use std::thread; use std::time::Instant; use std::time::{SystemTime, UNIX_EPOCH}; @@ -35,6 +39,9 @@ pub struct TrainingConfig { pub optimizer: AdamConfig, pub lr: f64, pub seed: Option, + pub num_workers: usize, + pub nn_batch_size: usize, + pub episodes_per_worker: usize, } pub fn train(training_config: TrainingConfig, device: impl Into) { @@ -98,8 +105,12 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { // Create RNG once and reuse it for sampling and shuffling // Seed from system time (platform default entropy may be unavailable in some contexts) - let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap(); - let seed = training_config.seed.unwrap_or(now.as_nanos() as u64); + let seed = training_config.seed.unwrap_or( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() as u64, + ); println!("using seed: {}", seed); let mut rng = SmallRng::seed_from_u64(seed); @@ -109,6 +120,65 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let infer_model = model.valid(); // Gen samples println!("Generating {} games...", training_config.num_episodes); + + let mut episodes_dispatched: usize = 0; + let mut episodes_recieved: usize = 0; + + 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(); + let mut nn_tx_map: HashMap> = HashMap::new(); + + for i in 0..training_config.num_workers { + let (tx, rx) = mpsc::channel::(); + nn_tx_map.insert(i, tx); + nn_channels.push_back(rx); + } + + let server_device = device.clone(); + thread::spawn(move || { + let server = NNServer::new( + server_cmd_rx, + server_request_rx, + nn_tx_map, + infer_model, + server_device, + training_config.nn_batch_size, + ); + + server.run_server() + }); + + for i in 0..training_config.num_workers { + if training_config.num_episodes > episodes_dispatched { + let (cmd_tx, cmd_rx) = mpsc::channel::(); + worker_map.insert(i, cmd_tx); + let episodes_to_dispatch = min( + training_config.episodes_per_worker, + training_config.num_episodes - episodes_dispatched, + ); + + let req_tx = server_request_tx.clone(); + let res_rx = nn_channels.pop_front().unwrap(); + let data_tx_thread = data_tx.clone(); + thread::spawn(move || { + let worker = Worker::new( + i, + episodes_to_dispatch, + req_tx, + res_rx, + cmd_rx, + data_tx_thread, + ); + + worker.run_worker() + }); + } + } + for episode in 0..training_config.num_episodes { println!("Episode: {}", episode); let mut board_state = BoardState::default(); @@ -120,7 +190,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { while board_state.status == BoardStateStatus::Ongoing { // let before = Instant::now(); - let results = mcts.search(&board_state, &infer_model, &device); + let results = mcts.search(&board_state, &infer_model, &device, true); // println!("{}", before.elapsed().as_millis()); episode_buffer.push(results); @@ -138,6 +208,8 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { board_state.apply_move(mv); ply_count += 1; + + // truncate long games if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing { board_state.status = BoardStateStatus::FiftyMove; } diff --git a/engine/src/worker.rs b/engine/src/worker.rs new file mode 100644 index 0000000..589ac19 --- /dev/null +++ b/engine/src/worker.rs @@ -0,0 +1,145 @@ +use crate::mcts::BoardState; +use crate::net::model::{ChessModel, TrainingSample}; +use burn::tensor::Device; +use burn::Tensor; +use std::collections::{HashMap, VecDeque}; +use std::sync::mpsc::{Receiver, Sender}; + +pub struct Worker { + id: usize, + games: Vec, + request_tx: Sender, + response_rx: Receiver, + cmd_rx: Receiver, + data_tx: Sender>, +} + +impl Worker { + pub fn new( + id: usize, + num_games: usize, + sender: Sender, + receiver: Receiver, + cmd_rx: Receiver, + data_tx: Sender>, + ) -> Worker { + Worker { + id, + games: (0..num_games) + .map(|_| BoardState::default()) + .collect::>(), + request_tx: sender, + response_rx: receiver, + cmd_rx, + data_tx, + } + } + + pub fn run_worker() {} +} + +pub enum NNServerCmd { + Shutdown, +} + +pub enum WorkerCmd { + Shutdown, +} + +pub struct NNServer { + command_rx: Receiver, + request_rx: Receiver, + response_txs: HashMap>, + model: ChessModel, + device: Device, + batch_size: usize, +} + +impl NNServer { + pub fn new( + command_rx: Receiver, + request_rx: Receiver, + response_txs: HashMap>, + model: ChessModel, + device: Device, + batch_size: usize, + ) -> NNServer { + NNServer { + command_rx, + request_rx, + response_txs, + model, + device, + batch_size, + } + } +} + +#[derive(Clone, Eq, Hash, PartialEq)] +struct NNJobData { + worker_id: usize, + request_id: u64, + game_id: u32, +} +pub struct NNRequest { + nn_job_data: NNJobData, + request_tensor: Tensor<4>, +} + +pub struct NNResponse { + nn_job_data: NNJobData, + response_tensors: (Tensor<2>, Tensor<2>), +} + +impl NNServer { + pub fn run_server(&self) { + let mut queue: VecDeque = VecDeque::new(); + + loop { + if queue.is_empty() { + let first = self + .request_rx + .recv() + .expect("failed to receive NN request"); + queue.push_back(first); + } + + while let Ok(req) = self.request_rx.try_recv() { + queue.push_back(req); + } + + let mut jobs_to_send: Vec = Vec::new(); + + for i in 0..self.batch_size { + if let Some(job) = queue.pop_front() { + jobs_to_send.push(job); + } else { + break; + } + } + + let batched_tensor = Tensor::cat( + jobs_to_send + .iter() + .map(|j| j.request_tensor.clone()) + .collect(), + 0, + ); + + let (p, v) = self.model.forward(batched_tensor); + let p_split = p.chunk(jobs_to_send.len(), 0); + let v_split = v.chunk(jobs_to_send.len(), 0); + + for (i, job) in jobs_to_send.iter().enumerate() { + let worker_tx = &self.response_txs[&job.nn_job_data.worker_id]; + + worker_tx + .send(NNResponse { + nn_job_data: job.nn_job_data.clone(), + response_tensors: (p_split[i].clone(), v_split[i].clone()), + }) + .expect("Expected response to be sent successfully"); + } + } + } +}