关键词:
深度神经网络
显存
张量
卸载
重计算
摘要:
深度神经网络(Deep Neural Network,DNN)是人工智能领域的关键技术,在多种应用领域展示出卓越的性能。近年来,为了提高DNN的准确性,训练数据集的规模和模型的复杂度均大幅增加。由于DNN训练必须保留大量的张量,随着模型规模的不断扩大,对GPU(Graphics Processing Unit,GPU)的显存资源需求显著增加。然而,由于硬件技术限制,GPU设备的显存扩容速度远低于DNN训练作业的发展速度。现有GPU显存的硬件限制制约了大规模DNN模型的发展。卸载和重计算是降低训练时显存资源需求的常用机制。卸载是将显存中的张量通过PCI-E总线迁移到大容量的CPU内存中,重用时再将其预取回GPU的显存,但卸载/预取机制可能带来训练过程阻塞的代价。重计算机制只保留部分张量作为检查点,通过从检查点重新计算来满足张量的重用需求,其缺点是引入了张量重复计算的代价。两种机制引入代价的原理不同,分别消耗PCI-E总线资源和GPU计算资源,存在结合优化的潜力。但不合理的结合策略会引入较高的阻塞和重计算代价,反而对训练性能造成严重的影响。本文主要研究如何在有限显存的条件下高效地结合卸载和重计算,以满足DNN训练的显存需求。论文的主要工作如下:(1)本文提出了混合显存优化统一代价模型,该模型以训练延迟为基准,准确描述卸载/预取和重计算机制的显存优化代价。基于该模型,本文提出了基于卸载和重计算混合机制的显存优化问题(Memory Optimization Problem for GPU Based on Hybrid Mechanisms,HGMOP)。该问题以张量计算的依赖关系和GPU显存容量为约束条件,在保证训练作业正确性的前提下,最小化显存优化带来的训练延迟代价。(2)本文设计了两种算法来求解基于卸载和重计算混合机制的最小化显存优化代价问题。本文设计了基于混合整数规划方法的显存优化算法MIP-MO(Memory Optimization Algorithm Based on Mixed Integer Programming,MIP-MO)。实验结果表明,与相关工作相比,MIP-MO算法能够平均提升DNN训练性能达23.5%。为了解决MIP-MO算法在处理大规模显存优化问题时求解时间长的问题,本文设计了一种基于深度强化学习的显存优化算法DRL-MO(Memory Optimization Algorithm Based on Deep Reinforcement Learning,DRL-MO)。DRL-MO采用了图同构神经网络来提取张量依赖关系,利用动作过滤模块来挑选满足张量依赖约束的操作,利用优势演员-评论家算法实现稳定的强化学习训练。实验结果表明,相比MIP-MO算法,DRL-MO算法开销平均减小79倍。与相关工作相比,DRL-MO算法能够平均提升DNN训练性能达18.1%。(3)现有训练框架未能提供完善的接口,在实现卸载和重计算的机制方面存在一定的限制,本文设计并实现了轻量的深度学习训练平台Ant-NN(Ant Neural Network,Ant-NN)。该平台是基于CUDA和CUDNN实现,支持张量卸载/预取和重计算等灵活的张量管理机制,支持DNN作业类型扩展和显存优化策略扩展。本文将所提出的两种显存优化算法集成于Ant-NN平台,实现了从全卸载、DRL-MO到MIP-MO的渐进式显存优化策略调整。该平台实现了对用户透明的显存优化策略实施,同时满足用户对训练作业快速启动和高效优化的需求。