性能达 1.5+ 倍！昇腾 AI 助力分子动力学模拟研究

本文分享自华为云社区《性能达1.5+倍！昇腾AI助力分子动力学模拟研究》，作者：昇腾 CANN。

分子动力学是一套分子模拟方法，主要依靠牛顿力学来模拟分子体系运动，是研究微观世界的有效手段。但传统研究手段的时间复杂度较高，仅限于研究数千量级原子的小型系统，在应用中存在普遍局限性。而深势科技提出的 DeePMD-kit 可将 AI 技术成功应用于分子动力学模拟，并实现了上亿原子体系的模拟，目前该项技术已经发展成为端到端的开源解决方案，以科学智能（AI for Science）为分子动力学的行业应用带来了更多可能。

而昇腾 AI 凭借超强算力及完整软件栈配套，为科学智能提供了天然土壤，通过昇腾 AI 硬件及异构计算架构 CANN，借助软硬件协同能力充分释放硬件算力，为 DeePMD 的大规模分子动力学模拟提供了高性能解决方案。

基于 CANN 自定义 DeePMD 高性能算子

对诸如 DeePMD 之类的神经网络模型进行加速的关键之一是对网络中的算子进行深度优化，从而发挥出硬件的澎湃算力，而作为专门面向 AI 场景的异构计算架构，CANN 正是搭起了上层深度学习框架和底层 AI 硬件的桥梁。

基于 CANN 的算子编程接口，开发人员在 DeePMD 网络中开发相关的自定义算子，其中涉及到数据排序、读取与存储等离线操作，也涉及矩阵、标量的计算，这些自定义算子不但对精度有较高的要求，而且也很大程度上决定了模型整体的性能。

CANN 能够高效协同昇腾不同异构单元自定义高性能算子，充分释放 AI Core、AI CPU 和 Vector Core 的异构算力。比如将离散的距离计算和排序部署在 AI CPU 上，将可以并行的矩阵、标量的计算部分部署在 AI Core 上，以发挥出每个计算单元的能力，充分发挥硬件计算性能。

高效协同昇腾不同异构单元自定义高性能算子

基于 CANN 的 DeePMD 融合规则及网络优化

算子融合也是网络性能优化重点之一，也是业界常用手段，在 DeePMD 原生实现中，算子融合就作为性能工程的主要手段之一。

而在 DeePMD 网络性能优化中，算子开发人员进一步结合网络特点，借助 CANN 设计新的融合规则，包括 PAD 算子支持动静合一、Mul 支持 NZ+ND 和 ND+NZ、MatMul + Add + TanhGrad 支持 Buffer 融合等，这些融合在整网性能提升中起到了关键作用。同时借助 CANN 的智能调优工具 AOE，自动化完成子图调优、算子调优，以及 TransData 消除等一系列优化，极大提升了模型调优效率。

DeePMD 网络融合示例

此外项目优化了整网混合精度推理流程，将半精度和单精度混合使用，有效减少内存占用，充分释放硬件算力，进一步加速模型执行，在确保分子模拟精度的前提下带来最大性能收益。在不同场景对分子模拟的物理性质与仿真结果对比测试中，以铜拉伸应力应变模拟实验为例可以看到，昇腾 AI 基础软硬件平台的实验结果已经非常近似实际数据。

昇腾 AI 基础软硬件平台上铜的拉伸应力应变模拟结果

昇腾 AI 的整体优化解决方案使 DeePMD-kit 工具在分子动力学模拟计算上取得 1.5+倍现有产品的性能提升成果，助力分子动力学服务成功商用。未来，昇腾 AI 将继续以超强算力和软硬协同能力助力科学计算进入科学智能新阶段，携手更多伙伴凝心聚力，共同向上发展，构筑科学智能领先格局，全面迈进数智时代。

点击关注，第一时间了解华为云新鲜技术~