使用异步执行实现批量 RPC 处理

原文：https://pytorch.org/tutorials/intermediate/rpc_async_execution.html

作者：Shen Li

先决条件：

本教程演示了如何使用@rpc.functions.async_execution装饰器来构建批量 RPC 应用，该装饰器通过减少阻止的 RPC 线程数和合并被调用方上的 CUDA 操作来帮助加快训练速度。 这使用 TorchServer 的相同想法进行批量推断。

注意

本教程需要 PyTorch v1.6.0 或更高版本。

基础知识

先前的教程显示了使用torch.distributed.rpc构建分布式训练应用的步骤，但并未详细说明在处理 RPC 请求时被调用方发生的情况。 从 PyTorch v1.5 开始，每个 RPC 请求都会在被调用方上阻塞一个线程，以在该请求中执行该函数，直到该函数返回为止。 这适用于许多用例，但有一个警告。 如果用户函数例如通过嵌套 RPC 调用在 IO 上阻塞，或者例如在等待其他 RPC 请求解除阻塞的信号时阻塞，则被调用方上的 RPC 线程将必须空闲，直到 IO 完成或发生信令事件为止。 结果，RPC 被调用者可能使用了不必要的更多线程。 造成此问题的原因是 RPC 将用户函数视为黑盒，并且几乎不了解该函数会发生什么。 为了允许用户函数产生和释放 RPC 线程，需要向 RPC 系统提供更多提示。

从 v1.6.0 开始，PyTorch 通过引入两个新概念来解决此问题：

  • torch.futures.Future 类型封装了异步执行，还支持安装回调函数。
  • 一个@rpc.functions.async_execution装饰器，允许应用告诉被调用方目标函数将返回将来的函数，并且在执行期间可以暂停并产生多次。

使用这两个工具，应用代码可以将用户函数分解为多个较小的函数，将它们作为Future对象上的回调链接在一起，然后返回包含最终结果的Future。 在被调用方，当获取Future对象时，它还将安装后续的 RPC 响应准备和通讯作为回调，这将在最终结果准备好时触发。 这样，被调用者不再需要阻塞一个线程并等待直到最终返回值准备就绪。 有关简单示例，请参考@rpc.functions.async_execution 的 API 文档。

除了减少被调用方上的空闲线程数之外，这些工具还有助于使批量 RPC 处理更容易，更快捷。 本教程的以下两节演示了如何使用@rpc.functions.async_execution装饰器来构建分布式批更新参数服务器和批量强化学习应用。

批量更新参数服务器

考虑具有一个参数服务器（PS）和多个训练器的同步参数服务器训练应用。 在此应用中，PS 保留参数并等待所有训练器报告坡度。 在每次迭代中，它都会等到收到所有训练者的梯度后，再一次更新所有参数。 下面的代码显示 PS 类的实现。 update_and_fetch_model方法是用@rpc.functions.async_execution装饰的，将由训练器调用。 每次调用都会返回一个Future对象，该对象将填充有更新的模型。 大多数训练器发起的调用仅将梯度累积到.grad字段，立即返回，并在 PS 上产生 RPC 线程。 最后到达的训练器将触发优化器步骤，并消耗所有先前报告的梯度。 然后，它使用更新的模型设置future_model，该模型又通过Future对象通知其他训练器的所有先前请求，并将更新后的模型发送给所有训练器。

  1. import threading
  2. import torchvision
  3. import torch
  4. import torch.distributed.rpc as rpc
  5. from torch import optim
  7. num_classes, batch_update_size = 30, 5
  9. class BatchUpdateParameterServer(object):
  10.     def __init__(self, batch_update_size=batch_update_size):
  11.         self.model = torchvision.models.resnet50(num_classes=num_classes)
  12.         self.lock = threading.Lock()
  13.         self.future_model = torch.futures.Future()
  14.         self.batch_update_size = batch_update_size
  15.         self.curr_update_size = 0
  16.         self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
  17.         for p in self.model.parameters():
  18.             p.grad = torch.zeros_like(p)
  20.     def get_model(self):
  21.         return self.model
  23.     @staticmethod
  24.     @rpc.functions.async_execution
  25.     def update_and_fetch_model(ps_rref, grads):
  26.         # Using the RRef to retrieve the local PS instance
  27.         self = ps_rref.local_value()
  28.         with self.lock:
  29.             self.curr_update_size += 1
  30.             # accumulate gradients into .grad field
  31.             for p, g in zip(self.model.parameters(), grads):
  32.                 p.grad += g
  34.             # Save the current future_model and return it to make sure the
  35.             # returned Future object holds the correct model even if another
  36.             # thread modifies future_model before this thread returns.
  37.             fut = self.future_model
  39.             if self.curr_update_size >= self.batch_update_size:
  40.                 # update the model
  41.                 for p in self.model.parameters():
  42.                     p.grad /= self.batch_update_size
  43.                 self.curr_update_size = 0
  44.                 self.optimizer.step()
  45.                 self.optimizer.zero_grad()
  46.                 # by settiing the result on the Future object, all previous
  47.                 # requests expecting this updated model will be notified and
  48.                 # the their responses will be sent accordingly.
  49.                 fut.set_result(self.model)
  50.                 self.future_model = torch.futures.Future()
  52.         return fut

