[论文] FastMPS：重新审视大规模矩阵乘积态采样中的数据并行

Source: arXiv - 2512.20064v1

概述

Fast‑MPS 通过将数据并行与一种沿键维度划分工作的全新“张量并行”层相结合，重新激活了对大规模矩阵乘积态（MPS）进行采样的数据并行。作者展示了这种两级并行能够打破以往 MPS 模拟器的内存和 I/O 瓶颈，使得能够进行前所未有的大规模量子采样实验（8 k+ 站点，χ = 10⁴），并实现超过 10× 的加速。

关键贡献

• Multi‑level parallel framework: 将经典的数据并行（跨独立样本）与一种新颖的张量并行相结合，后者在 MPS 键合维度上分配繁重的张量收缩。
• Memory‑ and I/O‑aware design: 引入即时压缩和通信/计算重叠，以保持每个进程的内存占用低并隐藏数据传输延迟。
• Scalable implementation: 基于 MPI + NCCL（或类似的高性能集合通信）构建，可在数千个 CPU/GPU 进程上高效运行。
• Benchmark on Gaussian Boson Sampling (GBS): 相比最先进的模拟器实现了超过 10 倍的加速，并将规模推进到 8,176 个格点，χ = 10⁴。
• Open‑source reference: 作者发布了一个原型库，可嵌入现有的张量网络工具包。

方法论

1. 问题分解

   • 数据并行：每个 MPI rank 独立生成总 MPS 样本的一个子集。
   • 张量并行：在同一 rank 内，MPS 张量在第二维度的进程之间拆分，这些进程共同持有键索引（χ 维度）的不同切片。

2. 压缩张量存储

   • 在分配之前，对每个张量进行量化/压缩（例如使用低秩近似或块稀疏编码），以降低内存占用且不牺牲采样保真度。

3. 通信与计算的重叠

   • 当张量的一个切片正在收缩时，下一个切片从远程 rank 流入。此流水线通过非阻塞 MPI 调用和 CUDA 流（在使用 GPU 时）进行调度。

4. 集合收缩引擎

   • 自定义的 all‑reduce / all‑gather 模式聚合张量并行组的部分结果，然后继续处理 MPS 链中的下一个位置。

5. 采样循环

   • 算法遍历 MPS 位点，执行一系列条件概率计算（通过张量收缩）来抽取每个物理索引，完全同标准 MPS 采样，只是现在实现了完全并行化。

整体工作流可视化为一个二维进程网格：行负责不同的样本，列负责同一样本的不同 χ 切片。

结果与发现

作者还展示了采样分布的统计特性（例如 GBS 中的光子数直方图）与精确 MPS 的统计特性在可忽略的误差范围内匹配，证实压缩不会降低科学正确性。

实际意义

• 量子采样研究：研究人员现在可以模拟更大规模的 GBS 实例，有助于验证近期量子光子设备并对量子优势主张进行基准测试。
• 张量网络库：Fast‑MPS 的两级并行可以抽象为流行 Python 包（例如 tensornetwork、quimb）的后端，为开发者的任何基于 MPS 的工作流提供即插即用的性能提升。
• 高性能机器学习：MPS 正在成为序列数据的紧凑模型；Fast‑MPS 使在现有 HPC 集群上对海量数据集进行训练/推理成为可能。
• 受限内存环境：压缩 + 重叠策略可适用于其他大张量工作负载（例如深度学习模型并行、科学模拟），在 I/O 成为瓶颈的情况下提升效率。
• 可扩展的云部署：由于该方法依赖标准的 MPI/NCCL 原语，可迁移到基于云的 GPU 农场（AWS、Azure），无需定制硬件，从而实现按需的大规模张量采样即服务。

限制与未来工作

• 压缩权衡: 虽然作者报告称对采样保真度的影响最小，但激进的压缩可能影响更为细微的量子态属性；仍需进行系统性的误差分析。
• 硬件异构性: 当前实现假设一个相当均匀的 CPU/GPU 集群；将其扩展到混合精度或异构节点配置可能并非易事。
• 超出 MPS 的通用性: Fast‑MPS 针对 MPS 的线性链结构进行优化；将相同思路应用于更高维张量网络（如 PEPS、MERA）将需要新的通信模式。
• 并行布局自动化: 目前在数据并行组和张量并行组之间的最佳划分依赖手动调优；自动调谐器可以使框架更易于使用。

总体而言，Fast‑MPS 为将 MPS 采样扩展到此前难以实现的问题规模提供了实用路径，其设计原则有望影响更广泛的高性能张量计算领域。

作者

• Yaojian Chen
• Si‑Qiu Gong
• Lin Gan
• Yanfei Liu
• An Yang
• Yinuo Wang
• Chao‑yang Lu
• Guangwen Yang

论文信息

• arXiv ID: 2512.20064v1
• 分类: cs.DC
• 出版时间: 2025年12月23日
• PDF: Download PDF