[论文解读] Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication

Grzegorz Kwaśniewski, Marko Kabić+4 more|arXiv (Cornell University)|Aug 26, 2019

Parallel Computing and Optimization Techniques参考文献 50被引用 10

一句话总结

COSMA 是一种近似通信最优的并行矩阵-矩阵乘法算法，通过红蓝 Pebble 游戏抽象推导出最优的顺序调度，并在并行化过程中保持 I/O 最优性。它在无需手动调优的情况下，相较于 ScaLAPACK、CARMA 和 CTF，最高可提升 12.8 倍（平均提升 2.2 倍），在 Piz Daint 上实现了高达 88% 的峰值性能。

ABSTRACT

We propose COSMA: a parallel matrix-matrix multiplication algorithm that is near communication-optimal for all combinations of matrix dimensions, processor counts, and memory sizes. The key idea behind COSMA is to derive an optimal (up to a factor of 0.03\% for 10MB of fast memory) sequential schedule and then parallelize it, preserving I/O optimality. To achieve this, we use the red-blue pebble game to precisely model MMM dependencies and derive a constructive and tight sequential and parallel I/O lower bound proofs. Compared to 2D or 3D algorithms, which fix processor decomposition upfront and then map it to the matrix dimensions, it reduces communication volume by up to $\sqrt{3}$ times. COSMA outperforms the established ScaLAPACK, CARMA, and CTF algorithms in all scenarios up to 12.8x (2.2x on average), achieving up to 88\% of Piz Daint's peak performance. Our work does not require any hand tuning and is maintained as an open source implementation.

研究动机与目标

• 解决现有并行矩阵-矩阵乘法算法在非方阵或任意处理器数量下通信效率低下的局限性。
• 开发一种统一的算法，适用于所有矩阵维度、处理器数量和内存大小组合，并实现 I/O 最优性。
• 通过推导一种构造性、近似最优的调度，消除高性能矩阵乘法中手动参数调优的需求。
• 在多种工作负载和硬件配置下，实现相对于 ScaLAPACK、CARMA 和 CTF 等成熟库的实际性能提升。

提出的方法

• 使用红蓝 Pebble 游戏建模矩阵-矩阵乘法中的数据依赖关系，并推导出紧致的顺序与并行 I/O 下界。
• 通过在内存约束下求解 Pebbling 问题，构建最优的顺序调度，确保在给定内存大小下的最小 I/O。
• 在保持 I/O 最优性的前提下并行化所推导的顺序调度，将计算映射到任意数量的处理器上，不受 2 的幂次限制。
• 应用算法优化技术，包括分块数据布局、通信-计算重叠以及单边通信，以提升性能。
• 利用 I/O 下界构造性证明，确保该算法在 10MB 快速内存下距离理论最小值仅差 0.03%，证实其通信成本近乎最优。
• 将该算法实现为开源库，以确保可复现性，并可在 HPC 系统中实际部署。

实验结果

研究问题

• RQ1 是否可以设计一种矩阵-矩阵乘法算法，使其在所有矩阵维度和处理器数量下均实现通信最优且具备通用性？
• RQ2 红蓝 Pebble 游戏在多大程度上可用于为并行算法推导出构造性、紧致的 I/O 下界？
• RQ3 在真实场景中，通信最优算法与 ScaLAPACK、CARMA 和 CTF 等先进库相比，性能表现如何？
• RQ4 是否可以实现一种近似最优的算法，而无需手动调整块大小或处理器网格配置？
• RQ5 渐近通信最优性对现代超级计算机上实际运行时间和性能的实际影响是什么？

主要发现

• COSMA 相较于 2D 和 2.5D 算法，通信量最多减少 √3 倍，尤其在非方阵和非 2 的幂次处理器数量下优势显著。
• 该算法在 Piz Daint 上最高实现 88% 的峰值性能，相较于所有测试配置下的 ScaLAPACK、CARMA 和 CTF 均有显著提升。
• 在某些场景下，COSMA 相较于现有性能最佳的库最高实现 12.8 倍加速，平均提升 2.2 倍，展现出实际优越性。
• 该算法在 10MB 快速内存下距离理论 I/O 下界仅差 0.03%，证实其通信成本近乎最优。
• 该方法可推广至其他线性代数内核（如 LU 分解和 Cholesky 分解），并支持多级内存。
• 开源实现消除了手动调优的需求，使其具备生产就绪性，可广泛部署于各类 HPC 环境。