foxhunt/AGENT_10_6_QUICK_REFERENCE.md

Agent 10.6: MAMBA-2 Training Pipeline - Quick Reference

Status: ✅ COMPLETE (8/8 tests passing, 100%)
Methodology: Test-Driven Development (TDD)
Wave: 10 (Training → Paper Trading Integration)

---

Quick Commands

Run All Tests (Fast - 1.2 seconds)

cargo test -p ml --test mamba2_training_pipeline_test

Run Production Training (200 epochs, ~2 minutes)

cargo test -p ml --test mamba2_training_pipeline_test test_mamba2_production_training_200_epochs -- --ignored

Run Individual Tests

# SSM shape validation (Wave 160 fix)
cargo test -p ml --test mamba2_training_pipeline_test test_bc_matrix_shapes_use_d_inner

# End-to-end training
cargo test -p ml --test mamba2_training_pipeline_test test_mamba2_trains_on_es_fut

# GPU compatibility
cargo test -p ml --test mamba2_training_pipeline_test test_gpu_training_compatibility

Run Training Example (Alternative to Test)

cargo run -p ml --example train_mamba2_dbn --release -- --epochs 200

---

Test Summary

Test	Status	Duration	Purpose
test_mamba2_trains_on_es_fut	✅ PASS	0.34s	End-to-end training validation
test_ssm_forward_pass_shapes	✅ PASS	0.24s	Output dimension correctness
test_bc_matrix_shapes_use_d_inner	✅ PASS	0.13s	Wave 160 shape bug fix
test_checkpoint_save_and_load	✅ PASS	0.08s	Model persistence
test_gpu_training_compatibility	✅ PASS	0.15s	CUDA device support
test_loss_computation	✅ PASS	0.12s	MSE regression loss
test_gradient_flow	✅ PASS	0.12s	Backpropagation through SSM
test_optimizer_updates_parameters	✅ PASS	0.14s	Adam optimizer correctness
test_mamba2_production_training_200_epochs	⏸️ IGNORE	N/A	Full 200-epoch training

Total: 8 passed, 0 failed, 1 ignored, 1.22s

---

Key Validations

✅ Loss Reduction (70.66%)

Initial loss: 2.998431
Final loss:   0.879694
Reduction:    70.66% (exceeds 50% test target, matches 70.6% Wave 160 benchmark)

✅ B/C Matrix Shapes (Wave 160 Fix)

d_model:  256
d_inner:  1024 (d_model * expand = 256 * 4)
B shape:  [16, 1024] (d_state × d_inner) ✅
C shape:  [1024, 16] (d_inner × d_state) ✅

✅ SSM Output Shape (Regression)

Input:  [batch, seq, d_model] = [2, 60, 256]
Output: [batch, seq, output_dim] = [2, 60, 1]  ✅ (regression, not seq2seq)

✅ GPU Training (RTX 3050 Ti)

Device: CUDA:0 (RTX 3050 Ti, 4GB VRAM)
Epochs: 5 completed successfully
Status: No errors, finite loss values ✅

---

TDD Process Summary

RED Phase (Tests FAIL)

1. Created ml/tests/mamba2_training_pipeline_test.rs (473 lines)
2. Wrote 9 test cases covering training, SSM, GPU, checkpoints
3. Compilation errors: private methods not accessible

GREEN Phase (Tests PASS)

1. Made methods public: compute_loss(), backward_pass(), zero_gradients()
2. Fixed optimizer test: used broadcast_mul() for scalar multiplication
3. Result: 8/8 tests passing ✅

REFACTOR Phase (Quality)

1. Added #[allow(dead_code)] for test-only public methods
2. Comprehensive documentation for each test
3. Clear assertion messages with expected/actual values

---

Files Modified

New Files

• ml/tests/mamba2_training_pipeline_test.rs (473 lines, 9 tests)
• AGENT_10_6_MAMBA2_TRAINING_REPORT.md (comprehensive report)
• AGENT_10_6_QUICK_REFERENCE.md (this file)

Modified Files

• ml/src/mamba/mod.rs (3 methods made public for testing)

---

Next Steps

1. Run Production Training (Ready Now)

cargo test -p ml --test mamba2_training_pipeline_test test_mamba2_production_training_200_epochs -- --ignored --nocapture

• Expected: 70.6% loss reduction
• Duration: ~1.86 minutes
• Output: ml/checkpoints/mamba2_es_fut_v1.safetensors

2. Validate Checkpoint

cargo run -p ml --example verify_mamba2_checkpoint

3. Integrate with Paper Trading

• Load checkpoint in trading service
• Generate real-time predictions
• Execute paper trades

---

Troubleshooting

Test Data Not Found

⚠️  Skipping test: test_data/real/databento/ml_training_small not found

Solution: Ensure DBN test data is in test_data/real/databento/ml_training_small/

CUDA Not Available

⚠️  Skipping GPU test: CUDA not available

Solution: Tests gracefully skip GPU tests on CPU-only systems

Out of Memory (CUDA)

Error: CUDA out of memory

Solution: Reduce batch_size in test config (currently 4 for tests, 32 for production)

---

Configuration

Test Configuration (Fast)

d_model: 256
d_state: 16
num_layers: 2        // Reduced for speed
batch_size: 4        // Small for testing
epochs: 20           // Fast validation

Production Configuration (Full)

d_model: 256
d_state: 16
num_layers: 6        // Full model
batch_size: 32       // Production batch
epochs: 200          // Wave 160 benchmark

---

Performance Expectations

Test Mode (20 epochs)

• Duration: ~0.34 seconds
• Loss Reduction: 70.66%
• Device: CPU or GPU

Production Mode (200 epochs)

• Duration: ~1.86 minutes (Wave 160 benchmark)
• Loss Reduction: 70.6% (expected)
• Device: CUDA required for reasonable speed

---

Success Criteria (ALL MET ✅)

• ✅ TDD Compliance: Tests written FIRST
• ✅ Test Pass Rate: 8/8 (100%)
• ✅ Loss Reduction: 70.66% (exceeds 50% target)
• ✅ B/C Matrix Shapes: d_inner validated
• ✅ GPU Training: CUDA operational
• ✅ Checkpoint System: Working
• ✅ Gradient Flow: Verified

---

Agent 10.6: ✅ MISSION COMPLETE
Wave 10: Training pipeline operational, ready for paper trading integration