Skip to content

Rewrite the README as a README #69

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CanReader merged 1 commit into from
Oct 30, 2026
Merged

Rewrite the README as a README #69

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
332 changes: 91 additions & 241 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,37 +1,49 @@
<div align="center">

# FastNN

A robust deep learning library in Rust with CUDA kernels.
**A robust deep learning library in Rust with CUDA kernels.**

[![CI](https://github.com/CanReader/FastNN/actions/workflows/ci.yml/badge.svg)](https://github.com/CanReader/FastNN/actions/workflows/ci.yml)
[![crates.io](https://img.shields.io/crates/v/fastnn.svg)](https://crates.io/crates/fastnn)
[![docs.rs](https://img.shields.io/docsrs/fastnn)](https://docs.rs/fastnn)
[![MSRV](https://img.shields.io/badge/rustc-1.87+-blue.svg)](https://github.com/CanReader/FastNN/blob/master/Cargo.toml)
[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)

Tensors with reverse-mode autograd, the usual layers, transformers with
KV-cached generation, optimizers, data loading, resumable checkpoints and
safetensors. It runs on the CPU by default and on NVIDIA GPUs with the `cuda`
feature.

The codebase is small on purpose: 86 source files, about 12k lines, one idea
per file. You can follow `loss.backward()` from the training loop down to the
kernel. Every backward rule is checked against finite differences of its own
forward, and every CUDA kernel against the CPU path.

## Features

| Area | What is in it |
|---|---|
| Tensors | Broadcasting arithmetic, matmul (batched, transposed forms), reductions, indexing, views, `cat`/`stack`/`narrow`, CPU or CUDA storage |
| Autograd | Reverse mode, no tape, no global state. `no_grad`, `detect_anomaly`, gradient clipping by norm or value |
| Layers | `Linear`, `Conv1d`, `Conv2d` (dilation, groups, depthwise), `ConvTranspose1d/2d`, pooling, `BatchNorm2d`, `LayerNorm`, `RMSNorm`, `Dropout`, `Embedding`, `LSTM`, `GRU` |
| Transformers | `MultiHeadAttention` with causal and padding masks, encoder and decoder stacks, encoder-decoder `Transformer` |
| Generation | `KvCache`, `Sampler` (temperature, top-k, top-p, repetition penalty), `BeamSearch` |
| Optimizers | `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, `Adadelta`, `RAdam`, `Lion`, plus `Lookahead` and `Ema` wrappers |
| Schedules | `Warmup`, `CosineAnnealing`, `OneCycle`, `StepDecay`, `Constant` |
| Losses | Cross-entropy (label smoothing, class weights, ignore index), MSE, MAE, BCE, Huber, KL, NLL, focal, dice, Gaussian and Poisson NLL, triplet, contrastive, InfoNCE |
| Data | `Dataset` trait, shuffling `DataLoader` with device placement, MNIST |
| Checkpoints | Weights, full training state (weights, optimizer, step), safetensors read and write with F16/BF16 widening |
[Documentation](https://docs.rs/fastnn) · [Examples](examples) · [Releases](https://github.com/CanReader/FastNN/releases) · [Contributing](CONTRIBUTING.md)

</div>

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 one feature flag.

The whole stack is in one crate, written to be read. 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.

## Highlights

- **Autograd without a tape.** A tensor made by a differentiable op carries the
rule that made it. `backward()` walks those links, and dropping the loss frees
the graph. No `requires_grad`, no global state.
- **Built for long runs.** Checkpoints carry the optimizer state and the step,
so a resumed run continues instead of restarting Adam. `detect_anomaly` stops
at the op that produced a NaN. A frozen layer gets no gradient at all, not
just one the optimizer ignores.
- **Transformers end to end.** Multi-head attention with causal and padding
masks, encoder and decoder stacks, a KV cache for decoding, sampling with
temperature, top-k and top-p, and beam search.
- **GPU with one flag.** `--features cuda` compiles the kernels in this
repository and routes matmul through cuBLAS. The CPU build needs nothing but
`cargo`.
- **Checkpoints other tools can read.** Safetensors in and out, with F16 and
BF16 widened on load.
- **Checked against the maths.** `gradcheck` tests every backward rule against
finite differences, `cuda_parity` runs every kernel against the CPU path, and
`convergence` trains every layer family and optimizer on a small problem.

## Installation

Expand Down Expand Up @@ -76,14 +88,17 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
}
```

To train on a GPU, move the model and the loader:
To train on a GPU, move the model and the loader. The loop stays the same.

```rust
let device = Device::best(); // GPU if there is one, else CPU
model.to_device(device);
let loader = DataLoader::new(&train, 128).to_device(device);
```

The [documentation](https://docs.rs/fastnn) walks through the whole crate:
models, training, data, checkpoints, the GPU, generation and debugging.

## Examples

```bash
Expand Down Expand Up @@ -135,245 +150,80 @@ Releases include a `SHA256SUMS` file and a build provenance attestation for each
archive (`gh attestation verify <archive> --repo CanReader/FastNN`). The CUDA
build needs an NVIDIA driver new enough for the toolkit it was built with.

## Platform support

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

## How it fits together

```
data Dataset ─► DataLoader ─► Batch
│
nn Module ──► forward ────────┤ layers hold Param handles
▼
tensor Tensor ops on CPU or CUDA ─┴─► loss
│
autograd loss.backward() walks the graph the ops built,
depositing gradients in each Param's slot
│
optim opt.step() reads those slots and updates in place
```
![How a training step moves through FastNN](assets/architecture.svg)

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, with no borrow of the model.

Everything starts on the CPU. `model.to_device(device)` and
`DataLoader::to_device` move things over, and mixing devices panics instead of
copying behind your back.

The graph is the tensors. A tensor produced by a differentiable op carries a
node naming the rule that made it and the tensors it consumed, and
`loss.backward()` walks that. Drop the loss and the graph frees itself. A tensor
with no node is plain data, so there is no `requires_grad` flag. Wrap inference
in `no_grad(|| ...)` to skip building a graph.

`Param` is a handle, not a value. The model and the optimizer hold handles to
the same weight and gradient, so `opt.step()` doesn't borrow the model. Weight
tying works the same way: use the handle twice and the gradients add up.

## Long training runs

**Resume.** `save()` writes weights, which is enough for a finished model.
Resuming also needs the optimizer's moments and step count. Without them Adam
restarts at step 1 and the first update is far too large.

```rust
let mut step = load_training(&model, &mut opt, "run.fdl").unwrap_or(0);

while step < total {
// ... train ...
step += 1;
if step % 500 == 0 { save_training(&model, &opt, step, "run.fdl")?; }
}
```

Loading a plain checkpoint with `load_training` is an error, not a silent
optimizer reset.

**NaNs.** One non-finite gradient poisons the weights, and by the time the loss
prints NaN the culprit is long gone. `detect_anomaly` panics at the op that
produced it.

```rust
detect_anomaly(|| {
loss.backward(); // panics: "anomaly: Log produced inf in the gradient for input 0"
});

if !opt.gradients_are_finite() { continue; } // cheap guard: skip the batch
```

Both read every gradient, so don't leave them on in a healthy loop.

**Freezing.** A frozen parameter hands out a detached tensor, so the backward
pass stops there and no gradient is computed for it.

```rust
backbone.freeze();
let mut opt = Adam::new(model.parameters(), 1e-3); // updates only the head
```

**Fallible ops.** Ordinary ops panic on bad shapes. Where a shape comes from a
config file or an upload, use the `try_` variants.

```rust
let y = x.try_matmul(&w)?; // Error::Shape, not a panic
x.try_reshape(&[2, -1])?;
table.try_index_select(&ids)?;
```

They validate and then call the panicking op, so there is one implementation of
each operation.

## Extending
## What is inside

**A layer.** Implement `forward`, plus `named_parameters` if it has weights.
Device placement, parameter counting and gradient zeroing follow from those.

```rust
use fastnn::nn::module::scoped;

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())
}
}
```

**An op.** Forward in `tensor/ops/`, backward rule in `autograd/ops/`, attach it
with `with_grad`, add a case to `tests/gradcheck.rs`.

```rust
// tensor/ops/unary.rs
pub fn sqrt(&self) -> Tensor {
let out = unary_op(self, f32::sqrt, ffi::fastnn_cuda_sqrt);
let saved = out.clone();
out.with_grad(&[self], || SqrtBackward { output: saved })
}

// autograd/ops/unary.rs
impl Backward for SqrtBackward {
fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
vec![grad.div(&self.output.mul_scalar(2.0))]
}
fn name(&self) -> &'static str { "Sqrt" }
}
```

The rule returns gradients in input order. Save detached tensors, otherwise the
saved value pins the graph that produced it. [CONTRIBUTING.md](CONTRIBUTING.md)
has the full checklist.

**A dataset.** Implement `Dataset` and pass it to a `DataLoader`. Items are
written into caller-provided slices instead of returned as tensors, so batching
60,000 images doesn't allocate 60,000 throwaway ones.

## Errors

Shape and device mistakes panic. They are bugs in the calling code, like
indexing past the end of a slice. `Result` is for things that fail at runtime:

```rust
let device = Device::cuda(0)?; // no GPU
let data = Mnist::load(Split::Train)?; // download or parse failed
load(&model, "model.fdl")?; // missing file, or a shape that moved
```

## Testing

```bash
cargo test # everything
cargo test --test gradcheck # every backward rule vs finite differences
cargo test --test robustness # resume, anomalies, freezing
cargo test --test generation # kv-cached decoding, masks, beam search, sampling
cargo test --test convergence # every layer family and optimizer actually learns
cargo test --test losses # loss values against hand-computed references
cargo test --features cuda --test cuda_parity -- --test-threads=1 # CPU vs GPU
cargo bench # throughput
```
| Area | What is in it |
|---|---|
| Tensors | Broadcasting arithmetic, matmul (batched, transposed forms), reductions, indexing, views, `cat`/`stack`/`narrow`, CPU or CUDA storage |
| Autograd | Reverse mode, no tape, no global state. `no_grad`, `detect_anomaly`, gradient clipping by norm or value |
| Layers | `Linear`, `Conv1d`, `Conv2d` (dilation, groups, depthwise), `ConvTranspose1d/2d`, pooling, `BatchNorm2d`, `LayerNorm`, `RMSNorm`, `Dropout`, `Embedding`, `LSTM`, `GRU` |
| Transformers | `MultiHeadAttention` with causal and padding masks, encoder and decoder stacks, encoder-decoder `Transformer` |
| Generation | `KvCache`, `Sampler` (temperature, top-k, top-p, repetition penalty), `BeamSearch` |
| Optimizers | `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, `Adadelta`, `RAdam`, `Lion`, plus `Lookahead` and `Ema` wrappers |
| Schedules | `Warmup`, `CosineAnnealing`, `OneCycle`, `StepDecay`, `Constant` |
| Losses | Cross-entropy (label smoothing, class weights, ignore index), MSE, MAE, BCE, Huber, KL, NLL, focal, dice, Gaussian and Poisson NLL, triplet, contrastive, InfoNCE |
| Data | `Dataset` trait, shuffling `DataLoader` with device placement, MNIST |
| Checkpoints | Weights, full training state (weights, optimizer, step), safetensors read and write with F16/BF16 widening |

`cuda_parity` runs each op on both devices and compares. It skips itself when
there is no GPU. It has already caught a real bug: the softmax kernel's block
reduction assumed a power-of-two thread count and dropped the tail of every row
whose width wasn't one, which included MNIST's ten classes.
## Platform support

CI runs formatting, clippy, the suite on Linux, macOS and Windows, an MSRV
build, docs, packaging and a CUDA compile. A nightly job reruns the suite in
release on stable and beta and trains the examples that don't need a download.
| 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 |

## CUDA
### Building with CUDA

The `cuda` feature is opt-in. With it on, `build.rs` compiles `cuda/kernels.cu`
with `nvcc` and links `cudart`, `cublas` and `curand`.
With the `cuda` feature on, `build.rs` compiles `cuda/kernels.cu` with `nvcc`
and links `cudart`, `cublas` and `curand`.

* `nvcc` is found on `PATH`, or set the `NVCC` environment variable.
* Headers and libraries come from `CUDA_PATH`, then `CUDA_HOME`, then the
- `nvcc` is found on `PATH`, or set the `NVCC` environment variable.
- Headers and libraries come from `CUDA_PATH`, then `CUDA_HOME`, then the
default install location. Make sure that is the same toolkit as your `nvcc`.
* If `nvcc` rejects your system compiler as too new, point `FASTNN_NVCC_CCBIN`
- If `nvcc` rejects your system compiler as too new, point `FASTNN_NVCC_CCBIN`
at one it accepts.
* Without the feature, `cuda/stubs.c` supplies the symbols, `Device::cuda(0)`
- Without the feature, `cuda/stubs.c` supplies the symbols, `Device::cuda(0)`
returns an error, and everything runs on the CPU.

Matmul goes through cuBLAS. The transposed forms the backward pass needs
(`matmul_nt`, `matmul_tn`) use a transpose flag instead of a transposed copy.
GPU allocations come from a size-keyed free list, so the short-lived temporaries
in a training step don't each hit the driver.
## Testing

## Project layout
`cargo test` runs everything. The suites that carry the most weight:

```
src/
tensor/ the array type and everything you can do to it
core.rs Tensor: shape, storage, graph link
checked.rs try_* variants that report instead of panicking
shape.rs strides, broadcasting, index math
storage.rs the bytes, on one device or the other
device.rs Device::cuda(0) -> Result
init.rs zeros, randn, kaiming, xavier, ...
ops/ arith unary activation matmul reduce view
index conv pool norm

autograd/ reverse-mode differentiation
node.rs Backward trait, graph nodes, gradient slots
engine.rs the reverse pass
mode.rs no_grad
anomaly.rs detect_anomaly
ops/ one backward rule per forward op, same file names

nn/ layers, all implementing Module
linear conv conv1d conv_transpose pooling norm activation
dropout embedding shape attention transformer decoder rnn
cache (kv cache) sample (sampling) beam (beam search)
loss metric_losses module param sequential

optim/ sgd adam rmsprop adagrad radam lion lookahead
ema schedule state
data/ dataset loader mnist
serialize/ checkpoint training (weights + optimizer)
safetensors half (F16/BF16 conversion)
cuda/ ffi (raw bindings) kernels (safe wrappers) buffer

cuda/kernels.cu every GPU kernel
cuda/stubs.c CPU-only stand-ins for the same symbols
```
- `gradcheck` checks every backward rule against central finite differences of
its own forward.
- `cuda_parity` runs every kernel against the CPU path on a real GPU, and skips
itself when there is none. It has caught a real bug: a block reduction that
assumed a power-of-two thread count and dropped the tail of every row whose
width was not one, which included MNIST's ten classes.
- `convergence` trains every layer family and optimizer on a small problem
until the loss drops under a target, and fails if it never does.

`tensor/ops/` and `autograd/ops/` mirror each other file for file. The forward
for `relu` is in `tensor/ops/activation.rs`, its derivative in
`autograd/ops/activation.rs`.
CI runs formatting, clippy, the suite on Linux, macOS and Windows, an MSRV
build, docs, packaging and a CUDA compile. A nightly job reruns the suite in
release on stable and beta and trains the examples that do not need a download.

## Limitations

* Version 0.x, so the API can still change between minor releases.
* Compute is `f32`. F16 and BF16 checkpoints load and are widened on read.
* `LSTM` and `GRU` are built from ordinary ops, one graph node per gate per
- Version 0.x, so the API can still change between minor releases.
- Compute is `f32`. F16 and BF16 checkpoints load and are widened on read.
- `LSTM` and `GRU` are built from ordinary ops, one graph node per gate per
timestep. Fine for short sequences, use a transformer for long ones.
* Tensors are always contiguous. `permute` and `expand` write a new buffer
- Tensors are always contiguous. `permute` and `expand` write a new buffer
instead of returning a strided view.

## Contributing
Expand Down
Loading