[论文解读] Addressing the Memory Bottleneck in AI Model Training

David Ojika, Bhavesh Patel+4 more|arXiv (Cornell University)|Mar 11, 2020

Advanced Neural Network Applications参考文献 7被引用 6

一句话总结

本论文展示了在单节点服务器上使用英特尔优化版TensorFlow与第二代英特尔至强可扩展处理器，首次成功训练了内存占用达1 TB的深度神经网络，该服务器配备大容量系统内存。该方法实现了内存密集型人工智能模型的纵向扩展训练，特别是在医学影像领域，克服了传统内存瓶颈，无需使用分布式系统。

ABSTRACT

Using medical imaging as case-study, we demonstrate how Intel-optimized TensorFlow on an x86-based server equipped with 2nd Generation Intel Xeon Scalable Processors with large system memory allows for the training of memory-intensive AI/deep-learning models in a scale-up server configuration. We believe our work represents the first training of a deep neural network having large memory footprint (~ 1 TB) on a single-node server. We recommend this configuration to scientists and researchers who wish to develop large, state-of-the-art AI models but are currently limited by memory.

研究动机与目标

• 解决训练大规模AI模型时的内存瓶颈问题，特别是在医学影像领域。
• 证明配备大容量系统内存的单节点服务器可训练以往需要分布式系统的模型。
• 为基础设施有限的研究人员提供一种实用且可扩展的分布式训练替代方案。
• 验证在单台服务器上训练具有极高内存需求的最先进模型的可行性。
• 为开发大型AI模型的科学家提供软硬件堆栈推荐。

提出的方法

• 在采用第二代英特尔至强可扩展处理器的x86架构服务器上使用英特尔优化版TensorFlow。
• 利用大容量系统内存（最高达1 TB）将整个模型及激活状态存储在内存中。
• 采用纵向扩展架构而非分布式训练，以避免通信开销。
• 优化TensorFlow内部的内存管理和数据流，以最大化可用RAM的利用率。
• 选择医学影像作为案例研究，因为3D卷积网络对内存需求极高。
• 通过充足的内存带宽和NUMA感知内存分配配置服务器，以维持高吞吐量。

实验结果

研究问题

• RQ1 配备1 TB系统内存的单节点服务器能否成功训练内存占用同样达1 TB的深度神经网络？
• RQ2 与分布式训练相比，单服务器上的纵向扩展训练在内存效率和训练吞吐量方面表现如何？
• RQ3 何种软硬件堆栈可实现无需模型并行化的内存密集型模型稳定训练？
• RQ4 英特尔优化版TensorFlow在大规模AI工作负载中能将内存利用率提升到何种程度？
• RQ5 在无集群基础设施的情况下，是否可行将最先进的医学影像模型训练于单台服务器上？

主要发现

• 作者成功在单节点服务器上训练了内存占用达1 TB的深度神经网络，这是首次实现此类演示。
• 该训练通过在配备大容量系统内存的第二代英特尔至强可扩展处理器上使用英特尔优化版TensorFlow实现。
• 纵向扩展配置消除了对模型并行或数据并行的需求，简化了训练设置。
• 该方法实现了复杂医学影像模型的稳定训练，无需分布式系统协调。
• 结果验证了当具备足够内存时，大型最先进的AI模型可在单台服务器上训练。
• 本研究为受基础设施限制的研究人员提供了切实可行的替代方案，提供了一条更简单、更高效的训练路径。