From 01192b9fb8b58aa54a078d0c647438bd0fc5640c Mon Sep 17 00:00:00 2001 From: DragonDuck24 Date: Wed, 9 Sep 2026 15:47:57 -0400 Subject: [PATCH] training fixes --- engine/src/main.rs | 10 ++-- engine/src/train.rs | 120 +++++++++++++++++++++++++++++++------------ engine/src/worker.rs | 7 ++- 3 files changed, 98 insertions(+), 39 deletions(-) diff --git a/engine/src/main.rs b/engine/src/main.rs index f11def3..1163813 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(200, 1.0, 0.05, 0.25, 128); + let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128); let adam_config = AdamConfig::new(); @@ -28,9 +28,11 @@ fn main() { load_model: true, hidden_channels: 48, num_blocks: 4, - batch_size: 1024, // num positions sampled to a batch (2048) - num_episodes: 96, // num games generated per iteration (5000) - buffer_max_size: 25_000, // max number of samples in the buffer + batch_size: 1024, // num positions sampled to a batch (2048) + 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 mcts_config, optimizer: adam_config, lr: 2e-4, diff --git a/engine/src/train.rs b/engine/src/train.rs index 7dd19c8..a947d52 100644 --- a/engine/src/train.rs +++ b/engine/src/train.rs @@ -33,6 +33,8 @@ pub struct TrainingConfig { pub hidden_channels: usize, pub num_blocks: usize, pub batch_size: usize, // num positions sampled to a batch (2048) + pub minibatch_size: usize, // num positions sampled to a batch (2048) + pub num_train_steps: usize, // number of training steps of size minibatch_size pub num_episodes: usize, // num games generated per iteration (5000) pub buffer_max_size: usize, // max number of samples in the buffer pub mcts_config: MctsConfig, @@ -241,49 +243,101 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len()); println!( - "Finished self-play! Training on {} samples (buffer size: {})...", + "Finished self-play! Training on {} samples {} times with a minibatch size of {} (buffer size: {})...", num_samples, + training_config.num_train_steps, + training_config.minibatch_size, replay_buffer.len() ); - let mut indices: Vec = (0..replay_buffer.len()).collect(); - indices.shuffle(&mut rng); - - let samples = indices - .into_iter() - .take(num_samples) - .map(|i| replay_buffer[i].clone()) - .collect::>(); let batcher = ChessBatcher {}; + // Instead of one step, do multiple epochs over the sampled batch + for step in 0..training_config.num_train_steps { + // Re-shuffle and create mini-batches within the sampled data + let mut indices: Vec = (0..replay_buffer.len()).collect(); + indices.shuffle(&mut rng); + indices.truncate(num_samples); - let batch = batcher.batch(samples, &autodiff_device); + let mut loss_val: f32 = 0.0; + let mut p_loss_val: f32 = 0.0; + let mut v_loss_val: f32 = 0.0; - let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets); + // Process in mini-batches of e.g. 256 + for chunk in indices.chunks(training_config.minibatch_size) { + let samples: Vec = + chunk.iter().map(|&i| replay_buffer[i].clone()).collect(); - let loss_val = output.loss.clone().into_data().to_vec::().unwrap()[0]; - let p_loss_val = output - .policy_loss - .clone() - .into_data() - .to_vec::() - .unwrap()[0]; - let v_loss_val = output - .value_loss - .clone() - .into_data() - .to_vec::() - .unwrap()[0]; + 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); + model = optim.step(training_config.lr, model, grads); - println!( - "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}", - iter, loss_val, p_loss_val, v_loss_val - ); - println!("Iteration took {}s", iter_start.elapsed().as_secs()); + if step == training_config.num_train_steps - 1 { + loss_val = output.loss.clone().into_data().to_vec::().unwrap()[0]; + p_loss_val = output + .policy_loss + .clone() + .into_data() + .to_vec::() + .unwrap()[0]; + v_loss_val = output + .value_loss + .clone() + .into_data() + .to_vec::() + .unwrap()[0]; + } + } + if step == training_config.num_train_steps - 1 { + println!( + "Iteration {}, last minibatch training loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}", + iter, loss_val, p_loss_val, v_loss_val + ); + println!("Iteration took {}s", iter_start.elapsed().as_secs()); + } + } - let grads = output.loss.backward(); - let grads = GradientsParams::from_grads(grads, &model); - - model = optim.step(training_config.lr, model, grads); + // let mut indices: Vec = (0..replay_buffer.len()).collect(); + // indices.shuffle(&mut rng); + // + // let samples = indices + // .into_iter() + // .take(num_samples) + // .map(|i| replay_buffer[i].clone()) + // .collect::>(); + // + // let batcher = ChessBatcher {}; + // + // let batch = batcher.batch(samples, &autodiff_device); + // + // let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets); + // + // let loss_val = output.loss.clone().into_data().to_vec::().unwrap()[0]; + // let p_loss_val = output + // .policy_loss + // .clone() + // .into_data() + // .to_vec::() + // .unwrap()[0]; + // let v_loss_val = output + // .value_loss + // .clone() + // .into_data() + // .to_vec::() + // .unwrap()[0]; + // + // println!( + // "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}", + // iter, loss_val, p_loss_val, v_loss_val + // ); + // println!("Iteration took {}s", iter_start.elapsed().as_secs()); + // + // let grads = output.loss.backward(); + // let grads = GradientsParams::from_grads(grads, &model); + // + // model = optim.step(training_config.lr, model, grads); iter += 1; diff --git a/engine/src/worker.rs b/engine/src/worker.rs index da43794..303dc65 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -84,7 +84,8 @@ impl Worker { // println!("{}", before.elapsed().as_millis()); episode_buffer.push(results); - let temp = if ply_count < 30 { 1.0 } else { 0.2 }; + 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 }; let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp); @@ -125,9 +126,11 @@ impl Worker { _ => 0.0, }; + let blended_target = 0.75 * outcome + 0.25 * result.value; + samples.push(TrainingSample::from_mcts_with_outcome( result.clone(), - outcome, + blended_target, )); }