[論文レビュー] 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、分散クラスタを含む異種システム上で、最小限のパフォーマンスオーバーヘッドで効率的な実行を可能にする。
• 正しく性質を保ちつつ、不変性と再現性を維持したまま、イン-place 操作とメモリ再利用によるメモリフットプリントの削減を実現する。
• 最小限のコード変更で複数マシンにまたがる分散トレーニングをシームレスにサポートし、高いスケーラビリティを達成する。

提案手法

• シンボリック計算グラフと命令型テンソル演算の両方をスケジューリングする統合エンジンを用い、すべての演算におけるデータおよび依存関係の流れを追跡する。
• すべての計算ユニット（NDArray、乱数生成器、メモリ領域）をタグ付きリソースとして表現することで、細粒度の依存関係追跡とスケジューリングを可能にする。
• ミューテーション操作を第一級リソースとしてサポートし、イン-place 更新を可能にし、効率的なパラメータ更新とメモリ再利用を実現する。
• 分散トレーニング用に二段階のパラメータサーバー・アーキテクチャ（KVStore）を実装し、マシン内およびマシン間の同期レイヤーを備えることで帯域幅を削減し、異なる一貫性モデルをサポートする。
• エンジンを介して通信と計算のスケジューリングを統合し、デバイスおよびマシン間でのデータの一貫性とシームレスな調整を保証する。
• イン-place 操作やコalesced メモリ割り当てを含む、能動的なメモリ最適化技術を適用することで、トレーニングおよび推論中のメモリ使用量を最小限に抑える。

実験結果

リサーチクエスチョン

• RQ1 シンボリックおよび命令型プログラミングパラダイムを、最適化の機会と開発者の柔軟性の両方を最大化する形で、1 つの機械学習システムに統合するにはどうすればよいか？
• RQ2 モバイルデバイスからマルチGPUクラスタまで、異種ハードウェア上でディープラーニングモデルの効率的実行を実現するための、システムレベルの設計パターンとは何か？
• RQ3 パフォーマンスや正しさを損なわずに、ディープラーニングトレーニングにおけるメモリフットプリントをどの程度まで削減できるか？
• RQ4 計算と通信のスケジューリング統合が、分散ディープラーニングのスケーラビリティおよび収束に与える影響はいかほどか？
• RQ5 統合エンジンは、シンボリック計算グラフと命令型テンソル演算の両方を管理しつつ、高いパフォーマンスと低オーバーヘッドを維持できるか？

主な発見

• MXNet は、1 GPU での実行において、Torch7 や Caffe と同等の原始的パフォーマンスを達成しており、TensorFlow は古い CUDNN バージョンを使用しているため約 2 倍遅い。
• イン-place 操作とコalesced メモリ割り当てにより、推論時で最大 4 倍、トレーニング時で 2 倍のメモリ使用量削減が達成され、VGGNet ではトレーニングに 16MB 未満の追加メモリで十分である。
• 10 台のマシンを用いた分散トレーニングでは、1 台のマシンと比較して平均的なデータパス時間が 10 倍速くなり（1 台で 14K 秒、10 台で 1.4K 秒）、ILSVRC12 データセットで超線形スケーラビリティを示している。
• 分散環境では、10 回のデータパス後に単一マシンでのトレーニングよりも収束が速く、効果的な通信および同期戦略が有効であることが示された。
• 統合エンジンにより、命令型演算とシンボリック計算のシームレスな統合が可能となり、パフォーマンスコストを負担せずに、同じプログラム内で両方のパラダイムを利用できるようになった。
• 二段階の KVStore アーキテクチャにより、マシン間通信の帯域幅が削減され、マシン内およびマシン間の同期に異なる一貫性モデルをサポートし、システムの効率性と信頼性が向上した。