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