brainwires-training

Model training and fine-tuning for the Brainwires Agent Framework — cloud fine-tuning and local LoRA/QLoRA/DoRA training.

Powered by Burn — Local training in this crate is built on the Burn deep learning framework, a comprehensive deep learning framework written in Rust. Burn provides the GPU-accelerated tensor operations, automatic differentiation, and neural network modules that make local LoRA/QLoRA/DoRA adapter training possible. We are grateful to the Burn team and contributors for their excellent work on the Rust ML ecosystem.

Overview

brainwires-training provides a unified TrainingManager that dispatches fine-tuning jobs to either cloud providers (OpenAI, Together, Fireworks, Anyscale, AWS Bedrock, Google Vertex AI) or a local Burn-based training backend supporting LoRA, QLoRA, and DoRA adapter methods with DPO/ORPO alignment losses.

Design principles:

• Dual-path — one API for cloud fine-tuning and local adapter training; switch with a feature flag
• Provider-agnostic — FineTuneProvider trait abstracts all cloud APIs behind submit, status, cancel, download
• Adapter-first — local training uses parameter-efficient adapters (LoRA/QLoRA/DoRA) rather than full fine-tuning
• Dataset-integrated — consumes brainwires-datasets types directly for seamless data → training pipelines
• Observable — TrainingProgress and TrainingMetrics provide real-time job monitoring

  ┌─────────────────────────────────────────────────────────────┐
  │                    brainwires-training                       │
  │                                                             │
  │                  ┌────────────────┐                         │
  │                  │TrainingManager │                         │
  │                  └───────┬────────┘                         │
  │                          │                                  │
  │              ┌───────────┴───────────┐                      │
  │              ▼                       ▼                      │
  │  ┌──────────────────┐   ┌───────────────────┐              │
  │  │  Cloud Providers  │   │  Local Burn Backend│             │
  │  │  ┌──────────────┐ │   │  ┌──────────────┐  │            │
  │  │  │   OpenAI     │ │   │  │   Adapters   │  │            │
  │  │  │   Together   │ │   │  │  LoRA/QLoRA  │  │            │
  │  │  │   Fireworks  │ │   │  │    DoRA      │  │            │
  │  │  │   Anyscale   │ │   │  └──────────────┘  │            │
  │  │  │   Bedrock    │ │   │  ┌──────────────┐  │            │
  │  │  │   Vertex     │ │   │  │  Alignment   │  │            │
  │  │  └──────────────┘ │   │  │  DPO / ORPO  │  │            │
  │  └──────────────────┘   │  └──────────────┘  │            │
  │                          │  ┌──────────────┐  │            │
  │                          │  │Checkpointing │  │            │
  │                          │  │   Export      │  │            │
  │                          │  └──────────────┘  │            │
  │                          └───────────────────┘             │
  │                                  │                         │
  │                                  ▼                         │
  │                        ┌─────────────────┐                 │
  │                        │ TrainedModel    │                 │
  │                        │  Artifact       │                 │
  │                        └─────────────────┘                 │
  └─────────────────────────────────────────────────────────────┘

  Flow: Dataset → TrainingManager → Cloud / Local → Artifacts

Quick Start

Add to your Cargo.toml:

[dependencies]
brainwires-training = "0.6"

Submit a cloud fine-tuning job:

use brainwires_training::{
    TrainingManager, CloudFineTuneConfig, FineTuneProviderFactory,
    TrainingHyperparams, TrainingJobStatus,
};

// Create a provider (OpenAI in this example)
let provider = FineTuneProviderFactory::create("openai", "your-api-key")?;

// Configure the job
let config = CloudFineTuneConfig {
    model: "gpt-4o-mini-2024-07-18".into(),
    training_file: "data/train.jsonl".into(),
    validation_file: Some("data/eval.jsonl".into()),
    hyperparams: TrainingHyperparams {
        epochs: Some(3),
        learning_rate: Some(1e-5),
        batch_size: Some(4),
        ..Default::default()
    },
    ..Default::default()
};

// Submit and monitor
let job_id = provider.submit(config).await?;
loop {
    let status = provider.status(&job_id).await?;
    match status {
        TrainingJobStatus::Completed => break,
        TrainingJobStatus::Failed(err) => return Err(err.into()),
        _ => tokio::time::sleep(std::time::Duration::from_secs(30)).await,
    }
}

Features

# Local training only (no cloud deps)
[dependencies]
brainwires-training = { version = "0.6", default-features = false, features = ["local"] }

# Full — cloud + local
[dependencies]
brainwires-training = { version = "0.6", features = ["full"] }

Architecture

TrainingManager

Central coordinator that dispatches to the correct backend based on configuration.

Cloud Providers

All providers implement the FineTuneProvider trait:

#[async_trait]
pub trait FineTuneProvider: Send + Sync {
    async fn submit(&self, config: CloudFineTuneConfig) -> Result<TrainingJobId>;
    async fn status(&self, job_id: &TrainingJobId) -> Result<TrainingJobStatus>;
    async fn cancel(&self, job_id: &TrainingJobId) -> Result<()>;
    async fn download(&self, job_id: &TrainingJobId) -> Result<TrainedModelArtifact>;
}

CloudFineTuneConfig

Local Training (Burn Backend)

Requires the local feature. Uses the Burn deep learning framework for GPU-accelerated adapter training.

Adapter Methods

LoraConfig

Alignment Methods

LocalTrainingConfig

LrScheduler

Job Tracking

Usage Examples

Local LoRA Training

