Best-checkpoint save (val Sharpe improvement, ~30% of epochs in convergent runs) blocked the epoch loop for 20-40s on each improvement: serialize_model() chained N (~26) per-tensor DtoH downloads via GpuTensor::to_host → memcpy_dtoh, each forcing an implicit stream sync on a busy training stream. The DtoH chain also violates feedback_no_htod_htoh_only_mapped_pinned.md (only cuMemHostAlloc DEVICEMAP allowed for CPU↔GPU paths). Plan B: - Introduce snapshot_model_to_pinned (mod.rs): allocate one MappedF32Buffer sized for all named weight slices concatenated, cuMemcpyDtoDAsync each slice into the buffer's device pointer (which aliases the host page), single stream sync, copy bytes out to a Send + 'static Vec<u8>. One sync per snapshot, replaces N. - serialize_snapshot_bytes (mod.rs): pure-CPU safetensors construction from CheckpointSnapshot. Static — callable without &self, so the worker can move the snapshot across thread boundary. - handle_epoch_checkpoints_and_early_stopping on val-Sharpe improvement: save_best_gpu_params (DtoD, fast) + snapshot to pinned + tokio::task::spawn_blocking the safetensors construction + checkpoint_callback invocation. JoinHandle parked on pending_checkpoint_handles. Training loop continues immediately. - await_pending_checkpoint_handles drains in-flight workers at training end (success branch + early-stop branches) and before any synchronous cold-path checkpoint write to keep disk ordering deterministic. - F bound on train / train_walk_forward / train_fold_from_slices gains + 'static so the callback can be moved into the worker. All public callers already use 'static-compatible move closures (test fixtures with shared mutable state migrate to Arc<Mutex<T>>). Internal pipeline uses CheckpointCallbackHandle = Arc<std::sync::Mutex<Box<dyn FnMut + Send + 'static>>> so the same callback flows through multi-fold walk-forward into every fold's worker. - serialize_model itself rewritten via the snapshot path: the no-DtoH rule now holds across ALL checkpoint paths (best, periodic, early-stop, plateau-exhausted). The pre-existing GpuTensor::to_host path is no longer reachable from the DQN trainer. The audit's spec called for an mpsc channel(1) drop-old worker, but the multi-fold + &mut F pre-existing API made the simpler fire-and-forget spawn_blocking pattern a cleaner fit (Mutex serialises any concurrent invocations; Vec<JoinHandle> drain at end guarantees disk writes complete before the trainer returns). Same overlap benefit (training rolls while serialize+disk run on a blocking thread); upper bound on in-flight work is one-per-improved- epoch which approximates the spec's depth=1 in realistic training runs. Per feedback_no_partial_refactor: every site that constructs a checkpoint payload migrated in lockstep — best-improvement uses the worker; periodic / plateau-exhausted / early-stop call the shared Arc<Mutex<F>> handle inline. All paths read params via snapshot_model_to_pinned, so the no-DtoH rule applies uniformly. Test fixtures (8 .rs files) updated for the + 'static bound (move closures + cloned PathBufs / Arc<Mutex<T>> for shared mutable state). Verified: SQLX_OFFLINE=true cargo check --workspace --tests clean (warnings unchanged from baseline). cargo test -p ml --lib --no-run clean. No fingerprint change. Wire-up audit entry extended with Plan B file:line edit sites (rides under the same Async-validation overlap section started by the companion Plan A commit). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
163 lines
5.1 KiB
Rust
163 lines
5.1 KiB
Rust
#![allow(
|
|
clippy::assertions_on_constants,
|
|
clippy::assertions_on_result_states,
|
|
clippy::clone_on_copy,
|
|
clippy::decimal_literal_representation,
|
|
clippy::doc_markdown,
|
|
clippy::empty_line_after_doc_comments,
|
|
clippy::field_reassign_with_default,
|
|
clippy::get_unwrap,
|
|
clippy::identity_op,
|
|
clippy::inconsistent_digit_grouping,
|
|
clippy::indexing_slicing,
|
|
clippy::integer_division,
|
|
clippy::len_zero,
|
|
clippy::let_underscore_must_use,
|
|
clippy::manual_div_ceil,
|
|
clippy::manual_let_else,
|
|
clippy::manual_range_contains,
|
|
clippy::modulo_arithmetic,
|
|
clippy::needless_range_loop,
|
|
clippy::non_ascii_literal,
|
|
clippy::redundant_clone,
|
|
clippy::shadow_reuse,
|
|
clippy::shadow_same,
|
|
clippy::shadow_unrelated,
|
|
clippy::single_match_else,
|
|
clippy::str_to_string,
|
|
clippy::string_slice,
|
|
clippy::tests_outside_test_module,
|
|
clippy::too_many_lines,
|
|
clippy::unnecessary_wraps,
|
|
clippy::unseparated_literal_suffix,
|
|
clippy::use_debug,
|
|
clippy::useless_vec,
|
|
clippy::wildcard_enum_match_arm,
|
|
clippy::else_if_without_else,
|
|
clippy::expect_used,
|
|
clippy::missing_const_for_fn,
|
|
clippy::similar_names,
|
|
clippy::type_complexity,
|
|
clippy::collapsible_else_if,
|
|
clippy::doc_lazy_continuation,
|
|
clippy::items_after_test_module,
|
|
clippy::map_clone,
|
|
clippy::multiple_unsafe_ops_per_block,
|
|
clippy::unwrap_or_default,
|
|
clippy::assign_op_pattern,
|
|
clippy::needless_borrow,
|
|
clippy::println_empty_string,
|
|
clippy::unnecessary_cast,
|
|
clippy::used_underscore_binding,
|
|
clippy::create_dir,
|
|
clippy::implicit_saturating_sub,
|
|
clippy::exit,
|
|
clippy::expect_fun_call,
|
|
clippy::too_many_arguments,
|
|
clippy::unnecessary_map_or,
|
|
clippy::unwrap_used,
|
|
dead_code,
|
|
unused_imports,
|
|
unused_variables,
|
|
clippy::cloned_ref_to_slice_refs,
|
|
clippy::neg_multiply,
|
|
clippy::while_let_loop,
|
|
clippy::bool_assert_comparison,
|
|
clippy::excessive_precision,
|
|
clippy::trivially_copy_pass_by_ref,
|
|
clippy::op_ref,
|
|
clippy::redundant_closure,
|
|
clippy::unnecessary_lazy_evaluations,
|
|
clippy::if_then_some_else_none,
|
|
clippy::unnecessary_to_owned,
|
|
clippy::single_component_path_imports,
|
|
)]
|
|
//! Tests for true gradient accumulation in DQN training.
|
|
//!
|
|
//! Verifies that `train_step_with_accumulation()` performs exactly 1 optimizer
|
|
//! step per call (not N steps), proving true accumulation works correctly.
|
|
|
|
#![allow(unused_crate_dependencies)]
|
|
|
|
use anyhow::{Context, Result};
|
|
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
|
use std::path::PathBuf;
|
|
use tracing::info;
|
|
use tracing::warn;
|
|
|
|
fn get_6e_fut_data_dir() -> Result<String> {
|
|
let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
|
.parent()
|
|
.context("Failed to get workspace root")?
|
|
.to_path_buf();
|
|
let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
|
|
if !data_dir.exists() {
|
|
anyhow::bail!("6E.FUT data not found: {}", data_dir.display());
|
|
}
|
|
Ok(data_dir.to_string_lossy().to_string())
|
|
}
|
|
|
|
/// Verify gradient accumulation with steps=4 produces training with finite losses.
|
|
#[tokio::test]
|
|
async fn test_accumulation_single_optimizer_step() -> Result<()> {
|
|
let data_dir = match get_6e_fut_data_dir() {
|
|
Ok(dir) => dir,
|
|
Err(e) => {
|
|
warn!(reason = %e, "Skipping test: data not available");
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
let checkpoint_dir = tempfile::tempdir()?;
|
|
let ckpt_path = checkpoint_dir.path().to_path_buf();
|
|
|
|
let mut hyperparams = DQNHyperparameters::conservative();
|
|
hyperparams.replay_buffer_vram_fraction = 0.0; // Disable AutoReplaySizer for test determinism
|
|
hyperparams.epochs = 5;
|
|
hyperparams.batch_size = 32;
|
|
hyperparams.learning_rate = 0.0001;
|
|
hyperparams.gradient_accumulation_steps = 4;
|
|
hyperparams.early_stopping_enabled = false;
|
|
hyperparams.checkpoint_frequency = 100;
|
|
|
|
let mut trainer = DQNTrainer::new(hyperparams)?;
|
|
let metrics = trainer
|
|
.train(&data_dir, "ES.FUT", move |_epoch, checkpoint_data, _is_best| {
|
|
let path = ckpt_path.join("accum_test.safetensors");
|
|
std::fs::write(&path, &checkpoint_data)?;
|
|
Ok(path.to_string_lossy().to_string())
|
|
})
|
|
.await?;
|
|
|
|
assert_eq!(
|
|
metrics.epochs_trained, 5,
|
|
"Should complete all 5 epochs"
|
|
);
|
|
|
|
let loss_history = trainer.loss_history();
|
|
assert!(
|
|
loss_history.len() >= 2,
|
|
"Need at least 2 epochs of loss history"
|
|
);
|
|
let initial_loss = loss_history.first().copied().unwrap_or(0.0);
|
|
let final_loss = loss_history.last().copied().unwrap_or(0.0);
|
|
assert!(
|
|
initial_loss.is_finite() && final_loss.is_finite(),
|
|
"Losses must be finite: initial={}, final={}",
|
|
initial_loss,
|
|
final_loss
|
|
);
|
|
|
|
info!(
|
|
epochs = metrics.epochs_trained,
|
|
accumulation_steps = 4,
|
|
initial_loss,
|
|
final_loss,
|
|
effective_batch = 32 * 4,
|
|
training_time_seconds = metrics.training_time_seconds,
|
|
"Gradient accumulation test results"
|
|
);
|
|
|
|
Ok(())
|
|
}
|