[论文解读] MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems

Tianqi Chen, Mu Li+8 more|arXiv (Cornell University)|Dec 3, 2015

Advanced Memory and Neural Computing参考文献 9被引用 1,917

一句话总结

MXNet 是一个轻量级、多语言的机器学习库，通过统一符号计算与命令式张量操作，实现异构系统（从移动设备到分布式GPU集群）上的高效执行。它通过统一的引擎调度、原地操作和优化通信，实现高性能与内存效率，在分布式训练中展现出超线性可扩展性，ImageNet 上达到10倍加速。

ABSTRACT

MXNet is a multi-language machine learning (ML) library to ease the development of ML algorithms, especially for deep neural networks. Embedded in the host language, it blends declarative symbolic expression with imperative tensor computation. It offers auto differentiation to derive gradients. MXNet is computation and memory efficient and runs on various heterogeneous systems, ranging from mobile devices to distributed GPU clusters. This paper describes both the API design and the system implementation of MXNet, and explains how embedding of both symbolic expression and tensor operation is handled in a unified fashion. Our preliminary experiments reveal promising results on large scale deep neural network applications using multiple GPU machines.

研究动机与目标

• 解决深度学习工作负载日益增长的复杂性，要求在多样化硬件上实现高效、可扩展的系统。
• 统一符号与命令式编程范式，平衡优化潜力与实现灵活性。
• 在异构系统（包括CPU、GPU与分布式集群）上实现高效执行，性能开销最小化。
• 通过原地操作与内存复用减少内存占用，同时保持正确性与可重现性。
• 通过极少的代码修改，实现跨多台机器的无缝分布式训练，达到高可扩展性。

提出的方法

• 使用统一引擎调度符号计算图与命令式张量操作，追踪所有操作之间的数据与依赖流动。
• 将所有计算单元——NDArray、随机数生成器与内存空间——表示为带标签的资源，以实现细粒度的依赖追踪与调度。
• 将变异操作作为一等资源支持，允许原地更新，从而实现高效的参数更新与内存复用。
• 实现两级参数服务器架构（KVStore），包含机内与机间同步层，以减少带宽并支持不同的一致性模型。
• 通过引擎集成通信与计算调度，确保设备与机器间的数据一致性与无缝协调。
• 应用激进的内存优化技术，包括原地操作与内存合并分配，以最小化训练与推理过程中的内存使用。

实验结果

研究问题

• RQ1 如何在单一机器学习系统中有效统一符号与命令式编程范式，以最大化优化机会与开发者灵活性？
• RQ2 哪些系统级设计模式能够实现从移动设备到多GPU集群的异构硬件上深度学习模型的高效执行？
• RQ3 在不牺牲性能或正确性的前提下，深度学习训练中的内存占用可减少到何种程度？
• RQ4 计算与通信调度的集成如何影响分布式深度学习的可扩展性与收敛性？
• RQ5 统一引擎能否在保持高性能与低开销的前提下，同时管理符号计算图与命令式张量操作？

主要发现

• 在单张GPU上，MXNet的原始性能与Torch7和Caffe相当，而TensorFlow由于使用旧版本CUDNN，性能约为其2倍慢。
• 通过原地操作与合并内存分配，推理阶段内存使用减少高达4倍，训练阶段减少2倍，VGGNet训练时额外内存需求低于16MB。
• 使用MXNet在10台机器上进行分布式训练，平均数据遍历时间从单机的14K秒减少至1.4K秒，实现10倍加速，表明在ILSVRC12数据集上具有超线性可扩展性。
• 在分布式设置中，经过10次数据遍历后，系统收敛速度优于单机训练，表明通信与同步策略有效。
• 统一引擎实现了命令式操作与符号计算的无缝集成，允许开发者在同一程序中使用两种范式，且无性能损失。
• 两级KVStore架构减少了机间通信带宽，支持机内与机间同步的不同一致性模型，提升了系统效率与可靠性。