数据并行训练中权重更新的自动跨副本分片
分布式、并行与集群计算
2020-05-05 v1 机器学习
机器学习
摘要
在深度神经网络的数据并行同步训练中，不同设备（副本）运行相同程序但处理训练批次的不同划分，而权重更新计算在所有副本上重复进行，因为权重没有可划分的批次维度。对于具有大权重的典型语言模型，以及具有小每副本批次大小（大规模训练中典型情况）的模型，这可能成为性能与可扩展性的瓶颈。本文提出一种利用静态分析和训练计算图上的变换，通过高效通信原语与数据格式化，在副本间自动分片权重更新计算的方法。我们展示该技术在 Cloud TPU 上对典型图像与语言模型实现了显著加速，且无需修改模型代码。该技术有助于缩小传统上昂贵的（ADAM）与廉价的（SGD）优化器之间的差距，因为它们仅占训练步时间的一小部分且峰值内存使用相似。它帮助我们在 Google 的 MLPerf 0.6 提交中实现了最先进的训练性能。
引用
@article{arxiv.2004.13336,
title = {Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training},
author = {Yuanzhong Xu and HyoukJoong Lee and Dehao Chen and Hongjun Choi and Blake Hechtman and Shibo Wang},
journal= {arXiv preprint arXiv:2004.13336},
year = {2020}
}
备注
12 pages, 23 figures, 1 table