A model training perf baseline and some optimizations
Intro
I’ve been building a practice of training some small language models. After experimenting with my local GPU and also some Google Colab accelerators, I’ve turned my focus to my local GPU. Correspondingly, it seemed like a good time to look at performance optimizations. I got some pretty quick wins for a 30% improvement in my baseline.
A baseline
I realized I needed a baseline of some sorts, and something much shorter in training time than the 4-47 hour training runs I’ve been doing. After some tests I decided to stick with a small gpt-neox model and a 25% subset of the small wikitext dataset I have setup.
The first training run took 27m 17s and had these performance specs:
100%|██████████| 898/898 [27:15<00:00, 1.82s/it]
----------------------------------------------------------------------
:: Run summary
----------------------------------------------------------------------
>> status: completed
steps: 898 / 898
tokens/step: 32,768
tokens trained: 29,425,664
train loss: 4.980603
runtime: 1,635.92s
steps/sec: 0.5490
tokens/sec: 17,987.2
TFLOP/s (est.): 13.264
MFU (est.): 32.7%
accelerator: NVIDIA GeForce RTX 3070 (bf16, world_size=1)
Some optimizations
There’s plenty to return to here, but to begin I just pointed AI at the nanochat repo and asked it to learn and then apply optimizations to my system. Here are some of the key changes:
After
The next run of the baseline was significantly faster at 18m 49s. Here are the specs:
100%|██████████| 898/898 [18:44<00:00, 1.25s/it]
>> Saving model and tokenizer...
Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 1.29it/s]
saved to /home/cmc/code/codycollier/alien-ink/output/baseline-perf-mist
----------------------------------------------------------------------
:: Run summary
----------------------------------------------------------------------
>> status: completed
steps: 898 / 898
tokens/step: 32,768
tokens trained: 29,425,664
train loss: 4.982572
runtime: 1,124.49s
steps/sec: 0.7990
tokens/sec: 26,168.1
TFLOP/s (est.): 19.296
MFU (est.): 47.5%
accelerator: NVIDIA GeForce RTX 3070 (bf16, world_size=1)
The GPU system differences
The GPU utilization and the data / network traffic were a little interesting.
Before:
After:
Background Vibes
While working this late afternoon, I’ve been watching the best
snowboard video to ever be made:
TB3 - Coming Down the Mountain (1993, Standard Films). A variation of my first snowboard (a Kemper) features in the
Tom Burt section. The Jim Rippey section is the best.