foxhunt/WAVE_1_AGENT_2_ML_TRAINING_ANALYSIS.md

WAVE 1 AGENT 2: ML Training Orchestration Analysis

Date: 2025-10-15 Agent: Claude Code Agent 2 Mission: Analyze ml/tests/ for training orchestration test failures Status: ANALYSIS COMPLETE - Test failures documented, missing implementations identified

---

Executive Summary

Test Status: 40/40 tests defined, 0/40 tests passing (100% failure rate) Root Cause: Missing UnifiedTrainable trait implementations for all 4 models Infrastructure: Job Queue and Redis integration COMPLETE, GPU management READY Deliverable: Comprehensive test failure matrix with remediation requirements

Critical Finding: The UnifiedTrainable trait exists and is well-designed, but NO models implement it yet. This is a missing glue layer problem, not an architectural issue.

---

Test Failure Matrix

1. MAMBA-2 Tests (10 tests, 0 passing)

Test Name	Status	Error Type	Required Fix
test_mamba2_trait_implementation	FAIL	No UnifiedTrainable impl	Implement trait for Mamba2SSM
test_mamba2_forward_pass	FAIL	Method call works, but no trait	Trait implementation
test_mamba2_backward_pass	FAIL	Method train_batch is private	Make public + trait impl
test_mamba2_optimizer_step	FAIL	No initialize_optimizer method	Add method + trait impl
test_mamba2_checkpoint_save	FAIL	Async checkpoint methods exist	Wrap in trait impl
test_mamba2_checkpoint_load	FAIL	Async checkpoint methods exist	Wrap in trait impl
test_mamba2_metrics_collection	FAIL	get_performance_metrics exists	Map to collect_metrics trait method
test_mamba2_training_step	FAIL	Method train_batch is private	Make public + trait impl
test_mamba2_device_transfer	FAIL	Device field exists	Expose via trait method
test_mamba2_nan_detection	FAIL	Training logic exists	Test via trait interface

MAMBA-2 Status: Infrastructure READY, only needs trait wrapper Implementation Effort: ~200 lines (trait impl + public method wrappers)

---

2. DQN Tests (10 tests, 0 passing)

Test Name	Status	Error Type	Required Fix
test_dqn_trait_implementation	FAIL	No UnifiedTrainable impl	Implement trait for WorkingDQN
test_dqn_forward_pass	FAIL	E0616: field q_network private	Add public forward method + trait impl
test_dqn_backward_pass	FAIL	E0616: field q_network private	Add backward pass method
test_dqn_optimizer_step	FAIL	E0599: no init_optimizer method	Add method + E0616: optimizer field private
test_dqn_checkpoint_save	FAIL	E0599: no save_checkpoint method	Add async save method
test_dqn_checkpoint_load	FAIL	E0599: no load_checkpoint method	Add async load method
test_dqn_metrics_collection	FAIL	E0599: no get_metrics method	Add metrics collection method
test_dqn_training_step	FAIL	E0061: train_step signature mismatch	Signature is train_step(batch: Option<Vec<Experience>>), test expects 5 args
test_dqn_device_transfer	FAIL	No device() method	Add device accessor
test_dqn_nan_detection	FAIL	Signature mismatch	Fix train_step signature

DQN Status: Core training works, missing public API + trait impl Implementation Effort: ~250 lines (trait impl + 7 new public methods)

DQN Config Issue: Tests use non-existent hidden_dim field (should be hidden_dims: Vec<usize>) Constructor Issue: Tests call WorkingDQN::new(config, device) but signature is new(config) (device hardcoded to CPU)

---

3. PPO Tests (10 tests, 0 passing)

Test Name	Status	Error Type	Required Fix
test_ppo_trait_implementation	FAIL	No UnifiedTrainable impl	Implement trait for WorkingPPO
test_ppo_forward_pass	FAIL	actor.action_probabilities exists	Wrap in trait forward method
test_ppo_backward_pass	FAIL	actor.forward exists	Add backward pass logic
test_ppo_optimizer_step	FAIL	init_optimizers method exists	Wrap in trait method
test_ppo_checkpoint_save	FAIL	No checkpoint methods	Implement async save (actor + critic)
test_ppo_checkpoint_load	FAIL	No checkpoint methods	Implement async load (actor + critic)
test_ppo_metrics_collection	FAIL	get_training_steps() exists	Add full metrics collection
test_ppo_training_step	FAIL	No batch training method	Add train_batch method
test_ppo_device_transfer	FAIL	actor.device() exists	Expose via trait method
test_ppo_nan_detection	FAIL	Forward pass works	Add validation in training

