Crate forge 

Forge — a WebGPU-native machine learning framework in Rust, scoped to GPT-2.

Production execution targets WebGPU via wgpu; the CPU backend is a mathematically identical reference used for testing and verification.

Optional features

One, off by default, so cargo add forge-ml gets the runtime alone.

• train — autograd and optim, the backward kernels, and Gpt2::loss/loss_grads. Off by default because Forge is an inference runtime that also happens to train; cargo add forge-ml should not compile a tape you never record. Construction and serialization (Gpt2::init_random, params, save_safetensors) stay core — they are not training, and the inference tests use them.

Re-exports

pub use device::Device;

pub use dtype::DType;

pub use error::ForgeError;

pub use error::Result;

pub use models::gpt2::AttnStep;

pub use models::gpt2::Gpt2;

pub use models::gpt2::Gpt2Config;

pub use models::gpt2::KvCache;

pub use models::gpt2::LayerDetail;

pub use models::gpt2::Sampler;

pub use models::gpt2::Sampling;

pub use models::gpt2::StepTrace;

pub use models::gpt2::Surprisal;

pub use models::gpt2::top_probs;

pub use shape::Shape;

pub use tensor::Tensor;

pub use tokenizer::AnyTokenizer;

pub use tokenizer::CharTokenizer;

pub use tokenizer::Gpt2Tokenizer;

pub use tokenizer::Tokenizer;

Modules

autogradtrain

Tape-based reverse-mode autograd (roadmap v4, Stage 8).

backend

device

dtype

error

models

nn

Minimal neural-network modules — exactly what GPT-2 needs.

ops

Backend-agnostic tensor operations. Each op validates shapes, dispatches to the CPU reference or the WGPU kernels, and returns a tensor on the same device as its inputs.

optimtrain

Optimizers (roadmap v4, Stage 9): AdamW with decoupled weight decay and optional global gradient-norm clipping. Backend-agnostic — state lives on the same device as the parameters and updates run through ops.

serialization

SafeTensors saving: named tensors -> file.

shape

tensor

tokenizer

Tokenizers: the GPT-2 byte-level BPE, and nanoGPT’s character-level vocab.