ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866929316841914368 |
|---|---|
| author | Zhu, Feiwen Nowaczynski, Arkadiusz Li, Rundong Xin, Jie Song, Yifei Marcinkiewicz, Michal Eryilmaz, Sukru Burc Yang, Jun Andersch, Michael |
| author_facet | Zhu, Feiwen Nowaczynski, Arkadiusz Li, Rundong Xin, Jie Song, Yifei Marcinkiewicz, Michal Eryilmaz, Sukru Burc Yang, Jun Andersch, Michael |
| contents | AlphaFold2 has been hailed as a breakthrough in protein folding. It can rapidly predict protein structures with lab-grade accuracy. However, its implementation does not include the necessary training code. OpenFold is the first trainable public reimplementation of AlphaFold. AlphaFold training procedure is prohibitively time-consuming, and gets diminishing benefits from scaling to more compute resources. In this work, we conducted a comprehensive analysis on the AlphaFold training procedure based on Openfold, identified that inefficient communications and overhead-dominated computations were the key factors that prevented the AlphaFold training from effective scaling. We introduced ScaleFold, a systematic training method that incorporated optimizations specifically for these factors. ScaleFold successfully scaled the AlphaFold training to 2080 NVIDIA H100 GPUs with high resource utilization. In the MLPerf HPC v3.0 benchmark, ScaleFold finished the OpenFold benchmark in 7.51 minutes, shown over $6\times$ speedup than the baseline. For training the AlphaFold model from scratch, ScaleFold completed the pretraining in 10 hours, a significant improvement over the seven days required by the original AlphaFold pretraining baseline. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2404_11068 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours Zhu, Feiwen Nowaczynski, Arkadiusz Li, Rundong Xin, Jie Song, Yifei Marcinkiewicz, Michal Eryilmaz, Sukru Burc Yang, Jun Andersch, Michael Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing Quantitative Methods AlphaFold2 has been hailed as a breakthrough in protein folding. It can rapidly predict protein structures with lab-grade accuracy. However, its implementation does not include the necessary training code. OpenFold is the first trainable public reimplementation of AlphaFold. AlphaFold training procedure is prohibitively time-consuming, and gets diminishing benefits from scaling to more compute resources. In this work, we conducted a comprehensive analysis on the AlphaFold training procedure based on Openfold, identified that inefficient communications and overhead-dominated computations were the key factors that prevented the AlphaFold training from effective scaling. We introduced ScaleFold, a systematic training method that incorporated optimizations specifically for these factors. ScaleFold successfully scaled the AlphaFold training to 2080 NVIDIA H100 GPUs with high resource utilization. In the MLPerf HPC v3.0 benchmark, ScaleFold finished the OpenFold benchmark in 7.51 minutes, shown over $6\times$ speedup than the baseline. For training the AlphaFold model from scratch, ScaleFold completed the pretraining in 10 hours, a significant improvement over the seven days required by the original AlphaFold pretraining baseline. |
| title | ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing Quantitative Methods |
| url | https://arxiv.org/abs/2404.11068 |