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

编辑 | ScienceAI真实化学体系包含大量的微观粒子,其精确的严格计算需要指数高的复杂度,对这些体系的模拟一直是材料、制药和催化等领域的难点和前沿。 为了解决这一问题,近日字节跳动 ByteDance Research 团队开发并开源了 ByteQC ——基于 GPU 加速的大规模量子化学计算工具集。 该工具集使用强大的 GPU 算力,大幅度加速了常见的量子化学算法,同时结合领域内前沿的量子嵌入方法实现了量子化学「黄金标准」精度下的大规模量子化学体系的模拟。

图片

编辑 | ScienceAI

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

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

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

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

图片

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

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

代码链接:https://github.com/bytedance/byteqc

摘要

在大规模体系中应用量子化学算法需要大量的计算资源,并且计算资源的需求随着体系规模和所需精度的提高而增长。字节团队开发并发展了开源项目 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 驱动科学大爆发!从蛋白质到数学证明,2024 年最值得关注的科技突破

编辑 | ScienceAI2024 年对于 AI for Science 而言,可谓硕果累累:两个诺贝尔奖再度聚焦人工智能与科学的先驱性结合。 其一是诺贝尔化学奖,颁发给了在蛋白质设计与蛋白质结构预测领域做出开创性贡献的 David Baker 博士、John Jumper 博士以及Demis Hassabis 博士;其二是诺贝尔物理学奖,授予了 John J. Hopfield 博士与 Geoffrey Hinton 博士,以表彰他们在人工神经网络及其机学习核心原理方面的奠基性工作。

LLM学习原子「结构语言」,生成未知化合物的晶体结构,登Nature子刊

编辑 | 萝卜皮生成合理的晶体结构通常是预测材料化学成分及其性质的第一步,但当前大多数预测方法计算成本高,制约了创新进程。 通过使用优质生成的候选结构来预测晶体结构,可以突破这一瓶颈。 在最新的研究中,英国雷丁大学(University of Reading)的研究人员介绍了 CrystaLLM,这是一种基于晶体学信息文件 (CIF) 格式的自回归大型语言建模 (LLM) 的多功能晶体结构生成方法。

AI赋能传统力场:字节跳动开发高精度通用小分子力场ByteFF

编辑 | ScienceAI小分子力场是药物发现中的重要工具,在计算机辅助药物设计中发挥关键作用。 化学空间覆盖广泛且高效精确的小分子力场将为药物发现奠定可靠的基础。 尽管基于机器学习的 MLFF(如 ANI-2x,MACE-OFF23 等)能够提供非常精确的小分子势能面预测,但它们的训练需要海量数据量,且推理速度较慢,还存在外推场景不确定度大等问题。