[论文解读] TF-Replicator: Distributed Machine Learning for Researchers

Peter Buchlovsky, David Budden+10 more|arXiv (Cornell University)|Feb 1, 2019

Adversarial Robustness in Machine Learning参考文献 47被引用 21

一句话总结

TF-Replicator 是一个基于 TensorFlow 的框架，使研究人员能够通过最小的代码修改，在 CPU、GPU 或 TPU 的分布式集群上轻松扩展深度学习模型，支持数据并行或模型并行训练。它在多种架构上实现了强大的可扩展性——在 ResNet-50、SN-GAN 和 D4PG 模型上得到验证——且无需具备分布式系统专业知识，未来将作为 TensorFlow 2.0 的一部分开源。

ABSTRACT

We describe TF-Replicator, a framework for distributed machine learning designed for DeepMind researchers and implemented as an abstraction over TensorFlow. TF-Replicator simplifies writing data-parallel and model-parallel research code. The same models can be effortlessly deployed to different cluster architectures (i.e. one or many machines containing CPUs, GPUs or TPU accelerators) using synchronous or asynchronous training regimes. To demonstrate the generality and scalability of TF-Replicator, we implement and benchmark three very different models: (1) A ResNet-50 for ImageNet classification, (2) a SN-GAN for class-conditional ImageNet image generation, and (3) a D4PG reinforcement learning agent for continuous control. Our results show strong scalability performance without demanding any distributed systems expertise of the user. The TF-Replicator programming model will be open-sourced as part of TensorFlow 2.0 (see https://github.com/tensorflow/community/pull/25).

研究动机与目标

• 解决深度学习研究中对可扩展、通用的分布式训练日益增长的需求。
• 通过抽象低层分布式系统问题，降低实现分布式训练的复杂性。
• 实现模型在不同硬件（CPU、GPU、TPU）和训练模式（同步/异步）之间的无缝部署。
• 支持复杂的非标准训练模式，如多损失优化和强化学习智能体。
• 通过允许研究人员在不学习新工具的情况下，将现有单机模型快速扩展到分布式集群，从而加速研究迭代。

提出的方法

• TF-Replicator 提供高层 API，抽象了 TensorFlow 的分布式执行机制，使用户只需定义一次模型，即可在异构集群上部署。
• 它采用基于副本的编程模型，每个副本在工作节点上运行模型的副本，并通过参数服务器或集体通信（如 NCCL）自动同步梯度。
• 该框架支持数据并行和模型并行训练，透明地处理设备分配，并在跨机器的设备之间自动插入通信操作。
• 它与 TensorFlow 2.0 的即时执行和 Keras 风格 API 集成，支持自然的控制流和灵活的训练循环。
• 通过可配置的设备分配和通信策略，支持混合并行性，允许用户结合数据并行和模型并行。
• 系统通过 ClusterSpec 自动处理集群配置，并将逻辑设备映射到物理硬件，实现单机与多机部署之间的可移植性。

实验结果

研究问题

• RQ1 高层抽象是否能在不依赖深层分布式系统知识的前提下，简化研究人员的分布式训练？
• RQ2 该框架在不同模型架构（如 CNN、GAN、RL 智能体）和硬件（GPU、TPU）上的可扩展性如何？
• RQ3 TF-Replicator 是否能够支持复杂的训练模式，如多损失优化和分层强化学习？
• RQ4 在可扩展性和训练速度方面，其性能与手工优化的分布式系统相比如何？
• RQ5 研究人员在多大程度上可以仅通过最小修改 reuse 单机代码实现分布式训练？

主要发现

• TF-Replicator 在 ImageNet 上的 ResNet-50 模型中实现了强大的可扩展性，跨 GPU 和 TPU 集群达到了具有竞争力的训练速度。
• 用于类别条件图像生成的 SN-GAN 模型在 TF-Replicator 上实现了有效扩展，证明了该框架对复杂生成模型的支持能力。
• D4PG 强化学习智能体在训练时间上实现了显著加速，8 块 V100 GPU 的性能与单个 TPUv2 设备（8 核）相当。
• 单个 TPUv2 设备的性能可与 8 块通过 NVLink 连接的 V100 GPU 相媲美，凸显了该框架在通信和设备管理方面的高效性。
• 该框架使研究人员能够以极小的代码修改，在多种集群架构上扩展模型，同时保持高吞吐量和低通信开销。
• TF-Replicator 在无需低层分布式系统编程的情况下，成功用于训练包含复杂训练循环（如多损失和分层 RL）的模型。