FlowAware：一种支持AI for Science任务的模型分布式自动并行方法

doi:10.11871/jfdc.issn.2096-742X.2025.05.006

数据与计算发展前沿 ›› 2025, Vol. 7 ›› Issue (5): 65-87.

CSTR: 32002.14.jfdc.CN10-1649/TP.2025.05.006

doi: 10.11871/jfdc.issn.2096-742X.2025.05.006

• 专刊：国产算力新力量，助力科学计算应用新发展 • 上一篇下一篇

FlowAware：一种支持AI for Science任务的模型分布式自动并行方法

曾艳¹(),吴宝福¹,易广政¹,黄成创¹,邱扬¹,陈越¹,万健^1,^2,^*(),胡帆³,金思聪¹,梁迦隽¹,李欣¹

1.杭州电子科技大学，浙江杭州 310018
2.浙江科技大学，浙江杭州 310023
3.浙江曙光信息技术有限公司，浙江杭州 310013

收稿日期:2025-02-28 出版日期:2025-10-20 发布日期:2025-10-23
通讯作者: 万健
作者简介:曾艳，博士，杭州电子科技大学副教授，研究方向包括分布式与并行计算、分布式机器学习及大数据处理。
本文中负责概念构思、指导监督及方法论制定。
ZENG Yan, Ph.D., is an Associate Professor at Hangzhou Dianzi University. Her research interests include distributed and parallel computing, distributed ma-chine learning, and big data.
In this paper, she is responsible for conceptualization, supervision, and methodology design. E-mail: yz@hdu.edu.cn|万健，博士，杭州电子科技大学教授，主要研究领域为网格计算、服务计算和云计算。
本文中负责项目管理及论文撰写、审阅与编辑。
WAN Jian, Ph.D., is a Professor at Hangzhou Dianzi University. His research interests include grid computing, service computing, and cloud computing.
In this paper, he is responsible for project administration, original draft writing, and review and editing of the manuscript.E-mail: wanjian@hdu.edu.cn

FlowAware: A Feature-Aware Automated Model Parallelization Method for AI-for-Science Tasks

ZENG Yan¹(),WU Baofu¹,YI Guangzheng¹,HUANG Chengchuang¹,QIU Yang¹,CHEN Yue¹,WAN Jian^1,^2,^*(),HU Fan³,JIN Sicong¹,LIANG Jiajun¹,LI Xin¹

1. Hangzhou Dianzi University, Hangzhou, Zhejiang 310018, China
2. Zhejiang University of Science and Technology, Hangzhou, Zhejiang 310023, China
3. Zhejiang Sugon Information Technology Co., Ltd, Hangzhou, Zhejiang 310013, China

Received:2025-02-28 Online:2025-10-20 Published:2025-10-23
Contact: WAN Jian
Supported by:
National Key Research and Development Program of China(2023YF-B3001501);National Natural Science Foundation of China (NSFC)(62302133);Key Research and Development Program of Zhejiang Province(2024C01026);Yangtze River Delta Project(2023Z-Y1068);Hangzhou Key Research Plan Project(2024SZD1A02);GHfund A(202302019816)

摘要/Abstract

摘要：

【目的】 本研究旨在解决AI for Science领域中深度学习模型分布式并行计算策略设计与实现困难等导致AI for Science任务计算低效的问题。【方法】 本文提出了一种支持 AI for Science任务的模型分布式自动并行方法FlowAware。该方法基于AI for Science框架 JAX，深入分析AI for Science任务特征以及深度学习模型的算子结构和数据流特性，结合集群拓扑结构，构建模型分布式并行计算策略搜索空间；在此基础上以负载均衡和通信最优化为目标，为AI模型搜索最优分布式并行计算策略。【结果】 在类似GPU的加速器集群和GPU集群上进行了对比实验，实验结果表明，相比Alpa，FlowAware的吞吐量最高可提升7.8倍。【结论】 FlowAware为AI for Science任务中AI模型提供了高效的分布式并行策略搜索方法，并显著加速了AI模型的计算性能。

关键词: AI for Science, 深度学习, 分布式并行计算

Abstract:

[Objective] This study aims to address the inefficiency of AI-for-Science tasks caused by the design and implementation challenges of applying the distributed parallel computing strategies to deep learning models, as well as their inefficient execution. [Methods] We propose an automatic distributed parallelization method for AI-for-Science tasks, called FlowAware. Based on the AI-for-Science framework JAX, this approach thoroughly analyzes task characteristics, operator structures, and data flow properties of deep learning models. By incorporating cluster topology information, it constructs a search space for distributed parallel computing strategies. Guided by load balancing and communication optimization objectives, FlowAware automatically identifies optimal distributed parallel computing strategies for AI models. [Results] Comparative experiments conducted on both GPU-like accelerator clusters and GPU clusters demonstrated that FlowAware achieves a throughput improvement of up to 7.8×compared to Alpa. [Conclusions] FlowAware effectively enhances the search efficiency of distributed parallel computing strategies for AI models in scientific computing tasks and significantly improves their computational performance.

Key words: AI for Science, deep learning, distributed parallel computing

