NanoGPT Training#

Overview#

Training GPT-2 scale models (~124M parameters) efficiently on a single GPU. It provides:

• GPT-124M Architecture: Standard transformer with RoPE and modern optimizations
• Tokenized Datasets: Loading pre-tokenized shards from HuggingFace Hub or local files
• Modern Optimizers: Muon optimizer with Newton-Schulz orthogonalization
• Mixed Precision: bfloat16 training on A100 for 2x speedup

Training options:

• Baseline GPT: Standard residual connections
• Experimental residual variants: Optional alternative residual schemes for stability/efficiency

Quick Reference#

TopicReferenceModel ArchitectureGPT ArchitectureData LoadingTokenized DataOptimizersOptimizersTraining LoopTraining LoopHyperparametersHyperparameters

Installation#

bash

pip install torch einops numpy huggingface_hub

Minimal Example#

python

import modal

app = modal.App("gpt-training")

image = modal.Image.debian_slim(python_version="3.11").pip_install(
    "torch", "einops", "numpy", "huggingface_hub"
)

@app.function(gpu="A100", image=image, timeout=3600)
def train():
    import torch
    from dataclasses import dataclass

    @dataclass
    class GPTConfig:
        block_size: int = 1024
        vocab_size: int = 50257
        n_layer: int = 12
        n_head: int = 12
        n_embd: int = 768
        dropout: float = 0.0
        bias: bool = False

    # Download data, build model, train
    # ... (see references for full implementation)

    return {"final_loss": final_loss}

@app.local_entrypoint()
def main():
    results = train.remote()
    print(results)

Common Imports#

python

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
from dataclasses import dataclass
from einops import rearrange, repeat, reduce
import numpy as np
import math

When to Use What#

ScenarioApproachStandard GPT trainingUse baseline model with standard residualsStability experimentsTry alternative residual variants or extra streamsSmall experimentsUse T4/A10G GPUFull trainingUse A100 with bfloat16Custom dataModify the dataset loader classDifferent model sizeAdjust GPTConfig parameters

Metrics to Monitor#

MetricTypical SignalNotesValidation lossSteady decreaseAbsolute value depends on dataset/tokenizerGrad normModerate, stable rangeLarge spikes indicate instabilityTraining stabilitySmooth curvesFrequent spikes suggest LR/batch issuesThroughputConsistent tokens/secUse for comparing configs

External Resources#

• nanoGPT: https://github.com/karpathy/nanoGPT
• build-nanogpt: https://github.com/karpathy/build-nanogpt
• modded-nanogpt: https://github.com/KellerJordan/modded-nanogpt
• FineWeb-Edu token shards: https://huggingface.co/datasets/karpathy/fineweb-edu-100B-gpt2-token-shards