[论文解读] Boosting the effective performance of massively parallel tensor network state algorithms on hybrid CPU-GPU based architectures via non-Abelian symmetries

Andor Menczer, Örs Legeza|arXiv (Cornell University)|Sep 23, 2023

Parallel Computing and Optimization TechniquesComputer Science被引用 3

一句话总结

本文提出了一种混合CPU-GPU实现的密度矩阵重整化群（DMRG）算法，利用非阿贝尔SU(2)对称性，将计算复杂度降低高达10倍，并在TFLOPS性能上相比先前的U(1)对称实现提升3–6倍。通过利用Wigner-Eckart定理将非阿贝尔张量代数与张量网络层解耦，并结合自适应缓冲和分步批处理操作优化内存访问，该方法在单节点八张A100 GPU上实现了110 TFLOPS的性能，且在同等精度下估计有效性能可达250–500 TFLOPS。

ABSTRACT

We present novel algorithmic solutions together with implementation details utilizing non-Abelian symmetries in order to boost the current limits of tensor network state algorithms on high performance computing infrastructure. In our in-house developed hybrid CPU-multiGPU solution scheduling is decentralized, threads are autonomous and inter-thread communications are solely limited to interactions with globally visible lock-free constructs. Our custom tailored virtual memory management ensures data is produced with high spatial locality, which together with the use of specific sequences of strided batched matrix operations translates to significantly higher overall throughput. In order to lower IO overhead, an adaptive buffering technique is used to dynamically match the level of data abstraction, at which cache repositories are built and reused, to system resources. The non-Abelian symmetry related tensor algebra based on Wigner-Eckhart theorem is fully detached from the conventional tensor network layer, thus massively parallel matrix and tensor operations can be performed without additional overheads. Altogether, we have achieved an order of magnitude increase in performance with respect to results reported in arXiv:2305.05581 in terms of computational complexity and at the same time a factor of three to six in the actual performance measured in TFLOPS. Benchmark results are presented on Hilbert space dimensions up to $2.88 imes10^{36}$ obtained via large-scale SU(2) spin adapted density matrix renormalization group simulations on selected strongly correlated molecular systems. These demonstrate the utilization of NVIDIA's highly specialized tensor cores, leading to performance around 110 TFLOPS on a single node supplied with eight NVIDIA A100 devices. In comparison to U(1) implementations with matching accuracy, our solution has an estimated effective performance of 250-500 TFLOPS.

研究动机与目标

• 通过在张量网络态算法中利用非阿贝尔对称性，克服量子系统模拟的指数级缩放问题。
• 在混合CPU-GPU高性能计算架构上实现高效、大规模并行的DMRG模拟，且对称性处理不引入性能开销。
• 通过自定义内存管理和自适应缓冲技术，实现更高的计算吞吐量和更低的I/O开销。
• 展示SU(2)自旋适配DMRG在高达2.88×10^36维希尔伯特空间中的可行性与性能提升。
• 为强关联量子系统在量子化学与凝聚态物理中的研究建立可扩展、高性能的计算框架。

提出的方法

• 采用去中心化、无锁的线程调度模型，并结合自主GPU内核，以最小化线程间通信。
• 基于Wigner-Eckart定理的非阿贝尔张量代数与张量网络层完全解耦，消除对称性相关的计算开销。
• 采用分步批处理矩阵运算，以最大化GPU张量核心上的空间局部性和吞吐量。
• 自适应缓冲动态匹配数据抽象层级与系统资源，降低I/O开销并提升缓存复用率。
• 自定义虚拟内存管理确保在多GPU节点间实现高数据局部性和高效的内存访问模式。
• 该实现基于混合CPU-多GPU架构，采用NVIDIA A100 GPU，并利用张量核心加速FP16和BF16运算。

实验结果

研究问题

• RQ1 非阿贝尔SU(2)对称性能否在混合CPU-GPU系统上的大规模并行张量网络态算法中被高效利用？
• RQ2 Wigner-Eckart定理如何被整合进张量网络算法中，而不会引入性能瓶颈？
• RQ3 在现代GPU架构上，结合非阿贝尔对称性与优化的内存访问及批处理操作，能实现多大的性能提升？
• RQ4 自适应缓冲在高对称性降维的大规模DMRG模拟中，能在多大程度上降低I/O开销？
• RQ5 在同等精度下，SU(2)对称DMRG相比U(1)对称实现，其可达到的有效性能是多少？

主要发现

• 与先前工作（arXiv:2305.05581）相比，该方法在计算复杂度上实现数量级的改进，将缩放从指数级降低至近多项式级。
• 在单节点八张NVIDIA A100 GPU上，该实现达到110 TFLOPS，充分展示了张量核心的有效利用。
• 在同等精度下，SU(2)实现的估计有效性能为250–500 TFLOPS，相比U(1)实现提升3–6倍。
• 该算法随GPU设备数量线性扩展，支持高效多节点部署，适用于千万亿次规模的模拟。
• 成功利用大规模SU(2)自旋适配DMRG对强关联分子体系实现了高达2.88×10^36维的希尔伯特空间模拟。
• 将非阿贝尔对称性代数与张量网络层解耦，实现了高吞吐量、低开销的张量操作并行化。