#![recursion_limit = "512"]
mod book;
mod csa;
mod diagnostics;
mod exporter;
mod positions;
mod scored;
mod teacher_cache;
mod trainer;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use sekirei_core::board::Board;
use sekirei_core::nnue::save_weights;
use csa::parse_csa;
use exporter::export_game;
use positions::load_positions;
use scored::load_scored;
use trainer::{ConflictMaskLayer, FreezeLayer, LrSchedule, ReplayComponent, Trainer};
static SIDECARE_WRITE_COUNTER: AtomicU64 = AtomicU64::new(0);
struct Args {
games_dir: Option<PathBuf>,
positions_path: Option<PathBuf>, output: PathBuf,
epochs: usize,
sample: usize, best_every: usize, min_rate: f32, quiet: bool, min_ply: usize, label_depth: u32, export: Option<PathBuf>, depths: Vec<u32>, build_book: Option<PathBuf>, book_max_ply: usize, book_min_count: u64, scored_path: Option<PathBuf>, min_stability: f32, stability_weighted: bool, label_threshold_cp: i32, phase_weights: HashMap<String, f32>, side_balance: bool, source_cap: usize, validation_ratio: f32, init_seed: u64, l2_bias_init: f32, split_seed: u64, shuffle_seed: Option<u64>,
cp_wdl_grad_trace: bool, wdl_target_scale: f32, sample_grad_trace: u64, trace_weights: bool, diagnostic_freeze_layer: Option<FreezeLayer>, diagnostic_freeze_from_position: u64, diagnostic_freeze_until_position: u64, diagnostic_ft_active_block: u64, diagnostic_ft_frozen_block: u64, diagnostic_ft_frozen_first: bool, diagnostic_ft_reactivate_from_position: u64, diagnostic_ft_reactivate_until_position: u64, diagnostic_ft_reactivate2_from_position: u64, diagnostic_ft_reactivate2_until_position: u64, diagnostic_replay_component: Option<ReplayComponent>, diagnostic_replay_from_position: u64, diagnostic_replay_until_position: u64, diagnostic_shadow_trace_from_position: u64, diagnostic_shadow_trace_until_position: u64, diagnostic_shadow_trace_probe_set: Option<PathBuf>, diagnostic_conflict_mask: Option<ConflictMaskLayer>, diagnostic_rate_matched_mask_count: u64, diagnostic_rate_matched_mask_total: u64, diagnostic_rate_matched_mask_seed: u64, checkpoint_dir: Option<PathBuf>, teacher_cache_path: Option<PathBuf>, reuse_teacher_cache: bool, wdl_lambda: Option<f32>, lr: f32, lr_schedule: LrSchedule, min_lr: f32, warmup_epochs: u32, lr_schedule_epochs: u32,
eval_only: Option<PathBuf>, grad_clip_norm: Option<f32>,
ft_clip_norm: Option<f32>,
l2_clip_norm: Option<f32>,
out_clip_norm: Option<f32>,
trace_positions: Vec<u64>,
}
fn parse_phase_weights(s: &str) -> HashMap<String, f32> {
s.split(',')
.filter_map(|pair| {
let (k, v) = pair.split_once('=')?;
let w: f32 = v.parse().ok()?;
Some((k.trim().to_string(), w))
})
.collect()
}
fn compute_side_weights(samples: &[positions::PositionSample]) -> HashMap<String, f32> {
let total = samples.len() as f32;
let black = samples.iter().filter(|s| s.side_to_move == "black").count() as f32;
let white = total - black;
[
(
"black".to_string(),
if black > 0.0 {
0.5 * total / black
} else {
1.0
},
),
(
"white".to_string(),
if white > 0.0 {
0.5 * total / white
} else {
1.0
},
),
]
.into_iter()
.collect()
}
fn parse_args() -> Result<Args, String> {
let argv: Vec<String> = std::env::args().skip(1).collect();
let mut games_dir = None;
let mut positions_path: Option<PathBuf> = None;
let mut output = PathBuf::from("weights.bin");
let mut epochs = 3usize;
let mut sample = 4usize;
let mut best_every = 0usize;
let mut min_rate = 1500.0f32;
let mut quiet = false;
let mut min_ply = 0usize;
let mut label_depth = 1u32;
let mut export: Option<PathBuf> = None;
let mut depths: Vec<u32> = vec![4, 6, 8];
let mut build_book: Option<PathBuf> = None;
let mut book_max_ply = 30usize;
let mut book_min_count = 20u64;
let mut scored_path: Option<PathBuf> = None;
let mut min_stability = 0.85f32;
let mut stability_weighted = false;
let mut label_threshold_cp = 120i32;
let mut phase_weights: HashMap<String, f32> = HashMap::new();
let mut side_balance = false;
let mut source_cap = 0usize;
let mut validation_ratio = 0.0f32;
let mut seed = 42u64;
let mut l2_bias_init = 0.5f32;
let mut trace_positions: Vec<u64> = Vec::new();
let mut shuffle_seed: Option<u64> = None;
let mut cp_wdl_grad_trace = false;
let mut wdl_target_scale = 1200.0f32;
let mut sample_grad_trace = 0u64;
let mut trace_weights = false;
let mut diagnostic_freeze_layer: Option<FreezeLayer> = None;
let mut diagnostic_freeze_from_position = 0u64;
let mut diagnostic_freeze_until_position = 0u64;
let mut diagnostic_ft_active_block = 0u64;
let mut diagnostic_ft_frozen_block = 0u64;
let mut diagnostic_ft_frozen_first = false;
let mut diagnostic_ft_reactivate_from_position = 0u64;
let mut diagnostic_ft_reactivate_until_position = 0u64;
let mut diagnostic_ft_reactivate2_from_position = 0u64;
let mut diagnostic_ft_reactivate2_until_position = 0u64;
let mut diagnostic_replay_component: Option<ReplayComponent> = None;
let mut diagnostic_replay_from_position = 0u64;
let mut diagnostic_replay_until_position = 0u64;
let mut diagnostic_shadow_trace_from_position = 0u64;
let mut diagnostic_shadow_trace_until_position = 0u64;
let mut diagnostic_shadow_trace_probe_set: Option<PathBuf> = None;
let mut diagnostic_conflict_mask: Option<ConflictMaskLayer> = None;
let mut diagnostic_rate_matched_mask_count = 0u64;
let mut diagnostic_rate_matched_mask_total = 0u64;
let mut diagnostic_rate_matched_mask_seed = 0u64;
let mut init_seed: Option<u64> = None;
let mut split_seed: Option<u64> = None;
let mut checkpoint_dir: Option<PathBuf> = None;
let mut teacher_cache_path: Option<PathBuf> = None;
let mut reuse_teacher_cache = false;
let mut wdl_lambda: Option<f32> = None;
let mut lr = 0.001f32;
let mut lr_schedule = LrSchedule::StepHalf;
let mut min_lr = 0.0f32;
let mut warmup_epochs = 0u32;
let mut lr_schedule_epochs: Option<u32> = None;
let mut eval_only: Option<PathBuf> = None;
let mut grad_clip_norm: Option<f32> = None;
let mut ft_clip_norm: Option<f32> = None;
let mut l2_clip_norm: Option<f32> = None;
let mut out_clip_norm: Option<f32> = None;
let mut i = 0;
while i < argv.len() {
match argv[i].as_str() {
"--games" => {
i += 1;
games_dir = argv.get(i).map(PathBuf::from);
}
"--positions" => {
i += 1;
positions_path = argv.get(i).map(PathBuf::from);
}
"--output" => {
i += 1;
if let Some(s) = argv.get(i) {
output = PathBuf::from(s);
}
}
"--epochs" => {
i += 1;
if let Some(s) = argv.get(i) {
epochs = s.parse().unwrap_or(3);
}
}
"--sample" => {
i += 1;
if let Some(s) = argv.get(i) {
sample = s.parse().unwrap_or(4);
}
}
"--best-every" => {
i += 1;
if let Some(s) = argv.get(i) {
best_every = s.parse().unwrap_or(0);
}
}
"--min-rate" => {
i += 1;
if let Some(s) = argv.get(i) {
min_rate = s.parse().unwrap_or(1500.0);
}
}
"--quiet" => {
quiet = true;
}
"--min-ply" => {
i += 1;
if let Some(s) = argv.get(i) {
min_ply = s.parse().unwrap_or(0);
}
}
"--label-depth" => {
i += 1;
if let Some(s) = argv.get(i) {
label_depth = s.parse().unwrap_or(1);
}
}
"--export" => {
i += 1;
export = argv.get(i).map(PathBuf::from);
}
"--depths" => {
i += 1;
if let Some(s) = argv.get(i) {
depths = s.split(',').filter_map(|d| d.parse().ok()).collect();
}
}
"--build-book" => {
i += 1;
build_book = argv.get(i).map(PathBuf::from);
}
"--book-max-ply" => {
i += 1;
if let Some(s) = argv.get(i) {
book_max_ply = s.parse().unwrap_or(book_max_ply);
}
}
"--book-min-count" => {
i += 1;
if let Some(s) = argv.get(i) {
book_min_count = s.parse().unwrap_or(book_min_count);
}
}
"--scored" => {
i += 1;
scored_path = argv.get(i).map(PathBuf::from);
}
"--min-stability" => {
i += 1;
if let Some(s) = argv.get(i) {
min_stability = s.parse().unwrap_or(0.85);
}
}
"--stability-weighted" => {
stability_weighted = true;
}
"--label-threshold-cp" => {
i += 1;
if let Some(s) = argv.get(i) {
label_threshold_cp = s.parse().unwrap_or(120);
}
}
"--wdl-lambda" => {
i += 1;
if let Some(s) = argv.get(i) {
wdl_lambda = s.parse().ok();
}
}
"--lr" => {
i += 1;
if let Some(s) = argv.get(i) {
lr = s.parse().unwrap_or(0.001);
}
}
"--lr-schedule" => {
i += 1;
if let Some(s) = argv.get(i) {
lr_schedule = LrSchedule::parse(s).unwrap_or(LrSchedule::StepHalf);
}
}
"--min-lr" => {
i += 1;
if let Some(s) = argv.get(i) {
min_lr = s.parse().unwrap_or(0.0);
}
}
"--warmup-epochs" => {
i += 1;
if let Some(s) = argv.get(i) {
warmup_epochs = s.parse().unwrap_or(0);
}
}
"--lr-schedule-epochs" => {
i += 1;
if let Some(s) = argv.get(i) {
lr_schedule_epochs = s.parse().ok();
}
}
"--grad-clip-norm" => {
i += 1;
if let Some(s) = argv.get(i) {
grad_clip_norm = s.parse().ok();
}
}
"--ft-clip-norm" => {
i += 1;
if let Some(s) = argv.get(i) {
ft_clip_norm = s.parse().ok();
}
}
"--l2-clip-norm" => {
i += 1;
if let Some(s) = argv.get(i) {
l2_clip_norm = s.parse().ok();
}
}
"--out-clip-norm" => {
i += 1;
if let Some(s) = argv.get(i) {
out_clip_norm = s.parse().ok();
}
}
"--l2-bias-init" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<f32>().ok()) {
l2_bias_init = v;
}
}
"--trace-positions" => {
i += 1;
if let Some(s) = argv.get(i) {
trace_positions = s.split(',').filter_map(|n| n.trim().parse().ok()).collect();
}
}
"--shuffle-seed" => {
i += 1;
shuffle_seed = argv.get(i).and_then(|s| s.parse().ok());
}
"--cp-wdl-grad-trace" => {
cp_wdl_grad_trace = true;
}
"--wdl-target-scale" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<f32>().ok()) {
wdl_target_scale = v;
}
}
"--sample-grad-trace" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
sample_grad_trace = v;
}
}
"--trace-weights" => {
trace_weights = true;
}
"--diagnostic-freeze-layer" => {
i += 1;
if let Some(s) = argv.get(i) {
diagnostic_freeze_layer = FreezeLayer::parse(s);
}
}
"--diagnostic-freeze-from-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_freeze_from_position = v;
}
}
"--diagnostic-freeze-until-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_freeze_until_position = v;
}
}
"--diagnostic-ft-active-block" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_active_block = v;
}
}
"--diagnostic-ft-frozen-block" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_frozen_block = v;
}
}
"--diagnostic-ft-frozen-first" => {
diagnostic_ft_frozen_first = true;
}
"--diagnostic-ft-reactivate-from-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_reactivate_from_position = v;
}
}
"--diagnostic-ft-reactivate-until-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_reactivate_until_position = v;
}
}
"--diagnostic-ft-reactivate2-from-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_reactivate2_from_position = v;
}
}
"--diagnostic-ft-reactivate2-until-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_ft_reactivate2_until_position = v;
}
}
"--diagnostic-replay-component" => {
i += 1;
if let Some(s) = argv.get(i) {
diagnostic_replay_component = ReplayComponent::parse(s);
}
}
"--diagnostic-replay-from-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_replay_from_position = v;
}
}
"--diagnostic-replay-until-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_replay_until_position = v;
}
}
"--diagnostic-shadow-trace-from-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_shadow_trace_from_position = v;
}
}
"--diagnostic-shadow-trace-until-position" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_shadow_trace_until_position = v;
}
}
"--diagnostic-shadow-trace-probe-set" => {
i += 1;
if let Some(s) = argv.get(i) {
diagnostic_shadow_trace_probe_set = Some(PathBuf::from(s));
}
}
"--diagnostic-conflict-mask" => {
i += 1;
if let Some(s) = argv.get(i) {
diagnostic_conflict_mask = ConflictMaskLayer::parse(s);
}
}
"--diagnostic-rate-matched-mask-count" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_rate_matched_mask_count = v;
}
}
"--diagnostic-rate-matched-mask-total" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_rate_matched_mask_total = v;
}
}
"--diagnostic-rate-matched-mask-seed" => {
i += 1;
if let Some(v) = argv.get(i).and_then(|s| s.parse::<u64>().ok()) {
diagnostic_rate_matched_mask_seed = v;
}
}
"--eval-only" => {
i += 1;
if let Some(s) = argv.get(i) {
eval_only = Some(PathBuf::from(s));
}
}
"--phase-weights" => {
i += 1;
if let Some(s) = argv.get(i) {
phase_weights = parse_phase_weights(s);
}
}
"--side-balance" => {
side_balance = true;
}
"--source-cap" => {
i += 1;
if let Some(s) = argv.get(i) {
source_cap = s.parse().unwrap_or(0);
}
}
"--validation-ratio" => {
i += 1;
if let Some(s) = argv.get(i) {
validation_ratio = s.parse().unwrap_or(0.0);
}
}
"--seed" => {
i += 1;
if let Some(s) = argv.get(i) {
seed = s.parse().unwrap_or(42);
}
}
"--init-seed" => {
i += 1;
if let Some(s) = argv.get(i) {
init_seed = s.parse().ok();
}
}
"--split-seed" => {
i += 1;
if let Some(s) = argv.get(i) {
split_seed = s.parse().ok();
}
}
"--checkpoint-dir" => {
i += 1;
checkpoint_dir = argv.get(i).map(PathBuf::from);
}
"--teacher-cache" => {
i += 1;
teacher_cache_path = argv.get(i).map(PathBuf::from);
}
"--reuse-teacher-cache" => {
reuse_teacher_cache = true;
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
_ => {}
}
i += 1;
}
if games_dir.is_none() && positions_path.is_none() {
return Err("either --games <dir> or --positions <jsonl> is required".to_string());
}
if games_dir.is_some() && positions_path.is_some() {
return Err("--games and --positions are mutually exclusive".to_string());
}
if wdl_lambda.is_some() && positions_path.is_some() {
return Err(
"--wdl-lambda requires --games (CSA path) -- shogiesa positions.jsonl carries no game_result yet".to_string(),
);
}
if eval_only.is_some() && positions_path.is_some() {
return Err("--eval-only requires --games (CSA path)".to_string());
}
let lr_schedule_epochs =
trainer::resolve_schedule_epochs(epochs as u32, lr_schedule_epochs, warmup_epochs)?;
Ok(Args {
games_dir,
positions_path,
output,
epochs,
sample,
best_every,
min_rate,
quiet,
min_ply,
label_depth,
export,
build_book,
book_max_ply,
book_min_count,
depths,
scored_path,
min_stability,
stability_weighted,
label_threshold_cp,
phase_weights,
side_balance,
source_cap,
validation_ratio,
init_seed: init_seed.unwrap_or(seed),
l2_bias_init,
split_seed: split_seed.unwrap_or(seed),
checkpoint_dir,
teacher_cache_path,
reuse_teacher_cache,
wdl_lambda,
lr,
lr_schedule,
min_lr,
warmup_epochs,
lr_schedule_epochs,
eval_only,
grad_clip_norm,
ft_clip_norm,
l2_clip_norm,
out_clip_norm,
trace_positions,
shuffle_seed,
cp_wdl_grad_trace,
wdl_target_scale,
sample_grad_trace,
trace_weights,
diagnostic_freeze_layer,
diagnostic_freeze_from_position,
diagnostic_freeze_until_position,
diagnostic_ft_active_block,
diagnostic_ft_frozen_block,
diagnostic_ft_frozen_first,
diagnostic_ft_reactivate_from_position,
diagnostic_ft_reactivate_until_position,
diagnostic_ft_reactivate2_from_position,
diagnostic_ft_reactivate2_until_position,
diagnostic_replay_component,
diagnostic_replay_from_position,
diagnostic_replay_until_position,
diagnostic_shadow_trace_from_position,
diagnostic_shadow_trace_until_position,
diagnostic_shadow_trace_probe_set,
diagnostic_conflict_mask,
diagnostic_rate_matched_mask_count,
diagnostic_rate_matched_mask_total,
diagnostic_rate_matched_mask_seed,
})
}
fn git_commit_hash() -> Option<String> {
std::process::Command::new("git")
.args(["rev-parse", "HEAD"])
.output()
.ok()
.filter(|o| o.status.success())
.and_then(|o| String::from_utf8(o.stdout).ok())
.map(|s| s.trim().to_string())
}
fn dataset_hash(paths: &[PathBuf]) -> u64 {
let mut sorted: Vec<&PathBuf> = paths.iter().collect();
sorted.sort();
let joined: String = sorted
.iter()
.map(|p| {
let len = fs::metadata(p).map(|m| m.len()).unwrap_or(0);
format!("{}:{}\0", p.display(), len)
})
.collect();
positions::sfen_hash(&joined, 0)
}
fn split_hash(keys: impl Iterator<Item = String>) -> u64 {
keys.fold(0u64, |acc, k| acc.wrapping_add(positions::sfen_hash(&k, 0)))
}
fn checkpoint_hash(bytes: &[u8]) -> u64 {
let mut h = 14695981039346656037u64;
for b in bytes {
h ^= *b as u64;
h = h.wrapping_mul(1099511628211);
}
h
}
fn eval_validation_set(
trainer: &mut Trainer,
games: &[csa::CsaGame],
valid_idxs: &[usize],
args: &Args,
cache: &mut HashMap<String, i32>,
) -> trainer::ValidStats {
let start = Instant::now();
let mut last_progress = Instant::now();
let mut acc = trainer::ValidStats::default();
for (i, &gi) in valid_idxs.iter().enumerate() {
acc = acc
+ trainer.eval_game(
&games[gi],
args.sample,
args.quiet,
args.min_ply,
args.label_depth,
args.wdl_lambda,
args.wdl_target_scale,
cache,
);
if last_progress.elapsed() >= Duration::from_secs(5) {
let elapsed = start.elapsed().as_secs_f64();
eprintln!(
" valid progress: game {}/{} positions={} elapsed={elapsed:.1}s cache_hits={} cache_misses={} search_time={:.1}s",
i + 1,
valid_idxs.len(),
acc.count,
trainer.cache_hits,
trainer.cache_misses,
trainer.search_time_ns as f64 / 1e9,
);
last_progress = Instant::now();
}
}
acc
}
fn build_diag(
trainer: &Trainer,
w: &sekirei_core::nnue::NnueWeights,
param_update_norm: Option<f32>,
) -> diagnostics::EpochDiagnostics {
let (output_mean, output_std) = diagnostics::mean_std(
trainer.output_sum,
trainer.output_sum_sq,
trainer.total_count,
);
let l2_activation_frequency_per_neuron = diagnostics::l2_activation_frequency_per_neuron(
&trainer.l2_zero_count,
trainer.l2_sample_count,
);
let l2_saturation_frequency_per_neuron = diagnostics::l2_saturation_frequency_per_neuron(
&trainer.l2_sat_count,
trainer.l2_sample_count,
);
let l2_activation_frequency_mean = if l2_activation_frequency_per_neuron.is_empty() {
0.0
} else {
l2_activation_frequency_per_neuron.iter().sum::<f32>()
/ l2_activation_frequency_per_neuron.len() as f32
};
let l2_saturation_frequency_mean = if l2_saturation_frequency_per_neuron.is_empty() {
0.0
} else {
l2_saturation_frequency_per_neuron.iter().sum::<f32>()
/ l2_saturation_frequency_per_neuron.len() as f32
};
let pooled_l2_values: Vec<f32> = trainer.l2_values.iter().flatten().copied().collect();
let p = diagnostics::percentiles(&pooled_l2_values, &[0.01, 0.10, 0.50, 0.90, 0.99]);
let (ft_grad_norm_mean, ft_grad_norm_std) = diagnostics::mean_std(
trainer.ft_grad_norm_sum,
trainer.ft_grad_norm_sum_sq,
trainer.total_count,
);
let (l2_grad_norm_mean, l2_grad_norm_std) = diagnostics::mean_std(
trainer.l2_grad_norm_sum,
trainer.l2_grad_norm_sum_sq,
trainer.total_count,
);
let (out_grad_norm_mean, out_grad_norm_std) = diagnostics::mean_std(
trainer.out_grad_norm_sum,
trainer.out_grad_norm_sum_sq,
trainer.total_count,
);
let gp = diagnostics::percentiles(&trainer.global_grad_norm_values, &[0.50, 0.90, 0.95, 0.99]);
let op = diagnostics::percentiles(&trainer.out_grad_norm_values, &[0.95, 0.99]);
let (out_grad_norm_after_mean, out_grad_norm_after_std) = diagnostics::mean_std(
trainer.out_grad_norm_after_sum,
trainer.out_grad_norm_after_sum_sq,
trainer.total_count,
);
let clip_rate = |count: u64| {
if trainer.total_count > 0 {
count as f64 / trainer.total_count as f64
} else {
0.0
}
};
let (ft_update_norm_mean, ft_update_norm_std) = diagnostics::mean_std(
trainer.ft_update_norm_sum,
trainer.ft_update_norm_sum_sq,
trainer.total_count,
);
let (l2_update_norm_mean, l2_update_norm_std) = diagnostics::mean_std(
trainer.l2_update_norm_sum,
trainer.l2_update_norm_sum_sq,
trainer.total_count,
);
let (out_update_norm_mean, out_update_norm_std) = diagnostics::mean_std(
trainer.out_update_norm_sum,
trainer.out_update_norm_sum_sq,
trainer.total_count,
);
let (target_mean, target_std) = diagnostics::mean_std(
trainer.target_sum,
trainer.target_sum_sq,
trainer.total_count,
);
let pred_eval_correlation = diagnostics::pearson_correlation(
trainer.total_count,
trainer.output_sum,
trainer.output_sum_sq,
trainer.eval_teacher_sum,
trainer.eval_teacher_sum_sq,
trainer.pred_eval_prod_sum,
);
let train_cp_component = if trainer.total_count > 0 {
trainer.cp_component_sum / trainer.total_count as f64
} else {
0.0
};
let train_wdl_component = if trainer.wdl_component_count > 0 {
Some(trainer.wdl_component_sum / trainer.wdl_component_count as f64)
} else {
None
};
diagnostics::EpochDiagnostics {
param_update_norm,
ft_active_ratio: diagnostics::ratio(&trainer.ft_ever_active),
ft_saturation_ratio: diagnostics::ratio(&trainer.ft_ever_saturated),
output_mean,
output_std,
quantized_ft_zero_ratio: diagnostics::quantized_ft_zero_ratio(w),
l2_ever_active_ratio: diagnostics::ratio(&trainer.l2_ever_active),
l2_ever_saturated_ratio: diagnostics::ratio(&trainer.l2_ever_saturated),
l2_dead_neurons: diagnostics::l2_dead_neurons(
&trainer.l2_zero_count,
trainer.l2_sample_count,
),
l2_activation_frequency_mean,
l2_saturation_frequency_mean,
l2_activation_frequency_per_neuron,
l2_saturation_frequency_per_neuron,
l2_preactivation_p01: p[0],
l2_preactivation_p10: p[1],
l2_preactivation_p50: p[2],
l2_preactivation_p90: p[3],
l2_preactivation_p99: p[4],
l2_bias_per_neuron: trainer.weights.l2_bias().to_vec(),
l2_row_weight_norm_per_neuron: diagnostics::l2_row_weight_norm_per_neuron(
trainer.weights.l2(),
2 * sekirei_core::nnue::L1,
sekirei_core::nnue::L2,
),
output_weight_norm: diagnostics::output_weight_norm(trainer.weights.out()),
output_bias: trainer.weights.out_bias(),
ft_grad_norm_mean,
ft_grad_norm_std,
l2_grad_norm_mean,
l2_grad_norm_std,
out_grad_norm_mean,
out_grad_norm_std,
global_grad_norm_p50: gp[0],
global_grad_norm_p90: gp[1],
global_grad_norm_p95: gp[2],
global_grad_norm_p99: gp[3],
ft_update_norm_mean,
ft_update_norm_std,
l2_update_norm_mean,
l2_update_norm_std,
out_update_norm_mean,
out_update_norm_std,
target_mean,
target_std,
pred_eval_correlation,
train_cp_component,
train_wdl_component,
grad_clip_count: trainer.grad_clip_count,
ft_clip_trigger_rate: clip_rate(trainer.ft_clip_count),
l2_clip_trigger_rate: clip_rate(trainer.l2_clip_count),
out_clip_trigger_rate: clip_rate(trainer.out_clip_count),
out_grad_norm_p95: op[0],
out_grad_norm_p99: op[1],
out_grad_norm_after_mean,
out_grad_norm_after_std,
masked_position_count: trainer.masked_position_count,
eligible_position_count: trainer.wdl_component_count,
conflict_group: diagnostics::build_conflict_group_summary(&trainer.conflict_group),
nonconflict_group: diagnostics::build_conflict_group_summary(&trainer.nonconflict_group),
ft_dead_neurons: diagnostics::l2_dead_neurons(
&trainer.ft_zero_count,
trainer.l2_sample_count,
),
ft_activation_frequency_mean: {
let f = diagnostics::l2_activation_frequency_per_neuron(
&trainer.ft_zero_count,
trainer.l2_sample_count,
);
if f.is_empty() {
0.0
} else {
f.iter().sum::<f32>() / f.len() as f32
}
},
}
}
fn print_grad_diag_lines(diag: &diagnostics::EpochDiagnostics, train_count: u64) {
let clip_rate = if train_count > 0 {
diag.grad_clip_count as f64 / train_count as f64
} else {
0.0
};
eprintln!(
" grad: ft={:.4}±{:.4} l2={:.4}±{:.4} out={:.5}±{:.5} (mean±std) global_p50={:.4} p90={:.4} p95={:.4} p99={:.4} clipped={}/{} ({:.1}%)",
diag.ft_grad_norm_mean,
diag.ft_grad_norm_std,
diag.l2_grad_norm_mean,
diag.l2_grad_norm_std,
diag.out_grad_norm_mean,
diag.out_grad_norm_std,
diag.global_grad_norm_p50,
diag.global_grad_norm_p90,
diag.global_grad_norm_p95,
diag.global_grad_norm_p99,
diag.grad_clip_count,
train_count,
clip_rate * 100.0,
);
eprintln!(
" layer_clip: ft={:.1}% l2={:.1}% out={:.1}% (independent per-layer trigger rates) out_grad_p95={:.4} p99={:.4} out_grad_norm={:.4}→{:.4} (before→after, mean)",
diag.ft_clip_trigger_rate * 100.0,
diag.l2_clip_trigger_rate * 100.0,
diag.out_clip_trigger_rate * 100.0,
diag.out_grad_norm_p95,
diag.out_grad_norm_p99,
diag.out_grad_norm_mean,
diag.out_grad_norm_after_mean,
);
eprintln!(
" update: ft={:.6}±{:.6} l2={:.6}±{:.6} out={:.6}±{:.6} (mean±std, applied Adam step)",
diag.ft_update_norm_mean,
diag.ft_update_norm_std,
diag.l2_update_norm_mean,
diag.l2_update_norm_std,
diag.out_update_norm_mean,
diag.out_update_norm_std,
);
eprintln!(
" target: mean={:.3} std={:.3} pred_eval_corr={:.4} cp_component={:.3} wdl_component={} out_weight_norm={:.4} out_bias={:.4}",
diag.target_mean,
diag.target_std,
diag.pred_eval_correlation,
diag.train_cp_component,
diag.train_wdl_component
.map(|w| format!("{w:.3}"))
.unwrap_or_else(|| "n/a".to_string()),
diag.output_weight_norm,
diag.output_bias,
);
}
fn save_trace_json(path: &Path, snapshots: &[diagnostics::TraceSnapshot]) -> std::io::Result<()> {
if snapshots.is_empty() {
return Ok(());
}
let json = serde_json::to_string_pretty(snapshots)?;
write_atomic(path, json.as_bytes())
}
fn save_sample_grad_jsonl(
path: &Path,
records: &[diagnostics::SampleGradRecord],
) -> std::io::Result<()> {
if records.is_empty() {
return Ok(());
}
let mut out = String::new();
for r in records {
out.push_str(&serde_json::to_string(r)?);
out.push('\n');
}
write_atomic(path, out.as_bytes())
}
fn load_shadow_trace_probe_boards(path: &Path) -> std::io::Result<Vec<Board>> {
let text = fs::read_to_string(path)?;
let mut boards = Vec::new();
for line in text.lines() {
let sfen = line.trim();
if sfen.is_empty() || sfen.starts_with('#') {
continue;
}
match Board::from_sfen(sfen) {
Ok(b) => boards.push(b),
Err(e) => eprintln!(" shadow-trace probe set: skip unparseable line ({e}): {sfen}"),
}
}
Ok(boards)
}
fn save_shadow_trace_jsonl(
path: &Path,
records: &[diagnostics::ShadowTraceRecord],
) -> std::io::Result<()> {
if records.is_empty() {
return Ok(());
}
let mut out = String::new();
for r in records {
out.push_str(&serde_json::to_string(r)?);
out.push('\n');
}
write_atomic(path, out.as_bytes())
}
fn save_weight_snapshots(checkpoint: &Path, snapshots: &[(u64, trainer::TrainWeights)]) {
for (pos, w) in snapshots {
let nn = w.to_nnue_weights();
let path = checkpoint.with_extension(format!("pos{pos}.bin"));
match sekirei_core::nnue::save_weights(&nn, &path) {
Ok(_) => eprintln!(" weights → {:?} (pos {pos})", path),
Err(e) => eprintln!(" weight snapshot save failed (pos {pos}): {e}"),
}
}
}
fn write_atomic_stream<F>(path: &Path, write: F) -> io::Result<()>
where
F: FnOnce(&mut BufWriter<File>) -> io::Result<()>,
{
let file_name = path.file_name().ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"checkpoint metadata path must name a file",
)
})?;
let temp_path = path.with_file_name(format!(
".{}.tmp-{}-{}",
file_name.to_string_lossy(),
std::process::id(),
SIDECARE_WRITE_COUNTER.fetch_add(1, Ordering::Relaxed)
));
let write_result = (|| -> io::Result<()> {
let file = File::create(&temp_path)?;
let mut out = BufWriter::new(file);
write(&mut out)?;
out.flush()?;
let file = out.into_inner().map_err(|error| error.into_error())?;
file.sync_all()
})();
if let Err(error) = write_result {
let _ = fs::remove_file(&temp_path);
return Err(error);
}
if let Err(error) = fs::rename(&temp_path, path) {
let _ = fs::remove_file(&temp_path);
return Err(error);
}
Ok(())
}
fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
write_atomic_stream(path, |out| out.write_all(contents))
}
fn copy_atomic(source: &Path, destination: &Path) -> io::Result<u64> {
let file_name = destination.file_name().ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"checkpoint destination path must name a file",
)
})?;
let temp_path = destination.with_file_name(format!(
".{}.tmp-{}-{}",
file_name.to_string_lossy(),
std::process::id(),
SIDECARE_WRITE_COUNTER.fetch_add(1, Ordering::Relaxed)
));
let result = (|| -> io::Result<u64> {
let copied = fs::copy(source, &temp_path)?;
File::open(&temp_path)?.sync_all()?;
fs::rename(&temp_path, destination)?;
Ok(copied)
})();
if result.is_err() {
let _ = fs::remove_file(&temp_path);
}
result
}
#[allow(clippy::too_many_arguments)]
fn save_checkpoint_meta(
path: &Path,
args: &Args,
epoch: usize,
train_count: u64,
valid_count: u64,
train_games: Option<u64>,
valid_games: Option<u64>,
split_hash: u64,
diag: &diagnostics::EpochDiagnostics,
git_commit: Option<&str>,
dataset_hash: u64,
checkpoint_hash: u64,
cache_hits: Option<u64>,
cache_misses: Option<u64>,
search_time_ns: Option<u64>,
valid_stats: Option<&trainer::ValidStats>,
) -> std::io::Result<()> {
let (
valid_cp_mse,
valid_wdl_loss,
valid_output_mean,
valid_output_std,
valid_output_min,
valid_output_max,
valid_output_range,
) = match valid_stats {
Some(s) => {
let cp_mse = if s.count > 0 {
Some(s.cp_mse_sum / s.count as f64)
} else {
None
};
let wdl_loss = if s.wdl_count > 0 {
Some(s.wdl_loss_sum / s.wdl_count as f64)
} else {
None
};
let (mean, std) = diagnostics::mean_std(s.output_sum, s.output_sum_sq, s.count);
let (min, max, range) = if s.count > 0 {
(
Some(s.output_min),
Some(s.output_max),
Some(s.output_max - s.output_min),
)
} else {
(None, None, None)
};
(cp_mse, wdl_loss, Some(mean), Some(std), min, max, range)
}
None => (None, None, None, None, None, None, None),
};
let meta = serde_json::json!({
"epoch": epoch,
"positions": args.positions_path,
"games_dir": args.games_dir,
"min_rate": args.min_rate,
"sample": args.sample,
"scored": args.scored_path,
"label_depth": args.label_depth,
"wdl_lambda": args.wdl_lambda,
"wdl_target_scale": args.wdl_target_scale,
"phase_weights": args.phase_weights,
"side_balance": args.side_balance,
"source_cap": args.source_cap,
"validation_ratio": args.validation_ratio,
"init_seed": args.init_seed,
"l2_bias_init": args.l2_bias_init,
"split_seed": args.split_seed,
"shuffle_seed": args.shuffle_seed,
"lr": args.lr,
"lr_schedule": format!("{:?}", args.lr_schedule),
"min_lr": args.min_lr,
"warmup_epochs": args.warmup_epochs,
"epochs": args.epochs,
"lr_schedule_epochs": args.lr_schedule_epochs,
"train_count": train_count,
"valid_count": valid_count,
"train_games": train_games,
"valid_games": valid_games,
"split_hash": split_hash,
"architecture": format!(
"INPUT={} L1={} L2={}",
sekirei_core::nnue::INPUT,
sekirei_core::nnue::L1,
sekirei_core::nnue::L2
),
"git_commit": git_commit,
"dataset_hash": dataset_hash,
"checkpoint_hash": format!("{checkpoint_hash:016x}"),
"cache_hits": cache_hits,
"cache_misses": cache_misses,
"search_time_secs": search_time_ns.map(|ns| ns as f64 / 1e9),
"param_update_norm": diag.param_update_norm,
"ft_active_ratio": diag.ft_active_ratio,
"ft_saturation_ratio": diag.ft_saturation_ratio,
"output_mean": diag.output_mean,
"output_std": diag.output_std,
"quantized_ft_zero_ratio": diag.quantized_ft_zero_ratio,
"l2_ever_active_ratio": diag.l2_ever_active_ratio,
"l2_ever_saturated_ratio": diag.l2_ever_saturated_ratio,
"l2_dead_neurons": diag.l2_dead_neurons,
"l2_activation_frequency_mean": diag.l2_activation_frequency_mean,
"l2_saturation_frequency_mean": diag.l2_saturation_frequency_mean,
"l2_activation_frequency_per_neuron": diag.l2_activation_frequency_per_neuron,
"l2_saturation_frequency_per_neuron": diag.l2_saturation_frequency_per_neuron,
"l2_preactivation_p01": diag.l2_preactivation_p01,
"l2_preactivation_p10": diag.l2_preactivation_p10,
"l2_preactivation_p50": diag.l2_preactivation_p50,
"l2_preactivation_p90": diag.l2_preactivation_p90,
"l2_preactivation_p99": diag.l2_preactivation_p99,
"l2_bias_per_neuron": diag.l2_bias_per_neuron,
"l2_row_weight_norm_per_neuron": diag.l2_row_weight_norm_per_neuron,
"output_weight_norm": diag.output_weight_norm,
"output_bias": diag.output_bias,
"ft_grad_norm_mean": diag.ft_grad_norm_mean,
"ft_grad_norm_std": diag.ft_grad_norm_std,
"l2_grad_norm_mean": diag.l2_grad_norm_mean,
"l2_grad_norm_std": diag.l2_grad_norm_std,
"out_grad_norm_mean": diag.out_grad_norm_mean,
"out_grad_norm_std": diag.out_grad_norm_std,
"global_grad_norm_p50": diag.global_grad_norm_p50,
"global_grad_norm_p90": diag.global_grad_norm_p90,
"global_grad_norm_p95": diag.global_grad_norm_p95,
"global_grad_norm_p99": diag.global_grad_norm_p99,
"ft_update_norm_mean": diag.ft_update_norm_mean,
"ft_update_norm_std": diag.ft_update_norm_std,
"l2_update_norm_mean": diag.l2_update_norm_mean,
"l2_update_norm_std": diag.l2_update_norm_std,
"out_update_norm_mean": diag.out_update_norm_mean,
"out_update_norm_std": diag.out_update_norm_std,
"target_mean": diag.target_mean,
"target_std": diag.target_std,
"pred_eval_correlation": diag.pred_eval_correlation,
"train_cp_component": diag.train_cp_component,
"train_wdl_component": diag.train_wdl_component,
"grad_clip_norm": args.grad_clip_norm,
"grad_clip_count": diag.grad_clip_count,
"ft_clip_norm": args.ft_clip_norm,
"l2_clip_norm": args.l2_clip_norm,
"out_clip_norm": args.out_clip_norm,
"diagnostic_freeze_layer": args.diagnostic_freeze_layer.map(|l| l.as_str()),
"diagnostic_freeze_from_position": args.diagnostic_freeze_from_position,
"diagnostic_freeze_until_position": args.diagnostic_freeze_until_position,
"diagnostic_ft_active_block": args.diagnostic_ft_active_block,
"diagnostic_ft_frozen_block": args.diagnostic_ft_frozen_block,
"diagnostic_ft_frozen_first": args.diagnostic_ft_frozen_first,
"diagnostic_ft_reactivate_from_position": args.diagnostic_ft_reactivate_from_position,
"diagnostic_ft_reactivate_until_position": args.diagnostic_ft_reactivate_until_position,
"diagnostic_ft_reactivate2_from_position": args.diagnostic_ft_reactivate2_from_position,
"diagnostic_ft_reactivate2_until_position": args.diagnostic_ft_reactivate2_until_position,
"diagnostic_replay_component": args.diagnostic_replay_component.map(|c| c.as_str()),
"diagnostic_replay_from_position": args.diagnostic_replay_from_position,
"diagnostic_replay_until_position": args.diagnostic_replay_until_position,
"diagnostic_shadow_trace_from_position": args.diagnostic_shadow_trace_from_position,
"diagnostic_shadow_trace_until_position": args.diagnostic_shadow_trace_until_position,
"diagnostic_shadow_trace_probe_set": args.diagnostic_shadow_trace_probe_set.as_ref().map(|p| p.display().to_string()),
"diagnostic_conflict_mask": args.diagnostic_conflict_mask.map(|c| c.as_str()),
"diagnostic_rate_matched_mask_count": args.diagnostic_rate_matched_mask_count,
"diagnostic_rate_matched_mask_total": args.diagnostic_rate_matched_mask_total,
"diagnostic_rate_matched_mask_seed": args.diagnostic_rate_matched_mask_seed,
"masked_position_count": diag.masked_position_count,
"eligible_position_count": diag.eligible_position_count,
"conflict_group": diag.conflict_group,
"nonconflict_group": diag.nonconflict_group,
"ft_dead_neurons": diag.ft_dead_neurons,
"ft_activation_frequency_mean": diag.ft_activation_frequency_mean,
"ft_clip_trigger_rate": diag.ft_clip_trigger_rate,
"l2_clip_trigger_rate": diag.l2_clip_trigger_rate,
"out_clip_trigger_rate": diag.out_clip_trigger_rate,
"out_grad_norm_p95": diag.out_grad_norm_p95,
"out_grad_norm_p99": diag.out_grad_norm_p99,
"out_grad_norm_after_mean": diag.out_grad_norm_after_mean,
"out_grad_norm_after_std": diag.out_grad_norm_after_std,
"valid_cp_mse": valid_cp_mse,
"valid_wdl_loss": valid_wdl_loss,
"valid_output_mean": valid_output_mean,
"valid_output_std": valid_output_std,
"valid_output_min": valid_output_min,
"valid_output_max": valid_output_max,
"valid_output_range": valid_output_range,
});
let json = serde_json::to_vec_pretty(&meta)
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
write_atomic(path, &json)
}
fn split_games_by_index(
n_games: usize,
validation_ratio: f32,
seed: u64,
) -> (Vec<usize>, Vec<usize>) {
let split_threshold = (validation_ratio.clamp(0.0, 1.0) * 1000.0) as u64;
(0..n_games)
.partition(|&i| positions::sfen_hash(&i.to_string(), seed) % 1000 >= split_threshold)
}
fn print_usage() {
eprintln!(
"Usage: train (--games <dir> | --positions <jsonl>) [--output weights.bin] [--epochs 3] [--sample 4]"
);
eprintln!();
eprintln!(" --games <dir> Directory containing .csa game files");
eprintln!(" --positions <jsonl> shogiesa positions.jsonl (alternative to --games)");
eprintln!(" --output <file> Output weight file (default: weights.bin)");
eprintln!(" --epochs <n> Training epochs (default: 3)");
eprintln!(" --sample <n> Sample every N plies per game (default: 4)");
eprintln!(" --best-every <n> Save best-loss checkpoint every N games (default: 0 = off)");
eprintln!(
" --min-rate <r> Minimum rating for both players (default: 1500, 0 = no filter)"
);
eprintln!(" --quiet Skip positions in check or where next move is a capture");
eprintln!(" --min-ply <n> Skip the first N plies per game (default: 0)");
eprintln!(" --label-depth <n> Search depth for teacher labels (default: 1)");
eprintln!(" --export <path> Export observations JSONL for quietset (skips training)");
eprintln!(" --depths <list> Comma-separated depths for export (default: 4,6,8)");
eprintln!(
" --build-book <path> Build a statistical opening book from --games (skips training;"
);
eprintln!(" reuses --min-rate to filter which games count)");
eprintln!(" --book-max-ply <n> Max ply to record into the book (default: 30)");
eprintln!(" --book-min-count <n> Minimum times a move must appear to be kept (default: 20)");
eprintln!(" --scored <path> quietset scored JSONL — train only stable samples");
eprintln!(" --min-stability <f> Minimum stability_score to include (default: 0.85)");
eprintln!(" --stability-weighted Weight loss by stability_score instead of binary keep/drop");
eprintln!(
" --label-threshold-cp <n> Score threshold for adv/equal/disadv label (default: 120)"
);
eprintln!(
" --wdl-lambda <f> Blend in game result (CSA path only): teacher = λ·eval + (1-λ)·wdl (default: unset = eval-only)"
);
eprintln!(" --lr <f> Base learning rate (default: 0.001)");
eprintln!(
" --lr-schedule <name> constant | step-half | cosine (default: step-half, today's original behavior)"
);
eprintln!(" --min-lr <f> Floor applied to every schedule (default: 0.0)");
eprintln!(
" --warmup-epochs <n> Linear ramp to base_lr over the first N epochs (default: 0 = off)"
);
eprintln!(
" --lr-schedule-epochs <n> Schedule horizon the LR curve is shaped for; may exceed --epochs to reproduce the first N epochs of a longer schedule (default: --epochs)"
);
eprintln!(
" --grad-clip-norm <f> Global gradient-norm clip threshold; scales all layers' gradients down together when exceeded (default: unset = no clipping)"
);
eprintln!(
" --ft-clip-norm <f> FT-layer-only gradient-norm clip threshold, independent of --grad-clip-norm and the other --*-clip-norm flags (default: unset)"
);
eprintln!(
" --l2-clip-norm <f> L2-layer-only gradient-norm clip threshold (default: unset)"
);
eprintln!(
" --out-clip-norm <f> Output-layer-only gradient-norm clip threshold (default: unset)"
);
eprintln!(" --l2-bias-init <f> L2 layer's bias value at initialization (default: 0.5)");
eprintln!(
" --trace-positions <n1,n2,...> Position-counts since epoch start to snapshot L2/FT's per-neuron state (e.g. 0,1,2,4,8,16,32,64); writes <output>.epochN.trace.json. Default: unset (off)"
);
eprintln!(
" --shuffle-seed <n> Reshuffle each epoch's training order, seeded (default: unset -- original file order, unchanged)"
);
eprintln!(
" --cp-wdl-grad-trace Decompose the blended gradient into CP-only/WDL-only per position (CSA path + --wdl-lambda only; adds two diagnostic backward passes per position, real compute cost, default: off)"
);
eprintln!(
" --wdl-target-scale <f> Native range of wdl_target, via (wdl - 0.5) * scale (default: 1200.0, giving +/-600)"
);
eprintln!(
" --sample-grad-trace <n> Record a per-position gradient-correlation trace (game_id/outcome/residuals/L2 grad cosine similarity/gate state) for the first n positions each epoch; writes <output>.epochN.sample_grad.jsonl. Never reorders training. Default: 0 (off)"
);
eprintln!(
" --trace-weights At each --trace-positions marker, also dump a full weights checkpoint (<output>.epochN.posM.bin) for offline Δz decomposition. Requires --trace-positions. Default: off"
);
eprintln!(
" --diagnostic-freeze-layer <ft|l2|out> Causal freeze probe: skip the named layer's own Adam update while --diagnostic-freeze-from-position <= l2_sample_count <= --diagnostic-freeze-until-position this epoch. Gradient still flows through it to the other layers (not stop-gradient). Diagnostic only. Default: unset (off)"
);
eprintln!(
" --diagnostic-freeze-from-position <n> Position count (since epoch start) the freeze above starts being active at; no-op without --diagnostic-freeze-layer. Default: 0 (freeze from the first position)"
);
eprintln!(
" --diagnostic-freeze-until-position <n> Position count (since epoch start) the freeze above stays active until; no-op without --diagnostic-freeze-layer. Default: 0"
);
eprintln!(
" --diagnostic-ft-active-block <n> FT-only: within the freeze window above, cycle <n> positions of normal FT updates then --diagnostic-ft-frozen-block positions frozen, repeating. No-op unless --diagnostic-freeze-layer ft and both block flags are set (nonzero). Default: 0 (off, plain single-window freeze)"
);
eprintln!(
" --diagnostic-ft-frozen-block <n> Paired with --diagnostic-ft-active-block, see above. Default: 0 (off)"
);
eprintln!(
" --diagnostic-ft-frozen-first Paired with --diagnostic-ft-active-block/--diagnostic-ft-frozen-block: start each cycle frozen instead of active (for equal block lengths, produces the exact complement pattern of the default). No-op unless periodic mode is on. Default: off (cycle starts active)"
);
eprintln!(
" --diagnostic-ft-reactivate-from-position <n> Carves out one additional active sub-window inside an otherwise-frozen --diagnostic-freeze-layer ft span, for single-block necessity/sufficiency screens (freeze the whole window, reactivate exactly one block). Independent of --diagnostic-ft-active-block cycling. Default: 0"
);
eprintln!(
" --diagnostic-ft-reactivate-until-position <n> Paired with --diagnostic-ft-reactivate-from-position, see above. Default: 0 (off, no reactivation window)"
);
eprintln!(
" --diagnostic-ft-reactivate2-from-position <n> Second, independent reactivation window (for reactivating two disjoint blocks at once, e.g. a block-interaction screen). Same semantics as --diagnostic-ft-reactivate-from-position. Default: 0"
);
eprintln!(
" --diagnostic-ft-reactivate2-until-position <n> Paired with --diagnostic-ft-reactivate2-from-position, see above. Default: 0 (off)"
);
eprintln!(
" --diagnostic-replay-component <cp|wdl> CSA path only: within --diagnostic-replay-from-position/--diagnostic-replay-until-position, replace the normal blended (teacher,weight) with just this component's contribution, still scaled by its own wdl_lambda coefficient (not renormalized). Mathematically exact counterfactual replay, diagnostic only. Default: unset (off)"
);
eprintln!(
" --diagnostic-replay-from-position <n> Paired with --diagnostic-replay-component, see above. Default: 0"
);
eprintln!(
" --diagnostic-replay-until-position <n> Paired with --diagnostic-replay-component, see above. Default: 0 (off, no window)"
);
eprintln!(
" --diagnostic-shadow-trace-probe-set <path> CSA path only: file of one SFEN board per line. Within --diagnostic-shadow-trace-from-position/-until-position, branches CP-only/WDL-only/Blended one-step counterfactual FT+L2 updates from each live position's own pre-update state, evaluates every branch on this probe set, then discards -- never perturbs real training. Requires --wdl-lambda. Default: unset (off)"
);
eprintln!(
" --diagnostic-shadow-trace-from-position <n> Paired with --diagnostic-shadow-trace-probe-set, see above. Default: 0"
);
eprintln!(
" --diagnostic-shadow-trace-until-position <n> Paired with --diagnostic-shadow-trace-probe-set, see above. Default: 0 (off, no window)"
);
eprintln!(
" --diagnostic-conflict-mask <ft|ft-l2> CSA path only: stop the named layer(s)' update at positions where (score-eval_teacher)*(score-wdl_target) < 0 (prediction sits between the two teachers). Not PCGrad/projection -- per-position CP/WDL gradients are always exact scalar multiples, so projection degenerates to this same masking. Default: unset (off)"
);
eprintln!(
" --diagnostic-rate-matched-mask-count <n> Rate-matched control: mask FT at exactly this many positions per epoch, chosen independent of teacher conflict (seeded, unbiased selection). Normally set from a prior --diagnostic-conflict-mask run's own masked_position_count. Default: 0 (off)"
);
eprintln!(
" --diagnostic-rate-matched-mask-total <n> Paired with --diagnostic-rate-matched-mask-count: total eligible (wdl_target-having) positions per epoch, from the same prior run's eligible_position_count. Default: 0"
);
eprintln!(
" --diagnostic-rate-matched-mask-seed <n> Paired with --diagnostic-rate-matched-mask-count: seed for the position-selection RNG. Default: 0"
);
eprintln!(
" --eval-only <ckpt.bin> CSA path only: load a checkpoint, run one validation pass with cp_mse/wdl_loss, print, exit (no training)"
);
eprintln!(
" --phase-weights <spec> Phase multipliers: opening=0.5,middlegame=1.0,endgame=1.2"
);
eprintln!(" --side-balance Equalise black/white sample weights");
eprintln!(" --source-cap <n> Max samples per source file (0 = unlimited)");
eprintln!(
" --validation-ratio <f> Hold-out fraction for valid_loss, both --games and --positions (default: 0.0 = off)"
);
eprintln!(
" --seed <n> Sets both --init-seed and --split-seed at once (default: 42)"
);
eprintln!(
" --init-seed <n> Weight-init seed only, overrides --seed (default: --seed's value)"
);
eprintln!(
" --split-seed <n> Validation-split/source_cap seed only, overrides --seed (default: --seed's value)"
);
eprintln!(" --checkpoint-dir <dir> Directory for epoch checkpoints");
eprintln!(" --teacher-cache <path> JSONL cache of teacher scores (sfen → score_cp)");
eprintln!(" --reuse-teacher-cache Load teacher cache; skip search on cache hits");
eprintln!();
eprintln!("Data: download floodgate archives from http://wdoor.c.u-tokyo.ac.jp/shogi/");
}
fn collect_csa_files(dir: &Path) -> Vec<PathBuf> {
let mut files = Vec::new();
collect_csa_recursive(dir, &mut files);
files.sort();
files
}
fn collect_csa_recursive(dir: &Path, out: &mut Vec<PathBuf>) {
if let Ok(entries) = fs::read_dir(dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
collect_csa_recursive(&path, out);
} else if path.extension().and_then(|e| e.to_str()) == Some("csa") {
out.push(path);
}
}
}
}
fn main() {
let args = match parse_args() {
Ok(a) => a,
Err(e) => {
eprintln!("error: {e}");
print_usage();
std::process::exit(1);
}
};
let git_commit = git_commit_hash();
if let Some(pos_path) = &args.positions_path {
eprintln!("Positions mode: loading {:?}", pos_path);
let ds_hash = dataset_hash(std::slice::from_ref(pos_path));
let raw_samples = load_positions(pos_path);
if raw_samples.is_empty() {
eprintln!("No valid positions loaded");
std::process::exit(1);
}
let all_samples = if args.source_cap > 0 {
let n_before = raw_samples.len();
let s = positions::apply_source_cap(raw_samples, args.source_cap, args.split_seed);
eprintln!(
"{} positions loaded, {} after source_cap={} (split_seed={})",
n_before,
s.len(),
args.source_cap,
args.split_seed
);
s
} else {
eprintln!("{} positions loaded", raw_samples.len());
raw_samples
};
let split_threshold = (args.validation_ratio.clamp(0.0, 1.0) * 1000.0) as u64;
let (train_samples, valid_samples): (Vec<_>, Vec<_>) =
all_samples.into_iter().partition(|s| {
let sfen = sekirei_core::sfen::board_to_sfen(&s.board);
positions::sfen_hash(&sfen, args.split_seed) % 1000 >= split_threshold
});
let split_h = split_hash(
valid_samples
.iter()
.map(|s| sekirei_core::sfen::board_to_sfen(&s.board)),
);
eprintln!(
" train={} valid={} (validation_ratio={:.2}, split_seed={})",
train_samples.len(),
valid_samples.len(),
args.validation_ratio,
args.split_seed
);
let scored: HashMap<String, f32> = match &args.scored_path {
Some(p) => load_scored(p, args.min_stability),
None => HashMap::new(),
};
let side_weights = if args.side_balance {
compute_side_weights(&train_samples)
} else {
HashMap::new()
};
if args.side_balance {
eprintln!(
" side_balance: black={:.3} white={:.3}",
side_weights.get("black").copied().unwrap_or(1.0),
side_weights.get("white").copied().unwrap_or(1.0)
);
}
let checkpoint_dir = args
.checkpoint_dir
.clone()
.unwrap_or_else(|| args.output.parent().unwrap_or(Path::new(".")).to_path_buf());
let output_stem = args
.output
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("weights")
.to_string();
let mut combined_cache: HashMap<String, i32> = if args.reuse_teacher_cache {
match &args.teacher_cache_path {
Some(p) => teacher_cache::load(p, args.label_depth),
None => {
eprintln!("error: --reuse-teacher-cache requires --teacher-cache <path>");
std::process::exit(1);
}
}
} else {
HashMap::new()
};
let mut trainer = Trainer::new(args.init_seed, args.l2_bias_init);
trainer.grad_clip_norm = args.grad_clip_norm;
trainer.ft_clip_norm = args.ft_clip_norm;
trainer.l2_clip_norm = args.l2_clip_norm;
trainer.out_clip_norm = args.out_clip_norm;
trainer.trace_positions = args.trace_positions.iter().copied().collect();
trainer.cp_wdl_grad_trace = args.cp_wdl_grad_trace;
trainer.sample_grad_trace_limit = args.sample_grad_trace;
trainer.weight_snapshot_trace = args.trace_weights;
trainer.diagnostic_freeze_layer = args.diagnostic_freeze_layer;
trainer.diagnostic_freeze_from_position = args.diagnostic_freeze_from_position;
trainer.diagnostic_freeze_until_position = args.diagnostic_freeze_until_position;
trainer.diagnostic_ft_active_block = args.diagnostic_ft_active_block;
trainer.diagnostic_ft_frozen_block = args.diagnostic_ft_frozen_block;
trainer.diagnostic_ft_frozen_first = args.diagnostic_ft_frozen_first;
trainer.diagnostic_ft_reactivate_from_position =
args.diagnostic_ft_reactivate_from_position;
trainer.diagnostic_ft_reactivate_until_position =
args.diagnostic_ft_reactivate_until_position;
trainer.diagnostic_ft_reactivate2_from_position =
args.diagnostic_ft_reactivate2_from_position;
trainer.diagnostic_ft_reactivate2_until_position =
args.diagnostic_ft_reactivate2_until_position;
trainer.diagnostic_replay_component = args.diagnostic_replay_component;
trainer.diagnostic_replay_from_position = args.diagnostic_replay_from_position;
trainer.diagnostic_replay_until_position = args.diagnostic_replay_until_position;
trainer.diagnostic_shadow_trace_from_position = args.diagnostic_shadow_trace_from_position;
trainer.diagnostic_shadow_trace_until_position =
args.diagnostic_shadow_trace_until_position;
trainer.diagnostic_shadow_trace_wdl_lambda = args.wdl_lambda.unwrap_or(0.0);
if let Some(probe_path) = &args.diagnostic_shadow_trace_probe_set {
match load_shadow_trace_probe_boards(probe_path) {
Ok(boards) => trainer.diagnostic_shadow_trace_probe_boards = boards,
Err(e) => eprintln!(" shadow-trace probe set load failed: {e}"),
}
}
trainer.diagnostic_conflict_mask = args.diagnostic_conflict_mask;
trainer.diagnostic_rate_matched_mask_count = args.diagnostic_rate_matched_mask_count;
trainer.diagnostic_rate_matched_mask_total = args.diagnostic_rate_matched_mask_total;
trainer.diagnostic_rate_matched_mask_seed = args.diagnostic_rate_matched_mask_seed;
let mut prev_snapshot: Option<Vec<f32>> = None;
let mut best_valid_loss = f64::MAX;
let mut best_valid_checkpoint: Option<PathBuf> = None;
for epoch in 1..=args.epochs {
trainer.lr = trainer::compute_lr(
args.lr_schedule,
args.lr,
args.min_lr,
epoch as u32,
args.lr_schedule_epochs,
args.warmup_epochs,
);
trainer.reset_epoch_stats();
eprintln!("Epoch {epoch}/{} — lr = {:.6}", args.epochs, trainer.lr);
let shuffled_samples: Vec<positions::PositionSample>;
let epoch_samples: &[positions::PositionSample] = if let Some(seed) = args.shuffle_seed
{
let order = trainer::shuffled_order(train_samples.len(), seed ^ epoch as u64);
shuffled_samples = order.iter().map(|&i| train_samples[i].clone()).collect();
&shuffled_samples
} else {
&train_samples
};
let mut new_entries: Vec<(String, i32)> = Vec::new();
trainer.train_positions(
epoch_samples,
args.label_depth,
&scored,
args.stability_weighted,
&args.phase_weights,
&side_weights,
&combined_cache,
&mut new_entries,
);
let mut new_val_entries: Vec<(String, i32)> = Vec::new();
let (vloss_raw, vloss_w, vcount) = if valid_samples.is_empty() {
(0.0, 0.0, 0)
} else {
trainer.eval_positions(
&valid_samples,
args.label_depth,
&args.phase_weights,
&side_weights,
&combined_cache,
&mut new_val_entries,
)
};
new_entries.extend(new_val_entries);
let cache_misses_epoch = new_entries.len() as u64;
let cache_hits_epoch =
(train_samples.len() + valid_samples.len()) as u64 - cache_misses_epoch;
if epoch == 1 && !new_entries.is_empty() {
let n = new_entries.len();
for (sfen, cp) in new_entries {
combined_cache.entry(sfen).or_insert(cp);
}
eprintln!(" teacher cache: {n} new entries computed");
if let Some(cache_path) = &args.teacher_cache_path {
match teacher_cache::write(cache_path, &combined_cache, args.label_depth) {
Ok(_) => eprintln!(" teacher cache written → {:?}", cache_path),
Err(e) => eprintln!(" teacher cache write failed: {e}"),
}
}
} else if epoch == 1 {
eprintln!(
" teacher cache: all {} entries from cache (no search)",
combined_cache.len()
);
}
if !scored.is_empty() {
let total_seen = trainer.total_count + trainer.dropped_missing;
let missing_rate = if total_seen > 0 {
trainer.dropped_missing as f64 / total_seen as f64
} else {
0.0
};
let avg_weight = if trainer.total_count > 0 {
trainer.total_weight / trainer.total_count as f64
} else {
1.0
};
eprintln!(
" quietset: entries={} matched={} dropped_missing={} missing_rate={:.1}% avg_weight={:.3}",
scored.len(),
trainer.total_count,
trainer.dropped_missing,
missing_rate * 100.0,
avg_weight,
);
if missing_rate > 0.5 {
eprintln!(
" warn: missing_rate={:.1}% — SFEN mismatch?",
missing_rate * 100.0
);
}
if trainer.total_count == 0 && trainer.dropped_missing > 0 {
eprintln!(
"error: scored file loaded ({} entries) but 0 positions matched.",
scored.len()
);
eprintln!(
"hint: --export and --scored must cover the same --games / --sample / --quiet / --min-ply / --min-rate."
);
eprintln!(
"hint: check `head -1 scored.jsonl` — sample_id or sfen must be a SFEN string."
);
std::process::exit(1);
}
}
let avg_final_weight = if trainer.total_count > 0 {
trainer.total_weight / trainer.total_count as f64
} else {
1.0
};
eprintln!(
" train: avg_loss={:.4} samples={} avg_final_weight={:.3}",
trainer.avg_loss(),
trainer.total_count,
avg_final_weight,
);
let valid_count = valid_samples.len() as u64;
if !valid_samples.is_empty() {
eprintln!(
" valid: loss_raw={:.4} loss_weighted={:.4} samples={}",
vloss_raw, vloss_w, vcount,
);
}
let checkpoint = checkpoint_dir.join(format!("{output_stem}.epoch{epoch}.bin"));
let w = trainer.weights.to_nnue_weights();
let mut ckpt_hash = 0u64;
match sekirei_core::nnue::save_weights(&w, &checkpoint) {
Ok(_) => {
eprintln!(" checkpoint → {:?}", checkpoint);
match fs::read(&checkpoint) {
Ok(bytes) => ckpt_hash = checkpoint_hash(&bytes),
Err(e) => eprintln!(" checkpoint hash read failed: {e}"),
}
}
Err(e) => eprintln!(" checkpoint save failed: {e}"),
}
let snapshot = trainer.weights.snapshot_params();
let param_update_norm = prev_snapshot
.as_ref()
.map(|prev| diagnostics::l2_diff_norm(prev, &snapshot));
prev_snapshot = Some(snapshot);
let diag = build_diag(&trainer, &w, param_update_norm);
eprintln!(
" diag: ft_active={:.3} l2_ever_active={:.3} ft_sat={:.3} l2_ever_sat={:.3} l2_dead={} l2_act_freq={:.3} l2_sat_freq={:.3} out_mean={:.3} out_std={:.3} ft_zero={:.3} update_norm={}",
diag.ft_active_ratio,
diag.l2_ever_active_ratio,
diag.ft_saturation_ratio,
diag.l2_ever_saturated_ratio,
diag.l2_dead_neurons,
diag.l2_activation_frequency_mean,
diag.l2_saturation_frequency_mean,
diag.output_mean,
diag.output_std,
diag.quantized_ft_zero_ratio,
diag.param_update_norm
.map(|n| format!("{n:.4}"))
.unwrap_or_else(|| "n/a".to_string()),
);
print_grad_diag_lines(&diag, trainer.total_count);
let meta_path = checkpoint.with_extension("meta.json");
if let Err(e) = save_checkpoint_meta(
&meta_path,
&args,
epoch,
trainer.total_count,
valid_count,
None, None,
split_h,
&diag,
git_commit.as_deref(),
ds_hash,
ckpt_hash,
Some(cache_hits_epoch),
Some(cache_misses_epoch),
None, None, ) {
eprintln!(" metadata save failed: {e}");
} else {
eprintln!(" metadata → {:?}", meta_path);
}
let trace_path = checkpoint.with_extension("trace.json");
if let Err(e) = save_trace_json(&trace_path, &trainer.trace_snapshots) {
eprintln!(" trace save failed: {e}");
} else if !trainer.trace_snapshots.is_empty() {
eprintln!(" trace → {:?}", trace_path);
}
let sample_grad_path = checkpoint.with_extension("sample_grad.jsonl");
if let Err(e) = save_sample_grad_jsonl(&sample_grad_path, &trainer.sample_grad_records)
{
eprintln!(" sample-grad trace save failed: {e}");
} else if !trainer.sample_grad_records.is_empty() {
eprintln!(" sample_grad → {:?}", sample_grad_path);
}
let shadow_trace_path = checkpoint.with_extension("shadow_trace.jsonl");
if let Err(e) =
save_shadow_trace_jsonl(&shadow_trace_path, &trainer.shadow_trace_records)
{
eprintln!(" shadow trace save failed: {e}");
} else if !trainer.shadow_trace_records.is_empty() {
eprintln!(" shadow_trace → {:?}", shadow_trace_path);
}
save_weight_snapshots(&checkpoint, &trainer.weight_snapshots);
if args.validation_ratio > 0.0 && vcount > 0 && vloss_raw < best_valid_loss {
best_valid_loss = vloss_raw;
best_valid_checkpoint = Some(checkpoint.clone());
}
}
if let Some(best_ckpt) = &best_valid_checkpoint {
let best_path = args.output.with_extension("best.bin");
match copy_atomic(best_ckpt, &best_path) {
Ok(_) => eprintln!(
" best (valid_loss={best_valid_loss:.4}) → {:?} (from {:?})",
best_path, best_ckpt
),
Err(e) => eprintln!(" best checkpoint copy failed: {e}"),
}
}
let w = trainer.weights.to_nnue_weights();
match sekirei_core::nnue::save_weights(&w, &args.output) {
Ok(_) => eprintln!("Final weights saved → {:?}", args.output),
Err(e) => {
eprintln!("Save failed: {e}");
std::process::exit(1);
}
}
return;
}
let files = collect_csa_files(args.games_dir.as_ref().unwrap());
if files.is_empty() {
eprintln!("No .csa files found in {:?}", args.games_dir);
std::process::exit(1);
}
eprintln!("Found {} CSA files in {:?}", files.len(), args.games_dir);
let ds_hash = dataset_hash(&files);
eprint!("Parsing games... ");
let games: Vec<_> = files
.iter()
.filter_map(|p| fs::read_to_string(p).ok())
.filter_map(|text| parse_csa(&text))
.filter(|g| {
if args.min_rate <= 0.0 {
return true;
}
g.black_rate.is_some_and(|r| r >= args.min_rate)
&& g.white_rate.is_some_and(|r| r >= args.min_rate)
})
.collect();
eprintln!("{} games loaded (min_rate={})", games.len(), args.min_rate);
if games.is_empty() {
eprintln!("No valid games parsed — check CSA format");
std::process::exit(1);
}
if let Some(book_path) = &args.build_book {
eprintln!(
"Book mode → {:?} max_ply={} min_count={}",
book_path, args.book_max_ply, args.book_min_count
);
if let Err(e) = write_atomic_stream(book_path, |out| {
book::build_book(&games, args.book_max_ply, args.book_min_count, out)
}) {
eprintln!("Cannot write book file: {e}");
std::process::exit(1);
}
eprintln!("Book done → {:?}", book_path);
return;
}
if let Some(export_path) = &args.export {
eprintln!("Export mode → {:?} depths={:?}", export_path, args.depths);
if let Err(e) = write_atomic_stream(export_path, |out| {
for game in &games {
export_game(
game,
args.sample,
args.quiet,
args.min_ply,
&args.depths,
args.label_threshold_cp,
out,
)?;
}
Ok(())
}) {
eprintln!("Cannot write export file: {e}");
std::process::exit(1);
}
eprintln!("Export done → {:?}", export_path);
return;
}
let scored: HashMap<String, f32> = match &args.scored_path {
Some(p) => load_scored(p, args.min_stability),
None => HashMap::new(),
};
let (train_idxs, valid_idxs) =
split_games_by_index(games.len(), args.validation_ratio, args.split_seed);
let split_h = split_hash(valid_idxs.iter().map(|i| i.to_string()));
eprintln!(
" train_games={} valid_games={} (validation_ratio={:.2}, split_seed={})",
train_idxs.len(),
valid_idxs.len(),
args.validation_ratio,
args.split_seed
);
let mut trainer = Trainer::new(args.init_seed, args.l2_bias_init);
trainer.grad_clip_norm = args.grad_clip_norm;
trainer.ft_clip_norm = args.ft_clip_norm;
trainer.l2_clip_norm = args.l2_clip_norm;
trainer.out_clip_norm = args.out_clip_norm;
trainer.trace_positions = args.trace_positions.iter().copied().collect();
trainer.cp_wdl_grad_trace = args.cp_wdl_grad_trace;
trainer.sample_grad_trace_limit = args.sample_grad_trace;
trainer.weight_snapshot_trace = args.trace_weights;
trainer.diagnostic_freeze_layer = args.diagnostic_freeze_layer;
trainer.diagnostic_freeze_from_position = args.diagnostic_freeze_from_position;
trainer.diagnostic_freeze_until_position = args.diagnostic_freeze_until_position;
trainer.diagnostic_ft_active_block = args.diagnostic_ft_active_block;
trainer.diagnostic_ft_frozen_block = args.diagnostic_ft_frozen_block;
trainer.diagnostic_ft_frozen_first = args.diagnostic_ft_frozen_first;
trainer.diagnostic_ft_reactivate_from_position = args.diagnostic_ft_reactivate_from_position;
trainer.diagnostic_ft_reactivate_until_position = args.diagnostic_ft_reactivate_until_position;
trainer.diagnostic_ft_reactivate2_from_position = args.diagnostic_ft_reactivate2_from_position;
trainer.diagnostic_ft_reactivate2_until_position =
args.diagnostic_ft_reactivate2_until_position;
trainer.diagnostic_replay_component = args.diagnostic_replay_component;
trainer.diagnostic_replay_from_position = args.diagnostic_replay_from_position;
trainer.diagnostic_replay_until_position = args.diagnostic_replay_until_position;
trainer.diagnostic_shadow_trace_from_position = args.diagnostic_shadow_trace_from_position;
trainer.diagnostic_shadow_trace_until_position = args.diagnostic_shadow_trace_until_position;
trainer.diagnostic_shadow_trace_wdl_lambda = args.wdl_lambda.unwrap_or(0.0);
if let Some(probe_path) = &args.diagnostic_shadow_trace_probe_set {
match load_shadow_trace_probe_boards(probe_path) {
Ok(boards) => trainer.diagnostic_shadow_trace_probe_boards = boards,
Err(e) => eprintln!(" shadow-trace probe set load failed: {e}"),
}
}
trainer.diagnostic_conflict_mask = args.diagnostic_conflict_mask;
trainer.diagnostic_rate_matched_mask_count = args.diagnostic_rate_matched_mask_count;
trainer.diagnostic_rate_matched_mask_total = args.diagnostic_rate_matched_mask_total;
trainer.diagnostic_rate_matched_mask_seed = args.diagnostic_rate_matched_mask_seed;
if let Some(eval_ckpt) = &args.eval_only {
let nn = match sekirei_core::nnue::read_weights(eval_ckpt) {
Ok(w) => w,
Err(e) => {
eprintln!("eval-only: failed to load {:?}: {e}", eval_ckpt);
std::process::exit(1);
}
};
trainer.weights = trainer::TrainWeights::from_nnue_weights(&nn);
let mut cache: HashMap<String, i32> = if args.reuse_teacher_cache {
match &args.teacher_cache_path {
Some(p) => teacher_cache::load(p, args.label_depth),
None => {
eprintln!("error: --reuse-teacher-cache requires --teacher-cache <path>");
std::process::exit(1);
}
}
} else {
HashMap::new()
};
let stats = eval_validation_set(&mut trainer, &games, &valid_idxs, &args, &mut cache);
let vloss = if stats.count > 0 {
stats.loss_sum / stats.count as f64
} else {
0.0
};
let cp_mse = if stats.count > 0 {
stats.cp_mse_sum / stats.count as f64
} else {
0.0
};
let wdl_loss = if stats.wdl_count > 0 {
stats.wdl_loss_sum / stats.wdl_count as f64
} else {
0.0
};
let (out_mean, out_std) =
diagnostics::mean_std(stats.output_sum, stats.output_sum_sq, stats.count);
let out_range = if stats.count > 0 {
stats.output_max - stats.output_min
} else {
0.0
};
println!(
"eval-only {:?}: valid_loss={vloss:.4} valid_cp_mse={cp_mse:.4} valid_wdl_loss={wdl_loss:.4} valid_output_mean={out_mean:.3} valid_output_std={out_std:.6} valid_output_range={out_range:.6} samples={} wdl_samples={}",
eval_ckpt, stats.count, stats.wdl_count,
);
return;
}
let mut best_loss = f64::MAX;
let mut prev_snapshot: Option<Vec<f32>> = None;
let mut best_valid_loss = f64::MAX;
let mut best_valid_checkpoint: Option<PathBuf> = None;
let mut teacher_cache: HashMap<String, i32> = if args.reuse_teacher_cache {
match &args.teacher_cache_path {
Some(p) => teacher_cache::load(p, args.label_depth),
None => {
eprintln!("error: --reuse-teacher-cache requires --teacher-cache <path>");
std::process::exit(1);
}
}
} else {
HashMap::new()
};
for epoch in 1..=args.epochs {
trainer.lr = trainer::compute_lr(
args.lr_schedule,
args.lr,
args.min_lr,
epoch as u32,
args.lr_schedule_epochs,
args.warmup_epochs,
);
trainer.reset_epoch_stats();
eprintln!("Epoch {epoch}/{} — lr = {:.6}", args.epochs, trainer.lr);
let epoch_train_idxs: Vec<usize> = if let Some(seed) = args.shuffle_seed {
let order = trainer::shuffled_order(train_idxs.len(), seed ^ epoch as u64);
order.iter().map(|&oi| train_idxs[oi]).collect()
} else {
train_idxs.clone()
};
let train_phase_start = Instant::now();
let mut last_progress = Instant::now();
let total_games = epoch_train_idxs.len();
for (i, &gi) in epoch_train_idxs.iter().enumerate() {
let game = &games[gi];
trainer.train_game(
gi as u64,
game,
args.sample,
args.quiet,
args.min_ply,
args.label_depth,
&scored,
args.stability_weighted,
args.wdl_lambda,
args.wdl_target_scale,
&mut teacher_cache,
);
let game_num = i + 1;
if last_progress.elapsed() >= Duration::from_secs(5) {
let elapsed = train_phase_start.elapsed().as_secs_f64();
let pos_per_sec = if elapsed > 0.0 {
trainer.total_count as f64 / elapsed
} else {
0.0
};
eprintln!(
" progress: game {game_num}/{total_games} (idx {gi}) positions={} elapsed={elapsed:.1}s pos/s={pos_per_sec:.1} cache_hits={} cache_misses={} searches={} search_time={:.1}s",
trainer.total_count,
trainer.cache_hits,
trainer.cache_misses,
trainer.cache_misses,
trainer.search_time_ns as f64 / 1e9,
);
last_progress = Instant::now();
}
if game_num % 10_000 == 0 {
let loss = trainer.avg_loss();
eprintln!(
" epoch {epoch} game {:>7} avg_loss = {:.4}",
game_num, loss
);
if args.best_every > 0 && game_num % args.best_every == 0 && loss < best_loss {
best_loss = loss;
let best_path = args.output.with_extension("best.bin");
let w = trainer.weights.to_nnue_weights();
match save_weights(&w, &best_path) {
Ok(_) => {
eprintln!(" *** best checkpoint (loss={loss:.4}) → {best_path:?}")
}
Err(e) => eprintln!(" best checkpoint save failed: {e}"),
}
}
}
}
let train_phase_secs = train_phase_start.elapsed().as_secs_f64();
let valid_phase_start = Instant::now();
let valid_stats =
eval_validation_set(&mut trainer, &games, &valid_idxs, &args, &mut teacher_cache);
let valid_phase_secs = valid_phase_start.elapsed().as_secs_f64();
let vloss_sum = valid_stats.loss_sum;
let vcount = valid_stats.count;
let valid_cp_mse = if vcount > 0 {
valid_stats.cp_mse_sum / vcount as f64
} else {
0.0
};
let valid_wdl_loss = if valid_stats.wdl_count > 0 {
valid_stats.wdl_loss_sum / valid_stats.wdl_count as f64
} else {
0.0
};
let (valid_output_mean, valid_output_std) =
diagnostics::mean_std(valid_stats.output_sum, valid_stats.output_sum_sq, vcount);
let valid_output_range = if vcount > 0 {
valid_stats.output_max - valid_stats.output_min
} else {
0.0
};
if !valid_idxs.is_empty() {
let vloss = if vcount > 0 {
vloss_sum / vcount as f64
} else {
0.0
};
eprintln!(
" valid: loss={vloss:.4} cp_mse={valid_cp_mse:.4} wdl_loss={valid_wdl_loss:.4} out_mean={valid_output_mean:.3} out_std={valid_output_std:.6} out_range={valid_output_range:.6} samples={vcount}"
);
}
if epoch == 1
&& let Some(cache_path) = &args.teacher_cache_path
{
match teacher_cache::write(cache_path, &teacher_cache, args.label_depth) {
Ok(_) => eprintln!(
" teacher cache written → {:?} ({} entries)",
cache_path,
teacher_cache.len()
),
Err(e) => eprintln!(" teacher cache write failed: {e}"),
}
}
if !scored.is_empty() {
let total_seen = trainer.total_count + trainer.dropped_missing;
let missing_rate = if total_seen > 0 {
trainer.dropped_missing as f64 / total_seen as f64
} else {
0.0
};
let avg_weight = if trainer.total_count > 0 {
trainer.total_weight / trainer.total_count as f64
} else {
1.0
};
eprintln!(
" quietset: entries={} matched={} dropped_missing={} missing_rate={:.1}% avg_weight={:.3}",
scored.len(),
trainer.total_count,
trainer.dropped_missing,
missing_rate * 100.0,
avg_weight,
);
if missing_rate > 0.5 {
eprintln!(
" warn: missing_rate={:.1}% is high — SFEN mismatch or incomplete scored file?",
missing_rate * 100.0
);
}
if trainer.total_count == 0 && trainer.dropped_missing > 0 {
eprintln!(
"error: scored file loaded ({} entries) but 0 positions matched.",
scored.len()
);
eprintln!("hint: scored.jsonl must cover the same games used for training.");
eprintln!(
"hint: check `head -1 scored.jsonl` — sample_id or sfen must be a SFEN string."
);
std::process::exit(1);
}
}
eprintln!(
"Epoch {epoch}/{}: avg_loss = {:.4} samples = {}",
args.epochs,
trainer.avg_loss(),
trainer.total_count,
);
let serialize_phase_start = Instant::now();
let checkpoint = args.output.with_extension(format!("epoch{epoch}.bin"));
let w = trainer.weights.to_nnue_weights();
let mut ckpt_hash = 0u64;
match save_weights(&w, &checkpoint) {
Ok(_) => {
eprintln!(" checkpoint saved → {:?}", checkpoint);
match fs::read(&checkpoint) {
Ok(bytes) => ckpt_hash = checkpoint_hash(&bytes),
Err(e) => eprintln!(" checkpoint hash read failed: {e}"),
}
}
Err(e) => eprintln!(" checkpoint save failed: {e}"),
}
let snapshot = trainer.weights.snapshot_params();
let param_update_norm = prev_snapshot
.as_ref()
.map(|prev| diagnostics::l2_diff_norm(prev, &snapshot));
prev_snapshot = Some(snapshot);
let diag = build_diag(&trainer, &w, param_update_norm);
eprintln!(
" diag: ft_active={:.3} l2_ever_active={:.3} ft_sat={:.3} l2_ever_sat={:.3} l2_dead={} l2_act_freq={:.3} l2_sat_freq={:.3} out_mean={:.3} out_std={:.3} ft_zero={:.3} update_norm={} cache_hit={} cache_miss={}",
diag.ft_active_ratio,
diag.l2_ever_active_ratio,
diag.ft_saturation_ratio,
diag.l2_ever_saturated_ratio,
diag.l2_dead_neurons,
diag.l2_activation_frequency_mean,
diag.l2_saturation_frequency_mean,
diag.output_mean,
diag.output_std,
diag.quantized_ft_zero_ratio,
diag.param_update_norm
.map(|n| format!("{n:.4}"))
.unwrap_or_else(|| "n/a".to_string()),
trainer.cache_hits,
trainer.cache_misses,
);
print_grad_diag_lines(&diag, trainer.total_count);
let meta_path = checkpoint.with_extension("meta.json");
if let Err(e) = save_checkpoint_meta(
&meta_path,
&args,
epoch,
trainer.total_count,
vcount,
Some(train_idxs.len() as u64),
Some(valid_idxs.len() as u64),
split_h,
&diag,
git_commit.as_deref(),
ds_hash,
ckpt_hash,
Some(trainer.cache_hits),
Some(trainer.cache_misses),
Some(trainer.search_time_ns),
Some(&valid_stats),
) {
eprintln!(" metadata save failed: {e}");
} else {
eprintln!(" metadata → {:?}", meta_path);
}
let serialize_phase_secs = serialize_phase_start.elapsed().as_secs_f64();
eprintln!(
" epoch {epoch} summary: games={} positions={} search_time={:.1}s train_phase={:.1}s valid_phase={:.1}s serialize_phase={:.1}s epoch_total={:.1}s cache_hits={} cache_misses={}",
total_games,
trainer.total_count,
trainer.search_time_ns as f64 / 1e9,
train_phase_secs,
valid_phase_secs,
serialize_phase_secs,
train_phase_secs + valid_phase_secs + serialize_phase_secs,
trainer.cache_hits,
trainer.cache_misses,
);
let trace_path = checkpoint.with_extension("trace.json");
if let Err(e) = save_trace_json(&trace_path, &trainer.trace_snapshots) {
eprintln!(" trace save failed: {e}");
} else if !trainer.trace_snapshots.is_empty() {
eprintln!(" trace → {:?}", trace_path);
}
let sample_grad_path = checkpoint.with_extension("sample_grad.jsonl");
if let Err(e) = save_sample_grad_jsonl(&sample_grad_path, &trainer.sample_grad_records) {
eprintln!(" sample-grad trace save failed: {e}");
} else if !trainer.sample_grad_records.is_empty() {
eprintln!(" sample_grad → {:?}", sample_grad_path);
}
let shadow_trace_path = checkpoint.with_extension("shadow_trace.jsonl");
if let Err(e) = save_shadow_trace_jsonl(&shadow_trace_path, &trainer.shadow_trace_records) {
eprintln!(" shadow trace save failed: {e}");
} else if !trainer.shadow_trace_records.is_empty() {
eprintln!(" shadow_trace → {:?}", shadow_trace_path);
}
save_weight_snapshots(&checkpoint, &trainer.weight_snapshots);
if args.validation_ratio > 0.0 && vcount > 0 {
let vloss = vloss_sum / vcount as f64;
if vloss < best_valid_loss {
best_valid_loss = vloss;
best_valid_checkpoint = Some(checkpoint.clone());
}
}
}
if let Some(best_ckpt) = &best_valid_checkpoint {
let best_path = args.output.with_extension("best.bin");
match copy_atomic(best_ckpt, &best_path) {
Ok(_) => eprintln!(
" best (valid_loss={best_valid_loss:.4}) → {:?} (from {:?})",
best_path, best_ckpt
),
Err(e) => eprintln!(" best checkpoint copy failed: {e}"),
}
}
let w = trainer.weights.to_nnue_weights();
match save_weights(&w, &args.output) {
Ok(_) => eprintln!("Final weights saved → {:?}", args.output),
Err(e) => {
eprintln!("Save failed: {e}");
std::process::exit(1);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
#[test]
fn split_games_by_index_partitions_every_index_exactly_once() {
let (train, valid) = split_games_by_index(500, 0.2, 42);
let mut combined: Vec<usize> = train.iter().chain(valid.iter()).copied().collect();
combined.sort_unstable();
let expected: Vec<usize> = (0..500).collect();
assert_eq!(combined, expected);
let train_set: HashSet<usize> = train.into_iter().collect();
let valid_set: HashSet<usize> = valid.into_iter().collect();
assert!(train_set.is_disjoint(&valid_set));
}
#[test]
fn split_games_by_index_zero_ratio_holds_out_nothing() {
let (train, valid) = split_games_by_index(200, 0.0, 42);
assert_eq!(train.len(), 200);
assert!(valid.is_empty());
}
#[test]
fn split_games_by_index_ratio_one_holds_out_everything() {
let (train, valid) = split_games_by_index(200, 1.0, 42);
assert!(train.is_empty());
assert_eq!(valid.len(), 200);
}
#[test]
fn split_games_by_index_is_deterministic_for_the_same_seed() {
let a = split_games_by_index(300, 0.3, 7);
let b = split_games_by_index(300, 0.3, 7);
assert_eq!(a, b);
}
#[test]
fn dataset_hash_is_deterministic_for_the_same_file_list() {
let dir = tempfile::tempdir().unwrap();
let a = dir.path().join("a.csa");
let b = dir.path().join("b.csa");
fs::write(&a, "hello").unwrap();
fs::write(&b, "worldworld").unwrap();
let h1 = dataset_hash(&[a.clone(), b.clone()]);
let h2 = dataset_hash(&[b, a]); assert_eq!(h1, h2);
}
#[test]
fn dataset_hash_changes_when_a_file_size_changes() {
let dir = tempfile::tempdir().unwrap();
let a = dir.path().join("a.csa");
fs::write(&a, "hello").unwrap();
let before = dataset_hash(std::slice::from_ref(&a));
fs::write(&a, "hello, much longer content now").unwrap();
let after = dataset_hash(&[a]);
assert_ne!(before, after);
}
#[test]
fn checkpoint_sidecar_write_is_atomic() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("checkpoint.meta.json");
write_atomic(&path, b"first").unwrap();
write_atomic(&path, b"second").unwrap();
assert_eq!(fs::read(&path).unwrap(), b"second");
assert!(
!path
.with_file_name(format!(".checkpoint.meta.json.tmp-{}", std::process::id()))
.exists()
);
}
#[test]
fn best_checkpoint_copy_is_atomic_and_cleans_up_temp_file() {
let dir = tempfile::tempdir().unwrap();
let source = dir.path().join("weights.epoch1.bin");
let destination = dir.path().join("weights.best.bin");
fs::write(&source, b"new checkpoint").unwrap();
fs::write(&destination, b"old checkpoint").unwrap();
assert_eq!(copy_atomic(&source, &destination).unwrap(), 14);
assert_eq!(fs::read(&destination).unwrap(), b"new checkpoint");
assert_eq!(
fs::read_dir(dir.path())
.unwrap()
.filter_map(Result::ok)
.filter(|entry| entry
.file_name()
.to_string_lossy()
.contains(".best.bin.tmp-"))
.count(),
0
);
}
#[test]
fn split_games_by_index_differs_across_well_separated_seeds() {
let a = split_games_by_index(300, 0.3, 1);
let b = split_games_by_index(300, 0.3, 999_983);
assert_ne!(a, b);
}
}