Forge

What is it

Forge is a runtime for training and running neural networks, aiming to be the efficient and portable one: a single Rust crate and one set of WGSL kernels, with no CUDA toolchain and no Python interpreter anywhere in the loop.

Why

Running a transformer usually means Python, a CUDA toolchain, and a dependency stack tied to one vendor's hardware. Forge is a single Rust crate that trains and infers on any GPU wgpu reaches — Vulkan, Metal, D3D12, and WebGPU in a browser tab — with no CUDA, no Python interpreter, and no server round trip. That's not just a portability trick. Portable and efficient is where this is going: one runtime on every device you own, down to the edge, and eventually across them — the floor a vendor-independent runtime has to stand on before it can carry real research, and eventually local, safety-focused AI work, on hardware you control end to end.

Demo

Two pages run the runtime in a browser tab on your own GPU, built from this repository and served locally — the deployed site still describes 0.2.0:

make surprise   # select any text; the model tints it by how surprised it was
make council    # four small models merging their hidden states, one character at a time

https://github.com/user-attachments/assets/d3487d8f-40a1-4c84-9f98-d4a7c5ce1550

Run it

# Generate — the 43 MB char-level Shakespeare model ships in the repo
cargo run --release --example generate -- --model assets/shakespeare_char --prompt "ROMEO:"

# Real GPT-2 124M weights — the parity run from tests/gpt2_e2e.rs
./scripts/download_gpt2.sh
cargo run --release --example generate -- --backend wgpu --prompt "Hello Forge!"

# Train from scratch
./scripts/download_shakespeare.sh
cargo run --release --example train_shakespeare -- --backend wgpu

# Terminal model browser + run dashboard  (tools/forge-top)
cargo run --release -p forge-top -- --path models/

# A council of four small models, deciding one character together  (tools/council)
cargo run --release -p forge-council --example council_demo -- --prompt "ROMEO:"

# The two browser pages, built and served locally  (tools/)
make surprise
make council

# Tests
cargo test --release

As a library:

# The crate publishes as `forge-ml` (the name `forge` was taken on crates.io in
# 2017); the library it builds is still `forge`, so imports read `use forge::…`.
[dependencies]
forge-ml = "0.3"

use forge::{AnyTokenizer, Device, Gpt2, Gpt2Config, Sampling};

// assets/shakespeare_char ships in the repo; swap in models/gpt2 after
// running download_gpt2.sh — AnyTokenizer picks char or BPE by what it finds.
let device = Device::wgpu()?; // or Device::Cpu
let config = Gpt2Config::from_json("assets/shakespeare_char/config.json")?;
let model = Gpt2::from_safetensors("assets/shakespeare_char/model.safetensors", config, &device)?;
let tok = AnyTokenizer::from_dir("assets/shakespeare_char")?;
let text = model.generate(&tok, "ROMEO:", 40, Sampling::Greedy)?;

The crate, and what is built on it

forge-ml is the runtime and nothing else: tensors, kernels, GPT-2, tokenizers, serialization, and the browser bindings. One optional feature, off by default:

Feature	What it adds	Why it's optional
`train`	reverse-mode autograd, AdamW, the nine backward kernels, and `Gpt2::loss` / `loss_grads`. Needed by `examples/train_shakespeare.rs` and `make train`.	Forge is an inference runtime that also happens to train. `cargo add forge-ml` should not compile a tape you never record, and `src/wasm.rs` exports no training at all. Construction and serialization — `Gpt2::init_random`, `params`, `save_safetensors` — are not gated: they are not training, and the inference tests use them.

forge-ml = { version = "0.3", features = ["train"] }

Everything else lives in tools/, downstream of the runtime — a crate or a page that depends on forge and adds no kernel, no dtype and no device work of its own:

Tool	What it is
`tools/council`	Four small GPT-2s run in parallel and merged in hidden space by an entropy router, plus the page that draws the vectors they exchange
`tools/forge-top`	A terminal model browser and run dashboard
`tools/surprise`	A page that tints text by how surprised the model was to read it

Through 0.1.0 the first two were council and tui features of this crate. They are separate crates from 0.3.0 because the test of what belongs in the runtime is whether a third-party crate could have written it against the public API — and all three could. The primitives they stand on stayed: Gpt2::hidden_step, logits_from_hidden, wte_host, surprisal_async, Sampler, top_probs.

cargo run --release --features train --example train_shakespeare -- --backend wgpu

On Linux, wgpu needs a Vulkan ICD — install mesa-vulkan-drivers for a software fallback, or run scripts/setup_nvidia_vulkan.sh for NVIDIA inside a container. Check with cargo run --release --example wgpu_probe.

Acknowledgement

Supported by the RACE Merit Allocation Scheme (RMAS), RMIT University.

License

GNU Affero General Public License v3.0 — see LICENSE.

forge-ml 0.3.0