Elamur Research⌕
☰
home / papers / 2024 / MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs

MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs

Paper · 2024 · Ziheng Jiang, Haibin Lin, Yinmin Zhong et al.

Summary

When an LLM training job takes over a whole cluster of tens of thousands of GPUs, two engineering problems dominate. The first is keeping the hardware busy, measured as Model FLOPs Utilization (MFU). The second is keeping that efficiency steady for weeks, because hardware failures and slow nodes ("stragglers") happen constantly at this scale. ByteDance and Peking University describe MegaScale, their production system for this setting.

For efficiency, they tune the model and the infrastructure together:

  • architecture and optimizer changes,
  • a mix of four parallelism types, with custom overlap of communication and computation,
  • a faster data pipeline and faster setup of communication groups,
  • network tuning.

For stability, they rely on deep monitoring, automatic fault detection and recovery, and fine-grained tracing tools.

On a 175B model with 12,288 GPUs, MegaScale reaches 55.2% MFU, 1.34× better than Megatron-LM. A multi-week production run recovered automatically more than 100 times. The paper matters because most LLM reports leave out the systems work behind them.

Key points

  • Headline result: 55.2% MFU training a 175B transformer on 12,288 GPUs, a 1.34× improvement over Megatron-LM.
  • Model-level changes: parallel transformer block, sliding window attention and the LAMB optimizer. The authors check that these changes do not harm convergence.
  • Parallelism:
    • Data, pipeline (interleaved 1F1B), tensor and sequence parallelism are combined, with ZeRO-2-style sharding.
    • Tensor-parallel traffic stays inside a node.
    • Data-parallel groups are placed first, to limit traffic between minipods.
  • Systems tuning:
    • Data pipeline: prefetching and tree-based loading.
    • Group setup: non-blocking initialization of collective groups without global barriers.
    • Network: custom topology, fewer ECMP hash conflicts, custom congestion control, tuned retransmit timeouts.
  • Production stability: an in-house, closed model with several hundred billion parameters was trained for several weeks on multi-trillion tokens, recovering automatically more than 100 times. Some components are being open-sourced in veScale on GitHub.

Interesting ideas

  • I1. Diagnose slowdowns by elimination. The team timed each training phase with CUDA events and saw that forward, backward and optimizer times stayed flat. They also found bandwidth was stable. That left only late-starting ranks in the gradient reduce-scatter as the cause. Splitting a slowdown into measured phases and ruling them out one by one works for any distributed pipeline, not just GPU training.
  • I2. Place traffic where the links can carry it. The chattiest operations (tensor parallelism) stay on the fastest links inside a node, and lighter traffic crosses node boundaries. Any multi-machine job can use the same rule: put tightly coupled work on the cheapest interconnect.
  • I3. Plan for failure from the start. MegaScale combines rich heartbeat messages, a suite of tests that pinpoint faulty nodes, and fast checkpoint/recovery, so recovery is routine rather than an emergency. This is useful for any long-running job, such as large corpus processing or crawling.
  • I4. Remove global synchronization points. Making setup steps asynchronous and dropping global barriers cut startup cost at scale. The same idea speeds up any system where many workers would otherwise wait for one another.
  • I5. Visualize dependencies between workers. A timeline, heat-map and 3D-parallel dependency view make it easy to spot which worker holds everyone else up. A lighter version of such tooling could help teams debug smaller multi-GPU fine-tuning runs.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Peking UniversityBeijing · China315 works · 43,268 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksInformation Systems · 1,730 worksHardware and Architecture · 167 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,894 worksGPU & parallel computing · 41 works

Cite

Cite as Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, Yulu Jia, Sun He, Hongmin Chen, Zhihao Bai, Qi Hou, Shipeng Yan, Ding Zhou, Yiyao Sheng, Zhuo Jiang, Haohan Xu, Haoran Wei, Zhang Zhang, Pengfei Nie, Leqi Zou, Sida Zhao, Liang Xiang, Zherui Liu, Zhe Li, Xiaoying Jia, Jianxi Ye, Xin Jin and Xin Liu. 2024. MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs. arXiv:2402.15627.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2402.15627 · arXiv 2402.15627