对于训练器，它们都使用来自 PS 的相同参数集进行初始化。 在每次迭代中，每位训练器首先进行前进和后退操作，以局部生成梯度。 然后，每个训练器都使用 RPC 向 PS 报告其梯度，并通过同一 RPC 请求的返回值取回更新的参数。 在训练器的实现中，目标函数是否标记有@rpc.functions.async_execution都没有关系。 训练器只需使用rpc_sync调用update_and_fetch_model，这会阻塞训练器，直到返回更新的模型。

  1. batch_size, image_w, image_h  = 20, 64, 64
  3. class Trainer(object):
  4.     def __init__(self, ps_rref):
  5.         self.ps_rref, self.loss_fn = ps_rref, torch.nn.MSELoss()
  6.         self.one_hot_indices = torch.LongTensor(batch_size) \
  7.                                     .random_(0, num_classes) \
  8.                                     .view(batch_size, 1)
  10.     def get_next_batch(self):
  11.         for _ in range(6):
  12.             inputs = torch.randn(batch_size, 3, image_w, image_h)
  13.             labels = torch.zeros(batch_size, num_classes) \
  14.                         .scatter_(1, self.one_hot_indices, 1)
  15.             yield inputs.cuda(), labels.cuda()
  17.     def train(self):
  18.         name = rpc.get_worker_info().name
  19.         # get initial model parameters
  20.         m = self.ps_rref.rpc_sync().get_model().cuda()
  21.         # start training
  22.         for inputs, labels in self.get_next_batch():
  23.             self.loss_fn(m(inputs), labels).backward()
  24.             m = rpc.rpc_sync(
  25.                 self.ps_rref.owner(),
  26.                 BatchUpdateParameterServer.update_and_fetch_model,
  27.                 args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]),
  28.             ).cuda()

在本教程中，我们将跳过启动多个进程的代码，有关完整实现，请参考示例回购。 请注意，可以在没有@rpc.functions.async_execution装饰器的情况下实现批量。 但是，这将需要在 PS 上阻塞更多的 RPC 线程，或者使用另一轮 RPC 来获取更新的模型，后者将增加代码的复杂性和通信开销。

本节使用一个简单的参数服务器训练示例来说明如何使用@rpc.functions.async_execution装饰器实现批量 RPC 应用。 在下一节中，我们将使用批量重新实现上一分布式 RPC 框架入门指南中的强化学习示例，并演示其对训练速度的影响。

批量 CartPole 求解器

