minGPT-Style Addition Example

This is a TorchLean-native version of the spirit of Karpathy's minGPT/projects/adder experiment. The original minGPT adder trains a compact GPT to complete digit strings of the form

digits(a) ++ digits(b) ++ reverseDigits(a+b).

For example, in the one-digit setting 8 + 7 = 15 is represented as the digit sequence 8 7 5 1. At inference time the model sees 8 7 and greedily generates the two result digits 5 1, which we reverse back to 15.

This is not a text chatbot. It is a controlled algorithmic sequence task for the CUDA GPT training loop:

• synthetic data is generated in Lean,
• the model is a GPT-style causal Transformer built from TorchLean layers,
• training is CUDA-only by default,
• optimizer choices follow the minGPT-style setup (adamw, adam, or sgd),
• evaluation greedily completes every one-digit addition problem.

Performance note: this uses the eager CUDA runtime, not a persistent CUDA graph. The heavy tensor operations run on the GPU, including fused attention when --fast-kernels is on, but each step still records a fresh autograd tape and synchronizes parameter refs through the current scalar training bridge. This is the correctness-facing example; full PyTorch-style throughput requires persistent device parameters plus compiled/fused graph execution.

The GPT-shaped architecture is constructed through the public TorchLean model constructor nn.models.CausalTransformerOneHot, so the example can stay focused on the adder task mechanics.

Reference: https://github.com/karpathy/minGPT/tree/master/projects/adder.

def NN.Examples.Models.Sequence.GptAdder.exeName : String

CLI subcommand label used by the shared model runner.

def NN.Examples.Models.Sequence.GptAdder.defaultLogJson : System.FilePath

Default JSON loss-curve path for this command.

def NN.Examples.Models.Sequence.GptAdder.ndigit : ℕ

Number of input digits per operand.

We start with the one-digit curriculum because it trains directly in the eager CUDA runtime while still including carry examples such as 8 + 7 = 15.

def NN.Examples.Models.Sequence.GptAdder.vocab : ℕ

Digit-only vocabulary, matching minGPT's adder task (0..9).

def NN.Examples.Models.Sequence.GptAdder.batch : ℕ

Full one-digit table batch size.

This is 100, not 1: scalar-sized GPU workloads underutilize the device. In all-pairs mode one optimizer step sees every one-digit addition problem, and evaluation completes the whole table with two batched greedy forward passes.

def NN.Examples.Models.Sequence.GptAdder.trainCount : ℕ

Karpathy's adder uses a held-out split; for one digit this is 80 train / 20 test.

def NN.Examples.Models.Sequence.GptAdder.testCount : ℕ

Held-out one-digit examples when --train-split is enabled.

def NN.Examples.Models.Sequence.GptAdder.seqLen : ℕ

GPT block size: a, b, and all but the final reversed result digit.

For ndigit = 1, full rendered examples have length 1 + 1 + 2 = 4; model inputs have length 3, exactly as in minGPT's get_block_size = 3 * ndigit + 1 - 1.

def NN.Examples.Models.Sequence.GptAdder.numHeads : ℕ

Number of attention heads.

Karpathy's minGPT default for the adder is gpt-nano (3 heads, width 48). TorchLean's eager CUDA trainer is tape-based, so we use a middle-sized model that is substantially larger than the original compact setup (1,050 params) while keeping torchlean gpt_adder practical to run.

def NN.Examples.Models.Sequence.GptAdder.headDim : ℕ

Per-head width.

def NN.Examples.Models.Sequence.GptAdder.dModel : ℕ

Transformer embedding width.

def NN.Examples.Models.Sequence.GptAdder.ffnHidden : ℕ

Feed-forward hidden width (4 * dModel, matching the common GPT MLP ratio).

def NN.Examples.Models.Sequence.GptAdder.layers : ℕ

Number of Transformer blocks.

def NN.Examples.Models.Sequence.GptAdder.activeTargetPositions : ℕ

Number of positions per row that contribute to the minGPT adder loss.

def NN.Examples.Models.Sequence.GptAdder.activeTargetCount : ℕ

Number of non-ignored next-token targets in the training batch.

