Implementation:LaurentMazare Tch rs Min GPT

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Deep Learning, Natural Language Processing, Transformer Models
Last Updated	2026-02-08 00:00 GMT

Overview

A minimal GPT (Generative Pre-trained Transformer) implementation in Rust featuring causal self-attention, positional embeddings, and weight decay groups for text generation on the tinyshakespeare dataset.

Description

This example is a Rust port of Karpathy's minGPT. It implements the full GPT architecture with the following components:

Config struct: Holds model hyperparameters including vocabulary size, embedding dimension (512), number of attention heads (8), number of transformer layers (8), block size (128), and dropout rates (0.1 for attention, residual, and embedding dropout).

Custom Linear layer: Implements weight and bias as separate parameter groups to enable selective weight decay. The weight matrix belongs to WEIGHT_DECAY_GROUP (group 1, decay 0.1) while biases belong to NO_WEIGHT_DECAY_GROUP (group 0, decay 0.0). A linear_no_bias variant creates a non-trainable zero bias.

Causal Self-Attention: Implements multi-head attention with key, query, and value projections. A lower-triangular causal mask prevents attending to future positions. The mask is applied by filling masked positions with f64::NEG_INFINITY before softmax. Attention dropout is applied during training.

Transformer Block: Combines layer normalization, causal self-attention, and a feed-forward network (linear -> GELU -> linear) with residual connections and dropout.

GPT model: Stacks token embeddings, learned positional embeddings, N transformer blocks, final layer normalization, and a linear head projecting back to vocabulary size.

Training uses the AdamW optimizer with per-group weight decay, cross-entropy loss, and gradient clipping at 0.5. The model supports both train and predict modes via command-line arguments. Model checkpoints are saved every 10,000 iterations.

Text generation is autoregressive: tokens are fed through the model, softmax is applied to the last position's logits, and the next token is sampled via multinomial. The context window slides by dropping the oldest token to maintain block size.

Usage

Use this example to understand how to implement transformer-based language models in Rust with tch-rs, including causal masking, multi-head attention, weight decay groups, and autoregressive text generation.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: examples/min-gpt/main.rs
• Lines: 1-212

Signature

#[derive(Debug, Copy, Clone)]
struct Config {
    vocab_size: i64,
    n_embd: i64,
    n_head: i64,
    n_layer: i64,
    block_size: i64,
    attn_pdrop: f64,
    resid_pdrop: f64,
    embd_pdrop: f64,
}

struct Linear {
    pub ws: Tensor,
    pub bs: Tensor,
}

impl nn::Module for Linear { fn forward(&self, xs: &Tensor) -> Tensor }

fn linear(vs: nn::Path, in_dim: i64, out_dim: i64) -> Linear
fn linear_no_bias(vs: nn::Path, in_dim: i64, out_dim: i64) -> Linear
fn causal_self_attention(p: &nn::Path, cfg: Config) -> impl ModuleT
fn block(p: &nn::Path, cfg: Config) -> impl ModuleT
fn gpt(p: nn::Path, cfg: Config) -> impl ModuleT
fn sample(data: &TextData, gpt: &impl ModuleT, input: Tensor) -> String
pub fn main() -> Result<()>

Import

// Standalone binary example. Run with:
// cargo run --example min-gpt -- train
// cargo run --example min-gpt -- predict weights.ot "seed text"
use anyhow::{bail, Result};
use tch::data::TextData;
use tch::nn::{ModuleT, OptimizerConfig};
use tch::{nn, Device, IndexOp, Kind, Tensor};

I/O Contract

Inputs

Name	Type	Required	Description
mode	CLI argument	Yes	Either "train" or "predict".
data/input.txt	File (text)	Yes	Training corpus (e.g., tinyshakespeare).
weights.ot	File (weights)	For predict	Pre-trained model weights file.
seqstart	CLI argument (String)	For predict	Seed text for text generation.

Outputs

Name	Type	Description
stdout	Text	Training loss, sampled text (4096 chars) during training, or predicted text in predict mode.
gptN.ot	File (weights)	Model checkpoint saved every 10,000 iterations during training.

Usage Examples

use anyhow::{bail, Result};
use tch::data::TextData;
use tch::nn::{ModuleT, OptimizerConfig};
use tch::{nn, Device, IndexOp, Kind, Tensor};

const BLOCK_SIZE: i64 = 128;
const BATCH_SIZE: i64 = 64;
const LEARNING_RATE: f64 = 0.0003;

// Configure the GPT model
let cfg = Config {
    vocab_size: labels,
    n_embd: 512,
    n_head: 8,
    n_layer: 8,
    block_size: BLOCK_SIZE,
    attn_pdrop: 0.1,
    resid_pdrop: 0.1,
    embd_pdrop: 0.1,
};

let mut vs = nn::VarStore::new(Device::cuda_if_available());
let gpt = gpt(vs.root() / "gpt", cfg);

// Use AdamW with per-group weight decay
let mut opt = nn::AdamW::default().build(&vs, LEARNING_RATE)?;
opt.set_weight_decay_group(0, 0.0);  // no decay for biases/embeddings
opt.set_weight_decay_group(1, 0.1);  // decay for weight matrices

// Training loop
for batch in data.iter_shuffle(BLOCK_SIZE + 1, BATCH_SIZE) {
    let xs = batch.narrow(1, 0, BLOCK_SIZE).to_kind(Kind::Int64).to_device(device);
    let ys = batch.narrow(1, 1, BLOCK_SIZE).to_kind(Kind::Int64).to_device(device);
    let logits = xs.apply_t(&gpt, true);
    let loss = logits
        .view([BATCH_SIZE * BLOCK_SIZE, labels])
        .cross_entropy_for_logits(&ys.view([BATCH_SIZE * BLOCK_SIZE]));
    opt.backward_step_clip(&loss, 0.5);
}

Related Pages

• Principle:LaurentMazare_Tch_rs_GPT_Text_Generation