[论文解读] AMPNet: Asynchronous Model-Parallel Training for Dynamic Neural Networks

Alexander L. Gaunt, Matthew Johnson+5 more|arXiv (Cornell University)|May 27, 2017

Advanced Neural Network Applications参考文献 30被引用 29

一句话总结

AMPNet 提出了一种用于具有实例相关控制流的动态神经网络的异步模型并行训练算法，通过允许在无同步的情况下进行梯度更新，实现了高硬件利用率和更快的收敛速度。它通过利用稀疏性和通过新型中间表示（IR）实现的动态计算图，在 QM9 数据集上相比优化后的 CPU/TensorFlow GPU 实现获得了 9 倍的加速。

ABSTRACT

New types of machine learning hardware in development and entering the market hold the promise of revolutionizing deep learning in a manner as profound as GPUs. However, existing software frameworks and training algorithms for deep learning have yet to evolve to fully leverage the capability of the new wave of silicon. We already see the limitations of existing algorithms for models that exploit structured input via complex and instance-dependent control flow, which prohibits minibatching. We present an asynchronous model-parallel (AMP) training algorithm that is specifically motivated by training on networks of interconnected devices. Through an implementation on multi-core CPUs, we show that AMP training converges to the same accuracy as conventional synchronous training algorithms in a similar number of epochs, but utilizes the available hardware more efficiently even for small minibatch sizes, resulting in significantly shorter overall training times. Our framework opens the door for scaling up a new class of deep learning models that cannot be efficiently trained today.

研究动机与目标

• 解决现有训练框架在处理具有不规则、实例相关控制流的动态神经网络时的低效问题。
• 提升由于结构多样性导致小批量处理不切实际的模型训练中的硬件利用率。
• 实现在 CPU 和 FPGA 等专用硬件等互联设备上的动态模型可扩展、分布式训练。
• 设计一种原生支持动态控制流和异步执行的中间表示（IR），且无需控制依赖。
• 证明异步模型并行训练可收敛至与同步方法相当的精度，同时显著减少训练时间。

提出的方法

• 提出一种新型中间表示（IR），包含显式的分支和汇合控制流构造，支持动态计算的静态图编码。
• 将算法状态封装在消息中，使节点能够无控制依赖或集中调度地异步执行。
• 通过在消息负载中嵌入循环计数器、边 ID 和其他控制变量，使用有状态节点管理这些变量。
• 通过支持流水线模型并行化，使不同实例的前向和反向传播可独立进行，并支持异步参数更新。
• 在 IR 中支持数据聚合和副本编码，以恢复批处理效果并支持数据并行性。
• 在多核 CPU 上部署运行时，可并行处理多个实例，最大限度减少空闲时间和 DRAM 带宽瓶颈。

实验结果

研究问题

• RQ1 异步模型并行训练能否在提升动态神经网络硬件利用率的同时，实现与同步方法相当的收敛精度？
• RQ2 如何设计一种中间表示，使其原生表达动态控制流并支持高效的分布式执行？
• RQ3 异步训练在多大程度上可以缓解动态模型中流水线填充度与参数更新频率之间的权衡？
• RQ4 所提出的 IR 和运行时能否在 CPU 和 FPGA 等专用硬件上，为具有不规则计算的模型实现显著加速？
• RQ5 AMPNet 在小批量尺寸场景下与成熟框架（如 TensorFlow）相比性能如何？

主要发现

• 在 QM9 数据集上，AMPNet 使用门控图序列神经网络（GGSNN）实现的训练速度比手动优化的 TensorFlow CPU 实现快 9 倍。
• 尽管 GPU 的峰值 FLOPS 更高，AMPNet 在 TitanX GPU 上仍比 TensorFlow GPU 实现快 2.1 倍，这是由于更好地利用了稀疏性。
• 该框架收敛至与同步训练相似的测试精度，但训练时间显著缩短，即使在小批量尺寸下也成立。
• 一个简单的估算表明，使用 1 TFLOPS FPGA 网络的 AMPNet 在 QM9 上可实现每秒 6,000 个图的处理速度，相比 CPU 运行时提升 10 倍。
• IR 支持训练和推理的无缝集成，并可通过高效的数据聚合与复制恢复批处理优势。
• 该系统表明，通过解耦计算与参数更新，异步模型并行化可有效扩展动态模型，同时保持收敛性。