This project demonstrates the end-to-end training of a GPT-2–style model (124M parameters) using PyTorch and Hugging Face Datasets, in both single-GPU and Distributed Data Parallel (DDP) setups.
The goal is to reproduce the core training dynamics of the original GPT-2 paper on two large-scale text corpora.
The resulting model is available on the Hugging Face Hub: 👉nikolina-p/gpt2base.
| Dataset | Tokens | Training Time (8×A100 80 GB) | Loss |
|---|---|---|---|
| Project Gutenberg | 3.6 B | ~45 min | 3.4709 |
| FineWeb-Edu 10BT | 10 B | ~1 h 50 min | 3.2248 |
For loss curves, gradient norms, learning rates, and sample text generations, see: 👉training_log_analysis.ipynb .
The project also includes a performance tuning exercise performance_tuning.ipynb that explores various optimization techniques, achieving up to 10× improvement in throughput, following best practices from PyTorch and NVIDIA documentation.
Two tokenized datasets were prepared and hosted on the Hugging Face Hub:
- Project Gutenberg nikolina-p/gutenberg_flat
Project Gutenberg is a library of over 70,000 free eBooks. For this project, the English subset containing 38,026 books of classical English literature was used.
Source: manu/project_gutenberg
- FineWeb-Edu nikolina-p/fineweb_10BT_tokenized
FineWeb-Edu dataset is a high-quality dataset containing curated educational web pages filtered from FineWeB dataset.
Source: HuggingFaceFW/fineweb-edu (10BT split)
🧹 Cleaning: Project Gutenberg dataset was cleaned by removing duplicate books, stripping boilerplate headers and footers and cleaning whitespace and excessive blank lines.
🔁 Tokenization: Both datasets were tokenized using OpenAI's tiktoken tokenizer ('gpt-2' encoding).
🧩 Structure: Datasets were organized according to Hugging Face’s recommended format for efficient streaming and multi-GPU training.
To ensure balanced workload distribution across processes:
-
The number of shards per split was made divisible by 8 (for 8-GPU training).
-
Each shard contains an identical number of tokens, ensuring uniformity and consistent throughput.
-
124M parameters (12 transformer blocks × 12 attention heads).
-
Initialization: Positional embeddings, token embeddings, and projection layer initialized following OpenAI GPT-2 distributions
-
Weight tying: Between token embeddings and output projection (as in GPT-2)
-
Attention: FlashAttention used for improved efficiency
Two training modes are available:
-
Single-GPU training – baseline setup for smaller experiments.
-
Distributed (DDP) – multi-GPU parallel training using
torch.distributed.
In both training modes, it is possible to start training a model from scratch or resume from a checkpoint.
Mini-evaluations are performed periodically (“cycles”), printing metrics to the console. The number of validation steps, as a ratio of train steps in each cycle, is configurable via the val_ratio parameter.
Training function supports saving checkpoints, logging losses, and JSON-formatted results.
-
Optimizer: Fused AdamW
-
Scheduler: Cosine LR decay with warmup (scaled per GPT-2 paper)
-
Batch size: ≈ 0.5M tokens (effective batch size achieved via gradient accumulation)
-
Mixed precission training with
torch.bfloat16
Install dependencies
pip install -r requirements.txt
Single-GPU Training
python train_single.py
Distributed Training (DDP)
torchrun --nproc_per_node=8 train_ddp.py
The notebook performance_tuning.ipynb demonstrates step-by-step optimization techniques that together yielded up to a 10× increase in training throughput, including:
| Tehnique | Speedup |
|---|---|
| Enabling Tensor Cores | ×2.5 |
| Mixed Precision | ×1.1 |
| torch.compile() | ×2.5 |
| FlashAttention | ×1.2 |
| Layer size tuning | ×1.15 |
| Parallelization | negligible improvement |
All experiments follow reproducible configurations with documented benchmarks.
- "Build a Large Language Model (From Scratch)" by Sebastian Raschka - initial model implementation is borrowed from this book
- "Performance Tunning Guide", PyTorch
- Performance Background, NVIDIA
- Transformer Case Study, NVIDIA
- Iterable vs Map-style datasets, Huggingface