真實(shí)化學(xué)體系包含大量的微觀粒子，其精確的嚴(yán)格計算需要指數(shù)高的復(fù)雜度，對這些體系的模擬一直是材料、制藥和催化等領(lǐng)域的難點(diǎn)和前沿。

為了解決這一問題，近日字節(jié)跳動 ByteDance Research 團(tuán)隊開發(fā)并開源了 ByteQC —— 基于 GPU 加速的大規(guī)模量子化學(xué)計算工具集。該工具集使用強(qiáng)大的 GPU 算力，大幅度加速了常見的量子化學(xué)算法，同時結(jié)合領(lǐng)域內(nèi)前沿的量子嵌入方法實(shí)現(xiàn)了量子化學(xué)「黃金標(biāo)準(zhǔn)」精度下的大規(guī)模量子化學(xué)體系的模擬。論文以大尺寸分子團(tuán)簇，表面吸附問題為例，展示了 ByteQC 在真實(shí)材料計算中的應(yīng)用潛力。

• 論文鏈接：https://arxiv.org/abs/2502.17963
• 代碼鏈接：https://github.com/bytedance/byteqc

該論文作者中還包括 NVIDIA 和北京大學(xué)的合作者。

摘要

在大規(guī)模體系中應(yīng)用量子化學(xué)算法需要大量的計算資源，并且計算資源的需求隨著體系規(guī)模和所需精度的提高而增長。字節(jié)團(tuán)隊開發(fā)并發(fā)展了開源項(xiàng)目 ByteQC（ByteDance Quantum Chemistry）。

在硬件層面，ByteQC 在現(xiàn)代 GPU 上高效實(shí)現(xiàn)了多種標(biāo)準(zhǔn)量子化學(xué)算法，包括平均場計算（Hartree-Fock 方法和密度泛函理論）以及后 Hartree-Fock 方法（如 M?ller-Plesset 微擾理論、隨機(jī)相位近似、耦合簇方法和量子蒙特卡洛方法）。

在算法層面，ByteQC 提供了一種量子嵌入方法，該方法在保持量子化學(xué)「黃金標(biāo)準(zhǔn)」精度的同時，顯著擴(kuò)展了可計算的體系規(guī)模。

圖 1. ByteQC 軟件架構(gòu)

方法

GPU 的顯存顯著小于 CPU 內(nèi)存，同時架構(gòu)的不同導(dǎo)致很多 CPU 可以高效實(shí)現(xiàn)的復(fù)雜邏輯在 GPU 上很難實(shí)現(xiàn)。為了解決這些問題 ByteQC 在開發(fā)過程中主要使用了以下方法：

1. 引入高效計算庫

張量縮并是量子化學(xué)計算的主要熱點(diǎn)之一，為此作者團(tuán)隊引入了 NVIDIA 提供的高效張量計算庫 cuTENSR/cuTENSORMG。該計算庫在最小占用顯存的前提下高效計算張量縮并。作者團(tuán)隊完善了相關(guān)的函數(shù)封裝，將其引入到了 Python / Cupy 的生態(tài)中。

2. 高效實(shí)現(xiàn)復(fù)雜計算邏輯

在周期性體系屏蔽計算中需要在 GPU 上實(shí)現(xiàn)高效的動態(tài)生產(chǎn)者 - 消費(fèi)者模型，作者團(tuán)隊提出使用動態(tài)的 warp 特例化高效實(shí)現(xiàn)。在平均場 Fock 矩陣構(gòu)建中，涉及相鄰任意多的線程競態(tài)求和的問題。CUDA 自帶求和函數(shù)并未針對該特殊情況優(yōu)化，作者團(tuán)隊使用 warp 內(nèi)的 shuffle 指令實(shí)現(xiàn)了高效地求和。

圖 2. 基于 warp 特例化的生產(chǎn)者-消費(fèi)者模型

圖 3. 基于 warp 同步原語的相鄰 7 個線程的競態(tài)求和

3. 優(yōu)化緩存和簡單高效的原位操作
ByteQC 的諸多代碼實(shí)現(xiàn)均進(jìn)行了詳細(xì)的緩存分析，最大限度地實(shí)現(xiàn)了緩存的復(fù)用，減少了顯存需求。此外大量地使用 Cupy 提供的 kernel 接口，通過 CUDA kernel 實(shí)現(xiàn)了原位操作，減少了顯存的占用。

結(jié)果

基準(zhǔn)測試表明相比于 100 核 CPU，ByteQC 的標(biāo)準(zhǔn)量子化學(xué)算法最高可實(shí)現(xiàn)單 A100 GPU 60 倍加速，大多數(shù)模塊的多卡標(biāo)度可達(dá)到線性加速。對應(yīng)可以單 GPU 計算的體系規(guī)模也大幅提升：

• 耦合簇單、雙激發(fā)（CCSD）：1,610 軌道
• 帶微擾三重激發(fā)（CCSD (T)）：1,380 軌道
• 二階 M?ller-Plesset 微擾理論（MP2）：11,040 軌道
• 開放邊界條件下的平均場計算：37,120 軌道
• 周期邊界條件下的平均場計算：超過 100,000 軌道

圖 4. ByteQC 的子模塊加速比（數(shù)據(jù)點(diǎn)）和計算規(guī)模（虛線）

此外，結(jié)合 ByteQC 中提供的量子嵌入功能，團(tuán)隊在 2,753 軌道的水團(tuán)簇問題和 3,929 軌道的氮化硼表面水吸附問題上均實(shí)現(xiàn)了 CCSD (T) 水平的「黃金標(biāo)準(zhǔn)」精度的計算。

圖 5. （左）水團(tuán)簇結(jié)構(gòu)和（右）氮化硼表面水吸附結(jié)構(gòu)

總結(jié)

字節(jié)跳動 ByteDance Research 團(tuán)隊開發(fā)并開源的 ByteQC 軟件包克服了 GPU 開發(fā)過程中顯存受限，復(fù)雜邏輯難以高效實(shí)現(xiàn)的問題，實(shí)現(xiàn)了量子化學(xué)方法的高效 GPU 化。

此外，結(jié)合量子嵌入方法，ByteQC 可以在保持 CCSD (T) 的精度的前提下，計算更大的規(guī)模。通過這些創(chuàng)新和優(yōu)化，ByteQC 有望成為推動量子化學(xué)領(lǐng)域發(fā)展的工具。