foxhunt/MAMBA2_PRODUCTION_TRAINING_GUIDE.md

MAMBA-2 Production Training Guide

Status: ✅ READY FOR EXECUTION Implementation: Complete end-to-end training pipeline Data: Real DBN market data (ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT) GPU: RTX 3050 Ti optimized (4GB VRAM)

---

Quick Start

1. Pilot Training (50 Epochs - 30-45 minutes)

cd /home/jgrusewski/Work/foxhunt
cargo run -p ml --example train_mamba2_dbn --release -- --epochs 50

Expected Results:

• Training time: 30-45 minutes
• GPU utilization: 60-70%
• VRAM usage: ~2GB
• Loss convergence: Should see 10-20% reduction
• Checkpoints: Every 10 epochs + best model

2. Full Production Training (200 Epochs - 2-3 hours)

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

Expected Results:

• Training time: 2-3 hours
• GPU utilization: 60-70%
• VRAM usage: ~2GB
• Loss convergence: 30-50% reduction
• Early stopping: Triggered if no improvement for 20 epochs

---

Configuration

Default Hyperparameters

Model Architecture:
  d_model: 256              # Feature embedding dimension
  n_layers: 6               # Number of MAMBA-2 layers
  state_size: 16            # SSM state dimension
  num_heads: 8              # Multi-head attention
  expand: 2                 # Expansion factor

Training:
  epochs: 200               # Total training epochs
  batch_size: 32            # Batch size (VRAM optimized)
  learning_rate: 0.0001     # Initial learning rate
  seq_len: 60               # Sequence length (timesteps)
  dropout: 0.1              # Dropout rate
  grad_clip: 1.0            # Gradient clipping threshold
  weight_decay: 0.0001      # L2 regularization
  warmup_steps: 1000        # Learning rate warmup

Data:
  train_split: 0.8          # 80% training, 20% validation
  features: 16              # Base features per timestep
  technical_indicators: 10  # Additional indicators
  data_dir: test_data/real/databento/ml_training_small

Optimization:
  optimizer: Adam           # Built into MAMBA-2
  early_stopping_patience: 20
  checkpoint_frequency: 10  # Save every 10 epochs

Memory Estimation

Configuration	VRAM Usage	Status
Default (d_model=256, n_layers=6)	~2GB	✅ Safe
Large (d_model=512, n_layers=8)	~3.5GB	⚠️ Tight
XL (d_model=1024, n_layers=12)	>4GB	❌ Too Large

---

Data Pipeline

1. DBN Data Loading

DbnSequenceLoader::new(seq_len, d_model)
    .load_sequences(data_dir, train_split)

Features:

• Loads real OHLCV bars from DBN files
• Automatic price scaling (1e-9 for DBN format)
• Feature normalization (z-score)
• Symbol mapping (6E.FUT, ZN.FUT, etc.)

2. Feature Extraction

Per timestep (26 features total):

• OHLCV (5): Open, High, Low, Close, Volume
• Derived (4): Range, Body, Upper Wick, Lower Wick
• Technical Indicators (10): RSI, MACD, Bollinger Bands, ATR, EMA, etc.
• Padding: Zero-padded to d_model=256 dimension

3. Sequence Creation

• Input: [seq_len, d_model] = [60, 256] per sequence
• Target: [1, d_model] = next timestep prediction
• Sliding Window: Overlapping sequences for maximum data utilization

---

Training Loop

Architecture

Input [batch, seq_len, d_model]
    ↓
MAMBA-2 Layers (6x)
├── SSD Layer (Structured State Duality)
├── Selective State Space
├── Hardware-Aware Optimization
└── Residual + LayerNorm
    ↓
Output Projection [batch, seq_len, 1]
    ↓
MSE Loss (Regression)

Per-Epoch Flow

1. Forward Pass: Process batch through MAMBA-2 layers
2. Loss Computation: MSE between prediction and target
3. Backward Pass: Compute gradients via automatic differentiation
4. Gradient Clipping: Threshold at 1.0 for stability
5. Optimizer Step: Adam update with learning rate schedule
6. SSM Projection: Ensure spectral radius < 1.0 for stability
7. Validation: Compute validation loss every epoch
8. Checkpointing: Save if validation loss improves

Monitoring

Real-Time Metrics:

• Training loss (per epoch)
• Validation loss (per epoch)
• Perplexity: exp(loss)
• Learning rate (with warmup + cosine decay)
• Throughput (epochs/minute)
• GPU memory usage

SSM State Statistics (every 10 epochs):

