ByteQC:通往大规模实用化量子化学计算的曙光

2025-03-05 12:56:00
刘大牛
转自文章
248

图片

编辑 | ScienceAI

真实化学体系包含大量的微观粒子,其精确的严格计算需要指数高的复杂度,对这些体系的模拟一直是材料、制药和催化等领域的难点和前沿。

为了解决这一问题,近日 字节跳动 ByteDance Research 团队开发并开源了 ByteQC ——基于 GPU 加速的大规模量子化学计算工具集。

该工具集使用强大的 GPU 算力,大幅度加速了常见的量子化学算法,同时结合领域内前沿的量子嵌入方法实现了量子化学「黄金标准」精度下的大规模量子化学体系的模拟。

论文以大尺寸分子团簇,表面吸附问题为例,展示了 ByteQC 在真实材料计算中的应用潜力。

图片

该论文作者中还包括 NVIDIA 和 北京大学 的合作者。

论文链接:https://arxiv.org/abs/2502.17963

代码链接:https://github.com/bytedance/ byteq c

摘要

在大规模体系中应用量子化学算法需要大量的计算资源,并且计算资源的需求随着体系规模和所需精度的提高而增长。字节团队开发并发展了开源项目 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 有望成为推 动量 子化学领域发展的工具。

理论 人工智能 科学 化学 量子 AI for Science
发表评论
评论通过审核后显示。
文章分类
联系我们
联系人: 透明七彩巨人
Email: weok168@gmail.com