在 MXNET DAG 中高效嵌入 MPI 集合通信以扩展深度学习
分布式、并行与集群计算
2018-02-21 v1 机器学习
摘要
GPU 和 CPU 集群等高性能计算基础设施的可用性推动了分布式学习系统的发展。深度学习框架将神经网络表示为 DAG，并在 GPU 等计算资源上执行这些 DAG。在本文中，我们提出将 MPI 集合操作高效嵌入数据并行 DAG 的设计方案。不正确的设计很容易导致死锁或程序崩溃。特别地，我们展示了三种设计：Funneled、并发通信以及使用 MPI 集合与 DAG 的依赖链。这些设计通过允许与其他任务并发执行，自动实现计算与通信的重叠。我们直接将这些设计实现到 MXNET 的 KVStore API 中。这使我们能够直接利用其余基础设施。使用 ImageNet 和 CIFAR 数据集，我们展示了我们设计的潜力。特别地，我们的设计可扩展到 256 个 GPU，对于 ImageNet 1K 数据集的 epoch 时间低至 50 秒。
引用
@article{arxiv.1802.06949,
title = {Efficient Embedding of MPI Collectives in MXNET DAGs for scaling Deep Learning},
author = {Amith R Mamidala},
journal= {arXiv preprint arXiv:1802.06949},
year = {2018}
}