Model Checkpointing, SafeTensors & Browser WebGPU Inference

The Big Idea

You have done something extraordinary: you built a complete Large Language Model from first principles, beginning with tokenization in Section 1, constructing attention and transformer blocks in Section 2 & 3, and executing the full training loop in Lesson 4.1! But right now, those millions of carefully tuned weights exist only in volatile GPU memory. If you close your terminal or turn off your computer, all that learned intelligence is erased forever.

PyTorch stores learned weights in a state_dict: an elegant Python dictionary mapping each layer name (e.g. blocks.0.attn.qkv_proj.weight) to its trained floating-point tensor. We save this dictionary along with hyperparameter metadata to disk so our model can be reloaded in milliseconds.

Traditional Python torch.save relies on pickle, which can execute arbitrary malicious code upon deserialization. Production systems use SafeTensors: a format that stores raw tensor byte buffers with zero execution risk and lightning-fast memory mapping, matching nanoGPT checkpointing standards. Because our architecture is compact and modular (~10M parameters), it can also be deployed directly in the browser via the W3C WebGPU API with zero cloud fees!

Architecture Blueprint & Tensor Mechanics

Below is the complete routine for saving and loading PyTorch checkpoints along with state inspection and inference verification.

import os
import torch

def save_checkpoint(model, optimizer, step, loss, config, filepath="models/gpt_checkpoint.pt"):
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    checkpoint = {
        'step': step,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
        'config': config
    }
    torch.save(checkpoint, filepath)
    print(f"Checkpoint saved to {filepath} at step {step} with loss {loss:.4f}")

def load_checkpoint(filepath, device='cpu'):
    print(f"Loading checkpoint from {filepath}...")
    checkpoint = torch.load(filepath, map_location=device)
    config = checkpoint['config']
    
    # Reconstruct architecture using saved hyperparameters
    model = GPT(
        vocab_size=config['vocab_size'],
        n_embd=config['n_embd'],
        block_size=config['block_size'],
        n_layer=config['n_layer'],
        n_head=config['n_head']
    )
    # Load learned weights into model
    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)
    model.eval() # Essential for inference
    print(f"Successfully loaded model from step {checkpoint['step']}")
    return model, config

# Checkpoint parameter inspection
# total_params = sum(p.numel() for p in model.parameters())
# print(f"Total model parameters: {total_params / 1e6:.2f} Million")