明党参

首页 » 常识 » 诊断 » 源码解析PyTorch分布式1
TUhjnbcbe - 2022/3/6 19:28:00

0x00摘要0x01前言1.1先决条件1.2基础知识1.3代码0x02启动2.1总体启动2.2启动参数服务器0x03参数服务器0x04Trainer0x05对比0xFF参考

0x00摘要

在前面的文章之中,我们已经学习了PyTorch分布式的基本模块,接下来我们通过几篇文章来看看如何把这些模块应用到实践之中,顺便把PyTorch分布式逻辑整体梳理一下。本文介绍如何使用异步执行操作来实现批处理RPC,大家可以学习到PyTorch对参数服务器一个新的实现方式。

本文以IMPLEMENTINGBATCHRPCPROCESSINGUSINGASYNCHRONOUSEXECUTIONS的翻译为基础,加入了自己的理解。

注:本文没有完全按照原文顺序进行翻译,而是按照自己理解的思路重新组织了文章。

0x01前言1.1先决条件

本文的先决条件如下:

PyTorch分布式概述

分布式RPC框架入门

使用分布式RPC框架实现参数服务器

RPC异步执行装饰器

本教程演示了如何使用

rpc.functions.async_execution装饰器构建批处理RPC应用程序,这有助于通过减少被阻塞的RPC线程的数量,并且在被调用方整合CUDA操作来加快训练速度。这与使用TorchServer进行批量推理的想法相同。BatchRPC有助于将动作整合到较少的CUDA操作中,从而摊销开销。

注意:本教程需要PyTorchv1.6.0或更高版本。

1.2基础知识

之前的教程已经展示了使用torch.distributed.rpc构建分布式训练应用程序的步骤,但他们没有详细说明在处理RPC请求时被调用方会发生什么。从PyTorchv1.5开始,针对每个RPC请求,被调用者都会启动一个线程来执行该请求中的函数,该线程会阻塞直到该函数返回。这适用于许多用例,但有一个问题:如果用户函数在IO上阻塞,例如使用嵌套的RPC调用或信号(例如等待不同的RPC请求来解除阻塞),则被调用者上的RPC线程将不得不空闲等待,直到IO完成或信号(signal)事件发生。因此,RPC被调用者使用的线程可能会使用比实际需要更多。造成这个问题的原因是RPC把用户函数当成黑盒,对函数中发生的事情知之甚少。为了让用户函数能够让出和释放RPC线程,需要向RPC系统提供更多的提示。

从v1.6.0开始,PyTorch通过引入两个新概念来解决这个问题:

torch.futures.Future封装了一个异步执行,同时也支持安装回调函数。

rpc.functions.async_execution装饰器,它允许应用程序告诉被调用者,本目标函数将返回一个future,并且可以在执行过程中多次暂停和yield。

使用这两个工具,应用程序代码可以将用户函数分解为多个较小的函数,将它们链接在一起作为Future对象的回调方法,并返回包含最终结果的Future给调用者。在被调用方,在获取Future对象时,它也会安装后续的RPC响应处理作为回调方法,这些回调会在最终结果准备好时被触发。这样,被调用者不再需要阻塞一个线程,只是等待最终返回值准备好就行。简单的例子请参考

rpc.functions.async_execution的API文档。

除了减少被调用者的空闲线程数量外,这些工具还使批处理RPC处理更容易、更快。本教程演示了如何使用

rpc.functions.async_execution装饰器构建分布式批量更新参数服务器和批量处理强化学习应用程序。

注:我们不考虑强化学习的领域,那样会影响我们的思路,牵扯精力。

1.3代码

因为原文主要是强化学习代码讲解,而我们只

1
查看完整版本: 源码解析PyTorch分布式1