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)

```bash
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)

```bash
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

```yaml
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

```rust
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**:
```rust
model.load_checkpoint("ml/checkpoints/mamba2_dbn/best_model_epoch_42.ckpt").await?;
```

### Training Metrics

**CSV**: `training_losses.csv`
```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`
```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**:
   ```bash
   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