arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2110.13005cs.LGcs.AIcs.CLcs.DCcs.PF
收藏

AxoNN:面向极致规模深度学习的一种异步、消息驱动的并行框架

AxoNN: An asynchronous, message-driven parallel framework for extreme-scale deep learning

University of Maryland(马里兰大学)

Siddharth Singh, Abhinav Bhatele

AI总结:

AxoNN通过异步消息驱动调度和CPU内存暂存,将GPU内存消耗降低四倍,在大型Transformer模型上实现每GPU吞吐量达理论峰值54.78%,训练时间缩短15-25%。

摘要:

在过去几年中,训练最先进的神经网络所需的内存容量已远远超过现代硬件加速器的DRAM容量。这促使人们开发高效的算法,以便在大型GPU集群上并行训练这些神经网络。由于在现代GPU上计算相对廉价,设计和实现在这些并行训练算法中极其高效的通信对于提取最大性能至关重要。本文介绍了AxoNN,一种并行深度学习框架,它利用异步和消息驱动执行来调度每个GPU上的神经网络操作,从而减少GPU空闲时间并最大化硬件效率。通过将CPU内存用作暂存空间,在训练过程中定期卸载数据,AxoNN能够将GPU内存消耗减少四倍。这使我们能够将每个GPU的参数数量增加四倍,从而减少通信量并将性能提升超过13%。当在48-384个NVIDIA Tesla V100 GPU上测试具有120亿至1000亿参数的大型Transformer模型时,AxoNN实现了每GPU吞吐量达到理论峰值的49.4%-54.78%,并将训练时间减少了22-37天(15%-25%的加速),与最先进方法相比。

英文摘要:

In the last few years, the memory requirements to train state-of-the-art neural networks have far exceeded the DRAM capacities of modern hardware accelerators. This has necessitated the development of efficient algorithms to train these neural networks in parallel on large-scale GPU-based clusters. Since computation is relatively inexpensive on modern GPUs, designing and implementing extremely efficient communication in these parallel training algorithms is critical for extracting the maximum performance. This paper presents AxoNN, a parallel deep learning framework that exploits asynchrony and message-driven execution to schedule neural network operations on each GPU, thereby reducing GPU idle time and maximizing hardware efficiency. By using the CPU memory as a scratch space for offloading data periodically during training, AxoNN is able to reduce GPU memory consumption by four times. This allows us to increase the number of parameters per GPU by four times, thus reducing the amount of communication and increasing performance by over 13%. When tested against large transformer models with 12-100 billion parameters on 48-384 NVIDIA Tesla V100 GPUs, AxoNN achieves a per-GPU throughput of 49.4-54.78% of theoretical peak and reduces the training time by 22-37 days (15-25% speedup) as compared to the state-of-the-art.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