本节以 OpenAI Gym 中的 CartPole-v1 为例，说明批量 RPC 的性能影响。 请注意，我们的目标是演示@rpc.functions.async_execution的用法，而不是构建最佳的 CartPole 求解器或解决大多数不同的 RL 问题，我们使用非常简单的策略和奖励计算策略，并将重点放在多观察者单智能体的批量 RPC 实现。 我们使用与前面的教程类似的Policy模型，如下所示。 与上一教程相比，不同之处在于其构造器使用了一个附加的batch参数来控制F.softmax的dim参数，因为进行批量时，forward函数中的x参数包含来自多个观察者的状态，因此尺寸需要适当更改。 其他所有内容保持不变。

  1. import argparse
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  5. parser = argparse.ArgumentParser(description='PyTorch RPC Batch RL example')
  6. parser.add_argument('--gamma', type=float, default=1.0, metavar='G',
  7.                     help='discount factor (default: 1.0)')
  8. parser.add_argument('--seed', type=int, default=543, metavar='N',
  9.                     help='random seed (default: 543)')
  10. parser.add_argument('--num-episode', type=int, default=10, metavar='E',
  11.                     help='number of episodes (default: 10)')
  12. args = parser.parse_args()
  14. torch.manual_seed(args.seed)
  16. class Policy(nn.Module):
  17.     def __init__(self, batch=True):
  18.         super(Policy, self).__init__()
  19.         self.affine1 = nn.Linear(4, 128)
  20.         self.dropout = nn.Dropout(p=0.6)
  21.         self.affine2 = nn.Linear(128, 2)
  22.         self.dim = 2 if batch else 1
  24.     def forward(self, x):
  25.         x = self.affine1(x)
  26.         x = self.dropout(x)
  27.         x = F.relu(x)
  28.         action_scores = self.affine2(x)
  29.         return F.softmax(action_scores, dim=self.dim)

Observer的构造器也会相应地进行调整。 它还带有batch参数，该参数控制用于选择动作的Agent函数。 在批量模式下，它将调用Agent上的select_action_batch函数，该函数将很快出现，并且该函数将以@rpc.functions.async_execution装饰。

  1. import gym
  2. import torch.distributed.rpc as rpc
  4. class Observer:
  5.     def __init__(self, batch=True):
  6.         self.id = rpc.get_worker_info().id - 1
  7.         self.env = gym.make('CartPole-v1')
  8.         self.env.seed(args.seed)
  9.         self.select_action = Agent.select_action_batch if batch else Agent.select_action

与之前的教程分布式 RPC 框架入门相比，观察者的行为略有不同。 它不会在环境停止时退出，而是始终在每个剧集中运行n_steps迭代。 当环境返回时，观察者只需重置环境并重新开始。 通过这种设计，智能体将从每个观察者那里收到固定数量的状态，因此可以将它们打包成固定大小的张量。 在每个步骤中，Observer使用 RPC 将其状态发送到Agent，并通过返回值获取操作。 在每个剧集的结尾，它将所有步骤的奖励返还给Agent。 注意，Agent将使用 RPC 调用此run_episode函数。 因此，此函数中的rpc_sync调用将是嵌套的 RPC 调用。 我们也可以将此函数标记为@rpc.functions.async_execution，以避免阻塞Observer上的一个线程。 但是，由于瓶颈是Agent而不是Observer，因此可以在Observer进程中阻塞一个线程。

  1. import torch
  3. class Observer:
  4.     ...
  6.     def run_episode(self, agent_rref, n_steps):
  7.         state, ep_reward = self.env.reset(), NUM_STEPS
  8.         rewards = torch.zeros(n_steps)
  9.         start_step = 0
  10.         for step in range(n_steps):
  11.             state = torch.from_numpy(state).float().unsqueeze(0)
  12.             # send the state to the agent to get an action
  13.             action = rpc.rpc_sync(
  14.                 agent_rref.owner(),
  15.                 self.select_action,
  16.                 args=(agent_rref, self.id, state)
  17.             )
  19.             # apply the action to the environment, and get the reward
  20.             state, reward, done, _ = self.env.step(action)
  21.             rewards[step] = reward
  23.             if done or step + 1 >= n_steps:
  24.                 curr_rewards = rewards[start_step:(step + 1)]
  25.                 R = 0
  26.                 for i in range(curr_rewards.numel() -1, -1, -1):
  27.                     R = curr_rewards[i] + args.gamma * R
  28.                     curr_rewards[i] = R
  29.                 state = self.env.reset()
  30.                 if start_step == 0:
  31.                     ep_reward = min(ep_reward, step - start_step + 1)
  32.                 start_step = step + 1
  34.         return [rewards, ep_reward]