PPO Status: Actor/Critic networks work, missing orchestration methods + trait impl Implementation Effort: ~300 lines (trait impl + checkpoint I/O + batch training)

---

4. TFT Tests (10 tests, 0 passing)

Test Name	Status	Error Type	Required Fix
test_tft_trait_implementation	FAIL	E0432: unresolved import TFTModel	Fix import path (likely ml::tft::TFTModel)
test_tft_forward_pass	FAIL	Module import issue	Fix import + implement trait
test_tft_backward_pass	FAIL	Module import issue	Fix import + implement trait
test_tft_optimizer_step	FAIL	Module import issue	Fix import + implement trait
test_tft_checkpoint_save	FAIL	Module import issue	Fix import + implement trait
test_tft_checkpoint_load	FAIL	Module import issue	Fix import + implement trait
test_tft_metrics_collection	FAIL	Module import issue	Fix import + implement trait
test_tft_training_step	FAIL	Module import issue	Fix import + implement trait
test_tft_device_transfer	FAIL	Module import issue	Fix import + implement trait
test_tft_nan_detection	FAIL	Module import issue	Fix import + implement trait

TFT Status: BLOCKED by module import issue, then needs full trait impl Implementation Effort: ~300 lines (trait impl + all orchestration methods)

Import Issue: Test imports ml::tft::TFTModel but module may export different name or not export at all. Need to check ml/src/tft/mod.rs for correct export.

---

5. Orchestrator Integration Tests (10 tests, 0 passing)

Test Name	Status	Error Type	Required Fix
test_orchestrator_creation	PASS (placeholder)	Placeholder (assert true)	Implement real test after trait impls
test_orchestrator_model_registration	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_training_loop	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_checkpoint_management	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_metrics_aggregation	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_early_stopping	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_learning_rate_scheduling	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_validation_loop	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_error_recovery	PASS (placeholder)	Placeholder	Implement real test
test_orchestrator_multi_model_coordination	PASS (placeholder)	Placeholder	Implement real test

Orchestrator Status: Tests are PLACEHOLDERS (all assert!(true)), will fail once real tests written Implementation Effort: ~500 lines (10 comprehensive integration tests using real models)

---

UnifiedTrainable Trait Analysis

Trait Definition (ml/src/training/unified_trainer.rs)

Status: COMPLETE and well-designed Methods: 15 required methods covering full training lifecycle

pub trait UnifiedTrainable {
    fn model_type(&self) -> &str;                     // Model identifier
    fn device(&self) -> &Device;                      // GPU/CPU device
    fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError>;
    fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError>;
    fn backward(&mut self, loss: &Tensor) -> Result<f64, MLError>;  // Returns grad norm
    fn optimizer_step(&mut self) -> Result<(), MLError>;
    fn zero_grad(&mut self) -> Result<(), MLError>;
    fn get_learning_rate(&self) -> f64;
    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError>;
    fn get_step(&self) -> usize;
    fn collect_metrics(&self) -> TrainingMetrics;
    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError>;
    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError>;
    fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result<f64, MLError>;
}

Key Features:

• Standardized checkpoint format (safetensors + JSON metadata)
• Gradient norm tracking for explosion detection
• Learning rate scheduling support
• Validation loop integration
• Model-agnostic metrics collection

Architectural Quality: EXCELLENT - Clean abstraction, no leaky implementation details

---

Missing Implementations Summary

Model-by-Model Gap Analysis

Model	Core Logic	Public API	Trait Impl	Checkpoint I/O	Metrics	Estimated LOC
MAMBA-2	✅ READY	⚠️ PARTIAL	❌ MISSING	✅ READY (async)	✅ READY	~200
DQN	✅ READY	❌ MISSING	❌ MISSING	❌ MISSING	⚠️ PARTIAL	~250
PPO	✅ READY	⚠️ PARTIAL	❌ MISSING	❌ MISSING	⚠️ PARTIAL	~300
TFT	⚠️ UNKNOWN	❌ MISSING	❌ MISSING	❌ MISSING	❌ MISSING	~300

Total Implementation Effort: ~1,050 lines of glue code across 4 models

