Crate fastnn 

FastNN

A robust deep learning library in Rust with CUDA kernels.

FastNN trains neural networks in Rust. It has tensors with automatic differentiation, the layers and optimizers behind MLPs, CNNs, RNNs and transformers, data loading, and checkpoints a run can resume from. It runs on the CPU out of the box and on NVIDIA GPUs with the cuda feature.

The whole stack is in this crate: Rust on the CPU, and with the cuda feature the kernels in the same repository. It is written to be read, so you can follow loss.backward() from your training loop down to the kernel. Every backward rule is tested against finite differences of its own forward, every CUDA kernel against the CPU path, and unsafe appears only in the CUDA bindings.

A batch goes through the model, the ops build a graph as they run, and the loss sits at its end. backward() walks that graph and leaves a gradient in each Param. The optimizer holds the same Param handles, so step() reads the gradients and updates the weights in place.

Quick start

Train a classifier on MNIST:

use fastnn::prelude::*;
use fastnn::data::{Mnist, Split};

let train = Mnist::load(Split::Train)?;
let loader = DataLoader::new(&train, 128).shuffle(true);

let model = Sequential::new()
    .add(Flatten::new())
    .add(Linear::new(784, 128))
    .add(ReLU)
    .add(Linear::new(128, 10));

let mut opt = Adam::new(model.parameters(), 1e-3);

for batch in loader.iter() {
    let loss = cross_entropy(&model.forward(&batch.inputs), &batch.labels());

    opt.zero_grad();
    loss.backward();
    opt.step();
}

save(&model, "mnist.fdl")?;

The examples are complete programs. Run one with cargo run --release --example <name>:

Example	What it shows
simple_mlp	MLP learning XOR, the smallest training loop
mnist_mlp	Dataset download, batching, evaluation, checkpointing
mnist_cnn	Convolutions, batch norm, pooling, an LR schedule
char_lm	Small GPT trained from scratch, resumes after Ctrl-C
finetune	Pretrain, freeze the backbone, train a new head
gan	Two networks trained against each other
vae	Variational autoencoder on 2-D data
reinforce	Policy gradient on a bandit
multi_task	Shared trunk, two heads, one combined loss

Use --release. Debug builds are about 50x slower.

A Tour of FastNN

A training project meets the crate in a fixed order: build a model, train it, feed it data, save it, move it to a GPU, maybe generate from it, and debug it when a run goes bad. This tour follows that order and shows the API for each step. Everything a script needs is in the prelude, and each module’s own page goes deeper.

Models

A model is anything that implements Module, and Sequential chains layers into one. A layer’s weights are Param handles, shared with the optimizer, so training never has to borrow the model.

use fastnn::prelude::*;

let model = Sequential::new()
    .add(Linear::new(16, 32))
    .add(ReLU)
    .add(Dropout::new(0.1))
    .add(Linear::new(32, 4));

let logits = model.forward(&Tensor::randn(&[8, 16]));
assert_eq!(logits.shape(), &[8, 4]);
assert_eq!(model.num_parameters(), 16 * 32 + 32 + 32 * 4 + 4);

Your own layer needs forward, plus named_parameters if it has weights. Device placement, parameter counting and gradient zeroing follow from those.

use fastnn::nn::module::scoped;
use fastnn::prelude::*;

struct Residual { inner: Linear }

impl Module for Residual {
    fn forward(&self, x: &Tensor) -> Tensor {
        x.add(&self.inner.forward(x).relu())
    }
    fn named_parameters(&self) -> Vec<(String, Param)> {
        scoped("inner", self.inner.named_parameters())
    }
}

nn has the rest: convolutions with dilation and groups, pooling, batch, layer and RMS normalization, embeddings, LSTM and GRU, multi-head attention, and transformer encoder and decoder stacks.

Training

An optimizer takes the model’s parameters and a learning rate. A LrSchedule maps the step to a rate, and set_lr applies it. The losses are plain functions in nn.

use fastnn::prelude::*;

let model = Sequential::new().add(Linear::new(4, 3));
let mut opt = AdamW::new(model.parameters(), 1e-3);
let schedule = Warmup::new(CosineAnnealing::new(1e-3, 1000), 100);

let x = Tensor::randn(&[8, 4]);
let targets = [0, 1, 2, 0, 1, 2, 0, 1];

for step in 0..3 {
    opt.set_lr(schedule.lr_at(step));
    let loss = cross_entropy(&model.forward(&x), &targets);

    opt.zero_grad();
    loss.backward();
    clip_grad_norm(opt.parameters(), 1.0);
    opt.step();
}

