本文是一篇计算机论文，本文主要针对面向AI for Science任务的训练与推理框架展开研究。设计AI for Science任务的统一建模与表达，并提出算子结构与数据流特征感知的模型自动并行方法FlowAware，实现面向AI for Science任务的分布式自动并行计算框架ParAI4S。
第1章 绪论、
1.1 研究背景与意义
AI for Science在科学计算领域取得了重大进展，尤其是在生物医学[1]、材料科学[2]和分子动力学[3]等领域。各领域广泛应用深度学习大模型的预训练和快速推理能力，实现了海量数据的高效计算与处理。例如，DeepMind推出的AlphaFold2[4]基于10万种已知蛋白质序列和结构的数据集，在128张TPU v3上进行训练，随机初始化后训练耗时1周，微调耗时4天，成功将蛋白质结构预测误差控制在原子级精度，与传统耗时数月的实验结果性能相当。UC Berkeley提出一个用于无机粉末固态合成的实验平台A-Lab[5]，其利用主动学习，结合深度学习数据驱动、从文献数据中学习启发式合成方法，在17天的连续运行中，从58个目标中获得了41个新化合物，发现新材料效率远超过人工实验方法。由此可见，深度学习模型能够显著提升科学任务的计算效率，并保证模型精度，对促进科学计算各领域的发展具有重要作用。
随着应用场景的复杂化和数据的爆炸增长，越来越多的任务需要深度学习模型能够处理海量的数据并提取出更加复杂和细致的特征。这直接推动了模型参数规模不断增长，以增强模型的表达能力和学习能力。在这种趋势下，国内外科技公司和研究机构陆续发布了M6[6]、盘古[7]、PaLM[8]、通义千问[9]、Llama3[10]、GPT-4[11]等深度学习大模型，参数规模均超过千亿级，甚至达到万亿级。如图 1.1所示，模型参数量呈指数增长，Llama3[10]最大可达到4050亿参数，GPT模型由2020年GPT-3[12]模型的1750亿参数增长到GPT-4[11]模型的17600亿参数，达到万亿级别。这些模型凭借其庞大的参数量，在处理复杂科学任务数据时表现出色。数据集规模和模型参数的大幅度增长对深度学习模型的执行设备提出了更高的要求，这远远超出了单一设备的硬件能力。
1.2 研究现状
随着人工智能的不断发展，深度学习逐渐融入科学计算领域，并推动AI for Science的发展。AI for Science任务结合了AI模型，在蛋白质折叠、分子动力学、材料学等领域都表现出极佳的性能。其中，面向AI for Science任务的计算框架、大模型分布式训练与推理技术对AI for Science的发展和应用落地非常重要。本节将对其研究现状进行分析。
1.2.1 面向AI for Science任务的计算框架
AI for Science任务融合了深度学习，面向AI for Science任务的计算框架也由深度学习框架演变而来，研究人员通过扩展深度学习框架的算子和模型支持AI for Science计算。早期的深度学习框架为了提高研发神经网络算法的效率，Cuda-convnet[21]，Theano[22]，Distbelief[23]等框架尝试在GPU上运行神经网络训练程序来加速复杂神经网络训练。这些早期工作定义了深度学习框架需要支持的基本功能，如：神经网络基本计算单 ，自动微分等。随着算法的复杂度的增长和工程化的需要，Caffe[24]、Mx et[25]等框架开始将深度神经网络模型抽象为一组连续堆叠的层，每一层注册模型的前向计算和梯度计算。这些工作进一步加强和验证了系统设计对神经网络算法模块化的抽象，推动了新网络处理单 和网络结构的进一步发展。基于以上系统设计经验，谷歌在2015年推出了开源深度学习框架TensorFlow[13]，其特点在于强大的分布式计算、训练和模型构建功能。TensorFlow主要基于静态计算图，为部署和优化多种AI任务提供了支持，这些任务包括自然语言处理、计算机视觉和生物信息学等广泛领域。随后，在2016年Facebook推出了更加灵活的计算框架PyTorch[14]，它支持静态和动态计算图，能够高效地在多GPU上执行分布式计算。
第2章 相关理论基础
2.1 深度学习理论与模型
自20世纪80年代以来，以神经网络为代表的模式识别方法一直是人们探索人工智能的重要方向。随着硬件性能的提升，Alex等人[53]在2012年提出深度卷积神经网络，将人工智能引向深度学习时代。深度学习算法在图像[54]、文本生成与翻译[55]等领域的任务中表现出显著优势。近年来，深度学习模型也逐渐被应用在生物学[4]、材料学[3,56,57]和空气动力学[58]等科学计算领域。本文主要研究面向AI for Science领域深度学习模型分布式训练与推理技术。因此，本节将首先概述深度学习的理论基础和模型结构，然后介绍当前典型的深度学习网络。
2.1.1 深度学习理论基础
深度学习模型以算子为基础单 ，多个算子组合为层结构，堆叠层形成计算模型。模型可以学习数据的特征表示，通过使用多次反向传播算法来优化内部参数，拟合大型数据集中的隐藏特征。深度卷积网络[53]带来了处理图像、视频等计算机视觉领域的突破，而循环神经网络（Recurrent eural etworks, R ）[59]则适用于文本和语音等连续数据形式，多应用于自然语言处理领域。自注意力机制（Self-Attention）和Transformer结构[33]为大语言模型和蛋白质结构预测带来了新的方法。
2.2 AI for Science分布式计算技术
2.2.1 AI for Science计算工作机理
AI for Science是由AI和HPC融合发展产生，HPC有大量传统基础软件，AI相关模型与算法则日新月异，此外还有大数据和分布式相关软件，这些领域软件栈差异较大，互通性弱，需要科学工作模式。Shantenu等人将当前AI模型在科学发现工作模式总结为三种耦合计算模式[69]：
（1）ML-in-HPC：使用AI模型用来替换部分或整个科学发现任务过程。Atlfast3[70]是高能物理领域利用ML-in-HPC模式的代表性方法，该方法将参数化方法与机器学习技术相结合，替代传统方法Geant4，加速模拟粒子在物质中输运的物理过程。DeeP MD-kit[15]是基于深度学习的原子间势能和力场模型软件包，可快速训练、测试和评估势能模型，并与传统高性能方法LAMMPS[71]对接，少量修改后即可进行能量、力等状态计算。图 2.7展示了使用AI模型拟合分子势能，求解分子系统状态的抽象过程。
（2）ML-about-HPC：AI模型作为传统计算任务补充，指导任务执行，提高整体结果和效率。Proxima[72]基于阈值动态代理传统目标函数，根据输入数据选择代理模型和高成本目标函数，以平衡任务执行时间和精度，提高总体性能。
（3）ML-out-HPC：基于深度学习算法，如主动学习或强化学习，用于动态控制算法活动，引导整个工作流程。EXARL[73]是一种分布式和可扩展强化学习组件，可与LAMMPS对接，以控制传统高性能计算任务。
第3章 AI for Science 任务的统一建模与表达 ................... 25
3.1 统一建模与表达原语定义 ........................... 25
3.1.1 基本结构与定义 ............................... 25
3.1.2 分布式结构设计 ..................... 29
第4章 算子结构与数据流特征感知的模型自动并行方法 ................ 38
4.1 问题分析 ........................... 38
4.2 方法概述 .................................. 39
4.3 基于算子结构与数据流特征的模型结构自动化分析 ................ 40
第5章 面向 AI for Science 任务的分布式自动并行计算框架 Par AI4S .................. 49
5.1 架构设计 ............................... 49
5.2 用户编程接口 ........................ 50
第6章 实验结果分析
6.1 实验环境
6.1.1 软硬件环境
本章实验分别在GPU集群和DCU集群完成。其中GPU集群硬件环境拥有20个核心的虚拟化Intel CPU，主频为2.20 GHz，型号为Intel(R) Xeon(R) CPU E5-2650 v4，内存大小为32GB。集群节点配备高性能GPU计算设备16块，型号为 VIDIA Tesla P100，显存带宽540GB s，显存大小为12GB。配置如表 6.1。GPU集群软件环境为基于Ubuntu 22.04 LTS的操作系统，其内核版本采用5.15.0-91-generic，显卡驱动版本为545.23.08。软件环境配置如表 6.2所示。
第7章 总结与展望
7.1 工作总结
本文主要针对AI for Science任务的训练与推理计算框架展开研究。首先通过调研国内外相关研究，梳理并分析现有的面向AI for Science任务的计算框架和大模型分布式并行与推理技术的研究现状。其次，结合AI模型和数值计算的基本特性，设计AI for Science任务的统一建模与表达。然后，提出算子结构与数据流特征感知的模型自动并行方法FlowAware，根据AI模型和科学计算不同结构特征为AI for Science任务搜索分布式并行策略。最后，在AI for Science任务的统一建模与表达、算子结构与数据流特征感知的模型自动并行方法基础上设计并实现面向AI for Science任务的自动并行计算框架ParAI4S，在DCU集群和GPU集群验证ParAI4S框架的数据吞吐量和扩展性。本文创新工作主要包括以下几点：
（1）针对现有基于计算图的任务建模与表达方法不能描述AI for Science任务结构和集群拓扑，导致AI for Science分布式并行计算实现低效，特别是AI for Science领域深度学习模型的分布式训练与推理实现低效问题，本文提出AI for Science任务的统一建模与表达方法。该方法首先定义操作符、张量、分布式多级子图、分布式通信等原语，准确描述模型的分布式调度。然后，结合AI模型和科学计算的结构特性，提出AI for Science任务特征感知的统一建模与表达，统一两者的表示方法，以支撑AI for Science分布式并行计算的高效实现，特别是AI for Science领域深度学习模型的分布式训练与推理策略的设计与实现。
（2）针对AI for Science领域中深度学习模型分布式训练与推理策略设计困难、执行低效的问题，本文提出算子结构与数据流特征感知的模型自动并行方法FlowAware。该方法自动分析深度神经网络模型结构算子，推导计算代价和存储代价，由数据流依赖关系推导通信代价，构建并优化集群拓扑感知的图搜索空间，在此基础上以负载均衡为核心，根据任务特征为AI for Science领域中的深度学习模型搜索最优执行策略。
参考文献（略）