DeepMind又放福利：开源了一个内部的分布式机器学习库TF-Replicator

今天，DeepMind又放福利：开源了一个内部的分布式机器学习库TF-Replicator，可以帮助研究人员将TensorFlow模型轻松部署到GPU、TPU，并实现不同类型加速器之间的无缝切换。

最近AI领域的突破，从AlphaFold到BigGAN再到AlphaStar，一个反复出现的主题是，对方便、可靠的可扩展性的需求。

研究人员已经能够获取越来越多计算能力，得以训练更大的神经网络，然而，将模型扩展到多个设备并不是一件容易的事情。

今天，DeepMind又将其内部一个秘密武器公之于众——TF-Replicator，一个可以帮助研究人员将他们的TensorFlow模型轻松部署到GPU、Cloud TPU的分布式机器学习框架，即使他们之前完全没有使用分布式系统的经验。

TF-Replicator由DeepMind的研究平台团队开发，初衷是为DeepMind的研究人员提供一个简单的接入TPU的API，现在，TF-Replicator已经是DeepMind内部最广泛使用的TPU编程接口。

TF-Replicator允许研究人员针对机器学习定位不同的硬件加速器进行，将工作负载扩展到许多设备，并在不同类型的加速器之间无缝切换。

虽然它最初是作为TensorFlow上面的一个库开发的，但目前TF-Replicator的API已经集成到TensorFlow 2.0新的tf.distribute.Strategy中，作为 tf.distribute.Strategy的一部分开源：

https://www.tensorflow.org/alpha/guide/distribute_strategy

团队还公开了相关论文：TF-Replicator: Distributed Machine Learning for Researchers，全面描述了这个新框架的技术细节。

https://arxiv.org/abs/1902.00465

接下来，我们将介绍TF-Replicator背后的想法和技术挑战。

构建一个分布式机器学习库

虽然TensorFlow为CPU、GPU和TPU设备都提供了直接支持，但是在目标之间切换需要用户付出大量的努力。这通常涉及为特定的硬件目标专门编写代码，将研究想法限制在该平台的功能上。

一些构建在TensorFlow之上的现有框架，例如Estimators，已经试图解决这个问题。然而，它们通常针对生产用例，缺乏快速迭代研究思路所需的表达性和灵活性。

我们开发TF-Replicator的初衷是为DeepMind的研究人员提供一个使用TPU的简单API。TPU为机器学习工作负载提供了可扩展性，实现了许多研究突破，例如使用我们的BigGAN模型实现了最先进的图像合成。

TensorFlow针对TPU的原生API与针对GPU的方式不同，这造成了使用TPU的障碍。TF-Replicator提供了一个更简单、更用户友好的API，隐藏了TensorFlow的TPU API的复杂性。此外，研究平台团队与不同机器学习领域的研究人员密切合作，开发了TF-Replicator API，以确保必要的灵活性和易用性。

TF-Replicator API

使用TF-Replicator编写的代码与使用TensorFlow中为单个设备编写的代码类似，允许用户自由定义自己的模型运行循环。用户只需要定义(1)一个公开数据集的输入函数，以及(2)一个定义其模型逻辑的step函数(例如，梯度下降的单个step):

# Deploying a model with TpuReplicator.repl = tf_replicator.TpuReplicator( num_workers=1, num_tpu_cores_per_worker=8)with repl.context(): model = resnet_model() base_optimizer = tf.train.AdamOptimizer() optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)train_op = tf.reduce_mean(per_replica_loss)with tf.train.MonitoredSession() as session: repl.init(session) for i in xrange(num_train_steps): session.run(train_op) repl.shutdown(session)

将计算扩展到多个设备需要设备之间进行通信。在训练机器学习模型的背景下，最常见的通信形式是累积梯度(accumulate gradients)以用于优化算法，如随机梯度下降。

因此，我们提供了一种方便的方法来封装TensorFlow Optimizers，以便在更新模型参数之前在设备之间累积梯度。对于更一般的通信模式，我们提供了类似于MPI的原语，如“all_reduce”和“broadcast”。这些使得实现诸如全局批标准化之类的操作变得非常简单，这是扩展BigGAN模型训练的关键技术。

输入数据从主机发送到各个GPU, GPU立即开始处理。当需要在GPU之间交换信息时，它们会在发送数据之前进行同步。

实现

对于多GPU计算，TF-Replicator依赖于“图内复制”(“in-graph replication)模式，其中每个设备的计算在同一个TensorFlow graph中复制。设备之间的通信是通过连接设备对应子图中的节点来实现的。在TF-Replicator中实现这一点很具挑战性，因为在TensorFlow graph中的任何位置都可能发生通信。因此，构造计算的顺序至关重要。

我们的第一个想法是在一个单独的Python线程中同时构建每个设备的子图。当遇到通信原语时，线程同步，主线程插入所需的跨设备计算。之后，每个线程将继续构建其设备的计算。

然而，在我们考虑这种方法时，TensorFlow的图形构建API不是线程安全的，这使得在不同线程中同时构建子图非常困难。相反，我们使用图形重写(graph rewriting)在所有设备的子图构建完成后插入通信。在构造子图时，占位符被插入到需要通信的位置。然后，我们跨设备收集所有匹配占位符，并用适当的跨设备计算替换它们。

当TF-Replicator构建一个in-graph replicated计算时，它首先独立地为每个设备构建计算，并将占位符留给用户指定的跨设备计算。构建好所有设备的子图之后，TF-Replicator通过用实际的跨设备计算替换占位符来连接它们。

为AI研究构建一个平台

通过在TF-Replicator的设计和实现过程中与研究人员密切合作，我们最终构建一个库，让用户能够轻松地跨多个硬件加速器进行大规模计算，同时让他们拥有进行前沿AI研究所需的控制和灵活性。

例如，在与研究人员讨论之后，我们添加了MPI风格的通信原语，如all-reduce。TF-Replicator和其他共享基础架构使我们能够在稳健的基础上构建越来越复杂的实验，并在整个DeepMind快速传播最佳实践。

在撰写本文时，TF-Replicator已经成为DeepMind应用最广泛的TPU编程接口。虽然这个库本身并不局限于训练神经网络，但它最常用来训练大量数据。例如，BigGAN模型是在一个512核的TPUv3 pod训练的，batch size为2048。

在采用分布式actor-learner设置的增强学习智能体中，例如我们的重要性加权actor-learner架构，可扩展性是通过让许多actor通过与环境的交互生成新的体验来实现的。然后，learner对这些数据进行处理，以改进agent的策略，表示为一个神经网络。为了应对越来越多的actor，TF-Replicator可以很轻松地将learner分布在多个硬件加速器上。

这些以及更多例子在我们的arXiv论文中有更详细的描述。

Blog:

https://deepmind.com/blog/tf-replicator-distributed-machine-learning/

Paper：

https://arxiv.org/abs/1902.00465