2df1ea92e1
BREAKING CHANGES: - Removed orphaned dqn.rs monolithic trainer (4,975 lines) - Removed orphaned dqn_ensemble.rs module (816 lines) - Removed orphaned tft.rs and tft_complete_int8_integration_test.rs - TFT trainer split into modular directory structure DQN Module Refactoring: - Split trainers/dqn.rs into modular structure (config.rs, statistics.rs, trainer.rs) - Fixed hyperopt 39D search space (continuous params only) - Boolean flags (use_dueling, use_double_dqn, use_per, use_noisy_nets) are now FIXED architectural decisions - use_distributional defaults to false (Candle BUG #36 - scatter_add gradient issues) Clean Module Structure: - ml/src/trainers/dqn/ directory with proper mod.rs exports - ml/src/trainers/tft/ directory with config.rs, types.rs, model.rs, trainer.rs, tests.rs - All P0 features validated: TD-error clamping, batch diversity, LR scheduler, priority staleness Documentation: - Added comprehensive docs in docs/codebase-cleanup/ - ADR-001 for DQN refactoring decisions - Rainbow DQN component matrix and quick reference guides Build Status: Compiles with zero errors 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
259 lines
21 KiB
Plaintext
259 lines
21 KiB
Plaintext
╔═══════════════════════════════════════════════════════════════════════════╗
|
|
║ TFT TRAINER SPLIT - VISUAL GUIDE ║
|
|
╚═══════════════════════════════════════════════════════════════════════════╝
|
|
|
|
┌───────────────────────────────────────────────────────────────────────────┐
|
|
│ ORIGINAL FILE: ml/src/trainers/tft.rs (2,915 lines) │
|
|
└───────────────────────────────────────────────────────────────────────────┘
|
|
|
|
Lines │ Content │ Destination
|
|
─────────┼────────────────────────────────┼─────────────────────────────
|
|
1-33 │ Header & Imports │ trainer.rs (modified)
|
|
34-100 │ QAT Metrics (5 structs) │ types.rs
|
|
101-202 │ TFTModel Trait │ model.rs
|
|
203-277 │ TFTTrainer Struct │ trainer.rs
|
|
278-332 │ TrainingState │ types.rs (pub(crate))
|
|
333-385 │ Progress Types │ types.rs
|
|
386-538 │ TFTTrainerConfig │ config.rs
|
|
539-2498 │ TFTTrainer impl (30 methods) │ trainer.rs
|
|
2499-2507│ ValidationMetrics │ types.rs (pub(crate))
|
|
2508-2548│ TrainingMetrics │ types.rs
|
|
2550-2915│ Test Module │ tests.rs
|
|
|
|
┌───────────────────────────────────────────────────────────────────────────┐
|
|
│ NEW STRUCTURE: ml/src/trainers/tft/ (6 files) │
|
|
└───────────────────────────────────────────────────────────────────────────┘
|
|
|
|
mod.rs (~50 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ //! Module documentation │
|
|
│ mod config; │
|
|
│ mod model; │
|
|
│ mod trainer; │
|
|
│ mod types; │
|
|
│ #[cfg(test)] mod tests; │
|
|
│ │
|
|
│ pub use config::TFTTrainerConfig; │
|
|
│ pub use model::TFTModel; │
|
|
│ pub use trainer::TFTTrainer; │
|
|
│ pub use types::{QATMetrics, TrainingProgress, ...}; │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
config.rs (~350 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ use serde::{Deserialize, Serialize}; │
|
|
│ use crate::tft::{TFTConfig, training::TFTTrainingConfig}; │
|
|
│ │
|
|
│ pub struct TFTTrainerConfig { │
|
|
│ pub epochs: usize, │
|
|
│ pub learning_rate: f64, │
|
|
│ pub batch_size: usize, │
|
|
│ // ... ~40 fields total │
|
|
│ } │
|
|
│ │
|
|
│ impl Default for TFTTrainerConfig { ... } │
|
|
│ │
|
|
│ impl TFTTrainerConfig { │
|
|
│ pub fn to_model_config(&self) -> TFTConfig { ... } │
|
|
│ pub fn to_training_config(&self) -> TFTTrainingConfig { ... } │
|
|
│ } │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
types.rs (~250 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ use std::collections::HashMap; │
|
|
│ use serde::{Deserialize, Serialize}; │
|
|
│ │
|
|
│ // PUBLIC: QAT Metrics (5 structs, ~67 lines) │
|
|
│ pub struct QATMetrics { ... } │
|
|
│ pub struct ScaleStatistics { ... } │
|
|
│ pub struct ZeroPointStatistics { ... } │
|
|
│ pub struct ObserverRangeStatistics { ... } │
|
|
│ pub struct LayerQuantizationMetrics { ... } │
|
|
│ │
|
|
│ // PUBLIC: Progress Types (~53 lines) │
|
|
│ pub struct TrainingProgress { ... } │
|
|
│ pub struct ResourceUsage { ... } │
|
|
│ │
|
|
│ // PUBLIC: Result Types (~41 lines) │
|
|
│ pub struct TrainingMetrics { ... } │
|
|
│ │
|
|
│ // PRIVATE: Internal State (~64 lines) │
|
|
│ pub(crate) struct TrainingState { ... } │
|
|
│ pub(crate) struct ValidationMetrics { ... } │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
model.rs (~100 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ use std::sync::Arc; │
|
|
│ use candle_core::{Device, Tensor}; │
|
|
│ use candle_nn::VarMap; │
|
|
│ use crate::tft::{TFTConfig, TemporalFusionTransformer}; │
|
|
│ use crate::MLError; │
|
|
│ │
|
|
│ pub trait TFTModel: Send + Sync { │
|
|
│ fn forward(...) -> Result<Tensor, MLError>; │
|
|
│ fn get_device(&self) -> &Device; │
|
|
│ fn get_config(&self) -> &TFTConfig; │
|
|
│ fn get_varmap(&self) -> Arc<VarMap>; │
|
|
│ fn clear_cache(&mut self); │
|
|
│ } │
|
|
│ │
|
|
│ impl TFTModel for TemporalFusionTransformer { ... } │
|
|
│ │
|
|
│ /* Future QAT implementation */ │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
trainer.rs (~2,000 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ use std::collections::HashMap; │
|
|
│ use std::sync::Arc; │
|
|
│ // ... external imports ... │
|
|
│ │
|
|
│ use super::config::TFTTrainerConfig; │
|
|
│ use super::model::TFTModel; │
|
|
│ use super::types::{QATMetrics, TrainingProgress, ...}; │
|
|
│ │
|
|
│ pub struct TFTTrainer { │
|
|
│ model_config: TFTConfig, │
|
|
│ training_config: TFTTrainingConfig, │
|
|
│ model: Box<dyn TFTModel>, │
|
|
│ optimizer: Option<crate::Adam>, │
|
|
│ // ... ~20 more fields │
|
|
│ } │
|
|
│ │
|
|
│ impl TFTTrainer { │
|
|
│ // Constructor │
|
|
│ pub fn new(...) -> MLResult<Self> { ... } │
|
|
│ │
|
|
│ // Public API (3 methods) │
|
|
│ pub fn set_progress_callback(...) { ... } │
|
|
│ pub async fn train(...) -> MLResult<TrainingMetrics> { ... } │
|
|
│ │
|
|
│ // Training Methods (6 methods) │
|
|
│ async fn train_epoch(...) -> MLResult<f64> { ... } │
|
|
│ async fn validate_epoch(...) -> MLResult<ValidationMetrics> { ... } │
|
|
│ fn batch_to_tensors(...) { ... } │
|
|
│ fn compute_quantile_loss(...) { ... } │
|
|
│ fn compute_rmse(...) { ... } │
|
|
│ fn extract_attention_entropy(...) { ... } │
|
|
│ │
|
|
│ // QAT Methods (5 methods) │
|
|
│ async fn run_qat_calibration(...) { ... } │
|
|
│ async fn qat_to_quantized_checkpoint(...) { ... } │
|
|
│ fn export_qat_metrics(...) { ... } │
|
|
│ fn apply_qat_lr_schedule(...) { ... } │
|
|
│ async fn quantize_and_save_int8_checkpoint(...) { ... } │
|
|
│ │
|
|
│ // Helper Methods (15+ methods) │
|
|
│ fn initialize_optimizer(...) { ... } │
|
|
│ fn check_early_stopping(...) { ... } │
|
|
│ async fn save_checkpoint(...) { ... } │
|
|
│ async fn send_progress_update(...) { ... } │
|
|
│ fn get_resource_usage(...) { ... } │
|
|
│ // ... more helpers ... │
|
|
│ } │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
tests.rs (~165 lines)
|
|
═══════════════════════════════════════════════════════════════════════════
|
|
┌─────────────────────────────────────────────────────────────────────────┐
|
|
│ use super::*; │
|
|
│ use crate::checkpoint::FileSystemStorage; │
|
|
│ use std::path::PathBuf; │
|
|
│ use tempfile::TempDir; │
|
|
│ │
|
|
│ #[tokio::test] │
|
|
│ async fn test_tft_trainer_creation() { ... } │
|
|
│ │
|
|
│ #[tokio::test] │
|
|
│ async fn test_training_config_conversion() { ... } │
|
|
│ │
|
|
│ #[tokio::test] │
|
|
│ async fn test_checkpoint_save_load() { ... } │
|
|
│ │
|
|
│ #[tokio::test] │
|
|
│ async fn test_oom_retry_exponential_backoff() { ... } │
|
|
│ │
|
|
│ #[tokio::test] │
|
|
│ async fn test_oom_retry_minimum_batch_size() { ... } │
|
|
└─────────────────────────────────────────────────────────────────────────┘
|
|
|
|
┌───────────────────────────────────────────────────────────────────────────┐
|
|
│ DEPENDENCY FLOW │
|
|
└───────────────────────────────────────────────────────────────────────────┘
|
|
|
|
┌─────────────┐
|
|
│ mod.rs │ (orchestrates all)
|
|
└──────┬──────┘
|
|
│
|
|
┌────────────────┼────────────────┐
|
|
│ │ │
|
|
▼ ▼ ▼
|
|
┌──────────┐ ┌──────────┐ ┌──────────┐
|
|
│config.rs │ │ types.rs │ │ model.rs │
|
|
└──────────┘ └──────────┘ └──────────┘
|
|
│ │ │
|
|
└────────────────┼────────────────┘
|
|
│
|
|
▼
|
|
┌─────────────┐
|
|
│ trainer.rs │ (uses all above)
|
|
└──────┬──────┘
|
|
│
|
|
▼
|
|
┌─────────────┐
|
|
│ tests.rs │ (tests trainer)
|
|
└─────────────┘
|
|
|
|
┌───────────────────────────────────────────────────────────────────────────┐
|
|
│ EXTRACTION CHECKLIST │
|
|
└───────────────────────────────────────────────────────────────────────────┘
|
|
|
|
Step 1: Create types.rs
|
|
✓ Extract lines 34-100 (QAT metrics)
|
|
✓ Extract lines 278-332 (TrainingState - make pub(crate))
|
|
✓ Extract lines 333-385 (Progress types)
|
|
✓ Extract lines 2,499-2,507 (ValidationMetrics - make pub(crate))
|
|
✓ Extract lines 2,508-2,548 (TrainingMetrics)
|
|
✓ Add module header and imports
|
|
|
|
Step 2: Create model.rs
|
|
✓ Extract lines 101-202 (TFTModel trait + impls)
|
|
✓ Add module header and imports
|
|
|
|
Step 3: Create config.rs
|
|
✓ Extract lines 386-538 (TFTTrainerConfig)
|
|
✓ Add module header and imports
|
|
|
|
Step 4: Create trainer.rs
|
|
✓ Copy header from lines 1-33 (modified)
|
|
✓ Extract lines 203-277 (TFTTrainer struct)
|
|
✓ Extract lines 539-2,498 (impl blocks)
|
|
✓ Update imports to use super::{config, model, types}
|
|
|
|
Step 5: Create tests.rs
|
|
✓ Extract lines 2,550-2,915 (test module)
|
|
✓ Remove #[cfg(test)] attribute
|
|
✓ Add use super::*;
|
|
|
|
Step 6: Create mod.rs
|
|
✓ Write new file with module declarations
|
|
✓ Add re-exports for public API
|
|
|
|
Step 7: Verify
|
|
✓ cargo check --package ml
|
|
✓ cargo test --package ml trainers::tft
|
|
✓ cargo clippy --package ml
|
|
|
|
╔═══════════════════════════════════════════════════════════════════════════╗
|
|
║ READY TO PROCEED ║
|
|
║ ║
|
|
║ All documentation complete. Implementation agent can begin extraction. ║
|
|
╚═══════════════════════════════════════════════════════════════════════════╝
|