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! :)