Elamur Research⌕
☰
home / papers / 2024 / FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion

FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion

Paper · 2024 · Li-Wen Chang, Wenlei Bao, Qi Hou et al.

Summary

Tensor parallelism splits individual layers of large models across GPUs. This lowers latency, but it adds AllGather and ReduceScatter traffic that can take a large share of runtime. Earlier overlap schemes, such as TransformerEngine, cut a GEMM and its collective into a few chunks, one per device, and schedule them on streams. On GPUs this works poorly for three reasons: kernel timing cannot be controlled precisely, extra add operations create dependencies between chunks, and small GEMMs leave SMs idle.

Flux takes a different approach. It breaks communication and computation into tile-level pieces and fuses them into a single GEMM kernel, so each thread block handles a dependent compute tile and its communication. It is built modularly on NVIDIA CUTLASS and auto-tuned for different GPU generations and interconnects. Experiments cover GPT-3 175B and Llama-2 70B on A100 (PCIe/NVLink) and H800 NVLink clusters. Flux beats Megatron-LM in training and vLLM in prefill and decoding, and it outperforms TransformerEngine by a wider margin. This makes tensor parallelism more scalable in practice, especially where communication dominates runtime.

Key points

  • Flux can hide up to 96% of communication inside a fused kernel.
  • Speedups over non-overlapping baselines:
    • Training on 128 GPUs: up to 1.24x over Megatron-LM.
    • Inference on 8 GPUs: up to 1.66x for prefill and 1.30x for decoding over vLLM.
  • Speedups over TransformerEngine: up to 1.38x (training), 2.06x (prefill) and 2.10x (decoding). At the operation level the gain reaches 3.25x on A100 PCIe and 1.51x on H800 NVLink.
  • Test setup:
    • Models: GPT-3 175B and Llama-2 70B.
    • Training layout: 2-way data, 8-way pipeline and 8-way tensor parallelism.
    • Inference layout: 8-way tensor parallelism.
  • Optimizations beyond fusion: tile coordinate swizzling, GPU instruction selection and choice of communication order. On A100 PCIe, Flux sometimes finished faster than the NCCL communication-only time.

Interesting ideas

  • I1. Fusing communication into the compute kernel at tile granularity. Coarse chunk-and-stream scheduling has two costs: unpredictable timing and underused hardware. Fusing at tile level avoids both, and the approach could apply to other dependent compute/transfer pairs.
  • I2. Effective Communication Time (ECT) as a metric. ECT is total time minus the time of the best unsplit GEMM. This gives a fair way to compare overlap methods that change how the computation itself runs, and the same "subtract the ideal compute" trick could work for any overlapped pipeline.
  • I3. Motivating the work by measuring the non-overlapped share first. Before optimizing, the authors measured how much communication remains exposed across hardware types and workload phases (training, prefill, decoding). This is a cheap way to tell where an optimization will actually pay off.
  • I4. Modular design plus auto-tuning. Building on a template library (CUTLASS) and auto-tuning per GPU and interconnect lets one codebase adapt to heterogeneous clusters. Tuning per environment can even beat a standard library like NCCL for some problem sizes.

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🏢ByteDanceBeijing · China47 works · 4,339 citations

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksHardware and Architecture · 167 worksComputer Networks and Communications · 100 works

🧭 Topics 1Research subjects and languages the work is about.GPU & parallel computing · 41 works

Cite

Cite as Li-Wen Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Chengji Yao, Ziheng Jiang, Haibin Lin, Xin Jin and Xin Liu. 2024. FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion. arXiv:2406.06858.

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

@misc{chang2024flux,
  title = {{FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion}},
  author = {Li-Wen Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Chengji Yao, Ziheng Jiang, Haibin Lin, Xin Jin and Xin Liu},
  year = {2024},
  doi = {10.48550/arxiv.2406.06858},
  eprint = {2406.06858},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2406.06858}
}