Agent的构造器还采用batch参数，该参数控制如何对动作概率进行批量。 在批量模式下，saved_log_probs包含一张张量列表，其中每个张量包含一个步骤中所有观察者的动作抢夺。 如果不进行批量，则saved_log_probs是字典，其中的键是观察者 ID，值是该观察者的动作概率列表。

  1. import threading
  2. from torch.distributed.rpc import RRef
  4. class Agent:
  5.     def __init__(self, world_size, batch=True):
  6.         self.ob_rrefs = []
  7.         self.agent_rref = RRef(self)
  8.         self.rewards = {}
  9.         self.policy = Policy(batch).cuda()
  10.         self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2)
  11.         self.running_reward = 0
  13.         for ob_rank in range(1, world_size):
  14.             ob_info = rpc.get_worker_info(OBSERVER_NAME.format(ob_rank))
  15.             self.ob_rrefs.append(rpc.remote(ob_info, Observer, args=(batch,)))
  16.             self.rewards[ob_info.id] = []
  18.         self.states = torch.zeros(len(self.ob_rrefs), 1, 4)
  19.         self.batch = batch
  20.         self.saved_log_probs = [] if batch else {k:[] for k in range(len(self.ob_rrefs))}
  21.         self.future_actions = torch.futures.Future()
  22.         self.lock = threading.Lock()
  23.         self.pending_states = len(self.ob_rrefs)

非批量select_acion只需运行状态抛出策略，保存动作概率，然后立即将动作返回给观察者。

  1. from torch.distributions import Categorical
  3. class Agent:
  4.     ...
  6.     @staticmethod
  7.     def select_action(agent_rref, ob_id, state):
  8.         self = agent_rref.local_value()
  9.         probs = self.policy(state.cuda())
  10.         m = Categorical(probs)
  11.         action = m.sample()
  12.         self.saved_log_probs[ob_id].append(m.log_prob(action))
  13.         return action.item()

使用批量时，状态以观察者 id 为行 ID 存储在 2D 张量self.states中。 然后，它通过将回调函数安装到批量生成的self.future_actions Future对象上来链接Future，该对象将使用使用该观察者 ID 索引的特定行进行填充。 最后到达的观察者一口气通过策略运行所有批量状态，并相应地设置self.future_actions。 发生这种情况时，将触发安装在self.future_actions上的所有回调函数，并使用它们的返回值来填充链接的Future对象，该对象进而通知Agent为所有先前的 RPC 请求准备和传达来自其他观察者的响应。

  1. class Agent:
  2.     ...
  4.     @staticmethod
  5.     @rpc.functions.async_execution
  6.     def select_action_batch(agent_rref, ob_id, state):
  7.         self = agent_rref.local_value()
  8.         self.states[ob_id].copy_(state)
  9.         future_action = self.future_actions.then(
  10.             lambda future_actions: future_actions.wait()[ob_id].item()
  11.         )
  13.         with self.lock:
  14.             self.pending_states -= 1
  15.             if self.pending_states == 0:
  16.                 self.pending_states = len(self.ob_rrefs)
  17.                 probs = self.policy(self.states.cuda())
  18.                 m = Categorical(probs)
  19.                 actions = m.sample()
  20.                 self.saved_log_probs.append(m.log_prob(actions).t()[0])
  21.                 future_actions = self.future_actions
  22.                 self.future_actions = torch.futures.Future()
  23.                 future_actions.set_result(actions.cpu())
  24.         return future_action

