Training neural networks using Vast.ai
The plan was to create an analysis of how fast you can train a GPT-like neural network and compare different GPUs and how quick each one of them trains a neural network like this one.
I’ve learned several lessons throughout this process so I’d like to share those. The lessons are primarily related to how to run a training script in a GPU and set it up on a server.
What did I train?
I used one of Andrej Karphathy’s transformer overviews he covers in his YouTube channel where he goes over the pretraining loop of the ChatGPT, his overview covers how write a Transformer class using PyTorch and how to run training using TinyShakespeare to then generate infinite Shakespeare-like text.
I have a repository with the code from the video, a setup.sh script I used to configure containers and I added WanDB to the training loop to track results. WanDB is another tool that’s very useful to track and log training results.
The total number of parameters this model has is 10,788,161 (~10 million). Each parameter is essentially just a numberic float value.
Pretraining with different graphic cards
If you’re looking to train a neural network you can use a GTX 5080 which is expensive (around $1,500 USD) and it will get you half of what an H100 graphics card would offer, which are deployed to data centers to offer compute to companies around the world.
The alternative is to rent GPUs which on vast.ai is way cheaper, since you can pay by the hour (anywhere from $0.05/hr to $8/hr) and you don’t have to worry about upgrading your graphics card ever, you can just rent them as needed, and you’ll have the most flexibility since you can rent any size and any number of graphics cards. I’m not sponsored by them, it’s just that I’ve found their web console very easy to use.
Here are the results that I got from training the Karpathy’s GPT version with the same configuration as what’s currently in the repository above:
| Card name | Training time |
|---|---|
| 1 x RTX 3060 Ti | 32min 22sec |
| 1 x RTX 4060 Ti | 21min 46sec |
| 1 x RTX 5060 Ti | 17min 32sec |
| 1 x RTX 5080 | 9 min 22 sec |
| 1 x H100 | 5 min 41 sec |
How hard was this?
It didn’t take too long to start running training on a remote computer which made the process enjoyable, renting a server was a one-click operation, adding ssh keys was done in seconds, shutting down a server was one click as well.
The time consuming part was installing the correct PyTorch drivers on the remote server and coming up with the setup script that worked seamlessly when renting a new server, I suspect others will run into issues due to the difference in GPU drivers on each server, I haven’t looked into how to make that process smooth but I imagine it wouldn’t take long to support any of the most common drivers.
Lessons
- Running an evaluation in the training loop takes up significant resources. It can double the time it takes to train a neural network if you configure evaluations to run frequently.
- Different servers you rent will have different drivers, this can interfere with setting up the server and will lead you to run a few commands to uninstall and install pytorch with GPU drivers:
pipenv run pip uninstall torch -y
pipenv run pip install torch --index-url https://download.pytorch.org/whl/cu124- Automating the ML training process is it’s own domain, testing and verifying a training run can be quick, deploying an ML system to production very likely requires CI workflows to automate parts of the training and deployment process.
Who should train neural networks using vast.ai and wandb.ai?
Anyone who already knows about deep learning and any startups training their own speciallized models. The case where I don’t see this setup being useful is if the network you are trying to train is too big (creating a frontier LLM for example), but in any other case (recommendation systems, trigger word detection), this setup should work great.
Thanks for reading and I hope you learned something from this quick article!