tLM - A tiny LLM from Scratch

An implementation of GPT-style LLM Model with 20M parameters from scratch using PyTorch in Python

Training Details

Insights:

• No. of parameters: 19.83 Million (~20 Million)
• Data Type: FP16
• Best Loss: 2.267 (Initial: 8.375)
• Total Data Size: 59.31 Million (Training: 53.29 M and Validation: 5.92M)
• No. of transformer heads: 7
• Context Window: 512 tokens
• Embedding Dimension: 384
• Tokenizer Vocab Size: 4096
• train_iters = 100000

Dataset:

Opensource Wikipedia Data

Train Config:

• Optimizer: AdamW (Adam with Weight Decay)
• Scheduler: CosineAnnealingLR

Tools:

• PyTorch (Deep Learning Framework)
• Python (Programming)
• Weights and Biases (Experiment Tracking)

Setup

This project uses uv for Python dependency and virtual-environment management. It installs PyTorch, SentencePiece, NumPy, tqdm, and Weights & Biases into .venv.

uv sync

Run project commands through the managed environment:

uv run python main.py

Distributed tLM experiments

minimal-ddp/ contains two ways to train this repository's tLM/GPT across multiple processes. Each rank owns a full 19.8M-parameter model replica, processes different token batches, averages gradients, and verifies that the replicas match after training.

Manual minimal DDP

minimal-ddp/benchmark.py exposes the mechanics of data parallelism. It can compare synchronization after backward() against parameter-gradient hooks, and it can use several all-reduce implementations:

• builtin — PyTorch backend collective; this is the useful performance baseline.
• reduce_broadcast — sum at rank 0, then broadcast.
• point_to_point — explicit send/receive reduce plus broadcast.
• reduce_scatter_all_gather — compose all-reduce from two collectives.
• ring — a transparent hand-written ring all-reduce.

The --hooks mode registers an autograd hook for each parameter. With the builtin collective it launches asynchronous all-reduce as each gradient becomes ready, allowing communication to overlap with the rest of backward computation. Without hooks, it waits until all gradients are produced and synchronizes them afterward.

uv run torchrun --nnodes=1 --nproc_per_node=2 \
  --master_addr=127.0.0.1 --master_port=29500 \
  minimal-ddp/benchmark.py --algorithm builtin --hooks \
  --batch-size 2 --sequence-length 128

Run the complete manual matrix (five algorithms × hook/no-hook):

NPROC=4 STEPS=50 uv run ./minimal-ddp/run_matrix.sh

Native PyTorch DDP

minimal-ddp/native_ddp.py wraps the same tLM model with torch.nn.parallel.DistributedDataParallel. Native DDP automatically registers gradient hooks, groups gradients into buckets, overlaps optimized collectives with backward computation, and handles reducer coordination.

uv run torchrun --nnodes=1 --nproc_per_node=2 \
  --master_addr=127.0.0.1 --master_port=29501 \
  minimal-ddp/native_ddp.py --static-graph --gradient-as-bucket-view \
  --batch-size 2 --sequence-length 128

Tune native DDP with --bucket-cap-mb, --static-graph, and --gradient-as-bucket-view. Run its default benchmark with:

NPROC=4 STEPS=50 uv run ./minimal-ddp/run_native.sh

Manual and native runs write CSV timings under minimal-ddp/. Compare only rows with the same hardware, world size, per-rank batch size, sequence length, warm-up, and precision.

Benchmark results

The following local CPU/Gloo run used two ranks, a per-rank batch size of one, a sequence length of 64, five warm-up steps, and 20 measured steps. These charts are a communication-behavior comparison rather than a GPU scaling claim.

Native DDP with static_graph and gradient_as_bucket_view was the fastest configuration at 66.8 ms/step. Manual builtin all-reduce with gradient hooks took 71.9 ms/step, versus 112.6 ms/step without hooks.

Next Steps:

• Complete Notes.md
• Distributed training using Deepseed
• Integrate interpretability
• Make it more advanced (latest attention mechanisms)
• Incorporate RL based alignment
• Optimization techniques + On-device deployment

Training Charts:

1. Train Loss Curve

2. Val Loss Curve