High Performance Training

Performance

Hardware: NVIDIA RTX 3090 · CUDA 13.0 · PyTorch 2.11.0

Each optimization is applied incrementally on top of the previous.

Optimization	Tokens/s	TFLOPS/s	Peak TFLOPS/s	MFU (%)	Speedup
Baseline (fp32)	10 455	9.54	35.58	26.8	1.0x
bf16 mixed precision	24 007	21.90	142.6	15.4	2.30x
Flash Attention	45 238	33.58	142.6	23.5	4.33x
torch.compile	47 285	35.10	142.6	24.6	4.52x
Vocab padding	48 178	35.77	142.6	25.1	4.61x
Liger CrossEntropy	52 032	38.63	142.6	27.1	4.98x
Larger micro-batch	53 945	40.05	142.6	28.1	5.16x

Installation

Requires uv.

uv sync               # install runtime dependencies
uv sync --group dev   # also install linting/formatting tools
uv sync --group prof  # also install profiling tools (nvtx, tensorboard)

The hpt CLI is installed as a script into the project virtualenv. Activate it or prefix commands with uv run.

Usage

1. Tokenize a dataset

Downloads a HuggingFace dataset and encodes it to a binary file for training.

hpt tokenize <dataset_id> <tokenizer_name> [--subset <glob>] [--text-column <col>]

Example — download a small slice of FineWeb-Edu and tokenize with the GPT-2 vocabulary:

hpt tokenize HuggingFaceFW/fineweb-edu gpt2 --subset "sample/10BT/*"

2. Train

hpt train <dataset_id> <tokenizer_name> [--subset <glob>]

Example:

hpt train HuggingFaceFW/fineweb-edu gpt2 --subset "sample/10BT/*"

Metrics (loss, tokens/s, TFLOPs/s, GPU utilization) are logged to Weights & Biases.

3. PyTorch Profiler

Requires uv sync --group prof. Outputs a TensorBoard trace to .profiles/.

hpt train <dataset_id> <tokenizer_name> --profile
tensorboard --logdir .profiles

4. Nsight Systems

Requires NVIDIA Nsight Systems and uv sync --group prof.

./scripts/nsys_profile.sh <dataset_id> <tokenizer_name> [extra hpt train args]