WeightLifting

A PyTorch char-level transformer trained on HuggingFace's megaGymDataset. Time to tune some weights!

---

Setup

python -m venv .venv
.\.venv\Scripts\Activate.ps1

# PyTorch with CUDA 12.8
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128

# Everything else
pip install -e ".[dev]"

Verify the GPU is visible:

python -c "import torch; print(torch.cuda.get_device_name(0))"
# NVIDIA GeForce GTX 1660 Ti

---

Training

Downloads megaGymDataset (590 kB) on first run, then trains for 5,000 steps with mixed-precision on the GPU.

python scripts/train.py

Expected output (one line per eval, every 500 steps):

11:25:20 | INFO     | __main__             | Device: cuda (NVIDIA GeForce GTX 1660 Ti, 6.0 GB)
11:25:21 | INFO     | weightlifting.data.dataset | Loaded 1368 Q&A pairs from data\raw\fitness_cleaned.csv
11:25:21 | INFO     | __main__             | Dataset: 515947 samples  |  vocab size = 79
11:25:21 | INFO     | __main__             | Model: 10,776,192 parameters
11:27:35 | INFO     | weightlifting.training.trainer | step     0 | train 3.6832 | val 3.6830 | 127.6s
12:02:21 | INFO     | weightlifting.training.trainer | step   500 | train 0.7358 | val 0.7396 | 2086.0s
12:37:30 | INFO     | weightlifting.training.trainer | step  1000 | train 0.3756 | val 0.3773 | 2109.0s
13:11:16 | INFO     | weightlifting.training.trainer | step  1500 | train 0.2175 | val 0.2182 | 2026.0s
13:44:55 | INFO     | weightlifting.training.trainer | step  2000 | train 0.1383 | val 0.1389 | 2018.9s
14:18:34 | INFO     | weightlifting.training.trainer | step  2500 | train 0.0997 | val 0.1000 | 2019.2s
14:52:13 | INFO     | weightlifting.training.trainer | step  3000 | train 0.0833 | val 0.0831 | 2018.8s
15:25:51 | INFO     | weightlifting.training.trainer | step  3500 | train 0.0734 | val 0.0739 | 2017.8s
15:59:28 | INFO     | weightlifting.training.trainer | step  4000 | train 0.0673 | val 0.0678 | 2016.6s
16:33:07 | INFO     | weightlifting.training.trainer | step  4500 | train 0.0646 | val 0.0651 | 2018.1s
17:06:52 | INFO     | weightlifting.training.trainer | step  4999 | train 0.0625 | val 0.0629 | 2025.2s
...

Checkpoints are written to outputs/checkpoints/ckpt_NNNNN.pt after every eval.

Resuming from a checkpoint

python scripts/train.py --resume outputs\checkpoints\ckpt_01000.pt

The trainer restores both model weights and optimizer state and continues from step 4001. The max_iters cap in TrainConfig still applies — raise it in scripts/train.py if you want to train beyond 5,000 steps.

Tuning (edit configs/train.yaml or scripts/train.py)

Parameter	Default	Notes
max_iters	5 000	~6h on 1660 Ti
batch_size	128	tuned for 6 GB
block_size	256	context length in tokens
n_embd / n_head / n_layer	384 / 6 / 6	~10 M params
learning_rate	3e-4	AdamW

---

Generation

Sample text from any checkpoint:

python scripts/generate.py outputs/checkpoints/ckpt_05000.pt

# With a seed prompt
python scripts/generate.py outputs/checkpoints/ckpt_05000.pt --prompt "What is Sumo deadlift"

# More options
python scripts/generate.py outputs/checkpoints/ckpt_05000.pt \
    --prompt "core exercise" \
    --tokens 1000 \
    --temperature 0.8 \
    --top-k 200

---

Tests

python -m pytest                  # 11 tests, ~6 s
python -m pytest -v --tb=short   # verbose

---

Project structure

WeightLifting/
├── src/weightlifting/
│   ├── models/
│   │   ├── base.py          # BaseModel — save/load/param-count
│   │   └── transformer.py   # CharTransformer (decoder-only, nanoGPT-style)
│   ├── data/
│   │   └── dataset.py       # CharDataset + megaGymDataset auto-download
│   ├── training/
│   │   └── trainer.py       # Trainer — AMP, grad-clip, checkpointing, optional W&B
│   └── utils/
│       ├── device.py        # get_device() — cuda / mps / cpu auto-select
│       └── logging.py       # setup_logging() — stdout + file
├── scripts/
│   ├── train.py             # training entry point
│   └── generate.py          # Sample from a checkpoint
├── configs/
│   └── train.yaml           # Default hyperparameters
├── tests/
│   ├── test_models.py
│   └── test_data.py
├── data/
│   └── raw/                 # megaGymDataset cached here on first run (git-ignored)
├── outputs/
│   ├── checkpoints/         # .pt files written here (git-ignored)
│   └── logs/                # run.log written here
└── pyproject.toml

---

Weights & Biases tracking

W&B is enabled by default in scripts/train.py and logs to jkarancsi-cern/WeightLifting.

If W&B is not yet authenticated on a new machine:

pip install wandb
wandb login

To disable tracking for a run:

# in scripts/train.py, set both fields to None:
train_cfg = TrainConfig(..., wandb_project="WeightLifting", wandb_entity="jkarancsi-cern")

---