曾艳,吴宝福,易广政,黄成创,邱扬,陈越,万健,胡帆,金思聪,梁迦隽,李欣. FlowAware：一种支持AI for Science任务的模型分布式自动并行方法[J]. 数据与计算发展前沿, 2025, 7(5): 65-87.

ZENG Yan,WU Baofu,YI Guangzheng,HUANG Chengchuang,QIU Yang,CHEN Yue,WAN Jian,HU Fan,JIN Sicong,LIANG Jiajun,LI Xin. FlowAware: A Feature-Aware Automated Model Parallelization Method for AI-for-Science Tasks[J]. Frontiers of Data and Computing, 2025, 7(5): 65-87, https://cstr.cn/32002.14.jfdc.CN10-1649/TP.2025.05.006.

图/表 18

参考文献 20

[1]	ABRAMSON J, ADLER J, DUNGER J, et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3[J]. Nature, 2024, 630(8016): 493-500.
[2]	DU Y, WANG Y, HUANG Y, et al. M2Hub: Unlocking the Potential of Machine Learning for Materials Discovery[C]// Proceedings of the 37th International Conference on Neural Information Processing Systems, 2024: 77359-77378.
[3]	LIAO Y L, SMIDT T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs[J]. arXiv preprint arXiv:2206.11990, 2022.
[4]	JUMPER J, EVANS R, PRITZEL A, et al. Highly accurate protein structure prediction with AlphaFold[J]. Nature, 2021, 596(7873): 583-589.
[5]	SZYMANSKI N J, RENDY B, FEI Y, et al. An autonomous laboratory for the accelerated synthesis of novel materials[J]. Nature, 2023, 624(7990): 86-91.
[6]	WANG H, ZHANG L, HAN J, et al. DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics[J]. Computer Physics Communications, 2018, 228: 178-184.
[7]	WANG H, DING Y, GU J, et al. QuantumNAS: Noiseadaptive search for robust quantum circuits[C]//The 28th IEEE international symposium on high-performance computer architecture (HPCA-28), 2022: 692-708.
[8]	BRADBURY J, FROSTIG R, HAWKINS P, et al. JAX: composable transformations of Python+NumPy programs[EB/OL]. 2018. http://github.com/google/jax.
[9]	ZHENG L, LI Z, ZHANG H, et al. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning[C]//16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 2022: 559-578.
[10]	HAN M, ZENG Y, SHU H, et al. Device placement using Laplacian PCA and graph attention networks[J]. The Computer Journal, 2025, 68(2): 175-186.
[11]	ZENG Y, HUANG C C, NI Y J, et al. An Auto-Parallel Method for Deep Learning Models Based on Genetic Algorithm[C]//2023 IEEE 29th International Conference on Parallel and Distributed Systems (ICPADS), IEEE, 2023: 230-235.
[12]	ZENG Y, YI G, YIN Y, et al. Aware: Adaptive distributed training with computation, communication and position awareness for deep learning model[C]// 2022 IEEE 24th Int Conf on High Performance Computing & Communications; 8th Int Conf on Data Science & Systems; 20th Int Conf on Smart City; 8th Int Conf on Dependability in Sensor, Cloud & Big Data Systems & Application (HPCC/DSS/SmartCity/DependSys), IEEE, 2022: 1299-1306.
[13]	MA Y, YU D, WU T, et al. PaddlePaddle: An open-source deep learning platform from industrial practice[J]. Frontiers of Data and Computing, 2019, 1(1): 105-115.
[14]	Huawei MindSpore AI Development Framework[M]//Huawei Technologies Co. Ltd. Artificial Intelligence Technology, Singapore: Springer Nature, 2023: 137-162.
[15]	SHOEYBI M, PATWARY M, PURI R, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism[J]. arXiv preprint arXiv:1909.08053, 2019.
[16]	XU Q, YOU Y. An Efficient 2D Method for Training Super-Large Deep Learning Models[C]//2023 IEEE International Parallel and Distributed Processing Symposium (IPDPS), 2023: 222-232.
[17]	WANG B, XU Q, BIAN Z, et al. Tesseract: parallelize the tensor parallelism efficiently[C]// Proceedings of the 51st International Conference on Parallel Processing, 2022: 1-11.
[18]	BIAN Z, XU Q, WANG B, et al. Maximizing Parallelism in Distributed Training for Huge Neural Networks[J]. arXiv preprint arXiv:2105.14450, 2021.
[19]	JIA Z, ZAHARIA M, AIKEN A. Beyond Data and Model Parallelism for Deep Neural Networks[C]// Proceedings of machine learning and systems: Vol. 1, 2018: 1-13.
[20]	JEON B, CAI L, SRIVASTAVA P, et al. Baechi: fast device placement of machine learning graphs[C]// Proceedings of the 11th ACM Symposium on Cloud Computing, Virtual Event USA: ACM, 2020: 416-430.

Device	Specification
GPU Cluster Hardware	CPU	Intel(R) Xeon(R) CPU E5-2650, 2.20GHz, 20 cores, 40 threads
	Memory	32GB
	GPU	16×NVIDIA Tesla P100, 12GB memory, 540GB/s bandwidth
