Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/itsubaki/gpt/llms.txt

Use this file to discover all available pages before exploring further.

Pre-training teaches the GPT model to predict the next token in a sequence using cross-entropy loss over a tokenized corpus. The optimizer is AdamW with gradient clipping, and the model is checkpointed periodically throughout the run. All hyperparameters — model architecture, optimizer settings, batch size, and iteration count — are configurable through CLI flags passed to cmd/pretrain.

​
CLI flags

FlagDefaultDescription
-vocab-size1000Vocabulary size (must match tokenizer)
-context-len256Maximum context length (sequence length)
-embed-dim192Token embedding dimension
-num-of-heads6Number of attention heads per block
-num-of-blocks6Number of transformer blocks
-theta10000Theta constant for RoPE positional encoding
-max-learning-rate3e-4Peak learning rate for AdamW
-beta10.9AdamW first moment decay
-beta20.999AdamW second moment decay
-weight-decay0.01L2 weight decay for AdamW
-clip1.0Gradient clipping norm threshold
-max-iters20000Total number of training iterations
-batch-size32Number of sequences per batch
-tokens-pathtestdata/tiny_codes.binPath to the tokenized corpus .bin file
-model-pathtestdata/model_gpt.gobPath to save the trained model
-min-loss1.0Loss threshold — saves model.min when beaten
-pproffalseEnable CPU profiling to cpu.prof

​
Running pre-training

% make pretrain
go run ./cmd/pretrain/main.go
Pre-Training 100%|██████████████████████████████| 20000/20000
The progress bar shows elapsed time, ETA, iteration speed, and a live loss and perplexity reading: 
Pre-Training  100%|██████████████████████████████| 20000/20000 [12.3m<0.0s, 27.1 it/s] loss=0.9213(ppl=2.5123)

​
Training loop

The training loop inside cmd/pretrain/main.go follows the standard forward → loss → backward → update pattern: 
for i := range maxIters {
    // batch: (B, C) input tokens and (B, C) target tokens
    x, y := loader.Batch()

    // forward pass → logits shape (B, C, V)
    logits := m.Forward(x)

    // cross-entropy loss: reshape to (B*C, V) and (B*C)
    loss := F.CrossEntropy(
        F.Reshape(-1, logits.Size(-1))(logits), // (B, C, V) -> (B*C, V)
        F.Reshape(-1)(y),                       // (B, C) -> (B*C)
    )

    // backward and AdamW update with gradient clipping
    m.Cleargrads()
    loss.Backward()
    o.Update(m)
}
The TokenDataset yields sliding-window pairs: for position i, the input window is tokens[i:i+contextLen] and the target window is tokens[i+1:i+contextLen+1], implementing next-token prediction across every position in the context simultaneously.

​
Checkpointing

The model is saved to disk at two points during training:
  • Every 100 iterations — writes model_gpt.gob regardless of loss.
  • When loss beats min-loss — writes model_gpt.gob.min, capturing the best checkpoint seen so far.
if i%100 == 0 {
    if err := m.Save(modelPath); err != nil { ... }
}

if loss.At() < minLoss {
    if err := m.Save(modelPath + ".min"); err != nil { ... }
    minLoss = loss.At()
}
The final model is also saved unconditionally when the loop ends.

​
Loss logging

Every iteration appends a row to loss.csv with the iteration index and the current loss value: 
0,6.9078
100,4.2310
200,3.7891
...
19999,0.9213
Use plot loss.csv (installed via make install) to render a loss curve image after training completes.

​
Default model size

With the default flags, the model has approximately 300,000 trainable parameters: 
model parameters:
 VocabSize    : 1000
 MaxContextLen: 256
 EmbedDim     : 192
 NumOfHeads   : 6
 NumOfBlocks  : 6
------------------------------
The architecture stacks 6 transformer blocks, each containing RMSNorm → Multi-Head Attention with RoPE → residual → RMSNorm → SwiGLU → residual, followed by a final RMSNorm and a linear projection to logits.

​
Data loading

The DataLoader wraps a TokenDataset and provides shuffled, batched iteration over the token sequence: 
loader := dataloader.DataLoader{
    BatchSize: batchSize,
    Shuffle:   true,
    Dataset: &dataloader.TokenDataset{
        Tokens:     dataloader.MustLoadTokens(tokensPath),
        ContextLen: contextLen,
    },
}
TokenDataset.Len() returns len(Tokens) - ContextLen, so every valid sliding-window position is accessible as an independent training example.
Pass -pprof to write a CPU profile to cpu.prof during training. Visualise it with make pprof, which opens go tool pprof -http=:8080 cpu.prof in your browser. This is useful for identifying bottlenecks in the forward or backward pass.
To skip training and use pre-built weights, run make testdata. This downloads merge_rules.gob, tiny_codes.bin, model_gpt.gob, and model_gpt_sft.gob from the repository’s gob branch in seconds.

Build docs developers (and LLMs) love

Get started for freeTalk to us