機器之心編輯部
真實化學體系包含大量的微觀粒子，其精確的嚴格計算需要指數高的複雜度，對這些體系的模擬一直是材料、製藥和催化等領域的難點和前沿。
為了解決這一問題，近日字節跳動 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 提供了一種量子嵌入方法，該方法在保持量子化學「黃金標準」精度的同時，顯著擴充套件了可計算的體系規模。
圖 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
轉載請聯絡本公眾號獲得授權
投稿或尋求報道：[email protected]