现在，让我们定义如何将不同的 RPC 函数结合在一起。 Agent控制每个剧集的执行。 它首先使用rpc_async在所有观察者上开始该剧集，并阻止将由观察者奖励填充的返还期货。 请注意，以下代码使用 RRef 帮助器ob_rref.rpc_async()在具有提供的参数的ob_rref RRef 的所有者上启动run_episode函数。 然后将保存的动作概率和返回的观察者奖励转换为期望的数据格式，并开始训练步骤。 最后，它将重置所有状态并返回当前剧集的奖励。 此函数是运行一集的入口。

  1. class Agent:
  2.     ...
  4.     def run_episode(self, n_steps=0):
  5.         futs = []
  6.         for ob_rref in self.ob_rrefs:
  7.             # make async RPC to kick off an episode on all observers
  8.             futs.append(ob_rref.rpc_async().run_episode(self.agent_rref, n_steps))
  10.         # wait until all obervers have finished this episode
  11.         rets = torch.futures.wait_all(futs)
  12.         rewards = torch.stack([ret[0] for ret in rets]).cuda().t()
  13.         ep_rewards = sum([ret[1] for ret in rets]) / len(rets)
  15.         # stack saved probs into one tensor
  16.         if self.batch:
  17.             probs = torch.stack(self.saved_log_probs)
  18.         else:
  19.             probs = [torch.stack(self.saved_log_probs[i]) for i in range(len(rets))]
  20.             probs = torch.stack(probs)
  22.         policy_loss = -probs * rewards / len(rets)
  23.         policy_loss.sum().backward()
  24.         self.optimizer.step()
  25.         self.optimizer.zero_grad()
  27.         # reset variables
  28.         self.saved_log_probs = [] if self.batch else {k:[] for k in range(len(self.ob_rrefs))}
  29.         self.states = torch.zeros(len(self.ob_rrefs), 1, 4)
  31.         # calculate running rewards
  32.         self.running_reward = 0.5 * ep_rewards + 0.5 * self.running_reward
  33.         return ep_rewards, self.running_reward

其余代码是正常的进程启动和日志记录，与其他 RPC 教程类似。 在本教程中，所有观察者都被动地等待来自智能体的命令。 有关完整的实现，请参考示例回购。

  1. def run_worker(rank, world_size, n_episode, batch, print_log=True):
  2.     os.environ['MASTER_ADDR'] = 'localhost'
  3.     os.environ['MASTER_PORT'] = '29500'
  4.     if rank == 0:
  5.         # rank0 is the agent
  6.         rpc.init_rpc(AGENT_NAME, rank=rank, world_size=world_size)
  8.         agent = Agent(world_size, batch)
  9.         for i_episode in range(n_episode):
  10.             last_reward, running_reward = agent.run_episode(n_steps=NUM_STEPS)
  12.             if print_log:
  13.                 print('Episode {}\tLast reward: {:.2f}\tAverage reward: {:.2f}'.format(
  14.                     i_episode, last_reward, running_reward))
  15.     else:
  16.         # other ranks are the observer
  17.         rpc.init_rpc(OBSERVER_NAME.format(rank), rank=rank, world_size=world_size)
  18.         # observers passively waiting for instructions from agents
  19.     rpc.shutdown()
  21. def main():
  22.     for world_size in range(2, 12):
  23.         delays = []
  24.         for batch in [True, False]:
  25.             tik = time.time()
  26.             mp.spawn(
  27.                 run_worker,
  28.                 args=(world_size, args.num_episode, batch),
  29.                 nprocs=world_size,
  30.                 join=True
  31.             )
  32.             tok = time.time()
  33.             delays.append(tok - tik)
  35.         print(f"{world_size}, {delays[0]}, {delays[1]}")
  37. if __name__ == '__main__':
  38.     main()

批量 RPC 有助于将操作推断合并为较少的 CUDA 操作，从而减少了摊销的开销。 上面的main函数使用不同数量的观察者（从 1 到 10）在批量和无批量模式下运行相同的代码。下图使用默认参数值绘制了不同世界大小的执行时间。 结果证实了我们的期望，即批量有助于加快训练速度。

使用异步执行实现批量 RPC 处理 - 图1

了解更多