Reproduce “Parallelizing Large-Scale Tensor Network Contraction on Multiple GPUs” on OpenResearch

About this paper

Feng Pan, Hanfeng Gu, Paul Springer, Xipeng Li

Researchers at Singapore University of Technology and Design and NVIDIA Corporation developed a framework for parallelizing large-scale Tensor Network (TN) contraction on multiple GPUs. The approach distributes intermediate tensors across devices to overcome single-GPU memory limitations and reduce redundant computation compared to traditional slicing, achieving substantial super-linear speedups on up to 1024 H100 GPUs.

Key insights

• Offline GEMM-oriented mode reordering ensures that all pairwise contractions are ready for efficient matrix multiplication kernels without runtime transpositions, and naturally places longer-lived modes in leading positions, benefiting distribution.
• Systematically distributing individual logical tensors across multiple GPUs, rather than just slicing, drastically reduces total floating-point operations by avoiding redundant computations, making previously memory-bound problems tractable.
• The effectiveness of distributed contraction, particularly the translation of compute reduction into wall-clock performance gains, is strongly correlated with the interconnect bandwidth, performing optimally on high-bandwidth NVLink within a node.

Summary by alphaXiv.