• Mean, Std Dev, Min, Max
• Spectral Radius (stability check)
• Compression Ratio

---

Outputs

Checkpoint Files

Location: ml/checkpoints/mamba2_dbn/

checkpoint_epoch_10.ckpt          # Periodic checkpoint
checkpoint_epoch_20.ckpt
...
best_model_epoch_42.ckpt          # Best validation loss
final_model.ckpt                  # Final epoch model

Usage:

model.load_checkpoint("ml/checkpoints/mamba2_dbn/best_model_epoch_42.ckpt").await?;

Training Metrics

CSV: training_losses.csv

epoch,train_loss,val_loss,learning_rate
0,0.456789,0.478901,0.0001
1,0.423456,0.445678,0.0001
...

JSON: training_metrics.json

{
  "total_epochs": 200,
  "best_val_loss": 0.123456,
  "best_epoch": 142,
  "training_duration_hours": 2.5,
  "final_perplexity": 1.131,
  "config": { ... }
}

---

Performance Benchmarks

Pilot Run (50 epochs)

Target Metrics:

• Duration: 30-45 minutes
• Loss reduction: 10-20%
• Perplexity: <2.0
• Convergence: Moderate (still learning)

GPU Usage:

• VRAM: ~2GB (50% of 4GB)
• Utilization: 60-70% (memory-bound)
• Temperature: 65-75°C

Full Training (200 epochs)

Target Metrics:

• Duration: 2-3 hours
• Loss reduction: 30-50%
• Perplexity: <1.5
• Convergence: Strong (low variance)

Expected Loss Curve:

Epoch   Loss    Perplexity
0       0.450   1.568
50      0.380   1.462
100     0.320   1.377
150     0.280   1.323
200     0.250   1.284

---

Troubleshooting

Out of Memory (OOM)

Symptoms: CUDA OOM error during training

Solutions:

1. Reduce batch size: batch_size: 16 (from 32)
2. Reduce model dimension: d_model: 128 (from 256)
3. Reduce sequence length: seq_len: 32 (from 60)
4. Use CPU: Will auto-fallback if CUDA fails

Slow Convergence

Symptoms: Loss not decreasing after many epochs

Solutions:

1. Increase learning rate: learning_rate: 0.0003 (from 0.0001)
2. Reduce weight decay: weight_decay: 1e-5 (from 1e-4)
3. Increase warmup steps: warmup_steps: 2000 (from 1000)
4. Check data quality (feature normalization)

SSM Instability

Symptoms: Warning "spectral radius >= 1.0"

Solutions:

1. Already handled automatically (SSM projection)
2. Reduce learning rate if persistent
3. Increase gradient clipping: grad_clip: 0.5 (from 1.0)

No Training Data

Symptoms: "No training data loaded!"

Solutions:

1. Check DBN files exist: ls test_data/real/databento/ml_training_small/*.dbn
2. Verify file permissions: chmod 644 test_data/real/databento/ml_training_small/*.dbn
3. Check symbol mapping in DbnSequenceLoader

---

Next Steps

After Pilot Training (50 epochs)

1. Analyze Checkpoints:

   cargo run -p ml --example analyze_mamba2_checkpoints
   

2. Review Metrics:

   • Check training_losses.csv for convergence
   • Verify perplexity is decreasing
   • Confirm GPU utilization ~60-70%

3. Decision:

   • If loss reducing >10%: Proceed to full 200 epochs
   • If loss flat: Tune hyperparameters (increase LR)
   • If unstable: Reduce LR, increase grad clipping

After Full Training (200 epochs)

1. Model Validation:

   • Load best checkpoint
   • Test on unseen data
   • Measure inference latency (<5μs target)

2. Integration:

   • Deploy to ML Training Service
   • Register in model registry
   • Enable gRPC inference endpoint

3. Production Deployment:

   • A/B test with baseline models (DQN, PPO)
   • Monitor Sharpe ratio, win rate, drawdown
   • Scale to multi-symbol (ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT)

---

References

• MAMBA-2 Paper: "Mamba-2: Linear-Time Sequence Modeling with Selective State Spaces"
• Implementation: /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs
• Training Script: /home/jgrusewski/Work/foxhunt/ml/examples/train_mamba2_dbn.rs
• Data Loader: /home/jgrusewski/Work/foxhunt/ml/src/data_loaders/dbn_sequence_loader.rs
• Wave Context: Agent 78 DQN training success, Agent 40 MAMBA-2 production setup

---

Last Updated: 2025-10-14 Status: ✅ PRODUCTION READY Next Milestone: Execute 50-epoch pilot run, analyze results, proceed to full training