gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(eval): shape-mismatch on checkpoint load — read arch from safetensors metadata (atomic)

Browse Source 
Smoke v1 (train-grfcw) evaluate phase failed with "Failed to load DQN
checkpoint" for fold 0 and fold 1. MinIO log inspection confirmed
checkpoints WERE saved (1431144 bytes each) — the failure was
eval-side shape mismatch.

Root cause:
- Training uses STATE_DIM=128 (ml_core::state_layout), num_actions=108
  (factored b0*b1*b2*b3=4*3*3*3), num_order_types=3,
  num_urgency_levels=3.
- evaluate_baseline CLI defaults: --feature-dim=54, --num-actions=5
  (legacy from pre-branching DQN era).
- Loading 128-state-dim 108-action checkpoint into 54-feature 5-action
  net → tensor shape mismatch → `load_from_safetensors` returned
  parse error → `with_context(...)` wrapped it as the generic "Failed
  to load DQN checkpoint" message, hiding the actual shape error.
- Both GPU and CPU eval paths hit the same root cause.

Fix:
Both eval paths now call `DQNConfig::from_safetensors_file(&ckpt_path)`
to read architecture-critical fields from the checkpoint's embedded
metadata (state_dim, num_actions, hidden_dims, num_order_types,
num_urgency_levels, dueling_hidden_dim, num_atoms, gamma). Eval-time
fields (LR, epsilon, buffer caps) overridden; hyperopt-derived gamma/
v_min/v_max applied if present in hyperopt config.

Older checkpoints without embedded metadata fall back to CLI-args-built
config + warn! log. All production SP21+ checkpoints embed metadata
via the existing DQNConfig::checkpoint_metadata path.

Files changed:
- crates/ml/examples/evaluate_baseline.rs: shape-aware config for both
  dqn_eval_gpu_path (line ~1238) and dqn_eval_cpu_path (line ~1029)
- docs/dqn-wire-up-audit.md: 2026-05-11 audit entry

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, in-flight on 2937da889) won't
have this fix; smoke v4 dispatch on this commit will validate
evaluate phase succeeds for all folds. Look for
"[DQN GPU] Architecture from checkpoint: ..." log line per fold.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-11 08:28:53 +02:00
parent 2937da8898
commit 62b5a50e8b
2 changed files with 275 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

217
crates/ml/examples/evaluate_baseline.rs
View File

@@ -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)")?;

116
docs/dqn-wire-up-audit.md
View File

@@ -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)