Dropout and BatchNorm2d behave differently at inference, so call model.eval() before validation and model.train() after.

To fine-tune, freeze what should not move. A frozen parameter hands out a detached tensor, so the backward pass stops there and no gradient is computed for it at all. The optimizer skips a parameter with no gradient.

let backbone = Linear::new(16, 16);
backbone.freeze();

let model = Sequential::new().add(backbone).add(ReLU).add(Linear::new(16, 2));
let mut opt = Adam::new(model.parameters(), 1e-3);   // only the head will move

optim has SGD, Adam, AdamW, RMSprop, Adagrad, Adadelta, RAdam and Lion, the Lookahead and Ema wrappers, and the Warmup, CosineAnnealing, OneCycle and StepDecay schedules. CrossEntropyLoss adds label smoothing, class weights and an ignore index to the plain function.

Data

A Dataset is anything indexable that can write one item into a slice. A DataLoader turns it into shuffled batches and can place them on a device. Items go into one buffer per batch, so batching 60,000 images does not allocate 60,000 tensors.

use fastnn::prelude::*;

let inputs = Tensor::randn(&[100, 8]);
let targets = Tensor::from_vec((0..100).map(|i| (i % 3) as f32).collect(), &[100]);
let dataset = TensorDataset::new(&inputs, &targets);

let loader = DataLoader::new(&dataset, 32).shuffle(true);
for batch in loader.iter() {
    let labels = batch.labels();   // targets as class ids
    assert_eq!(batch.inputs.dim(0), labels.len());
}

Mnist::load downloads the digits once and caches them. The data page shows a dataset written from scratch.

Checkpoints

save writes the weights, which is all a finished model needs. A run that has to survive a crash needs more: Adam’s moments and step count drive its bias correction, and restored without them it restarts at step 1 with an oversized first update. save_training writes the weights, the optimizer state and the step together, and load_training refuses a plain weights file rather than silently resetting the optimizer.

let mut step = load_training(&model, &mut opt, &path).unwrap_or(0);

while step < total {
    // one training step
    step += 1;
    if step % 500 == 0 {
        save_training(&model, &opt, step, &path)?;
    }
}

save_safetensors and load_safetensors exchange weights with other tools. F16 and BF16 files load and are widened to f32.

GPU

Everything starts on the CPU. Move the model and the loader, and the training loop stays the same. Mixing devices panics instead of copying behind your back.

let device = Device::best();   // the first GPU, else the CPU
model.to_device(device);
let loader = DataLoader::new(&dataset, 128).to_device(device);

Matmul goes through cuBLAS and everything else through the kernels in cuda/kernels.cu, each checked against the CPU path by the parity tests. The kernels use fast math, so GPU and CPU results differ in the last few digits.

Generation

A causal TransformerStack decodes through a StackCache, so each new token attends to the whole prefix without recomputing it. Sampler turns logits into a token, and BeamSearch finds the most probable sequence instead.

use fastnn::prelude::*;

let (vocab, width) = (50, 32);
let tokens = Embedding::new(vocab, width);
let positions = Embedding::new(64, width);
let stack = TransformerStack::causal(width, 4, 64, 2, 0.0);
let head = Linear::new(width, vocab);
let sampler = Sampler::new().temperature(0.8).top_k(10);

let mut ids = vec![1];
let mut cache = stack.new_cache();
no_grad(|| {
    for _ in 0..5 {
        let last = *ids.last().unwrap();
        let x = tokens.lookup(&[last]).add(&positions.lookup(&[cache.len()]));
        let hidden = stack.forward_cached(&x.reshape(&[1, 1, -1]), &mut cache);
        ids.push(sampler.sample(&head.forward(&hidden)));
    }
});
assert_eq!(ids.len(), 6);

Sampling draws from the thread’s random source, so manual_seed makes a generation reproducible.

Debugging

One non-finite gradient poisons the weights, and by the time the loss prints NaN the op at fault is long gone. detect_anomaly checks every gradient as it is produced and panics at the first bad one, naming the op:

let w = Param::new(Tensor::from_vec(vec![0.0], &[1]));
detect_anomaly(|| w.tensor().log().sum().backward());
// panics: anomaly: Log produced inf at element 0 of [1] in the gradient for input 0

It reads every gradient, so keep it for debugging. The cheap guard for a healthy loop is gradients_are_finite:

if !opt.gradients_are_finite() {
    opt.zero_grad();   // skip this batch instead of stepping into NaN
}

