GPT-2 Checkpoint Example

This is the load-and-sample half of the byte-level GPT example.

1. Train and save a model checkpoint:

scripts/lake.sh -Kcuda=true build torchlean:exe
scripts/lake.sh -Kcuda=true exe torchlean gpt2 --device cuda \
    --tiny-shakespeare --steps 1 --windows 1 \
  --prompt "First Citizen:" --generate 0 \
  --save-checkpoint data/examples/gpt2_shakespeare.state.json

2. Load the checkpoint and sample text without optimizer updates:

scripts/lake.sh -Kcuda=true exe torchlean gpt2_saved --device cuda \
  --checkpoint data/examples/gpt2_shakespeare.state.json \
  --prompt "First Citizen:" --generate 0

What A Checkpoint Is Here

This example uses the simplest TorchLean checkpoint format:

• a shape-indexed pack of model state tensors,
• stored as exact Float.toBits values in JSON, and
• checked against the model's state layout before inference starts.

So save/load is model-agnostic: if we can name the model, TorchLean can compute the expected state shapes and reject files with incompatible shapes. Matching shapes alone do not identify the architecture or parameter meanings.

Why This Is A Separate Example

The inference-only workflow is direct: load a checkpoint, convert it into runtime handles, and sample text without building a training loop.

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

CLI subcommand name used in terminal banners and error messages.

def NN.Examples.Models.Sequence.Gpt2Saved.usage : String

Help text for checkpoint-only GPT-2 sampling.

structure NN.Examples.Models.Sequence.Gpt2Saved.Options : Type

Command-local options for loading one checkpoint and sampling from it.

• generation : TorchLean.text.GenerationOptions

  Prompt and token-sampling policy.

• checkpointPath : System.FilePath

  Model checkpoint loaded before sampling starts.

@[instance_reducible]

instance NN.Examples.Models.Sequence.Gpt2Saved.instReprOptions : Repr Options

def NN.Examples.Models.Sequence.Gpt2Saved.instReprOptions.repr : Options → ℕ → Std.Format

def NN.Examples.Models.Sequence.Gpt2Saved.Options.parse (args : List String) (defaults : TorchLean.text.GenerationOptions) : Except String (Options × List String)

Parse the checkpoint path followed by the shared generation flags.

def NN.Examples.Models.Sequence.Gpt2Saved.sampleCheckpoint (load : Options) (runtime : TorchLean.Runtime.Config := { }) : IO String

Load model state from disk and run sampling with the fixed byte-level GPT architecture.

The checkpoint must match Gpt2.model's state shapes. If the model configuration in Gpt2.lean changes (heads, width, layers, etc.), mismatched checkpoints fail the shape check before sampling starts.

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

CLI entrypoint for checkpoint sampling.