FastMPS: Revisit Data Parallel in Large-scale Matrix Product State Sampling

TL;DR

FastMPS introduces a multi-level parallel framework that combines data and tensor parallelism to efficiently scale large-scale MPS sampling, significantly improving speed and scalability in tensor network applications.

Contribution

The paper proposes FastMPS, a novel parallel framework that overcomes memory and I/O limitations, enabling scalable and efficient large-scale MPS sampling.

Findings

Achieves over 10x speedup over existing simulators.

Scales to thousands of processes for large MPS.

Enables simulations with 8,176 sites and bond dimension 10^4.

Abstract

Matrix Product State (MPS) is a versatile tensor network representation widely applied in quantum physics, quantum chemistry, and machine learning, etc. MPS sampling serves as a critical fundamental operation in these fields. As the problems become more complex, the scale of MPS is rapidly increasing. Traditional data parallelism is limited by memory and heavy I/O in large-scale MPS. Model parallelism that can handle large-scale MPS imposes rigid process bindings and lacks scalability. This work proposes Fast-MPS, a multi-level parallel framework for scalable MPS sampling. Our design combines data parallelism across samples with tensor parallelism along bond dimensions. We eliminate memory and I/O pressure through compression and overlapping, and revive data parallel in large-scale MPS sampling. We evaluate our approach on Gaussian Boson Sampling, a representative and demanding application. Fast-MPS achieves over 10x speedup compared to existing simulators, scales to thousands of processes, and enables simulations with 8,176 sites and bond dimension chi = 10^4, significantly outperforming the state of the art. Fast-MPS has demonstrated great potential in high-performance tensor network applications.