Shape and device mistakes panic with a message naming both shapes, since they are bugs in the calling code. Where a shape comes from outside your code, a config file or an upload, the try_ variants return an Error instead:

let a = Tensor::randn(&[2, 3]);
assert!(a.try_matmul(&Tensor::randn(&[4, 4])).is_err());

Tensors

A Tensor is an immutable f32 array on the CPU or a GPU. Ops return new tensors and broadcast the way NumPy does.

use fastnn::tensor::Tensor;

let a = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let bias = Tensor::from_vec(vec![10.0, 20.0, 30.0], &[3]);

let shifted = a.add(&bias);            // [2, 3] + [3] broadcasts over the rows
assert_eq!(shifted.to_vec(), vec![11.0, 22.0, 33.0, 14.0, 25.0, 36.0]);

let gram = a.matmul(&a.transpose());   // [2, 3] x [3, 2]
assert_eq!(gram.shape(), &[2, 2]);

assert_eq!(a.sum_axis(1).to_vec(), vec![6.0, 15.0]);
assert_eq!(a.argmax(1), vec![2, 2]);
assert_eq!(a.narrow(0, 1, 1).to_vec(), vec![4.0, 5.0, 6.0]);
assert_eq!(a.reshape(&[3, -1]).shape(), &[3, 2]);

Tensor::zeros, ones, randn, arange and eye create tensors, and the Kaiming and Xavier initialisers are there for weights. to moves a tensor between devices. The tensor::ops page lists every operation by kind.

Autograd

A tensor made by a differentiable op carries a node: the rule that made it and the tensors it consumed. backward() walks those nodes and leaves each weight’s gradient in its Param. A tensor with no node is plain data, so there is no requires_grad flag and no tape. Drop the loss and the graph is freed with it.

use fastnn::prelude::*;

let w = Param::new(Tensor::from_vec(vec![1.0, 2.0], &[2]));
let x = Tensor::from_vec(vec![3.0, 4.0], &[2]);

let loss = w.tensor().mul(&x).sum();   // 1*3 + 2*4
loss.backward();

assert_eq!(w.grad().unwrap().to_vec(), vec![3.0, 4.0]);

Wrap inference in no_grad and no graph is built:

let y = no_grad(|| w.tensor().exp());
assert!(y.grad_fn().is_none());

Every op is two small functions: the forward in tensor/ops/ and its rule in autograd::ops, in files with the same names. The rule returns gradients in input order and saves only detached tensors, and tests/gradcheck.rs checks it against central finite differences of its own forward. Adding an op is adding those two functions and a test case.

Feature flags

• cuda compiles cuda/kernels.cu with nvcc and links the CUDA runtime, cuBLAS and cuRAND. It needs a CUDA toolkit and targets GPUs from Turing to Hopper (sm_75 to sm_90). It is off by default. Without it every op runs on the CPU and Device::cuda returns an error.

These docs are built without cuda, since docs.rs has no CUDA toolkit.

Platform support

FastNN needs Rust 1.87 or newer.

Platform	CPU	CUDA
Linux x86_64	Tested in CI	Built in CI, parity tested on real hardware
Linux ARM64	Built for releases	Not supported
macOS (Intel, Apple Silicon)	Tested in CI	Not available
Windows x86_64	Tested in CI	Should build, not tested

Modules

• tensor is the array type, its ops and devices.
• autograd is the graph and the backward rules, one file per op category.
• nn is layers, losses, sampling and beam search.
• optim is optimizers, schedules and gradient clipping.
• data is datasets, batching and MNIST.
• serialize is checkpoints and safetensors.
• cuda is GPU memory and the kernel bindings.

Re-exports

pub use error::Error;

pub use error::Result;

Modules

autograd

Reverse-mode autodiff. Each differentiable tensor carries a Node with its backward rule and inputs, so the graph is just the tensors and there is no tape.

cuda

GPU layer: device setup, memory (buffer), raw bindings (ffi) and safe wrappers (kernels).

data

Datasets, batching, and MNIST.

error

The error type for I/O, downloads, checkpoints, and CUDA.

nn

Layers, losses and decoding helpers. Everything implements Module, and a layer’s weights are shared Param handles so an optimizer can hold them too.

optim

Optimizers and learning-rate schedules.

prelude

The usual imports for a training script.

rng

The per-thread random source behind init, dropout, and shuffling.

serialize

Checkpoints: .fdl weights, .fdl training state, safetensors, and the 16-bit float conversions safetensors needs.

tensor

Tensors: shape, storage and the link into the autograd graph.