The adder loss below masks ignored prefix positions to all-zero targets, then computes summed one-hot cross-entropy divided by this count. That matches minGPT's ignore_index=-1 normalization: average over active next-token labels, not over every (batch, position, vocab) entry.

def NN.Examples.Models.Sequence.GptAdder.paramCountShapes : List Shape → ℕ

Count scalar entries across a list of parameter shapes.

theorem NN.Examples.Models.Sequence.GptAdder.instNeZeroNatSeqLen : NeZero seqLen

theorem NN.Examples.Models.Sequence.GptAdder.instNeZeroNatDModel : NeZero dModel

def NN.Examples.Models.Sequence.GptAdder.cfg : TorchLean.nn.models.CausalOneHotConfig

GPT configuration shared by the typed shapes and model constructor.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.GptAdder.σ : Shape

Input shape: batched one-hot digit sequences.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.GptAdder.τ : Shape

Output shape: one digit-logit row per input position.

def NN.Examples.Models.Sequence.GptAdder.model : TorchLean.nn.M (TorchLean.nn.Sequential σ τ)

Compact GPT-style causal Transformer for digit addition.

def NN.Examples.Models.Sequence.GptAdder.adderLoss {α : Type} [TorchLean.Runtime.TensorScalar α] [DecidableEq Shape] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (logits targetOneHot : Runtime.Autograd.TorchLean.RefTy m α τ) : m (Runtime.Autograd.TorchLean.RefTy m α TorchLean.Shape.scalar)

Cross-entropy summed over non-ignored adder targets, normalized like minGPT ignore_index.

