Transformer Pretraining

A decoder-only Transformer for language model pretraining, with FSDP distributed training, SwiGLU activations, RoPE position encoding, grouped-query attention (GQA), KV-cache generation, WSD learning rate schedule, and Chinchilla-optimal token budgeting. See report-transformer-pretraining.pdf for details.

Setup

Requires Python 3.10+ and PyTorch 2.0+.

# Install with uv (recommended)
uv sync

# Or with pip
pip install -e .

Usage

Train a model

# Default config (~30M params, 1000 steps)
transformer-pretraining

# With a YAML config
transformer-pretraining --config configs/small.yaml

# CLI overrides on top of defaults
transformer-pretraining --d_model 512 --num_layers 8 --train_steps 10000

Local testing (tiny model, wikitext-2, CPU-friendly)

transformer-pretraining --config configs/local-test.yaml

Distributed training

torchrun --nnodes=1 --nproc_per_node=2 --rdzv_backend=c10d \
    $(which transformer-pretraining) --config configs/small.yaml

See scripts/templates/ for SLURM cluster submission examples.

Generate text

# Generate with a trained checkpoint
transformer-pretraining \
    --prompt "The future of AI is" \
    --resume_from outputs/checkpoint.pt \
    --temperature 0.8 \
    --top_k 40

Configuration

All training hyperparameters are defined via YAML config files with optional CLI overrides. See configs/ for examples:

Config	Description
configs/small.yaml	~30M param model, 4600 steps (Chinchilla-optimal for C4)
configs/local-test.yaml	Tiny model (64-dim, 2 layers) on wikitext-2 for quick iteration

CLI arguments take precedence over YAML values, which take precedence over code defaults.

Reproducing the original training run

The original training (30M param model on C4, 4600 steps, 2× A100 GPUs, FSDP) can be reproduced with:

transformer-pretraining --config configs/small.yaml

The original SLURM submission scripts are preserved in scripts/original/.