TF-Replicator：研究人员的分布式机器学习
- 人工智能教育在 DeepMind，研究平台团队构建基础架构以增强和加速我们的 AI 研究。
在 DeepMind，研究平台团队构建基础架构以增强和加速我们的 AI 研究。今天，我们很高兴分享我们如何开发 TF-Replicator，这是一个软件库，可帮助研究人员以最小的努力和以前没有分布式系统的经验，在 GPU 和云 TPU 上部署他们的 TensorFlow 模型。 TF-Replicator 的编程模型现已作为 TensorFlow 的 tf.distribute.Strategy 的一部分开源。此博客文章概述了 TF-Replicator 背后的想法和技术挑战。有关更全面的描述，请阅读我们的 arXiv 论文。
近期 AI 突破的一个反复出现的主题 - 从 AlphaFold 到 BigGAN 再到 AlphaStar-- 需要轻松可靠的可扩展性。越来越多的计算能力使研究人员能够利用新功能培养更大的神经网络。为了解决这个问题，研究平台团队开发了 TF-Replicator，它允许研究人员针对机器学习定位不同的硬件加速器，将工作负载扩展到许多设备，并在不同类型的加速器之间无缝切换。虽然它最初是作为 TensorFlow 上的库开发的，但 TF-Replicator 的 API 已经集成到 TensorFlow 2.0 的新 tf.distribute.Strategy 中。
虽然 TensorFlow 为 CPU，GPU 和 TPU（Tensor Processing Unit）设备提供直接支持，但在目标之间切换需要用户付出巨大努力。这通常涉及为特定硬件目标专门化代码，将研究思想限制在该平台的功能上。一些现有的框架构建在 TensorFlow 之上，例如估算师，寻求解决这个问题。但是，它们通常针对生产用例，缺乏快速迭代研究思路所需的表现力和灵活性。
构建分布式机器学习库
我们开发 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）定义模型逻辑的步骤函数（例如梯度下降的单步）：
＃使用repl.context（）部署模型，使用TpuReplicator.repl = tf_replicator.TpuReplicator（num_workers = 1，num_tpu_cores_per_worker = 8）：model = resnet_model（）base_optimizer = tf。 train.AdamOptimizer（）optimizer = repl.wrap_optimizer（base_optimizer）＃...用tf.train.MonitoredSession定义副本input_fn和step_fn.per_replica_loss = repl.run（step_fn，input_fn）train_op = tf.reduce_mean（per_replica_loss）的代码）作为会话：在xrange中的i的repl.init（会话）（num_train_steps）：session.run（train_op）repl.shutdown（会话）
将计算扩展到多个设备需要设备彼此通信。在训练机器学习模型的背景下，最常见的通信形式是累积梯度以用于优化算法，例如随机梯度下降。因此，我们提供了一种方便的方法来包装 TensorFlow 优化器，以便在更新模型的参数之前在设备之间累积渐变。对于更一般的通信模式，我们提供类似 MPI 的原语，例如 `all_reduce` 和 `broadcast`。这些使得实现全局批量标准化等操作变得微不足道，这种技术对于扩展 BigGAN 模型的培训至关重要（参见本文第 3 节）。
全屏 fullscreen_mobile 输入数据从主机发送到每个 GPU，立即开始处理。当需要在 GPU 之间交换信息时，它们会在发送数据之前进行同步。
实现
对于多 GPU 计算，TF-Replicator 依赖于 “图中” 复制“模式，其中每个设备的计算在同一 TensorFlow 图中复制。通过连接来自设备的相应子图的节点来实现设备之间的通信。在 TF-Replicator 中实现这一点具有挑战性，因为通信可以在数据流图中的任何点进行。因此，计算构造的顺序至关重要。
我们的第一个想法是在一个单独的 Python 线程中同时构建每个设备的子图。遇到通信原语时，线程同步，主线程插入所需的跨设备计算。之后，每个线程将继续构建其设备的计算。但是，当我们考虑这种方法时，TensorFlow 的图形构建 API 不是线程安全的，这使得在不同线程中同时构建子图非常困难。相反，我们使用图形重写来在构建所有设备的子图之后插入通信。构建子图时，占位符将插入需要通信的位置。然后，我们跨设备收集所有匹配的占位符，并用适当的跨设备计算替换它们。
全屏 fullscreen_mobile 当 TF-Replicator 构建图中复制计算时，它首先独立地为每个设备构建计算，并留下占位符，用户指定了跨设备计算。一旦构建了所有设备的子图，TF-Replicator 就会通过用实际的跨设备计算替换占位符来连接它们。
在 DeepMind 上构建 AI 研究平台
通过在 TF-Replicator 的设计和实施过程中与研究人员密切合作，我们能够构建一个库，允许用户轻松地在多个硬件加速器上进行计算，同时使他们具备切割所需的控制和灵活性 - 人工智能研究。例如，我们在与研究人员讨论后添加了 MPI 风格的通信原语，例如 all-reduce。 TF-Replicator 和其他共享基础架构使我们能够在强大的基础上构建越来越复杂的实验，并在整个 DeepMind 中快速传播最佳实践。
在撰写本文时，TF-Replicator 是 DeepMind 中使用最广泛的 TPU 编程接口。虽然图书馆本身并不局限于训练神经网络，但它最常用于大批量数据的训练。例如，BigGAN 模型在尺寸为 2048 的批次上训练，最多可达到 TPUv3 吊舱的 512 个核心。在具有分布式演员 - 学习者设置的强化学习代理中，例如我们的重要性加权的演员 - 学习者架构，通过让许多演员通过与环境交互来产生新体验来实现可伸缩性。然后，学习者处理该数据以改进代理的策略，表示为神经网络。为了应对越来越多的参与者，TF-Replicator 可以用来轻松地将学习者分配到许多硬件加速器上。这些和其他示例在我们的 arXiv 论文中有更详细的描述。
TF-Replicator 只是 DeepMind 研究平台团队构建的众多有影响力技术的例子之一。从 AlphaGo 到 AlphaStar，DeepMind 在人工智能方面的许多突破都是由团队实现的。如果您分享我们的使命并对加速最先进的人工智能研究感到兴奋，请通过 https://deepmind.com/careers 了解研究平台中的开放式软件工程职位（机器学习经验是可选的这些角色）。
这项工作由 DeepMind 的研究平台团队完成。我们要感谢 Frederic Besse，Fabio Viola，John Aslanides，Andy Brock，Aidan Clark，SergioGómezColmenarejo，Karen Simonyan，Sander Dieleman，Lasse Espeholt，Akihiro Matsukawa，Tim Harley，Jean-Baptiste Lespiau，Koray Kavukcuoglu，Dan Belov 和 DeepMind 的许多其他人在 TF-Replicator 的整个开发过程中提供了宝贵的反馈。我们还要感谢 Priya Gupta，Jonathan Hseu，Josh Levenberg，Martin Wicke 和 Google 的其他人，他们将这些想法作为 tf.distribute.Strategy 的一部分提供给所有 TensorFlow 用户。
全文完