Expand description
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.
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::Sampling;pub use models::gpt2::StepTrace;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§
- autograd
- 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.
- optim
- 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.