def NN.Examples.Models.Sequence.GptAdder.adderLossProgram {α : Type} [TorchLean.Runtime.TensorScalar α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.Program α [τ, τ] TorchLean.Shape.scalar

Adder-specific scalar loss.

Ignored prefix positions are encoded by all-zero one-hot rows (maskAdderTargets), so they contribute exactly zero to one-hot cross entropy. We divide the summed loss by the number of active target positions, matching minGPT's ignore_index-style normalization rather than averaging over ignored prefix rows.

def NN.Examples.Models.Sequence.GptAdder.fixedDigits (width n : ℕ) : List ℕ

Render n as exactly width base-10 digits, most-significant first.

def NN.Examples.Models.Sequence.GptAdder.renderExample (a b : ℕ) : List ℕ

minGPT adder rendering.

For ndigit = 1, a = 8, b = 7 becomes [8, 7, 5, 1], i.e. the sum 15 is stored reversed as 5, 1. Reversing the output digits makes carry propagation local in left-to-right generation.

def NN.Examples.Models.Sequence.GptAdder.keepTargetPosition (t : ℕ) : Bool

Karpathy/minGPT masks the loss on the operand-prefix positions.

In projects/adder/adder.py, the target vector y is shifted by one token and then y[:ndigit*2-1] = -1, where -1 is PyTorch's "ignore index" for cross entropy. TorchLean's current one-hot cross entropy does not have an ignore-index target, so we represent the same idea by using an all-zero one-hot vector on ignored positions. Because the loss is -sum(y * log p), these positions contribute exactly zero gradient.

def NN.Examples.Models.Sequence.GptAdder.maskAdderTargets {α : Type} [Zero α] (y : TorchLean.Tensor.T α (TorchLean.Shape.mat seqLen vocab)) : TorchLean.Tensor.T α (TorchLean.Shape.mat seqLen vocab)

Apply the minGPT adder loss mask to a shifted one-hot target matrix.

def NN.Examples.Models.Sequence.GptAdder.mkRowSample (a b : ℕ) : TorchLean.SupervisedSample Float (TorchLean.Shape.mat seqLen vocab) (TorchLean.Shape.mat seqLen vocab)

Build one unbatched one-hot causal-LM sample for an addition row, then apply the minGPT-style ignored-prefix mask to its target matrix.

def NN.Examples.Models.Sequence.GptAdder.mkSample (a b : ℕ) : TorchLean.SupervisedSample Float σ τ

Build one supervised next-digit sample from an addition problem.

def NN.Examples.Models.Sequence.GptAdder.pairAt (i : ℕ) : ℕ × ℕ

Deterministic exhaustive one-digit dataset order.

def NN.Examples.Models.Sequence.GptAdder.trainPairAt (trainSplit : Bool) (i : ℕ) : ℕ × ℕ

Training row assignment. In split mode, rows repeat the first 80 train examples.

def NN.Examples.Models.Sequence.GptAdder.parseProbe? (s : String) : Option (ℕ × ℕ)

Parse a+b into a one-digit operand pair; returns none for malformed prompts.

def NN.Examples.Models.Sequence.GptAdder.parseProbeList (s : String) : Except String (List (ℕ × ℕ))

Comma-separated list of one-digit a+b checks.

def NN.Examples.Models.Sequence.GptAdder.mkTrainSample (trainSplit : Bool) : TorchLean.SupervisedSample Float σ τ

Build a batched supervised sample with one row per one-digit addition problem.

def NN.Examples.Models.Sequence.GptAdder.decodeResult (revDigits : List ℕ) : ℕ

Decode reversed generated result digits back into a natural number.

def NN.Examples.Models.Sequence.GptAdder.argmaxAt (logits : TorchLean.Tensor.T Float τ) (pos : ℕ) : ℕ

Argmax token id at sequence position pos.

def NN.Examples.Models.Sequence.GptAdder.argmaxAtBatch (logits : TorchLean.Tensor.T Float τ) (bi : Fin batch) (pos : ℕ) : ℕ

Argmax token id at a sequence position for a chosen batch row.

def NN.Examples.Models.Sequence.GptAdder.inputFromDigits (digits : List ℕ) : TorchLean.Tensor.T Float σ

Build a model input tensor from the current generated digit prefix.

def NN.Examples.Models.Sequence.GptAdder.inputFromRows (rows : Fin batch → List ℕ) : TorchLean.Tensor.T Float σ

Build a batched model input from one digit prefix per row.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.GptAdder.Predictor : Type

Fitted adder predictor returned by the public trainer handle.

def NN.Examples.Models.Sequence.GptAdder.generateResultDigits (predict : Predictor) (a b : ℕ) : IO (List ℕ)

Greedily complete ndigit + 1 result digits from the operand digits.

The key detail is that when the current prefix has length k, the next-token prediction lives at position k - 1, not always at the final padded position.

def NN.Examples.Models.Sequence.GptAdder.predictSum (predict : Predictor) (a b : ℕ) : IO ℕ

Predict a + b by greedy decoding and reversing the minGPT result digits.

def NN.Examples.Models.Sequence.GptAdder.evalAllSlow (predict : Predictor) : IO ℕ

Evaluate all 100 one-digit additions.

structure NN.Examples.Models.Sequence.GptAdder.EvalScore : Type

Exact-match counts for train/test/all one-digit addition rows.

• trainCorrect : ℕ

  Correct rows in the training split.

• testCorrect : ℕ

  Correct rows in the held-out split.

• allCorrect : ℕ

  Correct rows across all one-digit additions.

@[implicit_reducible]

instance NN.Examples.Models.Sequence.GptAdder.instReprEvalScore : Repr EvalScore

def NN.Examples.Models.Sequence.GptAdder.instReprEvalScore.repr : EvalScore → ℕ → Std.Format

def NN.Examples.Models.Sequence.GptAdder.evalBatched (predict : Predictor) : IO EvalScore

Evaluate all 100 additions with batched greedy decoding.

For one-digit operands, generation needs two result digits. We first predict the ones digit from rows [a,b], append it, and then predict the carry/tens digit from rows [a,b,pred₀].

def NN.Examples.Models.Sequence.GptAdder.evalAllBatched (predict : Predictor) : IO ℕ

Batched exact-match score over all one-digit additions.

def NN.Examples.Models.Sequence.GptAdder.printProbe (predict : Predictor) (a b : ℕ) : IO Unit

Print one addition check in the same digit convention used for training.

structure NN.Examples.Models.Sequence.GptAdder.AdderOptionsextends NN.API.text.InteractiveTrainOptions : Type

Adder-specific CLI options layered on top of the shared interactive text training flags.

• steps : ℕ
• batchSize : ℕ
• log : Runtime.Training.LogDestination
• logPath : System.FilePath
• cudaMemWatch : ℕ
• lr : Float
• interactive : Bool
• optim : TorchLean.optim.Kind

  Optimizer.

  adamw is closest to minGPT's adder recipe. adam and sgd are useful for debugging and comparisons.

• a : ℕ

  Operand a used by the highlighted addition check.

• b : ℕ

  Operand b used by the highlighted addition check.

• probes : List (ℕ × ℕ)

  Extra comma-separated addition checks, e.g. 0+0,4+5,9+9.

• trainSplit : Bool

  Train on an 80/20 train/test split instead of all 100 one-digit additions.

• overfitProbe : Bool

  Train only the selected pair, useful for checking that the CUDA GPT can overfit one addition.

@[implicit_reducible]

instance NN.Examples.Models.Sequence.GptAdder.instReprAdderOptions : Repr AdderOptions

def NN.Examples.Models.Sequence.GptAdder.instReprAdderOptions.repr : AdderOptions → ℕ → Std.Format

def NN.Examples.Models.Sequence.GptAdder.AdderOptions.defaultProbes : List (ℕ × ℕ)

Default extra addition checks shown after training when --probes is omitted.

def NN.Examples.Models.Sequence.GptAdder.AdderOptions.parse (args : List String) : Except String (AdderOptions × List String)

Parse adder-specific CLI options.

def NN.Examples.Models.Sequence.GptAdder.AdderOptions.logNotes (cfg : AdderOptions) (opts : TorchLean.Options) : Array String

Standard TrainLog notes for the adder training loop.

inductive NN.Examples.Models.Sequence.GptAdder.CurriculumMode : Type

Training/evaluation curriculum used by the adder runner.

• overfitPair : CurriculumMode
• trainSplit : CurriculumMode
• fullTable : CurriculumMode

@[implicit_reducible]

instance NN.Examples.Models.Sequence.GptAdder.instDecidableEqCurriculumMode : DecidableEq CurriculumMode

def NN.Examples.Models.Sequence.GptAdder.instReprCurriculumMode.repr : CurriculumMode → ℕ → Std.Format

@[implicit_reducible]

instance NN.Examples.Models.Sequence.GptAdder.instReprCurriculumMode : Repr CurriculumMode

def NN.Examples.Models.Sequence.GptAdder.CurriculumMode.ofOptions (cfg : AdderOptions) : CurriculumMode

Decide which curriculum the current adder options request.

def NN.Examples.Models.Sequence.GptAdder.CurriculumMode.intro (mode : CurriculumMode) (cfg : AdderOptions) : String

Startup note for the selected curriculum.

def NN.Examples.Models.Sequence.GptAdder.CurriculumMode.sample (mode : CurriculumMode) (cfg : AdderOptions) : TorchLean.SupervisedSample Float σ τ

Training sample corresponding to the selected curriculum.

def NN.Examples.Models.Sequence.GptAdder.CurriculumMode.progressLine (mode : CurriculumMode) (predict : Predictor) (cfg : AdderOptions) (done : ℕ) (lossVal : Float) : IO String

Per-step progress line for the selected curriculum.

def NN.Examples.Models.Sequence.GptAdder.CurriculumMode.finalLine? (mode : CurriculumMode) (predict : Predictor) : IO (Option String)

Final evaluation line for the selected curriculum, if any.

def NN.Examples.Models.Sequence.GptAdder.interactiveLoop (predict : Predictor) : IO Unit

Simple terminal REPL for the trained CUDA model.

partial def NN.Examples.Models.Sequence.GptAdder.interactiveLoop.loop (predict : Predictor) (stdin : IO.FS.Stream) : IO Unit

def NN.Examples.Models.Sequence.GptAdder.trainAdderFloat (opts : TorchLean.Options) (trainOpts : AdderOptions) : IO Unit

Train the minGPT-style adder from scratch and report exact addition accuracy.

def NN.Examples.Models.Sequence.GptAdder.main (args : List String) : IO UInt32

CLI entrypoint for the CUDA GPT adder command.