@misc{jiang2024megascale,
  title = {{MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs}},
  author = {Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, Yulu Jia, Sun He, Hongmin Chen, Zhihao Bai, Qi Hou, Shipeng Yan, Ding Zhou, Yiyao Sheng, Zhuo Jiang, Haohan Xu, Haoran Wei, Zhang Zhang, Pengfei Nie, Leqi Zou, Sida Zhao, Liang Xiang, Zherui Liu, Zhe Li, Xiaoying Jia, Jianxi Ye, Xin Jin and Xin Liu},
  year = {2024},
  doi = {10.48550/arxiv.2402.15627},
  eprint = {2402.15627},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.15627}
}

References

69 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al., “Language models are few-shot learners,” Advances in neural information processing systems, vol. 33, pp. 1877–1901, 2020.
  2. “Introducing chatgpt.” https://openai.com/blog/chatgpt, 2022. ↗
  3. J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei, “Scaling laws for neural language models,” 2020.
  4. L. Floridi and M. Chiriatti, “Gpt-3: Its nature, scope, limits, and consequences,” Minds and Machines, vol. 30, pp. 681–694, 2020. find
  5. A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al., “Palm: Scaling language modeling with pathways,” arXiv preprint arXiv:2204.02311, 2022.
  6. K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in IEEE Conference on Computer Vision and Pattern Recognition, 2016.
  7. D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley, M. Patwary, V. A. Korthikanti, D. Vainbrand, P. Kashinkunti, J. Bernauer, B. Catanzaro, A. Phanishayee, and M. Zaharia, “Efficient large-scale language model training on gpu clusters using megatron-lm,” 2021. find
  8. I. Beltagy, M. E. Peters, and A. Cohan, “Longformer: The long-document transformer,” 2020.
  9. Y. You, J. Li, S. Reddi, J. Hseu, S. Kumar, S. Bhojanapalli, X. Song, J. Demmel, K. Keutzer, and C.-J. Hsieh, “Large batch optimization for deep learning: Training bert in 76 minutes,” in International Conference on Learning Representations, 2020. find
  10. M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” 2020. find
  11. S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “Zero: Memory optimizations toward training trillion parameter models.” ArXiv, May 2020. find
  12. Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, et al., “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism ,” in NeurIPS, 2019. find
  13. D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “Pipedream: Generalized pipeline parallelism for dnn training,” in ACM SOSP, 2019. find
  14. B. Wang and A. Komatsuzaki, “GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model.” https://github.com/kingoflolz/mesh-transformer-jax, May 2021. ↗
  15. Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y. Hao, A. Mathews, and S. Li, “Pytorch fsdp: Experiences on scaling fully sharded data parallel,” 2023. find
  16. T. Dao, “Flashattention-2: Faster attention with better parallelism and work partitioning,” arXiv preprint arXiv:2307.08691, 2023.
  17. S. Li, Y. Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, J. Smith, B. Vaughan, P. Damania, and S. Chintala, “Pytorch distributed: Experiences on accelerating data parallel training,” 2020.
  18. I. Group, “Ieee 802.1 qbb - priority-based flow control.” https://1.ieee802.org/dcb/802-1qbb/, 2009. ↗
  19. Y. Zhu, H. Eran, D. Firestone, C. Guo, M. Lipshteyn, Y. Liron, J. Padhye, S. Raindel, M. H. Yahia, and M. Zhang, “Congestion Control for Large-scale RDMA Deployments,” ACM SIGCOMM Computer Communication Review, vol. 45, no. 4, pp. 523–536, 2015. find
  20. G. Kumar, N. Dukkipati, K. Jang, H. M. Wassel, X. Wu, B. Montazeri, Y. Wang, K. Springborn, C. Alfeld, M. Ryan, et al., “Swift: Delay is Simple and Effective for Congestion Control in the Datacenter,” in SIGCOMM, pp. 514–528, 2020. https://github.com/NVIDIA/ ↗
  21. “Megatron-LM.” Megatron-LM/tree/main, 2021. find
  22. V. A. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro, “Reducing activation recomputation in large transformer models,” Proceedings of Machine Learning and Systems, vol. 5, 2023.
  23. OpenAI, “Gpt-4 technical report,” 2023.
  24. D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen, “Gshard: Scaling giant models with conditional computation and automatic sharding,” 2020. find
  25. A. Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. Das-Sarma, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, J. Kernion, K. Ndousse, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, and J. Kaplan, “A general language assistant as a laborator… find
  26. J. Wei, M. Bosma, V. Y. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le, “Finetuned language models are zero-shot learners,” 2022.
  27. S. Smith, M. Patwary, B. Norick, P. LeGresley, S. Rajbhandari, J. Casper, Z. Liu, S. Prabhumoye, G. Zerveas, V. Korthikanti, E. Zhang, R. Child, R. Y. Aminabadi, J. Bernauer, X. Song, M. Shoeybi, Y. He, M. Houston, S. Tiwary, and B. Catanzaro, “Using deepspeed and megatron to train megatron-turing n… — Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model
  28. J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre, “Training compute-op… — Training Compute-Optimal Large Language Models
  29. H. Su, X. Zhou, H. Yu, X. Shen, Y. Chen, Z. Zhu, Y. Yu, and J. Zhou, “Welm: A well-read pre-trained language model for chinese,” 2023. find
  30. S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. ChenG, C. Dewan, M. Diab, X. Li, X. V. Lin, T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and L. Zettlemoyer, “Opt: Open pre-trained transformer language models,” 2022.
  31. T. L. Scao, A. Fan, C. Akiki, E. Pavlick, S. Ilić, D. Hesslow, R. Castagné, A. S. Luccioni, F. Yvon, M. Gallé, et al., “Bloom: A 176b-parameter open-access multilingual language model,” arXiv preprint arXiv:2211.05100, 2022. find
  32. H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample, “Llama: Open and efficient foundation language models,” 2023.
  33. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan,… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  34. R. Taori, I. Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto, “Stanford alpaca: An instruction-following llama model.” https://github.com/tatsu-lab/stanford_alpaca, 2023. ↗
  35. W.-L. Chiang, Z. Li, Z. Lin, Y. Sheng, Z. Wu, H. Zhang, L. Zheng, S. Zhuang, Y. Zhuang, J. E. Gonzalez, I. Stoica, and E. P. Xing, “Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality,” 2023. find
  36. X. Geng, A. Gudibande, H. Liu, E. Wallace, P. Abbeel, S. Levine, and D. Song, “Koala: A dialogue model for academic research.” Blog post, April 2023. find
  37. Y. Ji, Y. Deng, Y. Gong, Y. Peng, Q. Niu, B. Ma, and X. Li, “Belle: Be everyone’s large language model engine.” https://github.com/LianjiaTech/BELLE, 2023. ↗
  38. Z. Li, E. Wallace, S. Shen, K. Lin, K. Keutzer, D. Klein, and J. Gonzalez, “Train big, then compress: Rethinking model size for efficient training and inference of transformers,” in International Conference on Machine Learning (ICML), 2020. find
  39. G. Xiao, J. Lin, M. Seznec, J. Demouth, and S. Han, “Smoothquant: Accurate and efficient post-training quantization for large language models,” arXiv, 2022. find
  40. E. Frantar, S. Ashkboos, T. Hoefler, and D. Alistarh, “Gptq: Accurate post-training quantization for generative pre-trained transformers,” arXiv, 2022. find
  41. T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer, “Llm. int8 (): 8-bit matrix multiplication for transformers at scale,” arXiv, 2022. find
  42. E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen, “Lora: Low-rank adaptation of large language models,” 2021.
  43. R. Child, S. Gray, A. Radford, and I. Sutskever, “Generating long sequences with sparse transformers,” 2019.
  44. A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret, “Transformers are rnns: Fast autoregressive transformers with linear attention,” 2020. find
  45. C. Zhu, W. Ping, C. Xiao, M. Shoeybi, T. Goldstein, A. Anandkumar, and B. Catanzaro, “Long-short transformer: Efficient transformers for language and vision,” 2021. find
  46. B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, H. Cao, X. Cheng, M. Chung, M. Grella, K. K. GV, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, K. S. I. Mantri, F. Mom, A. Saito, X. Tang, B. Wang, J. S. Wind, S. Wozniak, R. Zhang, Z. Zhang, Q. Zhao, P. Zhou, J. Zhu, an… find
  47. Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei, “Retentive network: A successor to transformer for large language models,” 2023.
  48. D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic, “QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding,” in NeurIPS, 2017.
  49. P. Micikevicius, S. Narang, J. Alben, G. F. Diamos, E. Elsen, D. García, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu, “Mixed Precision Training,” in ICLR, 2018. find
  50. M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, M. Isard, et al., “Tensorflow: A system for large-scale machine learning,” in OSDI, 2016.
  51. A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al., “PyTorch: An Imperative Style, High-Performance Deep Learning Library ,” in NeurIPS, 2019.
  52. A. Jayarajan, J. Wei, G. Gibson, A. Fedorova, and G. Pekhimenko, “Priority-based Parameter Propagation for Distributed DNN Training ,” in MLSys, 2019. find
  53. S. H. Hashemi, S. Abdu Jyothi, and R. Campbell, “Tic-Tac: Accelerating Distributed Deep Learning with Communication Scheduling,” in MLSys, 2019. find
  54. Y. Peng, Y. Zhu, Y. Chen, Y. Bao, B. Yi, C. Lan, C. Wu, and C. Guo, “A generic communication scheduler for distributed DNN training acceleration,” in SOSP, 2019. find
  55. Y. Bao, Y. Peng, Y. Chen, and C. Wu, “Preemptive All-reduce Scheduling for Expediting Distributed DNN Training,” in INFOCOM, 2020. find
  56. Y. Li, M. Yu, S. Li, S. Avestimehr, N. S. Kim, and A. Schwing, “Pipe-SGD: A Decentralized Pipelined SGD Framework for Distributed Deep Net Training,” in NeurIPS, 2018. find
  57. Y. Chen, C. Xie, M. Ma, J. Gu, Y. Peng, H. Lin, C. Wu, and Y. Zhu, “Sapipe: Staleness-aware pipeline for data parallel dnn training,” Advances in Neural Information Processing Systems, vol. 35, pp. 17981–17993, 2022. find
  58. C. Guo, L. Yuan, D. Xiang, Y. Dang, R. Huang, D. Maltz, Z. Liu, V. Wang, B. Pang, H. Chen, Z.-W. Lin, and V. Kurien, “Pingmesh: A large-scale system for data center network latency measurement and analysis,” SIG-COMM Comput. Commun. Rev., vol. 45, p. 139–152, aug 2015. find
  59. Y. Zhu, N. Kang, J. Cao, A. Greenberg, G. Lu, R. Mahajan, D. Maltz, L. Yuan, M. Zhang, B. Y. Zhao, and H. Zheng, “Packet-level telemetry in large datacenter networks,” SIGCOMM Comput. Commun. Rev., vol. 45, p. 479–491, aug 2015. find
  60. Y. Li, R. Miao, C. Kim, and M. Yu, “Lossradar: Fast detection of lost packets in data center networks,” in Proceedings of the 12th International on Conference on Emerging Networking EXperiments and Technologies, CoNEXT ’16, (New York, NY, USA), p. 481–495, Association for Computing Machinery, 2016. find
  61. C. Tan, Z. Jin, C. Guo, T. Zhang, H. Wu, K. Deng, D. Bi, and D. Xiang, “Netbouncer: Active device and link failure localization in data center networks,” in Proceedings of the 16th USENIX Conference on Networked Systems Design and Implementation, NSDI’19, (USA), p. 599–613, USENIX Association, 2019. find
  62. K. Liu, Z. Jiang, J. Zhang, H. Wei, X. Zhong, L. Tan, T. Pan, and T. Huang, “Hostping: Diagnosing intra-host network bottlenecks in RDMA servers,” in 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23), (Boston, MA), pp. 15–29, USENIX Association, April 2023. find
  63. S. Haider, N. R. Ansari, M. Akbar, and M. R. Perwez, “Fault tolerance in distributed paradigms,” 2011. find
  64. Y. Peng, Y. Bao, Y. Chen, C. Wu, and C. Guo, “Optimus: an efficient dynamic resource scheduler for deep learning clusters,” in Proceedings of the Thirteenth EuroSys Conference, pp. 1–14, 2018. find
  65. A. S. Tanenbaum, Distributed systems principles and paradigms. 2007. find
  66. S. Chakravorty, C. L. Mendes, and L. V. Kalé, “Proactive fault tolerance in mpi applications via task migration,” in International Conference on High-Performance Computing, pp. 485–496, Springer, 2006. find
  67. S. Chakravorty, C. Mendes, and L. V. Kale, “Proactive fault tolerance in large systems,” in HPCRI Workshop in conjunction with HPCA, vol. 2005, pp. 1–7, Citeseer, 2005. find
  68. Y. Chen, Y. Peng, Y. Bao, C. Wu, Y. Zhu, and C. Guo, “Elastic parameter server load distribution in deep learning clusters,” in Proceedings of the 11th ACM Symposium on Cloud Computing, pp. 507–521, 2020. find
  69. I. Behera and C. R. Tripathy, “Performance modelling and analysis of mobile grid computing systems,” International Journal of Grid and Utility Computing, vol. 5, no. 1, pp. 11–20, 2014. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/jiang-2024-megascale-scaling-large-language-model/