#[cfg(feature = "local")]
{
    use brainwires_training::{
        TrainingManager, LocalTrainingConfig, LoraConfig,
        AdapterMethod, ComputeDevice, TrainingHyperparams, LrScheduler,
    };

    let config = LocalTrainingConfig {
        adapter: AdapterMethod::LoRA,
        device: ComputeDevice::Wgpu,
        lora: LoraConfig {
            rank: 16,
            alpha: 32.0,
            dropout: 0.05,
            target_modules: vec!["q_proj".into(), "k_proj".into(), "v_proj".into()],
        },
        hyperparams: TrainingHyperparams {
            epochs: Some(3),
            learning_rate: Some(2e-4),
            batch_size: Some(8),
            lr_scheduler: Some(LrScheduler::Cosine),
            ..Default::default()
        },
        checkpoint_dir: Some("checkpoints/".into()),
        ..Default::default()
    };

    let manager = TrainingManager::new();
    let job_id = manager.submit_local(config, &dataset).await?;

    // Monitor progress
    let progress = manager.status(&job_id).await?;
    println!("Step {}, loss: {:.4}", progress.step, progress.loss);
}

DPO Alignment Training

#[cfg(feature = "local")]
{
    use brainwires_training::{
        LocalTrainingConfig, AdapterMethod, AlignmentMethod,
        ComputeDevice, TrainingHyperparams,
    };
    use brainwires_datasets::PreferenceDataset;

    let config = LocalTrainingConfig {
        adapter: AdapterMethod::LoRA,
        alignment: Some(AlignmentMethod::Dpo),
        device: ComputeDevice::Cuda(0),
        hyperparams: TrainingHyperparams {
            epochs: Some(1),
            learning_rate: Some(5e-5),
            batch_size: Some(4),
            ..Default::default()
        },
        ..Default::default()
    };

    let preference_data = PreferenceDataset::from_file("data/preferences.jsonl")?;
    let job_id = manager.submit_local(config, &preference_data).await?;
}

Cloud Provider Selection

#[cfg(feature = "cloud")]
{
    use brainwires_training::{FineTuneProviderFactory, CloudFineTuneConfig};

    // OpenAI
    let openai = FineTuneProviderFactory::create("openai", "sk-...")?;

    // Together AI
    let together = FineTuneProviderFactory::create("together", "tok-...")?;

    // AWS Bedrock
    let bedrock = FineTuneProviderFactory::create("bedrock", "aws-credentials")?;

    // Same config works with any provider
    let config = CloudFineTuneConfig {
        model: "meta-llama/Llama-3-8b".into(),
        training_file: "data/train.jsonl".into(),
        ..Default::default()
    };
    let job_id = together.submit(config).await?;
}

Monitoring Job Progress

use brainwires_training::{TrainingManager, TrainingJobStatus};

let manager = TrainingManager::new();
let status = manager.status(&job_id).await?;

match status {
    TrainingJobStatus::Running => {
        let progress = manager.progress(&job_id).await?;
        println!(
            "Epoch {}/{}, step {}, loss: {:.4}, lr: {:.2e}",
            progress.epoch, progress.total_epochs,
            progress.step, progress.loss, progress.learning_rate
        );
    }
    TrainingJobStatus::Completed => {
        let metrics = manager.metrics(&job_id).await?;
        println!(
            "Done! Train loss: {:.4}, eval loss: {:.4}, duration: {:?}",
            metrics.train_loss, metrics.eval_loss, metrics.duration
        );
    }
    TrainingJobStatus::Failed(err) => eprintln!("Failed: {err}"),
    _ => {}
}

Integration with Brainwires

Use via the brainwires facade crate:

[dependencies]
brainwires = { version = "0.6", features = ["training"] }

Or depend on brainwires-training directly for standalone training capabilities. The crate depends on brainwires-datasets for data types, so both are pulled in together.

References

Papers

• DPO: Direct Preference Optimization (2023) — the alignment method behind AlignmentMethod::Dpo
• ORPO: Monolithic Preference Optimization (2024) — single-stage alignment, AlignmentMethod::Orpo
• Cramming: Training on a Single GPU in One Day (2023) — efficient small-model training strategies
• Ring Attention: Near-Infinite Context (ICLR 2024) — long-context training
• Evolution Strategies for Billion-Parameter Fine-Tuning (Sept 2025)
• IR-Tuning: Efficient Layer-wise Fine-tuning (Sept 2025)
• QTHA: Quantum-Enhanced Fine Tuning (March 2025)
• Complexity-aware Fine-tuning (June 2025)
• LUNE: Efficient LLM Unlearning via LoRA (Dec 2025)
• Mamba-3: State Space Models (2025)
• Revisiting Chinchilla Scaling Laws (ACL 2025)

Technical Blogs & Guides

• LoRA/QLoRA/DoRA Production Guide (Dec 2025) — comprehensive adapter comparison
• DoRA — NVIDIA Technical Blog — weight-decomposed low-rank adaptation
• How to Align LLMs in 2025 with DPO — practical DPO guide
• Small LLM Training Guide 2026 — end-to-end training walkthrough
• Mixed Precision: FP8 vs BF16 — precision trade-offs
• NVIDIA FP8 Training Blog
• FSDP vs DeepSpeed — HuggingFace — distributed training comparison

Rust ML Ecosystem

• Burn Framework — the local training backend used by this crate
• Candle — HuggingFace — Rust inference framework
• tch-rs: PyTorch Rust Bindings
• PyO3 for AI/Data Science (2025)
• Rust + CUDA for ML (2025)

License

Licensed under the MIT License. See LICENSE for details.