hybrid-predict-trainer-rs

A hybridized predictive training implementation for deep learning in Rust.

Overview

hybrid-predict-trainer-rs implements a novel training paradigm that achieves 5-10x training speedup by intelligently predicting multiple training steps instead of computing full forward/backward passes for every iteration.

The Core Idea

Traditional training: [Forward → Backward → Update] × N steps

Hybrid predictive training:

[Warmup: Collect Statistics] → 
[Full Train: Learn Dynamics] → 
[Predict: Skip Backprop] → 
[Correct: Apply Residuals] → 
Loop back to Full Train

By learning the training dynamics during full training phases, we can predict multiple steps ahead without computing gradients, then apply corrections based on accumulated residuals.

Architecture

┌──────────────────────────────────────────────────────────────────┐
│                        HybridTrainer                             │
├──────────────────────────────────────────────────────────────────┤
│                                                                  │
│  ┌─────────┐    ┌─────────┐    ┌─────────┐    ┌─────────┐       │
│  │ WARMUP  │───▶│  FULL   │───▶│ PREDICT │───▶│ CORRECT │       │
│  │         │    │ TRAIN   │    │         │    │         │       │
│  └─────────┘    └────┬────┘    └─────────┘    └────┬────┘       │
│                      │                             │             │
│                      └─────────────────────────────┘             │
│                             (cycle repeats)                      │
│                                                                  │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │                     Supporting Systems                       │ │
│  ├─────────────────────────────────────────────────────────────┤ │
│  │  State Encoder  │  RSSM Dynamics  │  Divergence Monitor    │ │
│  │  Residual Store │  Bandit Selector │  Metrics Collector    │ │
│  └─────────────────────────────────────────────────────────────┘ │
│                                                                  │
└──────────────────────────────────────────────────────────────────┘

Phases

1. Warmup Phase

• Collects baseline statistics (loss distribution, gradient norms)
• Establishes training dynamics baseline
• Typically 100-500 steps

2. Full Training Phase

• Standard forward/backward training
• Trains the dynamics predictor from observations
• Extracts residuals for future corrections

3. Predictive Phase

• Uses learned dynamics model to predict Y steps
• No backward passes - dramatic speedup
• Monitors for divergence; falls back if confidence drops

4. Correction Phase

• Applies accumulated residuals to predictions
• Uses similarity-based weighting from residual history
• Online learning improves correction over time

Features

• Multiple predictor architectures: Linear, MLP, RSSM-lite
• Adaptive phase lengths: Bandit-based selection
• Multi-signal divergence detection: Loss, gradient, oscillation
• GPU acceleration: CubeCL + Burn for CUDA support
• Comprehensive metrics: JSON export, console summaries
• Comprehensive benchmarking: Criterion.rs performance analysis

Performance

Benchmark Results

Comprehensive performance benchmarks using Criterion.rs on all critical paths:

RSSM Prediction Performance

Scaling: Linear at ~24 µs per prediction step.

Component Overhead

Speedup Analysis

Overhead Comparison:

• Full training step: ~1.37 ms overhead (RSSM training + state)
• Predict step: ~0.41 ms overhead (RSSM prediction + state)
• Overhead reduction: 70% (predict vs full)

Expected Speedups (for typical training configurations):

Note: Actual speedup depends on model architecture, batch size, and prediction horizon configuration.

Running Benchmarks

# Run all benchmarks
cargo bench

# Run specific benchmark group
cargo bench --bench hybrid_trainer_benchmarks -- rssm_prediction

# Generate HTML reports (target/criterion/report/index.html)
cargo bench --bench hybrid_trainer_benchmarks

For detailed performance analysis, see PHASE_2_BENCHMARKING_REPORT.md.

Memory Management

Current Status

The hybrid trainer implements automatic VRAM management to handle Burn's functional API model copy behavior:

Short runs (0-200 steps): Stable memory usage (~3 GB for GPT-2 Small)

Medium runs (200-1000 steps): Automatic cleanup every 10 steps maintains ~3-6 GB

Long runs (1000+ steps): Gradual accumulation to ~10-14 GB over 1000+ steps

Automatic Mitigations

The trainer includes multiple layers of VRAM protection:

1. Periodic cleanup: Forces CUDA synchronization every 10 steps
2. Phase transition logging: Monitors VRAM usage at Warmup→Full→Predict→Correct transitions
3. Emergency checkpoints: Automatically saves when VRAM exceeds 14 GB
4. Adaptive defaults: Reduced max_predict_steps from 80→15 to minimize copies
5. Checkpoint-based recovery: Frequent saves (every 50 steps) enable reload for cleanup

Recommended Configurations

For different GPU memory sizes:

// 8 GB GPU (aggressive cleanup)
HybridTrainerConfig::builder()
    .max_predict_steps(10)
    .checkpoint_config(CheckpointConfig {
        save_interval: 25,
        ..Default::default()
    })
    .build()

