2025-03-05 18:03 四川
GPU加速助力量子化学计算。
机器之心发布
机器之心编辑部
真实化学体系包含大量的微观粒子，其精确的严格计算需要指数高的复杂度，对这些体系的模拟一直是材料、制药和催化等领域的难点和前沿。
为了解决这一问题，近日字节跳动 ByteDance Research 团队开发并开源了 ByteQC —— 基于 GPU 加速的大规模量子化学计算工具集。该工具集使用强大的 GPU 算力，大幅度加速了常见的量子化学算法，同时结合领域内前沿的量子嵌入方法实现了量子化学「黄金标准」精度下的大规模量子化学体系的模拟。论文以大尺寸分子团簇，表面吸附问题为例，展示了 ByteQC 在真实材料计算中的应用潜力。
该论文作者中还包括 NVIDIA 和北京大学的合作者。
摘要
在大规模体系中应用量子化学算法需要大量的计算资源，并且计算资源的需求随着体系规模和所需精度的提高而增长。字节团队开发并发展了开源项目 ByteQC（ByteDance Quantum Chemistry）。
在硬件层面，ByteQC 在现代 GPU 上高效实现了多种标准量子化学算法，包括平均场计算（Hartree-Fock 方法和密度泛函理论）以及后 Hartree-Fock 方法（如 Møller-Plesset 微扰理论、随机相位近似、耦合簇方法和量子蒙特卡洛方法）。
在算法层面，ByteQC 提供了一种量子嵌入方法，该方法在保持量子化学「黄金标准」精度的同时，显著扩展了可计算的体系规模。
图 1. ByteQC 软件架构
方法
GPU 的显存显著小于 CPU 内存，同时架构的不同导致很多 CPU 可以高效实现的复杂逻辑在 GPU 上很难实现。为了解决这些问题 ByteQC 在开发过程中主要使用了以下方法：
1. 引入高效计算库
张量缩并是量子化学计算的主要热点之一，为此作者团队引入了 NVIDIA 提供的高效张量计算库 cuTENSR/cuTENSORMG。该计算库在最小占用显存的前提下高效计算张量缩并。作者团队完善了相关的函数封装，将其引入到了 Python / Cupy 的生态中。
2. 高效实现复杂计算逻辑
在周期性体系屏蔽计算中需要在 GPU 上实现高效的动态生产者 - 消费者模型，作者团队提出使用动态的 warp 特例化高效实现。在平均场 Fock 矩阵构建中，涉及相邻任意多的线程竞态求和的问题。CUDA 自带求和函数并未针对该特殊情况优化，作者团队使用 warp 内的 shuffle 指令实现了高效地求和。
图 2. 基于 warp 特例化的生产者-消费者模型
图 3. 基于 warp 同步原语的相邻 7 个线程的竞态求和
3. 优化缓存和简单高效的原位操作
ByteQC 的诸多代码实现均进行了详细的缓存分析，最大限度地实现了缓存的复用，减少了显存需求。此外大量地使用 Cupy 提供的 kernel 接口，通过 CUDA kernel 实现了原位操作，减少了显存的占用。
结果
基准测试表明相比于 100 核 CPU，ByteQC 的标准量子化学算法最高可实现单 A100 GPU 60 倍加速，大多数模块的多卡标度可达到线性加速。对应可以单 GPU 计算的体系规模也大幅提升：
耦合簇单、双激发（CCSD）：1,610 轨道
带微扰三重激发（CCSD (T)）：1,380 轨道
二阶 Møller-Plesset 微扰理论（MP2）：11,040 轨道
开放边界条件下的平均场计算：37,120 轨道
周期边界条件下的平均场计算：超过 100,000 轨道
图 4. ByteQC 的子模块加速比（数据点）和计算规模（虚线）
此外，结合 ByteQC 中提供的量子嵌入功能，团队在 2,753 轨道的水团簇问题和 3,929 轨道的氮化硼表面水吸附问题上均实现了 CCSD (T) 水平的「黄金标准」精度的计算。
图 5. （左）水团簇结构和（右）氮化硼表面水吸附结构
总结
字节跳动 ByteDance Research 团队开发并开源的 ByteQC 软件包克服了 GPU 开发过程中显存受限，复杂逻辑难以高效实现的问题，实现了量子化学方法的高效 GPU 化。
此外，结合量子嵌入方法，ByteQC 可以在保持 CCSD (T) 的精度的前提下，计算更大的规模。通过这些创新和优化，ByteQC 有望成为推动量子化学领域发展的工具。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
