model.GPT — Construct, Train, and Serialize a GPT Model

The model package provides the GPT struct which assembles all transformer layers into a trainable and serializable language model. It implements the optimizer.Model interface from github.com/itsubaki/autograd, making it compatible with AdamW and other autograd optimizers out of the box.

GPT Struct

GPT holds the architecture hyperparameters as exported fields alongside an embedded model.Model from the autograd library. The embedded model manages the named layer registry and the full parameter map.

type GPT struct {
    VocabSize     int
    MaxContextLen int
    EmbedDim      int
    NumOfHeads    int
    NumOfBlocks   int
    Theta         float64
    model.Model   // embedded autograd Model (layer registry + Params/Cleargrads)
}

Field	Description
`VocabSize`	Number of tokens in the vocabulary
`MaxContextLen`	Maximum sequence length for KV cache allocation and RoPE precomputation
`EmbedDim`	Hidden and embedding dimension shared across all layers
`NumOfHeads`	Number of multi-head attention heads; head dimension = `EmbedDim / NumOfHeads`
`NumOfBlocks`	Number of transformer blocks stacked in sequence
`Theta`	RoPE base frequency (standard default: `10000`)

NewGPT()

NewGPT constructs and initializes a GPT model with randomly initialized weights.

func NewGPT(
    vocabSize int,
    maxContextLen int,
    embedDim int,
    numOfHeads int,
    numOfBlocks int,
    theta float64,
    useCache ...bool,
) *GPT

vocabSize

int

required

Size of the token vocabulary. Must match the tokenizer’s vocabulary. Example: 1000.

maxContextLen

int

required

Maximum sequence length used for KV cache allocation and RoPE frequency precomputation. Example: 256.

embedDim

int

required

Embedding and hidden dimension used consistently across all layers. Example: 192.

numOfHeads

int

required

Number of attention heads. The per-head dimension is computed as embedDim / numOfHeads. Example: 6.

numOfBlocks

int

required

Number of transformer blocks stacked in the model. Example: 6.

theta

float64

required

Base frequency for Rotary Position Embedding (RoPE). The conventional default is 10000.

useCache

...bool

Optional variadic flag. Pass true to enable KV caching in all attention layers. Required for efficient autoregressive generation; typically omitted during training.

Layer composition

NewGPT registers the following named layers in order:

Key	Type	Description
`"embed"`	`Embeddings`	Token embedding lookup table — shape `(VocabSize, EmbedDim)`
`"block[i]"`	`BlockT`	Transformer block `i` containing `norm1` (RMSNorm) → `attn` (MultiHeadAttention with RoPE) → residual → `norm2` (RMSNorm) → `ffn` (SwiGLU) → residual
`"norm"`	`RMSNorm`	Final layer normalization before unembedding
`"unembed"`	`Linear`	Linear projection from `EmbedDim` → `VocabSize` (no bias)

RoPE frequencies are precomputed once via function.RoPE(theta, embedDim, maxContextLen) and shared across all blocks.

The embedDim value in the test suite uses 394 (not 192) to produce integer head dimensions with numOfHeads=6. Always ensure embedDim is evenly divisible by numOfHeads.

NewGPTFrom()

NewGPTFrom loads a gob-encoded model checkpoint from disk, reconstructs the architecture from the saved hyperparameters, and copies all saved parameter values into the freshly built model.

func NewGPTFrom(path string, useCache ...bool) (*GPT, error)

path

string

required

Filesystem path to a .gob file previously written by GPT.Save().

useCache

...bool

Optional variadic flag. Pass true to enable KV caching in the reconstructed model, typically required for inference.

Returns (*GPT, error) — a fully initialized model on success, or a non-nil error if the file cannot be opened, decoded, or if the parameter shapes do not match.

m, err := model.NewGPTFrom("testdata/model_gpt.gob", true)
if err != nil {
    log.Fatal(err)
}

fmt.Println(m.VocabSize, m.MaxContextLen, m.EmbedDim)
// 1000 256 394

NewGPTFrom internally calls NewGPT with the hyperparameters decoded from the gob file, then calls Load to copy parameter data. The useCache flag you pass here is applied to the freshly constructed model — it does not need to match the flag used at training time.

Forward()

Forward runs a full forward pass through the model and returns raw logits.

func (m *GPT) Forward(ids *variable.Variable) *variable.Variable

ids

*variable.Variable

required

Token ID tensor with shape (B, C) where B is the batch size and C is the sequence length. Token values must be integers in the range [0, VocabSize), stored as float64.

Returns *variable.Variable — logit tensor with shape (B, C, VocabSize). Each position holds unnormalized log-probabilities over the full vocabulary. Data flow:

ids (B, C)
  → embed          → (B, C, EmbedDim)
  → block[0..N-1]  → (B, C, EmbedDim)   # residual stream through N blocks
  → norm           → (B, C, EmbedDim)   # final RMSNorm
  → unembed        → (B, C, VocabSize)  # logits

When using KV cache for token-by-token generation, call Forward with a shape-(1, 1) input for each new token. The attention layers read cached keys and values from all prior steps automatically.

During training, Forward is called with the full (B, C) sequence and gradients are tracked. During generation, wrap calls inside variable.Nograd() (done automatically by GenerateChan) to skip gradient allocation.

Save() and Load()

Save()

Save encodes the entire GPT struct — hyperparameters and all parameter tensors — to a gob file at the given path.

func (m *GPT) Save(path string) error

path

string

required

Destination file path. The file is created (or truncated) by os.Create.

Returns error — non-nil if the file cannot be created or if gob encoding fails.

Load()

Load copies parameter data from a layer.Parameters map (keyed by dotted path strings such as "block[0].attn.Wq.w") into the live model’s parameters in-place.

func (m *GPT) Load(params layer.Parameters) error

params

layer.Parameters

required

A map[string]*layer.Parameter typically obtained from a previously decoded GPT via savedModel.Params(). Every key in params must exist in the current model; an unknown key returns an error.

Returns error — non-nil if any key in params is not found in the model.

Load is called automatically by NewGPTFrom. You only need to call it directly if you are implementing a custom checkpoint-loading workflow (e.g., loading weights converted from another framework).

Named parameter keys follow the pattern <layer>.<sublayer>.<weight>. For a 6-block model the full parameter set includes:

block[0].attn.Wk.w    block[0].attn.Wq.w
block[0].attn.Wv.w    block[0].attn.Wo.w
block[0].ffn.W.w      block[0].ffn.V.w    block[0].ffn.O.w
block[0].norm1.gamma  block[0].norm2.gamma
...                   (repeated for block[1] through block[5])
embed.w
norm.gamma
unembed.w

ClearCache()

ClearCache resets the KV cache stored in every transformer block’s attention layer.

func (m *GPT) ClearCache()

You must call ClearCache() before each new generation sequence. Stale cache entries from a previous sequence will corrupt attention outputs and produce incoherent text. GenerateChan and GenerateText call ClearCache automatically — only call it manually if you drive Forward directly.

Complete usage example

The following snippet shows a typical training workflow: constructing a model, setting up an AdamW optimizer, running a training loop, and saving a checkpoint.

package main

import (
    "fmt"

    F "github.com/itsubaki/autograd/function"
    "github.com/itsubaki/autograd/hook"
    "github.com/itsubaki/autograd/optimizer"
    "github.com/itsubaki/gpt/model"
)

func main() {
    // Construct model
    m := model.NewGPT(1000, 256, 192, 6, 6, 10000)

    // AdamW optimizer with gradient clipping
    o := optimizer.AdamW{
        Alpha:       3e-4,
        Beta1:       0.9,
        Beta2:       0.999,
        WeightDecay: 0.01,
        Hook: []optimizer.Hook{
            hook.ClipGrad(1.0),
        },
    }

    // Training step (x, y are (B, C) token ID variables)
    logits := m.Forward(x) // (B, C, VocabSize)
    loss := F.CrossEntropy(
        F.Reshape(-1, logits.Size(-1))(logits), // (B*C, VocabSize)
        F.Reshape(-1)(y),                       // (B*C,)
    )

    m.Cleargrads()
    loss.Backward()
    o.Update(m)

    fmt.Printf("loss: %.4f\n", loss.At())

    // Save checkpoint
    if err := m.Save("model_gpt.gob"); err != nil {
        panic(err)
    }
}

Save checkpoints periodically during training (e.g., every 100 iterations) rather than only at the end. Pass the saved path to model.NewGPTFrom with useCache=true to resume generation without retraining.

Model

Tokenizer

Layers

Data & Utilities

model.GPT — Construct, Train, and Serialize a GPT Model

GPT Struct

NewGPT()

Layer composition

NewGPTFrom()

Forward()

Save() and Load()

Save()

Load()

ClearCache()

Complete usage example

Build docs developers (and LLMs) love

Model

Tokenizer

Layers

Data & Utilities

Documentation Index

​GPT Struct

​NewGPT()

​Layer composition

​NewGPTFrom()

​Forward()

​Save() and Load()

​Save()

​Load()

​ClearCache()

​Complete usage example

Build docs developers (and LLMs) love

GPT Struct

NewGPT()

Layer composition

NewGPTFrom()

Forward()

Save() and Load()

Save()

Load()

ClearCache()

Complete usage example