Large Scale Distributed Training with TorchX and Ray

Large-scale model training has generally been out of reach for people in open source because it requires an engineer to learn how to set up an infrastructure, how to build composable software systems, and how to set up robust machine learning scripts.

To that end, we’ve built the TorchX Ray scheduler which leverages the newly created Ray Job API to allow scientists to focus on writing their scripts and making infrastructure and systems setup relatively easy.

1. Setting up a multi GPU setup on any cloud provider is as easy as calling ray up the cluster.yaml
2. TorchX embraces a component-based approach to designing systems that makes your ops workflows composable
3. Running a distributed PyTorch script is then as simple as calling torchx run

In this session, we’ll go through a practical live demo of how to train multi GPU models, set up the infrastructure live, and provide some tips and best practices to productionize such workflows

Transcript

1. Large scale training with torchX and Ray Mark Saroufim
2. About me • Maintain pytorch/serve • Contribute to pytorch/torchx, pytorch/pytorch • Charter is open source production story • twitter.com/marksaroufim
3. Key problems 1. Setting up an infrastructure 2. Submitting jobs against that infrastructure 3. Getting logs and job status 4. Deploying an end to end system on that infrastructure Data scientist != Infra engineer
4. Acknowledgements Meta • Can Balioglu for developing the original interface for the Ray scheduler • Tristian Rice for unblocking CI issues • Aliaksandr Ivanou for rigorous code reviews and fixing last few bugs • Geeta Chauhan, Kiuk Chung & Diamond Bishop for leadership guidance Anyscale • Amog Kamsetty for many Ray coaching sessions • Jiao Dong for building a great Ray Job API • Jules Damji for writing most of the blog post • Richard Liaw for leadership guidance
5. Setting up the infrastructure aws configure ray up -y ray_cluster.yaml Typical YAML file to • Determine cloud provider • Machine types • Docker images
6. Follow along in Google Colab shorturl.at/wJS19
7. Submitting a job
8. TorchX
9. Install torchx pip install torchx-nightly
10. Submitting a job torchx run -s ray -cfg dashboard_address=$public_ip:20002, working_dir=. ./component.py:trainer
11. What’s a component? • Python Application Definition ◦ Entrypoint ◦ Resources associated with a job ◦ Environment variables ◦ Docker image • What about a trainer.py?
13. Job Status or Logs Launched app: ray://torchx/54.214.124.247:20002-raysubmit_ntquG1dDV6CtFUC5 torchx describe ray://torchx/54.214.124.247:20002-raysubmit_ntquG1dDV6CtFUC5 torchx log ray://torchx/54.214.124.247:20002-raysubmit_ntquG1dDV6CtFUC5
14. Torchx describe
15. Torchx log
16. Under the hood - Ray Job SDK
17. Ray Job status
18. ray_driver.py
19. ray_driver.py
20. ray_driver.py
21. Torchx & Ray Run torchX components on Ray • Model serving (torch serve) • Elastic job launcher script • Hyperparameter optimization • Setting env variables • Metric logging • Configuration management • Pytorch Lightning Training loop with callbacks
22. Torchx apps/pipelines store running on Ray
23. Give us a star ⭐ https://github.com/pytorch/torchx https://github.com/ray-project/ray