深度学习编译器模型训练负载均衡优化方法.pdfVIP

计算机科学与探索1673-9418/2024/18(01)-0111-16

JournalofFrontiersofComputerScienceandTechnologydoi:10.3778/j.issn.1673-9418.2209026

深度学习编译器模型训练负载均衡优化方法

王丽，高开，赵雅倩，李仁刚，曹芳，郭振华+

浪潮电子信息产业股份有限公司高效能服务器与存储技术国家重点实验室，济南250000

+通信作者E-mail:guozhenhua@

摘要：对于计算密集型的人工智能（AI）训练应用，其计算图网络结构更加复杂，数据加载、计算图的任务划

分以及任务调度的负载均衡性都会成为影响计算性能的关键因素。为了使深度学习编译器中模型训练应用

的任务调度达到负载均衡的状态，提出了三种计算图负载均衡优化方法：第一，通过自动建立数据加载与模型

训练的高效流水实现中央处理器和后端计算设备的负载均衡，提高了系统整体能效；第二，通过计算图的分层

优化技术，实现计算图在后端设备执行调度时的负载均衡；最后，通过自动建立层间的高效流水提高后端设备

的资源利用率。实验结果表明，计算图负载均衡优化方法实现了训练任务到底层硬件设备自动映射过程中系

统的负载均衡，与Tensorflow、nGraph等传统的深度学习框架和编译器相比，在不同模型训练中通过任务调度

负载均衡优化技术分别获得了2%~10%的性能提升，同时能够使系统整体的能耗降低10%以上。

关键词：模型训练；编译器优化；负载均衡；分层调度；自动流水

文献标志码：A中图分类号：TP183

DeepLearningCompilerLoadBalancingOptimizationMethodforModelTraining

WANGLi,GAOKai,ZHAOYaqian,LIRengang,CAOFang,GUOZhenhua+

StateKeyLaboratoryofHigh-EndStorageTechnology,InspurElectronicInformationIndustryCo.,Ltd.,Jinan

250000,China

Abstract:Forcomputing-intensiveartificialintelligence(AI)trainingtasks,thecomputationalgraphismorecom-

plex,anddataloading,taskdivisionofthecomputationalgraph,andloadbalancingoftaskschedulinghavebecome

thekeyfactorsaffectingthecomputingperformance.Thispaperproposesthreeoptimizationmethodstomakethe

taskschedulingofmodeltrainingindeeplearningcompilersreachtheloadbalancestate.Firstly,theloadbalance

betweenCPUandback-endcomputingdevicesisrealizedbyautomaticallyestablishinganefficientpipelinefordata

loadingandmodeltraining,whichimprovestheoverallenergyefficiencyofthesystem.Secondly,thelayeredopti-

mizationtechnologyofcomputationalgraphisusedtorealizetheloadbalanceofcomputationalgraphwhenthe

back-enddevicesarescheduling.Finally,thispaperimprovesthe