diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index c9b92abd1..62db92f41 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -1025,37 +1025,77 @@ fn evaluate_dqn_fold( // Create DQN with same config as training — all architecture-affecting params // must match exactly, otherwise checkpoint loading fails (tensor shape mismatch). + // + // **Shape-mismatch fix (2026-05-11)** — same as the GPU eval path above: + // pull architecture-critical fields from checkpoint metadata via + // `DQNConfig::from_safetensors_file`. The CLI defaults (`num_actions=5`, + // `feature_dim=54`) are legacy and would silently shape-mismatch a + // post-SP21 checkpoint (108 actions, STATE_DIM=128). #[allow(clippy::integer_division)] - let config = DQNConfig { - num_actions: args.num_actions, - hidden_dims: { - let base = hp_usize(hp, "hidden_dim_base").unwrap_or(256); - // Align to 8 for tensor cores (matches training) - let align = |x: usize| -> usize { x.div_ceil(8) * 8 }; - vec![align(base), align(base / 2), align(base / 4)] - }, - learning_rate: 1e-4, - gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32, - epsilon_start: 0.0, // No exploration during evaluation - epsilon_end: 0.0, - epsilon_decay: 1.0, - replay_buffer_capacity: 100, // Minimal buffer, not used for eval - batch_size: 64, - min_replay_size: 64, - target_update_freq: 500, - warmup_steps: 0, - // Architecture params — must match training checkpoint shapes - dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), - num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), - v_min: hp_f64(hp, "v_min").unwrap_or_else(|| { - let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); - -(10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) - }) as f32, - v_max: hp_f64(hp, "v_max").unwrap_or_else(|| { - let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); - (10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) - }) as f32, - ..DQNConfig::default() + let config = match DQNConfig::from_safetensors_file(&ckpt_path) { + Ok(mut cfg) => { + cfg.learning_rate = 1e-4; + cfg.epsilon_start = 0.0; + cfg.epsilon_end = 0.0; + cfg.epsilon_decay = 1.0; + cfg.replay_buffer_capacity = 100; + cfg.batch_size = 64; + cfg.min_replay_size = 64; + cfg.target_update_freq = 500; + cfg.warmup_steps = 0; + if let Some(g) = hp_f64(hp, "gamma") { + cfg.gamma = g as f32; + } + if let Some(vmin) = hp_f64(hp, "v_min") { + cfg.v_min = vmin as f32; + } + if let Some(vmax) = hp_f64(hp, "v_max") { + cfg.v_max = vmax as f32; + } + info!( + " [DQN CPU] Architecture from checkpoint: num_actions={}, hidden_dims={:?}, \ + num_atoms={}, num_order_types={}, num_urgency_levels={}, dueling_hidden_dim={}", + cfg.num_actions, cfg.hidden_dims, cfg.num_atoms, + cfg.num_order_types, cfg.num_urgency_levels, cfg.dueling_hidden_dim, + ); + cfg + } + Err(e) => { + warn!( + " [DQN CPU] No architecture metadata in checkpoint ({}); falling back to \ + CLI args (num_actions={}, may shape-mismatch): {}", + ckpt_path.display(), args.num_actions, e, + ); + DQNConfig { + num_actions: args.num_actions, + hidden_dims: { + let base = hp_usize(hp, "hidden_dim_base").unwrap_or(256); + let align = |x: usize| -> usize { x.div_ceil(8) * 8 }; + vec![align(base), align(base / 2), align(base / 4)] + }, + learning_rate: 1e-4, + gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32, + epsilon_start: 0.0, + epsilon_end: 0.0, + epsilon_decay: 1.0, + replay_buffer_capacity: 100, + batch_size: 64, + min_replay_size: 64, + target_update_freq: 500, + warmup_steps: 0, + dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), + num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), + v_min: hp_f64(hp, "v_min").unwrap_or_else(|| { + let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); + -(10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) + }) as f32, + v_max: hp_f64(hp, "v_max").unwrap_or_else(|| { + let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); + (10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) + }) as f32, + ..DQNConfig::default() + } + } }; let mut dqn = DQN::new(config).context("Failed to create DQN model")?; @@ -1234,35 +1274,96 @@ fn evaluate_dqn_fold_gpu( } }; + // **Shape-mismatch fix (2026-05-11)** — read architecture-critical + // config (num_actions, hidden_dims, num_atoms, num_order_types, + // num_urgency_levels, dueling_hidden_dim) from the checkpoint's + // safetensors metadata instead of CLI args. Defaults + // `args.num_actions=5` and `args.feature_dim=54` are legacy from + // the pre-branching DQN era; current training uses factored 108 + // actions over `STATE_DIM=128`. Loading a 128-state-dim + // 108-action checkpoint into a 5-action net silently shape- + // mismatches → load fails with a tensor-shape error (the + // smoke v1 failure mode). + // + // `from_safetensors_file` parses architecture-critical fields + // and validates `state_dim == STATE_DIM` (compile-time constant + // = 128). Falls back to a CLI-args-built config if metadata is + // missing (older checkpoints pre-`arch_hash` embed). #[allow(clippy::integer_division)] - let config = DQNConfig { - num_actions: args.num_actions, - hidden_dims: { - let base = hp_usize(hp, "hidden_dim_base").unwrap_or(256); - let align = |x: usize| -> usize { x.div_ceil(8) * 8 }; - vec![align(base), align(base / 2), align(base / 4)] - }, - learning_rate: 1e-4, - gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32, - epsilon_start: 0.0, - epsilon_end: 0.0, - epsilon_decay: 1.0, - replay_buffer_capacity: 100, - batch_size: 64, - min_replay_size: 64, - target_update_freq: 500, - warmup_steps: 0, - dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), - num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), - v_min: hp_f64(hp, "v_min").unwrap_or_else(|| { - let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); - -(10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) - }) as f32, - v_max: hp_f64(hp, "v_max").unwrap_or_else(|| { - let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); - (10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) - }) as f32, - ..DQNConfig::default() + let config = match DQNConfig::from_safetensors_file(&ckpt_path) { + Ok(mut cfg) => { + // Override eval-time fields (LR, epsilon, buffer + // capacity, batch sizes) — these don't affect inference + // but the checkpoint reader returns defaults that aren't + // tuned for eval. Override hyperopt-derived v_min/v_max + // if hp is present; otherwise the checkpoint metadata's + // values are already correct for inference. + cfg.learning_rate = 1e-4; + cfg.epsilon_start = 0.0; + cfg.epsilon_end = 0.0; + cfg.epsilon_decay = 1.0; + cfg.replay_buffer_capacity = 100; + cfg.batch_size = 64; + cfg.min_replay_size = 64; + cfg.target_update_freq = 500; + cfg.warmup_steps = 0; + if let Some(g) = hp_f64(hp, "gamma") { + cfg.gamma = g as f32; + } + if let Some(vmin) = hp_f64(hp, "v_min") { + cfg.v_min = vmin as f32; + } + if let Some(vmax) = hp_f64(hp, "v_max") { + cfg.v_max = vmax as f32; + } + info!( + " [DQN GPU] Architecture from checkpoint: num_actions={}, hidden_dims={:?}, \ + num_atoms={}, num_order_types={}, num_urgency_levels={}, dueling_hidden_dim={}", + cfg.num_actions, cfg.hidden_dims, cfg.num_atoms, + cfg.num_order_types, cfg.num_urgency_levels, cfg.dueling_hidden_dim, + ); + cfg + } + Err(e) => { + // Older checkpoint without metadata — fall back to + // CLI-args-built config and warn. Production + // checkpoints from SP21 onward DO embed metadata + // (see DQN::checkpoint_metadata). + warn!( + " [DQN GPU] No architecture metadata in checkpoint ({}); falling back to \ + CLI args (num_actions={}, may shape-mismatch): {}", + ckpt_path.display(), args.num_actions, e, + ); + DQNConfig { + num_actions: args.num_actions, + hidden_dims: { + let base = hp_usize(hp, "hidden_dim_base").unwrap_or(256); + let align = |x: usize| -> usize { x.div_ceil(8) * 8 }; + vec![align(base), align(base / 2), align(base / 4)] + }, + learning_rate: 1e-4, + gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32, + epsilon_start: 0.0, + epsilon_end: 0.0, + epsilon_decay: 1.0, + replay_buffer_capacity: 100, + batch_size: 64, + min_replay_size: 64, + target_update_freq: 500, + warmup_steps: 0, + dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), + num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), + v_min: hp_f64(hp, "v_min").unwrap_or_else(|| { + let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); + -(10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) + }) as f32, + v_max: hp_f64(hp, "v_max").unwrap_or_else(|| { + let gamma = hp_f64(hp, "gamma").unwrap_or(0.95); + (10.0 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0) + }) as f32, + ..DQNConfig::default() + } + } }; let mut dqn = DQN::new(config).context("Failed to create DQN model (GPU path)")?; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 684a4852c..cf37e28b0 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -15098,6 +15098,122 @@ will validate both fixes: - evaluate phase should find `dqn_fold{0,1,2}_best.safetensors` for each fold and proceed without "Failed to load DQN checkpoint". +## 2026-05-11 — evaluate_baseline shape-mismatch fix: load config from checkpoint metadata (atomic) + +### Scope (atomic single commit) + +Smoke v1 (train-grfcw) evaluate phase failed with: +``` +[DQN GPU] Fold 0 GPU eval failed: Failed to load DQN checkpoint + (GPU path): /workspace/output/dqn_fold0_best.safetensors. + Falling back to CPU path. +[DQN] Fold 0 evaluation failed: Failed to load DQN checkpoint: + /workspace/output/dqn_fold0_best.safetensors +``` + +Investigation via MinIO logs (mc cp argo-logs/train-grfcw/...) confirmed: +- Checkpoints WERE saved successfully: + `Best model saved to: /workspace/output/dqn_fold{0,1,2}_best.safetensors + (off-thread, 1431144 bytes)` — all three folds. +- Load failed in `evaluate_baseline.rs` because eval used CLI-default + config: `args.feature_dim=54`, `args.num_actions=5`. +- Production training uses `STATE_DIM=128` (per + `ml_core::state_layout`), `num_actions=108` (factored + `b0*b1*b2*b3 = 4*3*3*3`), `num_order_types=3`, + `num_urgency_levels=3`. +- Loading 128-state-dim 108-action checkpoint into 54-feature-dim + 5-action net → tensor shape mismatch → `load_from_safetensors` + returned `Failed to parse safetensors` (silently swallowed by + `with_context(...)` wrap), GPU path warned + fell back to CPU + path, CPU path failed with the same root cause. + +### Fix + +Both eval paths (`dqn_eval_gpu_path` line ~1238 and `dqn_eval_cpu_path` +line ~1029) now use `DQNConfig::from_safetensors_file(&ckpt_path)` +to read architecture-critical config directly from the checkpoint's +embedded safetensors metadata: + +```rust +let config = match DQNConfig::from_safetensors_file(&ckpt_path) { + Ok(mut cfg) => { + // Override eval-time fields (LR, epsilon, buffer caps don't + // affect inference) and apply hyperopt-derived overrides + // (gamma, v_min, v_max) if present. + cfg.learning_rate = 1e-4; + cfg.epsilon_start = 0.0; cfg.epsilon_end = 0.0; + cfg.epsilon_decay = 1.0; + cfg.replay_buffer_capacity = 100; + cfg.batch_size = 64; cfg.min_replay_size = 64; + cfg.target_update_freq = 500; cfg.warmup_steps = 0; + if let Some(g) = hp_f64(hp, "gamma") { cfg.gamma = g as f32; } + if let Some(vmin) = hp_f64(hp, "v_min") { cfg.v_min = vmin as f32; } + if let Some(vmax) = hp_f64(hp, "v_max") { cfg.v_max = vmax as f32; } + info!("Architecture from checkpoint: num_actions={}, hidden_dims={:?}, ...", ...); + cfg + } + Err(e) => { + // Older checkpoint without metadata — fall back to CLI args. + warn!("No architecture metadata in checkpoint; falling back to CLI args ..."); + // (CLI-args-built config preserved for backwards compat) + } +}; +``` + +`DQNConfig::from_safetensors_file` parses the metadata block +embedded by `DQNConfig::checkpoint_metadata()` at save time +(`dqn.state_dim`, `dqn.num_actions`, `dqn.hidden_dims`, +`dqn.num_order_types`, `dqn.num_urgency_levels`, +`dqn.dueling_hidden_dim`, `dqn.num_atoms`, `dqn.iqn_lambda`, +`dqn.gamma`) and validates `saved_state_dim == STATE_DIM` +(compile-time constant). Architecture-critical fields are now +GUARANTEED to match training even when the eval invocation +uses stale CLI defaults. + +### Backwards compatibility + +Older checkpoints without metadata fall back to CLI-args-built +config + `warn!` log. Production checkpoints from SP21 onward +ALL embed metadata via the existing +`save_to_safetensors`/`checkpoint_metadata` paths. No +behavioral regression for legacy checkpoints; new checkpoints +work without CLI tuning. + +### Files changed + +| File | Status | Purpose | +|------|--------|---------| +| `crates/ml/examples/evaluate_baseline.rs` | Shape-aware config | Both `dqn_eval_gpu_path` and `dqn_eval_cpu_path` use `DQNConfig::from_safetensors_file` for architecture; CLI args become fallback for legacy checkpoints | +| `docs/dqn-wire-up-audit.md` | This entry | 2026-05-11 audit log | + +### Pearls + invariants honoured + +- `feedback_no_partial_refactor` — both eval paths (GPU + CPU) get + the same fix atomically. +- `feedback_no_stubs` — fallback CLI-args path retained with + explicit `warn!` log instead of silent CLI-default use. +- `feedback_trust_code_not_docs` — the v1 eval failure looked + like a "checkpoint not saved" bug per the error message, but + MinIO log inspection (1431144 bytes confirmed) showed + checkpoints WERE saved. The actual root cause was eval-side + config mismatch — the error message was misleading because + `with_context(...)` wraps don't print the underlying + shape-mismatch detail by default. Lesson: when a "fix" doesn't + match the actual data, trust the data. + +### Verification + +``` +cargo check -p ml --examples --features cuda: 0 errors +cargo test -p ml --lib financials: 7/7 (unchanged) +cargo test -p ml --lib sp21_isv_slots: 4/4 (unchanged) +``` + +Behavioral gate: smoke v3 (train-psf86) evaluate phase should +succeed for all 3 folds. Look for `[DQN GPU] Architecture from +checkpoint: ...` log line confirming metadata read succeeded +for each fold. + ## 2026-05-11 — SP21 T2.2 Phase 7.5: E8 curriculum_weights → per-segment PER insert priority boost (atomic) ### Scope (atomic single commit)