// 16 GB GPU (balanced, default)
HybridTrainerConfig::default() // max_predict_steps=15, save_interval=50

// 24+ GB GPU (relaxed)
HybridTrainerConfig::builder()
    .max_predict_steps(30)
    .checkpoint_config(CheckpointConfig {
        save_interval: 100,
        ..Default::default()
    })
    .build()

Future Improvements

Planned optimizations for long training runs:

1. In-place parameter updates: Eliminate model.map() copies entirely
2. Burn PR upstream: Contribute mutable ModuleMapper to Burn framework
3. Explicit CUDA memory management: Direct cudarc integration for aggressive cleanup
4. Gradient checkpointing: Trade compute for memory on forward passes

For detailed technical analysis, see docs/PHASE_2B_FINAL_SUMMARY.md.

Installation

Add to your Cargo.toml:

[dependencies]
hybrid-predict-trainer-rs = "0.1"

With CUDA acceleration:

[dependencies]
hybrid-predict-trainer-rs = { version = "0.1", features = ["cuda"] }

Quick Start

use hybrid_predict_trainer_rs::{HybridTrainer, HybridTrainerConfig, Phase};

// Configure the trainer
let config = HybridTrainerConfig::builder()
    .warmup_steps(200)
    .max_predict_steps(60)
    .confidence_threshold(0.85)
    .build();

// Create trainer with your model and optimizer
let mut trainer = HybridTrainer::new(model, optimizer, config)?;

// Training loop
for batch in data_loader {
    let result = trainer.step(&batch)?;

    println!(
        "Step {} | Phase: {:?} | Loss: {:.4} | Predicted: {}",
        trainer.current_step(),
        result.phase,
        result.loss,
        result.was_predicted
    );
}

// Get training summary
let stats = trainer.statistics();
println!("Backward reduction: {:.1}%", stats.backward_reduction_pct);

Configuration

use hybrid_predict_trainer_rs::config::{
    HybridTrainerConfig, PredictorConfig, DivergenceConfig, CheckpointConfig
};

// Using builder pattern
let config = HybridTrainerConfig::builder()
    // Phase configuration
    .warmup_steps(200)              // Steps before enabling prediction
    .full_steps(20)                 // Full training steps per cycle
    .max_predict_steps(80)          // Maximum prediction horizon
    // Predictor settings
    .predictor_config(PredictorConfig::RSSM {
        deterministic_dim: 256,
        stochastic_dim: 32,
        num_categoricals: 32,
        ensemble_size: 3,
    })
    // Confidence and quality
    .confidence_threshold(0.85)
    .divergence_threshold(3.0)
    // Metrics collection
    .collect_metrics(true)
    .build();

Validation Results

End-to-end validation on real models:

Validation Infrastructure:

• 227 comprehensive tests (218 unit + 9 integration)
• Automated VRAM monitoring (validate_vram.sh)
• Criterion.rs benchmark suite (6 groups, 16 scenarios)
• All tests passing on Rust 1.92+

Larger model benchmarks (1B+ parameters) planned for future releases.

Roadmap

v0.2.0 (Current Release)

• Core training loop implementation
• RSSM dynamics model (RSSM-lite) integration
• GRU cell with forward pass and training
• Multi-signal divergence detection
• LinUCB bandit for phase selection
• Residual correction framework
• Comprehensive metrics collection
• 227 unit and integration tests
• VRAM management system (5-layer protection)
• Comprehensive Criterion.rs benchmarks
• GPT-2 Small validation (124M params)
• Intra-horizon micro-corrections
• Checkpoint automation

v0.3.0 (Planned)

• CubeCL CUDA kernels for state encoding
• CubeCL CUDA kernel for RSSM forward pass
• 1B+ parameter model validation
• Integration examples (candle, tch-rs)
• Advanced optimizer support (AdamW, LAMB)

v0.4.0+ (Future)

• Distributed training support
• Mixed precision support (fp16, bf16)
• Multi-GPU training
• Advanced residual compression techniques

Research Background

This implementation is based on research exploring predictive training methods:

• DreamerV3 (Hafner et al., 2023): RSSM architecture for world models
• PowerSGD (Vogels et al., 2019): Low-rank gradient compression
• Lookahead Optimizer (Zhang et al., 2019): Multi-step optimization

The key insight is that training dynamics are often predictable enough to skip expensive backward passes while maintaining convergence quality.

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

Development Setup

# Clone the repository
git clone https://github.com/tzervas/hybrid-predict-trainer-rs.git
cd hybrid-predict-trainer-rs

# Build with all features
cargo build --all-features

# Run tests
cargo test

# Run benchmarks
cargo bench

License

Licensed under the MIT License. See LICENSE-MIT for details.

Copyright (c) 2026 Tyler Zervas

Acknowledgments

• The Burn team for their excellent deep learning framework
• The CubeCL team for GPU compute capabilities
• The Candle team for tensor operations inspiration