Show HN: MLbot – Train ML models in your cloud with a single command
github.com
github.com
I just released the initial version of MLbot: a new open-source tool that I’ve been working on for running ML training jobs in your cloud, with a single command.
In short, it allows you to run your training script in the cloud by simply swapping “python” for “mlbot run”.
For example, if “python train.py …” can run your training script locally, then “mlbot run —instance-type p3dn.24xlarge —num-nodes 2 train.py …” should be able to run your code in the cloud across 2 GPU machines.
Since this tool runs entirely inside your cloud environment, you don’t have to transfer your training data to a 3rd party, while having full observability into the underlying infrastructure.
Why I built this:
In a recent ML project I was working on (with a 10+ TB training dataset), I found that the tools to run distributed training jobs in the cloud were fairly painful & complex to use for a single developer like myself.
Current solutions like Kubeflow felt way too bloated, while typical 3rd-party hosted solutions didn’t work for my use-case either, since it often meant transforming & duplicating my large training dataset into a service-specific format (which would be both very expensive & impractical due to data security).
After unsuccessfully trying numerous solutions (including months w/ AWS Sagemaker), I finally ended up using AWS Elastic Kubernetes Service combined with PyTorch Elastic + PyTorch Lightning, which turned out to work really well for scalable, distributed training.
At the time, I wrote some scripts that automated the process of packaging my raw training code and launching it as a distributed training job on EKS, and I found that this significantly improved my speed of iterating & launching new distributed training jobs.
So, I thought it would be nice to package my scripts and open-source it as a single tool, in case others find it useful.
The current code definitely needs to be refactored into a much better structure, but I wanted to first get the current version out there, so that I can learn if this solves a real pain-point that others also face.
Finally, I’d love to hear your feedback, and any thoughts/experience/pain-points you might have related to running distributed ML jobs in the cloud!
Thanks! :)
So, one possible way to train a model on ImageNet using this tool could be to:
1. Store the dataset on S3
2. Have the training script download the necessary data from S3 at the beginning (for example, each process could only download the shard of data that it needs, etc.)
3. Use MLbot to run your local training script on AWS as a distributed training job with X number of nodes in your EKS cluster
In this case, when you execute "mlbot run ..." from the command-line, MLbot will package your code into a docker image and then create X number of pods within EKS (where X = the number of training nodes).
Once all of the pods are running, etcd + PyTorch Elastic begin handling all of the communication / synchronization between the various pods & distributed training happens.
So in this example, your code would be transferred to your EKS cluster (where the training would happen) and the training pods would then fetch the ImageNet data from S3.
And most importantly, all of this happens inside of your cloud infrastructure.
Now that I think about it, I think it might be helpful to add an ImageNet-based example to the repo.
Yes, that would be helpful. I’m not sure if S3 is fast enough to serve Imagenet one batch at a time - might become a bottleneck.