---

GPU Resource Management Requirements

Job Queue GPU Semaphore (services/ml_training_service/src/job_queue.rs)

Status: PRODUCTION-READY ✅ Implementation: Tokio Semaphore for GPU slot management

pub struct JobQueue {
    gpu_semaphore: Arc<Semaphore>,           // Limits concurrent GPU jobs
    total_gpu_slots: usize,                  // Typically 1 for RTX 3050 Ti
    // ... other fields
}

// GPU resource acquisition
pub async fn acquire_gpu(&self) -> Result<SemaphorePermit<'_>> {
    self.gpu_semaphore.acquire().await
        .context("Failed to acquire GPU resource")
}

Features:

• Priority-based queue (DQN/PPO = High, MAMBA-2/TFT = Medium, TLOB/LIQUID = Low)
• Semaphore prevents GPU oversubscription
• FIFO within same priority level
• Crash recovery via Redis persistence

GPU Configuration:

• RTX 3050 Ti: 1 concurrent job (4GB VRAM limit)
• Future A100: 2-4 concurrent jobs (80GB VRAM)
• CPU-only mode: No semaphore limits

---

Device Management in Training

Requirement: Models must support Device::cuda_if_available(0) pattern

// CORRECT pattern (auto-fallback)
let device = Device::cuda_if_available(0)?;
let model = Model::new(config, device)?;

// WRONG pattern (hardcoded CPU)
let device = Device::Cpu;  // ❌ Ignores GPU even if available

Current Status:

• MAMBA-2: ✅ Device field exists, needs trait exposure
• DQN: ❌ Hardcoded CPU in constructor (line 279: let device = Device::Cpu;)
• PPO: ⚠️ Device managed per-network (actor/critic), needs consolidation
• TFT: ⚠️ Unknown, needs verification

GPU Memory Estimation (from GPU_TRAINING_BENCHMARK.md):

• DQN: 50-150MB (small networks)
• PPO: 50-200MB (actor + critic networks)
• MAMBA-2: 150-500MB (SSM state + complex layers)
• TFT: 1.5-2.5GB (multi-head attention + time series length)

Batch Size Constraints:

• RTX 3050 Ti (4GB VRAM): Requires small batches (32-64 samples)
• Gradient accumulation used for larger effective batch sizes

---

MinIO/Redis Integration Points

1. Redis Integration (Job Queue Persistence)

File: services/ml_training_service/src/job_queue.rs Status: PRODUCTION-READY ✅

Use Cases:

• Job queue crash recovery
• Distributed queue coordination (future multi-GPU setup)
• Job status persistence

pub async fn with_redis(capacity: usize, gpu_slots: usize, redis_url: &str) -> Result<Self> {
    let redis_client = RedisClient::open(redis_url).context("Failed to connect to Redis")?;
    // ... persist queue state to Redis
}

Redis Keys (namespace: ml_training_queue:):

• ml_training_queue:jobs - Hash of job_id -> serialized QueuedJob
• ml_training_queue:processing - Set of currently processing job IDs
• ml_training_queue:metrics - Queue performance metrics

Recovery Protocol:

1. On startup, load ml_training_queue:jobs from Redis
2. Check ml_training_queue:processing for crashed jobs
3. Re-enqueue crashed jobs based on priority
4. Resume processing from queue state

---

2. S3/MinIO Integration (Checkpoint Storage)

Files:

• ml/src/checkpoint/storage.rs - Storage abstraction layer
• ml/src/model_registry.rs - Model versioning with S3

Status: ARCHITECTURE READY, MinIO not yet integrated ⚠️

Checkpoint Storage Backends:

pub trait CheckpointStorage {
    async fn save_checkpoint(&self, filename: &str, data: &[u8], metadata: &CheckpointMetadata) -> Result<()>;
    async fn load_checkpoint(&self, filename: &str) -> Result<Vec<u8>>;
    async fn delete_checkpoint(&self, filename: &str) -> Result<()>;
    async fn list_all_checkpoints(&self) -> Result<Vec<CheckpointMetadata>>;
}

Available Implementations:

1. FileSystemStorage - Local disk (currently used) ✅
2. InMemoryStorage - Testing only ✅
3. S3Storage - AWS S3 integration ⚠️ (feature flag: s3-storage)

MinIO Integration Gap:

• S3Storage implementation exists but not tested with MinIO
• MinIO compatibility assumed (S3 API-compatible)
• Need to configure MinIO credentials in ServiceConfig
• Need to test checkpoint save/load with MinIO

Model Registry (ml/src/model_registry.rs):

• PostgreSQL storage for model metadata (version, hyperparameters, metrics)
• S3 URLs for model artifacts (e.g., s3://foxhunt-ml-models/dqn/1.0.0/)
• Checksums (SHA-256) for integrity verification

Integration Requirements:

1. Configure MinIO credentials in Vault
2. Test S3Storage with MinIO endpoint (e.g., http://minio:9000)
3. Implement automatic checkpoint upload to MinIO after training
4. Add MinIO fallback for local checkpoint storage

---

3. Integration Flow (Training → MinIO → Registry)

┌────────────────┐
│ Training Loop  │
│ (Orchestrator) │
└───────┬────────┘
        │
        ▼ save_checkpoint()
┌────────────────┐
│ UnifiedTrainable│
│ (Model Impl)   │
└───────┬────────┘
        │
        ▼ CheckpointStorage::save_checkpoint()
┌────────────────┐    ┌────────────────┐
│ FileSystem     │ OR │ S3Storage      │
│ Storage (dev)  │    │ (MinIO, prod)  │
└───────┬────────┘    └───────┬────────┘
        │                     │
        └──────────┬──────────┘
                   │
                   ▼ ModelRegistry::register_version()
           ┌────────────────┐
           │ PostgreSQL +   │
           │ MinIO URLs     │
           └────────────────┘

Current Gap: Models save checkpoints to FileSystemStorage, but don't register with ModelRegistry or upload to MinIO automatically.

Required Work:

1. Update UnifiedTrainingOrchestrator to use configurable CheckpointStorage
2. Add post-training step to register models with ModelRegistry
3. Test MinIO connectivity and S3Storage compatibility
4. Add Prometheus metrics for checkpoint upload success/failure

---

Architecture Quality Assessment

Strengths ✅

1. Clean Abstraction: UnifiedTrainable trait is well-designed, model-agnostic
2. GPU Management: Job queue semaphore prevents oversubscription
3. Redis Persistence: Crash recovery for training jobs
4. Checkpoint Format: Standardized safetensors + JSON metadata
5. Learning Rate Scheduling: Built into orchestrator (warmup, cosine annealing, step decay)
6. Gradient Monitoring: Norm tracking for explosion detection
7. Early Stopping: Configurable patience for validation loss

Weaknesses ⚠️

1. Missing Glue Code: NO models implement UnifiedTrainable yet (~1,050 LOC needed)
2. DQN Device Hardcoded: CPU-only, ignores CUDA availability
3. TFT Import Issue: Module export incorrect or missing
4. MinIO Not Integrated: S3Storage exists but not tested/configured
5. No Automatic Registration: Models don't auto-register with ModelRegistry after training
6. Test Placeholders: Orchestrator tests are all assert!(true) stubs

Architectural Risks 🔴

1. Training Pipeline Blocked: Cannot train ANY model until trait implementations done
2. GPU Underutilization: DQN hardcoded to CPU, wasting RTX 3050 Ti
3. Checkpoint Loss Risk: Local-only storage, no MinIO backup
4. Test Coverage Gap: 40 tests defined, 0 real tests passing (orchestrator tests are stubs)

---

Remediation Roadmap

Phase 1: Unblock Training (Priority: CRITICAL)

Goal: Get 1 model training end-to-end via orchestrator Duration: 1-2 days

1. Fix DQN Constructor (1 hour)

   • Change Device::Cpu to Device::cuda_if_available(0)?
   • Make device a constructor parameter
   • Update tests to use correct WorkingDQNConfig fields

2. Implement UnifiedTrainable for DQN (4 hours)

   • Add 7 missing public methods (forward, backward, optimizer_step, etc.)
   • Implement checkpoint save/load using FileSystemStorage
   • Add metrics collection mapping

3. Fix Test Compilation (2 hours)

   • Fix DQN config field names (hidden_dims not hidden_dim)
   • Fix train_step test signatures
   • Verify all 10 DQN tests compile

4. End-to-End DQN Training Test (3 hours)

   • Create integration test: train DQN for 10 epochs via orchestrator
   • Verify checkpoint save/load works
   • Verify metrics collection works
   • Document results

Deliverable: 1 model (DQN) fully working with orchestrator, 10/40 tests passing

---

Phase 2: Complete Model Coverage (Priority: HIGH)

Goal: All 4 models implement UnifiedTrainable Duration: 2-3 days

1. MAMBA-2 Trait Implementation (3 hours)

   • Wrap existing async methods in sync trait methods
   • Make train_batch public
   • Add initialize_optimizer method
   • Test all 10 MAMBA-2 tests

2. PPO Trait Implementation (4 hours)

   • Add dual-checkpoint save/load (actor + critic)
   • Implement batch training method
   • Consolidate device management
   • Test all 10 PPO tests

3. Fix TFT Import Issue (1 hour)

   • Check ml/src/tft/mod.rs for correct export
   • Fix test import statement
   • Verify TFTModel struct exists

4. TFT Trait Implementation (4 hours)

   • Implement all 15 trait methods
   • Add checkpoint I/O
   • Test all 10 TFT tests

Deliverable: 4/4 models working, 40/40 tests passing (excluding orchestrator placeholders)

---

Phase 3: MinIO Integration (Priority: MEDIUM)

Goal: Production checkpoint storage Duration: 1-2 days

1. Configure MinIO (2 hours)

   • Add MinIO credentials to Vault
   • Update ServiceConfig with S3Storage config
   • Test MinIO connectivity

2. Test S3Storage with MinIO (3 hours)

   • Unit tests for checkpoint upload/download
   • Performance benchmarks (upload time, bandwidth)
   • Error handling (network failures, auth errors)

3. Integrate with Orchestrator (2 hours)

   • Add CheckpointStorage parameter to orchestrator config
   • Update checkpoint save logic to use MinIO
   • Add Prometheus metrics for upload success/failure

4. Model Registry Integration (3 hours)

   • Add post-training hook to register models
   • Implement SHA-256 checksum calculation
   • Test version query API

Deliverable: Automatic checkpoint backup to MinIO, model registry populated

---

Phase 4: Orchestrator Tests (Priority: MEDIUM)

Goal: Replace placeholder tests with real integration tests Duration: 2 days

1. Training Loop Test (3 hours)

   • Train DQN for 5 epochs, verify loss decreases
   • Check checkpoint files created
   • Validate metrics logged

2. Early Stopping Test (2 hours)

   • Configure patience=3
   • Inject increasing validation loss
   • Verify training stops after 3 epochs

3. Learning Rate Scheduling Test (2 hours)

   • Test warmup + cosine annealing
   • Verify LR changes per epoch
   • Check final LR matches expected value

4. Checkpoint Management Test (2 hours)

   • Train for 10 epochs with checkpoint_frequency=5
   • Verify 2 checkpoints + 1 best checkpoint saved
   • Test checkpoint loading and resumption

5. Multi-Model Coordination Test (3 hours)

   • Enqueue 2 DQN + 2 PPO jobs
   • Verify priority ordering (High priority jobs first)
   • Check GPU semaphore prevents concurrent execution

Deliverable: 10/10 orchestrator tests passing, full E2E coverage

---

Test Execution Plan

Prerequisites

1. ✅ Rust toolchain (already installed)
2. ✅ Redis running (docker-compose up redis)
3. ✅ PostgreSQL running (docker-compose up postgres)
4. ⚠️ MinIO configured (needed for Phase 3)
5. ⚠️ Trait implementations complete (Phases 1-2)

Command Sequence

# Phase 1: DQN only
cargo test -p ml --test unified_training_tests test_dqn

# Phase 2: All models
cargo test -p ml --test unified_training_tests

# Phase 3: Job queue integration
cargo test -p ml_training_service --test job_queue_tests

# Phase 4: Full orchestrator
cargo test -p ml_training_service --test integration_tests

Success Criteria

• 40/40 unified training tests passing
• 10/10 orchestrator integration tests passing (after Phase 4)
• 100% trait implementation coverage (4/4 models)
• MinIO checkpoint upload working
• Model registry populated with versions

---

Critical Dependencies

External Services

Service	Required For	Status	Priority
Redis	Job queue persistence	✅ Running	HIGH
PostgreSQL	Model registry	✅ Running	HIGH
MinIO	Checkpoint backup	⚠️ Not configured	MEDIUM
Prometheus	Training metrics	✅ Running	LOW

Internal Dependencies

Component	Depends On	Status	Blocker?
UnifiedTrainingOrchestrator	UnifiedTrainable impls	❌ Missing	YES
Job Queue	Redis	✅ Ready	NO
Checkpoint Storage	Filesystem/S3 backend	⚠️ Partial	NO
Model Registry	PostgreSQL + S3 URLs	⚠️ Partial	NO

---

Estimated Timeline

Total Implementation: 6-9 days

Phase	Duration	Blockers	Risk
Phase 1 (DQN trait)	1-2 days	None	LOW
Phase 2 (All models)	2-3 days	Phase 1	LOW
Phase 3 (MinIO)	1-2 days	MinIO setup	MEDIUM
Phase 4 (Tests)	2 days	Phases 1-2	LOW

Critical Path: Phase 1 → Phase 2 → Phase 4 Parallel Work: Phase 3 (MinIO) can proceed alongside Phases 1-2

---

Recommendations

Immediate Actions (Priority: CRITICAL)

1. Fix DQN Device Management - Change hardcoded CPU to CUDA auto-detect (15 min)
2. Implement UnifiedTrainable for DQN - Get 1 model working end-to-end (4 hours)
3. Run First Integration Test - Train DQN via orchestrator (1 hour)

Short-term (Priority: HIGH)

1. Complete Trait Implementations - All 4 models (2-3 days)
2. Fix TFT Import - Unblock TFT tests (1 hour)
3. Test All Models - Verify 40/40 tests passing (1 day)

Medium-term (Priority: MEDIUM)

1. Integrate MinIO - Production checkpoint storage (1-2 days)
2. Write Real Orchestrator Tests - Replace placeholders (2 days)
3. Model Registry Integration - Automatic version tracking (3 hours)

Long-term (Priority: LOW)

1. Distributed Training - Multi-GPU coordination via Redis (future work)
2. Checkpoint Compression - Reduce MinIO storage costs (future work)
3. Training Telemetry - Detailed Prometheus metrics (future work)

---

Conclusion

Status: Training orchestration infrastructure is PRODUCTION-READY, but BLOCKED by missing trait implementations.

Key Finding: This is a glue code problem, not an architectural problem. The UnifiedTrainable trait is well-designed, the orchestrator is complete, and GPU management works. We just need ~1,050 lines of trait implementations to connect existing model logic to the orchestration framework.

Next Steps:

1. Implement UnifiedTrainable for DQN (Phase 1, 4 hours)
2. Verify end-to-end DQN training works
3. Expand to remaining 3 models (Phase 2, 2-3 days)
4. Integrate MinIO for production storage (Phase 3, 1-2 days)

Estimated Time to Production: 6-9 days of focused implementation work.

---

Appendix A: Test File Reference

File: /home/jgrusewski/Work/foxhunt/ml/tests/unified_training_tests.rs Lines: 819 lines Tests: 40 tests (10 per model × 4 models) Coverage: Forward/backward, checkpointing, metrics, device transfer, NaN detection

Test Categories:

1. Trait implementation (type check)
2. Forward pass (inference)
3. Backward pass (gradient computation)
4. Optimizer step (parameter updates)
5. Checkpoint save/load
6. Metrics collection
7. Training step (batch processing)
8. Device transfer (CPU/CUDA)
9. NaN detection (numerical stability)
10. Integration (orchestrator coordination)

---

Appendix B: UnifiedTrainable Trait Methods

Method	Purpose	Return Type	Complexity
model_type()	Identifier	&str	Trivial
device()	GPU/CPU device	&Device	Trivial
forward()	Inference	Result<Tensor>	Wrap existing
compute_loss()	Loss calculation	Result<Tensor>	Simple
backward()	Gradient computation	Result<f64>	Wrap existing
optimizer_step()	Parameter update	Result<()>	Wrap existing
zero_grad()	Clear gradients	Result<()>	Simple
get_learning_rate()	LR accessor	f64	Trivial
set_learning_rate()	LR setter	Result<()>	Simple
get_step()	Training step count	usize	Trivial
collect_metrics()	Gather stats	TrainingMetrics	Moderate
save_checkpoint()	Persist model	Result<String>	Complex
load_checkpoint()	Restore model	Result<Metadata>	Complex
validate()	Val loop	Result<f64>	Moderate

Implementation Effort: 50-75 lines per model (excluding checkpoint I/O complexity)

---

End of Analysis