GPU-like Accelerator Cluster Hardware	CPU	32-core processor
	Memory	64GB
	GPU Accelerator Cards	32 cards, 16GB VRAM each

Networks	Application Domain	Datasets	Task Type
ResNet	Machine Vision	cifar10	Training
PNasNet	Machine Vision	MNIST	Training
BERT	Natural Language Processing	WikiText-2	Training
T5	Natural Language Processing	C4	Training
Alphafold2	Protein Structure Prediction	bfd	Inference

Model	Nnmber of Backbone Layers				Operator Number
Model	1	2	3	4	Operator Number
ResNet-50	3	4	6	3	3467
ResNet-152	3	8	36	3	10131
ResNet-230	5	16	50	5	15227
ResNet-302	5	16	72	7	19931

Model	Number of Operators	Framework	Search Time
ResNet-152	10131	Alpa	127.7
ResNet-152	10131	FlowAware	38.5
ResNet-230	15227	Alpa	199.3
ResNet-230	15227	FlowAware	58.2
ResNet-302	19931	Alpa	264.1
ResNet-302	19931	FlowAware	85.5

Accelerator Card	Number of Accelerator Cards	Parameter Scale (Billion)	Scalability Efficiency (%)
GPU like accelerator card	4	7.5	—
	8	15	100.00%
	16	17.5	58.33%
	32	18.75	31.25%
GPU	8	7.5	—
GPU	16	10	66.67%

FlowAware：一种支持AI for Science任务的模型分布式自动并行方法

FlowAware: A Feature-Aware Automated Model Parallelization Method for AI-for-Science Tasks

RichHTML

PDF

可视化

摘要/Abstract

引用本文

使用本文

图/表 18

参考文献 20

相关文章 15

编辑推荐

Metrics

本文评价

Models	Strategy	Inference Accuracy (%)
ResNet-50	JAX	61.09%
ResNet-50	FlowAware	61.09%
ResNet-152	JAX	60.15%
ResNet-152	FlowAware	60.15%

[1]	辛昱杭,王琦祎,孙婧,赵春燕,刘雨佳,梁雪,陈杰. 基于雷达回波外推的短临降水预报模型TrajCast在国产加速器上的实践[J]. 数据与计算发展前沿, 2025, 7(5): 113-122.
[2]	王鹏,杨霄凤,何忠臣,杜君. 基于浅层-深层卷积循环神经网络的多光谱遥感图像全色锐化方法[J]. 数据与计算发展前沿, 2025, 7(5): 138-152.
[3]	贾子昂. 基于多源半监督学习的牙齿结构分割[J]. 数据与计算发展前沿, 2025, 7(2): 175-185.
[4]	马秋平, 张琪, 赵晓凡. 图表问答研究综述[J]. 数据与计算发展前沿, 2025, 7(1): 19-37.
[5]	水映懿, 张琪, 李根, 张士豪, 吴尚. 基于多类特征的社交网络影响力预测研究综述[J]. 数据与计算发展前沿, 2025, 7(1): 2-18.
[6]	金家立, 高思远, 高满达, 王文彬, 柳绍祯, 孙哲南. 基于生成对抗网络和扩散模型的人脸年龄编辑综述[J]. 数据与计算发展前沿, 2025, 7(1): 38-55.
[7]	卢成浩,陈秀宏. 基于隐式分区学习深度特征融合重建曲面网络[J]. 数据与计算发展前沿, 2024, 6(6): 19-31.
[8]	韦一金,樊景超. 基于改进的BERT-BiGRU-Attention的农业科技政策分类模型[J]. 数据与计算发展前沿, 2024, 6(6): 53-61.
[9]	何文通,罗泽. 基于联邦学习的野生动物红外相机图像目标检测[J]. 数据与计算发展前沿, 2024, 6(6): 85-96.
[10]	晏直誉, 茹一伟, 孙福鹏, 孙哲南. 基于主动感知机制的视频行为识别方法研究[J]. 数据与计算发展前沿, 2024, 6(5): 66-79.
[11]	廖立波, 王书栋, 宋维民, 张兆领, 李刚, 黄永盛. CEPC上基于DeepSets模型的喷注标记算法研究[J]. 数据与计算发展前沿, 2024, 6(3): 108-115.
[12]	严瑾, 董科军, 李洪涛. 融合语义和共现特征的Web跟踪器深度识别方法[J]. 数据与计算发展前沿, 2024, 6(3): 127-138.
[13]	寇大治. 基于深度学习的口腔全景片牙齿自动分割方法[J]. 数据与计算发展前沿, 2024, 6(3): 162-172.
[14]	蔡程飞, 李军, 焦一平, 王向学, 郭冠辰, 徐军. 基于深度学习的医学多模态数据融合方法在肿瘤学中的进展和挑战[J]. 数据与计算发展前沿, 2024, 6(3): 3-14.
[15]	郑懿诺, 孙沐毅, 张虹云, 张婧, 邓天政, 刘倩. 深度学习在口腔种植影像学中的应用：研究进展与挑战[J]. 数据与计算发展前沿, 2024, 6(3): 41-49.