借助亚马逊云科技，中科院孟金涛团队实现AI效率新突破

概 述

中国科学院深圳先进技术研究院是中国科学院在粤港澳大湾区布局建设的国家战略科技力量，作为深圳首个国立科研机构，构建了以科研为主的集科研、教育、产业、资本为一体的微型协同创新体系，高效推进科学研究与产业发展一体设计、一体推进。基于Amazon Graviton处理器，中国科学院深圳先进技术研究院高性能计算技术研究中心孟金涛团队（简称“孟金涛团队”）开发并优化了AutoGEMM计算库，首次实现了将ARM芯片上的矩阵乘法计算效率提升到98%以上，相比主流计算库实现1.3-2.0倍加速，为AI训练与科学计算提供了高效、低功耗的新路径。

目前，孟金涛团队使用的亚马逊云科技产品与解决方案包括：Amazon Graviton、Amazon EC2等。

机会｜构建全新算法

提升ARM架构计算效率

高性能计算技术研究中心是中科院深圳先进技术研究院下属单位，汇集了新世纪百千万人才工程国家级人选、中国科学院百人计划、中国科学院特聘研究员等人才，聚焦高性能计算、数据分析、优化算法取得了多项研究成果，已广泛应用于气象预测、智慧交通、新媒体、生物医药等领域。

作为高性能计算技术研究中心的中坚力量之一，孟金涛潜心于高性能计算、深度学习、生物信息学、图计算等方向的研究，其团队研究内容涵盖ARM指令集下的矩阵乘法（General Matrix Multiply，GEMM）微内核设计、JIT（Just-In-Time，即时生成类）代码自动生成、软硬件协同指令流水线优化以及TVM（Tensor Virtual Machine）框架级自动调优等技术。过去十年中，孟金涛团队主导研发了多个开源线性代数与深度学习加速库，已在国际超级计算大会SC、IEEE TPDS（IEEE Transactions on Parallel and Distributed Systems）、IPDPS等顶级学术会议、期刊上累计发表论文30余篇。

如今，使用RISC（reduced instruction set computing，精简指令集计算）的ARM架构处理器凭借低功耗与高并行性等特点，越来越多地成为数据中心与边缘计算的选择，如何释放ARM芯片在AI场景下的性能潜力正成为业界关注的焦点。在AI计算中，矩阵乘法（GEMM）作为深度学习和科学计算的核心算子，占据了模型训练与推理的绝大部分计算量，其性能表现直接决定了ARM芯片在AI场景中的应用效果。然而，现有主流计算库主要针对规则方阵进行优化，而深度学习实际应用中大量涌现的不规则矩阵计算效率尚不理想，限制了ARM芯片在AI领域的竞争力。

面对AI模型规模不断增长的趋势，孟金涛团队将研究重心转向了ARM架构下不规则GEMM的极限性能优化，希望通过自研AutoGEMM算法进一步提升ARM芯片的GEMM计算效率，实现与x86等使用CISC（Complex Instruction Set Computing，复杂指令集计算）的架构平齐的性能，以便更好地支持大模型训练推理、AIGC生成及生物信息学模拟等前沿应用。为此，团队需要一个具有代表性且灵活高效的ARM验证系统，来支持算法的开发与性能调优。

解决方案｜构建云端弹性算力体系

加速算法优化与验证

基于上述考虑，孟金涛团队选择了基于Amazon Graviton的Amazon EC2实例作为AutoGEMM的核心验证系统。作为亚马逊云科技自研的数据中心级ARM处理器，Amazon Graviton以其全面的技术能力和灵活的算力配置，为AutoGEMM的研发提供了理想的支撑环境。

在技术能力层面，Amazon Graviton是最早支持ARMv9指令集的商业芯片之一，代表了ARM架构在高性能计算领域的前沿水平。它无缝兼容主流Linux系统，具有完善的编译优化工具链、先进的向量扩展（SVE2）与性能分析工具，同时提供增强的安全特性，确保了AutoGEMM从调试、优化到应用的技术需求得到满足，为算法提供了贴近实际生产环境的验证条件。

在算力资源层面，AutoGEMM的开发需要针对不同矩阵形状和硬件配置进行大量参数调优。从Amazon Graviton3实例（如C7g、Hpc7g，单节点最高64 vCPU）到最新Amazon Graviton4实例（如C8g、R8g、M8g，可扩展至192 vCPU），丰富的算力配置使团队能够根据需求灵活调配资源。借助Amazon EC2的弹性扩展与按需付费模式，结合Amazon Graviton的高性价比优势，团队能够在可控预算内快速迭代，高效完成从微内核优化到大规模并行测试的全流程验证，显著提升了AutoGEMM的研发效率。

业务成果｜从算法突破到产业应用

AutoGEMM开启

ARM高性能计算新篇章

在亚马逊云科技支持下，孟金涛团队高效完成了AutoGEMM计算库的开发与全面性能验证。针对ARM架构芯片在处理不规则矩阵乘法（GEMM）计算时存在的效率瓶颈，AutoGEMM通过JIT代码自动生成、软硬件协同指令流水线优化以及框架级参数化调优等技术，在小矩阵运算中首次将ARM芯片的GEMM计算效率提升至98%以上，在不规则矩阵运算中相比主流开源库实现了1.3-2.0倍的性能提升，使RISC架构芯片的计算效率达到CISC架构（x86）芯片水平，证实了RISC架构不仅节能同时也能高效。

目前，AutoGEMM已成功在日本理化研究所（RIKEN）的富岳超级计算机上实现部署，支撑大模型预训练与AIGC内容生成等高算力任务。与此同时，上海智峪–先进院联合实验室也已将该技术应用于AI模型加速计算，并在第15届CASP（Critical Assessment of Structure Prediction，蛋白质结构预测关键评估）国际大赛中获得核酸建模冠军、蛋白–配体建模亚军。

未来，孟金涛团队将继续与亚马逊云科技深化合作，围绕基于Amazon Graviton、Amazon Trainium等芯片的异构高性能计算、生成式AI模型训练优化等方向持续深入研发，探索更高能效、更智能化的算力调度与算法共优化体系，推动AutoGEMM及相关技术在科研计算、AI训练与产业级HPC场景中的持续落地与演进。

参考文献：

[1]https://ieeexplore.ieee.org/document/10793222

[2]https://ieeexplore.ieee.org/document/9695263

[3]https://dl.acm.org/doi/10.1145/3712285.3759779