References

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

  1. 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.
  2. 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.
  3. A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-… — PaLM: Scaling Language Modeling with Pathways
  4. 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
  5. A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever, “Zero-shot text-to-image generation,” 2021.
  6. M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, R. Jenatton, L. Beyer, M. Tschannen, A. Arnab, X. Wang, C. Riquelme, M. Minderer, J. Puigcerver, U. Evci, M. Kumar, S. van Steenkiste, G. F. Elsayed, A. Mahendran, F. Yu, A. Ol… find
  7. A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever, “Robust speech recognition via large-scale weak supervision,” 2022.
  8. Y. Zhang, D. S. Park, W. Han, J. Qin, A. Gulati, J. Shor, A. Jansen, Y. Xu, Y. Huang, S. Wang, Z. Zhou, B. Li, M. Ma, W. Chan, J. Yu, Y. Wang, L. Cao, K. C. Sim, B. Ramabhadran, T. N. Sainath, F. Beaufays, Z. Chen, Q. V. Le, C.-C. Chiu, R. Pang, and Y. Wu, “Bigssl: Exploring the frontier of large-sc… find
  9. X. Liu, D. McDuff, G. Kovacs, I. Galatzer-Levy, J. Sunshine, J. Zhan, M.-Z. Poh, S. Liao, P. D. Achille, and S. Patel, “Large language models are few-shot health learners,” 2023. find
  10. S. Wu, O. Irsoy, S. Lu, V. Dabravolski, M. Dredze, S. Gehrmann, P. Kambadur, D. Rosenberg, and G. Mann, “Bloomberggpt: A large language model for finance,” 2023.
  11. M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, P. Tillet, F. P. Such,… — Evaluating Large Language Models Trained on Code
  12. A. Jangda, J. Huang, G. Liu, A. H. N. Sabet, S. Maleki, Y. Miao, M. Musuvathi, T. Mytkowicz, and O. Saarikivi, “Breaking the computation and communication abstraction barrier in distributed machine learning workloads,” in Proceedings of the 27th ACM International Conference on Architectural Support… find
  13. S. Wang, J. Wei, A. Sabne, A. Davis, B. Ilbeyi, B. Hechtman, D. Chen, K. S. Murthy, M. Maggioni, Q. Zhang, et al., “Overlap communication with dependent computation via decomposition in large deep learning models,” in Proceedings of the 28th ACM International Conference on Architectural Support for… find
  14. NVIDIA, “TransformerEngine.” https://github.com/NVIDIA/TransformerEngine, 2022. ↗
  15. J. Lamy-Poirier, “Breadth-first pipeline parallelism,” Proceedings of Machine Learning and Systems, vol. 5, 2023. find
  16. D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley, M. Patwary, V. Korthikanti, D. Vainbrand, P. Kashinkunti, J. Bernauer, B. Catanzaro, et al., “Efficient large-scale language model training on gpu clusters using megatron-lm,” in Proceedings of the International Conference for High Performance Compu… find
  17. F. Li, S. Zhao, Y. Qing, X. Chen, X. Guan, S. Wang, G. Zhang, and H. Cui, “Fold3d: Rethinking and parallelizing computational and communicational tasks in the training of large dnn models,” IEEE Transactions on Parallel and Distributed Systems, vol. 34, no. 5, pp. 1432–1449, 2023. find
  18. A. Shah, V. Chidambaram, M. Cowan, S. Maleki, M. Musuvathi, T. Mytkowicz, J. Nelson, O. Saarikivi, and R. Singh, “{TACCL}: Guiding collective algorithm synthesis using communication sketches,” in 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23), pp. 593–612, 2023. find
  19. M. Cowan, S. Maleki, M. Musuvathi, O. Saarikivi, and Y. Xiong, “Mscclang: Microsoft collective communication language,” in Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pp. 502–514, 2023. find
  20. G. Wang, H. Qin, S. A. Jacobs, C. Holmes, S. Rajbhandari, O. Ruwase, F. Yan, L. Yang, and Y. He, “Zero++: Extremely efficient collective communication for giant model training,” arXiv preprint arXiv:2306.10209, 2023. find
  21. V. Thakkar, P. Ramani, C. Cecka, A. Shivam, H. Lu, E. Yan, J. Kosaian, M. Hoemmen, H. Wu, A. Kerr, M. Nicely, D. Merrill, D. Blasig, F. Qiao, P. Majcher, P. Springer, M. Hohnerbach, J. Wang, and M. Gupta, “CUTLASS.” https://github.com/NVIDIA/cutlass, 2024. ↗
  22. NVIDIA, “Megatron-LM.” https://github.com/NVIDIA/Megatron-LM, 2021. ↗
  23. W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica, “Efficient memory management for large language model serving with pagedattention,” in Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
  24. M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” arXiv preprint arXiv:1909.08053, 2019. find
  25. 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.
  26. R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, J. Heek, K. Xiao, S. Agrawal, and J. Dean, “Efficiently scaling transformer inference,” Proceedings of Machine Learning and Systems, vol. 5, 2023. find
  27. S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “Zero: Memory optimizations toward training trillion parameter models,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–16, IEEE, 2020. find
  28. J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He, “{ZeRO-Offload}: Democratizing {Billion-Scale} model training,” in 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551–564, 2021.
  29. S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, and Y. He, “Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning,” in SC21: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–15, IEEE, 2021.
  30. Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, et al., “Pytorch fsdp: experiences on scaling fully sharded data parallel,” arXiv preprint arXiv:2304.11277, 2023. find
  31. Z. Jiang, H. Lin, Y. Zhong, Q. Huang, Y. Chen, Z. Zhang, Y. Peng, X. Li, C. Xie, S. Nong, et al., “Megascale: Scaling large language model training to more than 10,000 gpus,” in 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24’), 2024.
  32. NVIDIA, “NCCL.” https://github.com/NVIDIA/nccl, 2016. ↗
  33. NVIDIA, “NVSHMEM.” https://developer.nvidia.com/nvshmem, 2020. ↗
  34. Z. Chen, A. Kerr, R. Cai, J. Kosaian, H. Wu, Y. Ding, and Y. Xie, “EVT: Accelerating deep learning training with epilogue visitor tree,” in Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024 (in press). find
  35. M. Osama, D. Merrill, C. Cecka, M. Garland, and J. D. Owens, “Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu,” in Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, pp. 429–431, 2023. find
  36. Alibaba, “Megatron-LLaMA.” https://github.com/alibaba/Megatron-LLaMA, 2023. ↗
  37. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al., “Llama 2: Open foundation and fine-tuned chat models,” arXiv preprint arXiv:2307.09288, 2023.
  38. A. Danalis, K.-Y. Kim, L. Pollock, and M. Swany, “Transformations to parallel codes for communication-computation overlap,” in SC’05: Proceedings of the 2005 ACM/IEEE conference on Supercomputing, pp. 58–58, IEEE, 2005. find
  39. A. Danalis, L. Pollock, M. Swany, and J. Cavazos, “Mpi-aware compiler optimizations for improving communication-computation overlap,” in Proceedings of the 23rd international conference on Supercomputing, pp. 316–325, 2009. find
  40. V. Marjanović, J. Labarta, E. Ayguadé, and M. Valero, “Overlapping communication and computation by using a hybrid mpi/smpss approach,” in Proceedings of the 24th acm International Conference on Supercomputing, pp. 5–16, 2010. find
  41. T. Gysi, J. Bär, and T. Hoefler, “dcuda: hardware supported overlap of computation and communication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 609–620, IEEE, 2016. find
  42. 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,” Advances in neural information processing systems, vol. 32, 2019. find
  43. 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 Proceedings of the 27th ACM Symposium on Operating Systems Principles, pp. 1–15, 2019. find
  44. B. Yang, J. Zhang, J. Li, C. Ré, C. Aberger, and C. De Sa, “Pipemare: Asynchronous pipeline parallel dnn training,” Proceedings of Machine Learning and Systems, vol. 3, pp. 269–296, 2021. find
  45. S. Rashidi, M. Denton, S. Sridharan, S. Srinivasan, A. Suresh, J. Nie, and T. Krishna, “Enabling compute-communication overlap in distributed deep learning training platforms,” in 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), pp. 540–553, IEEE, 2021. find
  46. Microsoft, “MSCCL++.” https://github.com/microsoft/mscclpp, 2023. ↗
  47. N. Dryden, T. Moon, S. A. Jacobs, and B. Van Essen, “Communication quantization for data-parallel training of deep neural networks,” in 2016 2nd Workshop on Machine Learning in HPC Environments (MLHPC), pp. 1–8, IEEE, 2016. find
  48. D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic, “Qsgd: Communication-efficient sgd via gradient quantization and encoding,” Advances in neural information processing systems, vol. 30, 2017.
  49. W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li, “Terngrad: Ternary gradients to reduce communication in distributed deep learning,” Advances in neural information processing systems, vol. 30, 2017. find
  50. H. Xu, K. Kostopoulou, A. Dutta, X. Li, A. Ntoulas, and P. Kalnis, “Deepreduce: A sparse-tensor communication framework for federated deep learning,” Advances in Neural Information Processing Systems, vol. 34, pp. 21150– 21163, 2021. find
  51. H. Xu, C.-Y. Ho, A. M. Abdelmoniem, A. Dutta, E. H. Bergou, K. Karatsenidis, M. Canini, and P. Kalnis, “Grace: A compressed communication framework for distributed machine learning,” in 2021 IEEE 41st international conference on distributed computing systems (ICDCS), pp. 561–572, IEEE, 2021. find
  52. J. Fei, C.-Y. Ho, A. N. Sahu, M. Canini, and A. Sapio, “Efficient sparse collective communication and its application to accelerate distributed deep learning,” in Proceedings of the 2021 ACM SIGCOMM 2021 Conference, pp. 676–691, 2021. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/chang-2024-flux-fast-software-based-communication/