Source-linked AI summary
MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, Zheng Zhang
TL;DR
MXNet addresses the challenge of building increasingly large and complex machine-learning systems by combining declarative symbolic expressions with imperative tensor computation. It unifies these interfaces through a dependency-aware backend, memory-management strategies, and distributed synchronization, with encouraging preliminary experiments on distributed training.
Problem
The increasing scale and complexity of machine-learning algorithms make efficient system design and implementation increasingly challenging.
Method
MXNet combines symbolic computation graphs, imperative tensor operations, dependency-aware scheduling, memory reuse, and distributed KVStore synchronization in a multi-language library.
Results
Distributed training initially converges more slowly but outperforms single-machine training after 10 data passes, with a reported super-linear speedup.
Takeaways & Limitations
MXNet combines symbolic expression with tensor computation while supporting lightweight, multi-language, distributed execution and encouraging experimental results.
Abstract
from arXiv · showhide
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.
1 Introduction
MXNet addresses the growing scale and complexity of machine learning by combining declarative computation structure with imperative flexibility across multiple host languages. Its unified backend also targets efficient execution, memory use, and distributed operation.
- Deep neural networks increasingly require billions of floating-point operations per sample, creating challenges for ML system design and implementation.
- Declarative programming exposes global computation graphs for optimization, while imperative programming provides flexibility for parameter updates and interactive debugging.
- MXNet embeds its interface in C++, Python, R, Go, and Julia.
- A unified backend fuses symbolic graphs and imperative operations, tracks dependencies, schedules them jointly, and reuses memory through in-place updates.
- MXNet combines declarative optimization with imperative tensor operations while supporting GPU clusters and a lightweight, multi-language implementation.
2 Programming Interface
MXNet exposes symbolic graphs, imperative tensors, and distributed synchronization through interfaces embedded in host languages. Lazy evaluation lets the backend coordinate these components while preserving efficient execution.
- 2.1 Symbol: Declarative Symbolic Expressions: Symbols define computation graphs from composable operators, including matrix operations and neural-network layers, with free or derived variables.
- 2.1 Symbol: Declarative Symbolic Expressions: Bound symbols support forward evaluation, automatic symbolic differentiation, and utilities for loading, saving, memory estimation, and visualization.
- 2.2 NDArray: Imperative Tensor Computation: NDArray provides imperative tensor computation that bridges declarative symbolic expressions and host-language operations, including GPU matrix multiplication.
- 2.2 NDArray: Imperative Tensor Computation: Lazy evaluation lets the backend resolve dependencies between NDArray operations and symbolic expressions, making mixed implementations as efficient as a single symbolic expression.
- 2.3 KVStore: Data Synchronization Over Devices: KVStore synchronizes data across devices through push and pull primitives, user-defined updates, and sequential or eventual consistency.
- 2.3 KVStore: Data Synchronization Over Devices: Lazy evaluation schedules distributed data pushes and pulls through the backend engine, giving the mixed implementation the same performance as a single declarative program.
3 Implementation
MXNet combines computation graphs, imperative operations, dependency-aware scheduling, memory reuse, and distributed communication in one backend engine. Its implementation uses graph optimizations, linear-time allocation heuristics, multi-resource scheduling, mutation tracking, and hierarchical synchronization.
- MXNet transforms bound computation graphs before evaluation to optimize execution and allocate internal-variable memory.
- Graph Optimization: Graph optimization prunes unnecessary subgraphs, groups operators, and uses manually optimized large operations.
- Memory Allocation: An ideal memory-allocation strategy requires O(n^2) time, where n is the number of variables.
- Memory Allocation: MXNet proposes inplace and co-share memory-allocation heuristics with linear time complexity.
- Dependency Engine: The engine schedules computation and communication operations across CPUs, GPUs, and memory or PCIe resources once dependencies resolve.
- Dependency Engine: The engine tracks mutations as resource writes, enabling imperative array updates, parameter-memory reuse, and reproducible random operations.
- Data Communication: KVStore uses level-1 servers for within-machine synchronization and level-2 servers for inter-machine synchronization.
- Data Communication: Aggregating outbound level-1 data reduces bandwidth requirements, while intra- and inter-machine synchronization may use different consistency models.
4 Evaluation
MXNet is evaluated for raw performance, memory usage, and distributed scalability. It matches comparable systems on a single GPU, reduces internal memory usage through allocation strategies, and achieves super-linear distributed speedup.
- Raw performance: MXNet has performance similar to Torch7 and Caffe on single-GPU forward-backward benchmarks, while TensorFlow is consistently 2x slower.Experiments used batch size 32 on one Nvidia GTX 980 card.
- Memory usage: Combining inplace and co-share reduces memory footprint 2x during training and up to 4x during prediction.Even VGG training requires less than 16MB of extra memory.
- Scalability: Distributed GoogLeNet training converges more slowly initially than single-machine training but outperforms it after 10 data passes.
- Scalability: The average cost of a data pass is 14K seconds on one machine and 1.4K seconds on 10 machines, revealing super-linear speedup.
5 Conclusion
MXNet combines symbolic expression with tensor computation in a lightweight, multi-language library that supports distributed settings. The authors report encouraging preliminary experiments and release the code.
- MXNet combines symbolic expression with tensor computation to maximize efficiency and flexibility.
- The library is lightweight, embeds in multiple host languages, and runs in distributed settings.
- The authors describe the experimental results as encouraging and state that the code is available online.