GMI-DRL: Empowering Multi-GPU Deep Reinforcement Learning with GPU Spatial Multiplexing

作者/机构: Yuke Wang, Boyuan Feng, Zheng Wang, †Tong Geng, †Ang Li, and Yufei Ding (University of California, Santa Barbara); †Tong Geng, †Ang Li (Pacific Northwest National Laboratory)

A1 主要贡献

核心问题: 深度强化学习（DRL）具有异构的计算工作负载（物理模拟 vs. GEMM-based神经网络操作）和交错执行的模式，这导致在现代强大的GPU平台上计算效率低下。如最先进的NVIDIA Isaac Gym在A100 GPU上进行PPO训练时，GPU的整体利用率持续低于50%（平均为32%）。此外，随着DRL应用变得更加复杂，需要大量的训练样本，这驱动了在多GPU平台上扩展DRL模拟和策略训练的需求。然而，现有的DRL工作要么只关注单GPU设置，要么简单地将单GPU设计扩展到多GPU平台，缺乏专门的分布式设计和优化，难以充分发挥多GPU系统的潜力。


图 1: (a) 基础DRL计算流程。s: 环境状态向量; a: Agent的动作选择; r: 奖励值; w: 策略模型的权重参数; (b) Isaac Gym在单张A100 GPU上进行PPO训练时的GPU利用率。

研究目标: 本文的核心思想是通过GPU空间多路复用技术（如MPS和MIG），将DRL执行过程中闲置的GPU资源充分利用起来，从而提升性能。GPU空间多路复用技术可以将一个大型GPU划分为一组小型子GPU，用于处理不同任务，从而为计算和通信的设计与优化提供了新的空间。然而，该技术也带来了挑战，例如子GPU之间的“内存屏障”会阻碍DRL中频繁的任务间通信，以及如何为异构的DRL工作负载匹配最合适的子GPU资源配置。

创新点: 为解决上述问题，本文提出了GMI-DRL，一个通过GPU空间多路复用技术加速多GPU深度强化学习的系统性设计。
1. 引入GPU多路复用实例（GMI）概念：GMI-DRL提出了一个统一的、资源可调的子GPU设计，称为GPU多路复用实例（GMI）。这使得能够以更细粒度和资源高效的方式利用GPU，并为设计和优化扩展了空间。
2. 通信高效的设计：GMI-DRL为跨GMI的通信提供了专门支持。它涵盖了多种通信模式，包括为同步DRL训练设计的延迟优化的跨GMI梯度归约，以及为异步DRL训练设计的吞吐量优化的跨GMI经验共享。
3. 计算自适应的管理策略：GMI-DRL具有自适应的GMI管理能力。这包括多个GMI布局模板，以最大化DRL中的计算和通信效率；以及一种新的启发式方法，用于确定GMI的资源预算（如SM和设备内存），以适应异构的DRL工作负载。
4. 工作负载和硬件感知：GMI-DRL提供了用户友好的API来捕获DRL工作负载的关键特征，以指导运行时配置搜索。它还以后端形式支持MPS和MIG，并能根据DRL应用的实际需求（如通信）和硬件规格（如GPU架构）自动决定使用哪一种。
5. 首次探索：GMI-DRL是首个通过结合MPS和MIG的优势来探索GPU空间多路复用在处理具有异构计算和通信混合工作负载（如DRL）方面的巨大潜力的工作。其封装的通用编程、计算和通信范式不仅适用于DRL，还可扩展到其他应用领域。


图 2: GMI-DRL 概览。

A3 背景知识

多进程服务 (MPS)：MPS【26, NVIDIA. Nvidia multi-process service.】是CUDA API的一种替代性、二进制兼容的实现，它可以透明地支持协作式多进程CUDA应用程序。MPS主要由三个组件构成：1) 控制守护进程：管理MPS服务器的启动和停止，并协调MPS客户端与服务器之间的连接。2) 客户端运行时：MPS客户端运行时内置于CUDA驱动库中，任何CUDA应用程序都可以透明地使用。用户定义的应用程序创建的每个进程在打算使用GPU时都会由一个MPS客户端处理。3) 服务器进程：服务器是客户端共享GPU的连接，并提供客户端之间的并发性。为了让应用程序从MPS中受益，必须满足几个基本要求：(i) 应用程序应以多进程并行的方式执行，以及 (ii) 单个进程对GPU资源的利用率应该较低。

多实例GPU (MIG)：MIG【25, Nvidia. Nvidia multi-instance gpu.】是NVIDIA Ampere GPU（如A100）中引入的一种新的GPU空间多路复用解决方案，相比MPS，它能提供更高硬件资源级别的隔离。MIG的主要设计目标是将一个大型GPU物理上划分为多个部分，供不同用户运行多个相互独立的作业/应用程序（如数据处理、神经网络推理和训练），且它们之间没有任何交互。MIG提供多样的GPU实例组合，以资源高效的方式利用GPU。如图3所示，A100 GPU的计算资源被划分为8个分区（7个可用，1个保留）。这有助于轻松地使GPU适应不同的应用目的，例如在同一GPU上共置神经网络训练和推理。


图 3: A100上可用的MIG实例组合。注意，灰色框表示保留的计算/内存空间。“2g.10gb”代表一个拥有GPU总SM数中的2个和10GB设备内存的MIG实例。

MPS与MIG的比较与挑战：表1定性分析了MPS和MIG之间的差异。现有关于MIG和MPS的工作主要集中在处理许多独立作业/应用程序的场景，例如DNN模型推理/服务【5, Seungbeom Choi, et al. Multi-model machine learning inference serving with gpu spatial partitioning. arXiv preprint arXiv:2109.01611, 2021.】【34, Cheng Tan, et al. Serving dnn models with multi-instance gpus: A case of the reconfigurable machine scheduling problem. arXiv preprint arXiv:2109.11067, 2021.】。本文的工作是首次探索利用GPU空间多路复用技术的潜力，以提高单个作业/应用程序（如DRL训练）的性能，该应用程序包含相互紧密交互的异构任务（如模拟、动作预测和策略训练）。本文指出需要解决几个关键问题：1) 如何在GPU实例之间有效共享数据以最小化通信开销；2) 如何有效划分GPU资源，以满足应用程序/作业中异构任务的需求（例如策略推理的延迟和策略训练的吞吐量），并最大化整体系统性能。

表 1: MIG 和 MPS 的比较。注意：“Iso.”：隔离；“QoS”：服务质量；“Com”：不同实例间的通信。

A2 方法细节

GMI-DRL 概览

基于进程的GMI编程范式：GMI-DRL采用基于进程的编程范式构建，以提供更好的服务级别性能保证和减少计算干扰【25, Nvidia. Nvidia multi-instance gpu.】【26, NVIDIA. Nvidia multi-process service.】。不同的进程拥有相互独立的配置（如环境变量）和内存地址空间。

DRL_role基类设计：为方便DRL编程，我们引入了DRL_role基类，如代码清单1所示，其包含几个主要函数。初始化函数__init__会根据GMI_id设置GMI，并将其注册到全局GMI管理器中。它还会根据提供的GPU ID将GMI附加到特定的GPU上。注册后，如果需要进行GMI间通信，GMI还会被分配到一个GMI组。此外，还包括GMI的主要执行例程GMI_run和通信原语（如集体通信和点对点通信），这些函数需要根据DRL_role的具体目的（如模拟器、训练器）进行实现。

 import GMI_DRL # import other packages 
 GMI_DRL.config.num_GPUs=3 # GPU configurations. 
 GMI_DRL.config.GPU_arch="sm_80" # sm_80(A100),sm_70(V100) 
 GMI_DRL.config.enable_NCCL=True # NCCL availability. 
 GMI_DRL.config.GMI_backend="MPS" # GMI backends. 
 # define class for DRL tasks. 
 class DRL_role(object): 
    # Initialize the base environment. 
    def __init__(self, GMI_id, role, gpu_id): 
        self.GMI_id = GMI_id 
        self.role = role 
        self.mgr = GMI_DRL.GMI_manager.add_GMI(GMI_id) 
        self.mgr.set_GPU(GMI_id, gpu_id) 
        self.group = self.mgr.get_group(GMI_id, GMI_id) 
    def GMI_run(self, param1, param2, ...): 
        # major routine of DRL tasks, 
        # e.g., simulators, agents, and trainers. 
    def GMI_collective(self, data): 
        # some data processing work 
        proc_data = proc_fun(data) 
        # allreduce data within a group of GMIs. 
        self.mgr.allreduce(proc_data, self.group) 
    def GMI_send(self, data, dst_GMI_id): 
        # some data processing work 
        proc_data = proc_fun(data) 
        # Asynchronously send data to another GMI. 
        self.mgr.send(proc_data, dst_GMI_id) 
    def GMI_recv(self, src_GMI_id): 
        # Synchronously receive data from another GMI. 
        data = self.mgr.recv(src_GMI_id)

代码清单 1

GMI后端选择：我们根据DRL任务的实际需求和GPU架构来确定GMI的后端（MIG或MPS）。在DRL训练中，我们选择MPS以获得更好的通信效率；而在DRL服务中，我们选择MIG以获得更好的计算性能。在计算能力为70≤sm<80的GPU（如V100）上，只有MPS可用。而在sm=80的GPU（如A100）上，MPS和MIG均可用。更多关于GMI后端的比较将在§6.2中详述。

DRL环境模拟器角色：环境模拟器是生成环境状态和动作奖励最关键的组件。每个模拟器可以容纳成百上千个并发运行的环境，并与Agent进行交互。要在GMI中创建一个DRL模拟器，用户应在init函数中创建并初始化模拟引擎，并指定模拟配置（如环境数量）。模拟器的run函数是一个具有指定上限（如总训练迭代次数）的循环。在每个服务或训练迭代中，模拟器通过接收来自Agent的动作和旧状态向量，并向Agent发送新状态和奖励，来与Agent进行交互。通常，出于性能考虑，我们将环境模拟器与Agent共置在同一个GMI上以便于数据共享（详见§5），因此，跨GMI的通信原语（如集体通信）不会应用于DRL模拟器。

DRL Agent角色：Agent是根据环境状态做出动作决策的关键组件。与模拟器的实现类似，用户将实现一个GMI_run函数，以迭代方式与模拟器进行交互。根据DRL应用的需求，Agent将负责以两种不同方式与其他GMI通信。在同步训练设置中（例如PPO算法【30, John Schulman, et al. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.】），Agent和Trainer共置于同一个GMI中（§5.1），Agent将与Trainer共享收集到的经验向量（由状态、动作和奖励组成），然后等待模型更新完成。在异步DRL训练中，Agent收集的经验将与不同GMI上的Trainer共享。这时会实现send_p2p通信函数，用于经验数据的移动和策略模型的同步。

DRL Trainer角色：Trainer是根据从Agent收集的经验数据来更新策略模型的主要组件。其run函数将是一个循环迭代过程，接收经验并迭代更新模型。在训练期间，Trainer将经验数据输入策略神经网络模型进行前向和后向计算以生成策略梯度。接着，Trainer将依赖集体通信原语（如AllReduce）与其他Trainer同步策略梯度，以维护一个一致的策略模型。最后，每个Trainer将使用同步后的策略梯度来更新策略模型。根据DRL Trainer的底层GMI映射（在同一GPU上或跨GPU），系统会自动确定不同的选项以实现高效的策略模型同步（§4.1）。

专业的GMI通信

通信挑战与解决方案：DRL是一个通信密集型应用，需要在不同组件之间频繁交换信息。然而，基于GMI的设计引入了内存屏障，使通信变得复杂。为此，我们引入了一种专门的GMI通信设计来简化通信工作。我们主要解决DRL训练中的两种主要通信类型：1) 用于同步DRL训练（如PPO【30, John Schulman, et al. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.】）中不同Trainer之间延迟优化的策略梯度归约；2) 用于异步DRL训练（如A3C【19, Volodymyr Mnih, et al. Asynchronous methods for deep reinforcement learning. International conference on machine learning (ICML), 2016.】）中Agent和Trainer之间吞吐量优化的经验共享。由于模拟器和Agent出于性能考虑（§5）被置于同一GMI内，它们之间的通信简单且发生在GMI内部，因此在此省略。异步训练中Agent和Trainer之间的策略参数共享，根据研究其性能影响相较于经验共享非常小（<5%），因此也予以忽略。

布局感知的梯度归约

现有通信原语的局限性：在同步DRL训练期间，为了维护策略模型的一致视图，我们需要一种低延迟开销的高效通信方式。在多GPU平台上的数据并行DNN训练中，已有一系列高度优化的通信原语被开发出来（如NCCL【22, Nvidia. Nvidia collective communication library (nccl). developer.nvidia.com/nccl.】，NVIDIA IPC【24, NVIDIA. Nvidia inter-process communication. docs.nvidia.com/cuda/cuda-samples/index. html#simpleipc.】）。然而，这些原语是为GPU粒度（每个GPU一个进程）的GPU间数据移动量身定制的，不能直接应用于子GPU粒度（每个GPU多个进程）的GMI间通信。这是因为现有的GMI后端设计，如MPS【26, NVIDIA. Nvidia multi-process service.】和MIG【25, Nvidia. Nvidia multi-instance gpu.】，强制执行进程间内存隔离以避免内存安全问题，例如侧信道攻击。因此，即使GMI并发共享同一物理GPU，数据（如张量）也无法在它们之间共享。为此，我们引入了一种新颖的布局感知梯度归约（LGR）策略，包括三种潜在的解决方案。


图 4: 跨GMI策略梯度归约的探索：(a) 多进程归约；(b) 多环归约；(c) 分层归约。为简化起见，图中省略了将结果广播到所有GMI的最后一步。

多进程归约 (MPR)：该方案将GMI间通信视为纯粹的进程间同步问题，如图4(a)所示。它需要首先将策略梯度从GMI移动到CPU，然后在CPU上跨不同进程应用梯度归约。这是处理GMI间梯度归约的基本解决方案，也是处理任何GMI布局的最通用方案，无论其底层GPU布局如何（例如在同一GPU上或跨多个GPU）。然而，在更复杂的GMI布局中（例如跨多个GPU的GMI），此策略有几个弱点：1) 在CPU主机和GPU设备内存之间产生过多的数据流量；2) 闲置了高速的GPU间硬件基础设施，如NVLink和NVSwitch；3) 依赖较慢的CPU进行归约计算。

多环归约 (MRR)：对于具有跨GPU归约需求的GMI布局，利用现有的GPU间通信基础设施（如NCCL通过NVLink进行基于环的归约）将是一个有前景的选择。值得注意的是，虽然GMI阻止了在同一GPU上通过NCCL原语进行同步，但它允许分布在不同GPU上的多个GMI通过NCCL进行通信。如图4(b)所示，不同GPU上的GMI形成了多个不相交的NCCL环进行NVLink通信（步骤1中的红、蓝、紫色曲线）。一旦所有环的归约完成，我们启动另一个NCCL环归约来同步前一阶段多环归约的所有部分结果（步骤2中的绿线）。虽然该策略可以利用高速的GPU间通信，但它受限于每个GPU的GMI数量小于或等于GPU总数的设置。否则，最后的同步NCCL环必须覆盖同一GPU上的多个GMI（即多个环的端点），这将在NCCL中触发“多CUDA流错误”。

分层归约 (HAR)：为了平衡前两种方案的适应性和性能，我们引入了一种分层归约（HAR）技术，以高效地处理基于GMI的同步DRL训练（图4(c)）。其核心思想是在不同通信层面（如GMI间或GPU间）系统地结合不同类型的通信策略（如多进程和多环），以(i)最大化通信并行性，(ii)最小化不必要的数据移动，以及(iii)利用强大的硬件互连和通信库的优势。HGR设计包括两个主要步骤：步骤1是在同一GPU上的Trainer GMI之间同步策略模型梯度。步骤2是跨GPU梯度同步，它将在属于不同GPU的GMI之间执行AllReduce操作。此步骤的关键是从每个GPU中选择一个GMI（“领导者”GMI）参与此GPU间同步。我们通过其ID满足GMI_id % M = t的条件来识别领导者GMI，其中0 ≤ t < M，M是每个GPU的GMI数量。

通信策略选择算法：在实践中，我们根据DRL Trainer的GMI布局来决定使用上述三种选项中的哪一种，如算法1所示。该算法以GMI及其所在GPU的列表作为输入。例如，输入列表MPL = [[0, 1, 2], [3, 4, 5], [6, 7, 8], ..]可解释为ID为{0,1,2}的GMI位于GPU=0上。
算法 1: 通信策略选择

输入: GMI到GPU的映射列表 MPL = [[0, 1, 2], [3, 4, 5], [6, 7, 8], .., ].
输出: 选择的通信策略.
1 GMIperGPU_set = set();
2 /* 当所有GMI在同一个GPU上时 */
3 if len(MPL) ≤ 1 then
4    return “Multi-Process Reduction”;
5 end
6 /* 遍历所有GPU上的所有GMI */
7 for GMI_li in MPL do
8    GMIperGPU_set.add(len(GMI_li));
9 end
10 /* 不同GPU拥有不同数量的GMI */
11 if GMIperGPU_set.size() > 1 then
12    return “Hierarchical Reduction”;
13 end
14 /* 每个GPU的GMI数量大于GPU总数 */
15 if GMIperGPU_set.pop() > len(MPL) then
16    return “Hierarchical Reduction”;
17 end
18 return “Multi-Ring Reduction”;

时间复杂度比较：我们进一步对这三种方法进行了理论时间复杂度比较，如表2所示。其中，t是每个GPU的GMI数量；g是GPU数量；B1是通过进程间通信的GMI间带宽；Mp是策略模型参数的大小；B2是通过NCCL的带宽。

表 2: MPR, MRR, 和 HAR 的时间复杂度。

梯度广播：通过上述三种策略之一完成归约后，我们将通过梯度广播将同步后的梯度刷回到所有GMI。与之前的策略梯度同步步骤（通信和计算交错）相比，梯度广播原语的成本较低，因为它没有操作依赖性，可以轻松地并行化。

基于通道的经验共享

异步训练的通信需求：作为异步DRL训练（例如A3C【19, Volodymyr Mnih, et al. Asynchronous methods for deep reinforcement learning. International conference on machine learning (ICML), 2016.】）的关键组成部分，跨GMI的经验共享连接了Agent GMI和Trainer GMI。与同步DRL训练中的梯度归约相比，经验数据量更大，需要一个以吞吐量为优化的连接和传输策略来实现更高的训练吞吐量性能。我们引入了多通道解决方案来系统地解决经验传输瓶颈。图5(a)(b)(c)中的简要图示从概念上展示了我们设计的优势，例如操作节省、批处理和传输效率。我们的核心设计思想是，经验数据在收集、传输和训练阶段的粒度可以不同，并且可以为最大化每个步骤的性能而进行专门优化。为实现这种可变传输粒度，我们整合了几个组件（图5）。


图 5: 多通道通信。(a) DRL训练所需经验（“Exp_”）的主要组成部分；(b) 单通道/多通道设计中数据通信模式的比较；(c) 单通道/多通道设计中数据通信工作流的比较；(d) 基于通道的经验共享示例（从A2到T0）。注意，“GMI-A”是DRL Agent的GMI，“GMI-T*”是DRL Trainer的GMI；“E0...Ek”是为不同通道（CH）分类的经验数据。

经验分发器 (DP)：该组件作为异步后端服务附加到每个Agent上。其主要任务是对经验数据进行分类，并将其推送到相应的输出通道，例如状态通道和奖励通道。为不同类型的数据创建不同通道的主要原因有：1) 在处理不同粒度和传输速率的数据收集与传输时具有更大的灵活性；2) 更有效地处理同一批经验中不同大小的数据。

经验压缩器 (CP)：这是一个系统级服务，控制所有Agent GMI每个通道的经验共享。它将管理每次数据传输的大小，以便优化不同GMI之间的数据移动，从而实现高吞吐量。CP的一个典型操作是将来自同一通道的多个张量连接起来，以增加每次数据移动的大小。

经验迁移器 (MG)：这是一个系统级服务，管理从Agent到Trainer的经验路由。如果Agent和Trainer GMI位于同一GPU上，MG将直接按通道将数据转发给Trainer。如果Agent和Trainer GMI位于不同GPU上，MG将首先按通道收集GMI间的经验数据，然后按通道将经验分发给工作负载最轻的Trainer。

经验批处理器 (BT)：该组件附加到每个Trainer GMI上。它负责数据准备和模型训练。数据准备中的关键操作包括切片（如果需要小批量）和堆叠（如果需要大批量）。BT可以根据DRL任务的实际需求针对不同目标进行优化。例如，对于复杂的机器人模拟，使用较小的批量以获得更高的模型更新频率。而在其他设置中，如局部运动模拟，可以使用较大的批量来减少数据噪声的影响。

自适应GMI管理

本节将详细介绍我们用于DRL计算的自适应GMI管理策略，包括用于将DRL任务粗粒度映射到GMI的任务感知GMI映射，以及用于微调DRL任务和GMI配置以最大化DRL执行效率的工作负载感知GMI选择。

表 3: 术语及其含义。

任务感知的GMI映射

映射目标：为了找到DRL任务在GMI上的最佳布局，我们引入了一种任务感知的GMI映射方法，该方法仔细考虑了各种DRL任务的特性（如计算和通信）以及硬件平台（如GPU和GPU互连）的特点。我们将这些设计思想融入到我们为DRL服务（计算主导）和DRL训练（计算-通信混合）设计的高效GMI任务模板中。

DRL服务 (DRL Serving)：DRL服务的目标是通过持续的环境-Agent交互来收集足够的经验。它由两个组件组成：环境模拟器和Agent。优化DRL服务的关键是实现环境模拟器和Agent之间每一轮交互的低延迟，以便快速收集更多经验。环境模拟器和Agent之间的交互需要频繁且细粒度的信息交换以实现高服务性能。然而，在基于GMI的设计中，由于现有GPU多路复用技术【25, Nvidia. Nvidia multi-instance gpu.】【26, NVIDIA. Nvidia multi-process service.】强制执行内存隔离，跨GMI的通信会产生不可忽略的成本。因此，不同GMI之间的数据移动（即使它们属于同一GPU）也必须绕道CPU主机内存。这一限制使得为环境模拟器和Agent使用专用GMI的设计方案（即任务专用GMI，TDG）效率低下，因为它在GMI边界间持续共享状态、动作和奖励数据时开销很高。基于此观察，我们提出了DRL服务块（即任务共置GMI，TCG）的概念，其中每个GMI管理环境模拟器和Agent的顺序执行。在这种设计中，环境模拟器生成的状态数据可以通过低成本的GMI内内存访问轻松地与后续的Agent共享以进行动作预测。同样，Agent的动作向量也可以被模拟器轻松检索用于下一轮交互，而无需GMI间的数据移动（COM = 0）。

TCG与TDG的成本分析：为了比较TCG和TDG，我们从资源和通信成本方面进行了分析。我们首先定义了术语（表3），然后在表4中详细比较。首先，我们使用以下公式确定$R_s$, $R_a$和$R_t$的I值：

其中$R_{SM}$和$R_{Mem}$是在单个DRL服务/训练进程独占一个GPU运行时测得的计算和内存资源消耗。在大多数情况下，$R_{SM}$通常是主导因素。我们的经验研究表明，$\alpha$约为0.2，$R_s \approx 10R_a$，$T_s \approx 6T_a$，因此，总的服务吞吐量可以计算为：

TCG的优势：考虑到表4和基于大量性能剖析得出的$COM_{BW} \approx 2 \cdot (T_s + T_a)$，我们可以估计出我们的TCG方案的整体服务吞吐量将比TDG高（约2.5倍）。这是因为模拟器-Agent共置在吞吐量上的资源惩罚（即并行度下降）是次要的（$(T_s + T_a) \cdot \max\{R_s, R_a\}/(T_s \cdot R_s + T_a \cdot \alpha \cdot R_a) - 1 \approx 0.16 \times$），相比之下，通过节省GMI间数据移动带来的吞吐量收益（约3倍）更为显著。

表 4: DRL服务中TDG和TCG的比较。

同步DRL训练：同步DRL训练是主流范式（例如PPO【30, John Schulman, et al. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.】）。它包含三个关键组件：环境模拟器、Agent和Trainer。主要有两个阶段：(i) 经验收集（环境模拟器和Agent之间进行多轮交互，例如32或64轮）和 (ii) 模型训练（基于经验更新策略模型）。这两个阶段在每个训练迭代中顺序发生，训练性能对两个阶段的延迟都非常敏感。优化同步DRL训练的主要目标是最大化DRL计算效率并最小化通信开销（例如模型参数的同步）。为实现此目标，需要同时实现高性能的经验收集和模型更新。虽然DRL服务为经验收集提供了一个高效的模板，我们仍需仔细考虑Trainer在同步DRL训练中的恰当放置。


图 6: 同步和异步DRL训练的GMI映射。

同步训练的GMI映射方案：第一个设计选项是扩展任务专用GMI（TDGEX）设计策略，为DRL Trainer设置专用的GMI。它将首先从所有服务GMI（包含环境模拟器和Agent）收集经验，用于Trainer端的策略训练。一旦策略训练完成，所有Trainer将推送最新的参数以更新Agent上的策略模型。这种设计使得Trainer GMI可以服务于多个服务GMI。然而，延迟敏感的同步DRL训练难以承受跨GMI传输经验的高开销。此外，经验数据对通信不友好，因为它包含异构类型的数据（如状态、动作和奖励），这些数据是分开生成且形状各异的。一个更有前景的方向是将环境、Agent和Trainer对齐在同一个GMI中，以实现完全基于GMI的数据并行训练。这可以看作是扩展的任务共置GMI（TCGEX）。不同的GMI将维护一个包含所有DRL训练任务的完整流水线，我们称之为“整体训练GMI”。如图6(a)所示，一个整体训练GMI引入了第三个阶段 (iii) 全局策略同步，以同步不同GMI中的策略参数。完成后，每个GMI的本地参数将相应更新。

TCGEX与TDGEX的成本分析：我们也对TCG_EX和TDG_EX进行了比较。我们的经验研究表明，$\alpha$约为0.2，$\beta$约为0.3，$R_s \approx 10R_a \approx 5R_t$，$T_s \approx 6T_a \approx 3T_t$，整体系统吞吐量计算如下：

TCGEX的优势：考虑到表5和基于我们性能剖析得出的$COM_{BW} \approx 7 \cdot (T_s + T_a + T_t)$，我们可以推断出我们的TCGEX方案的整体系统吞吐量将比TDGEX显著增加（约5倍）。这是因为吞吐量上的资源惩罚（$(T_s + T_a + T_t) \cdot \max\{R_s, R_a, R_t\}/(T_s \cdot R_s + T_a \cdot \alpha \cdot R_a + T_t \cdot \beta \cdot R_t) - 1 \approx 0.5\times$）可以被避免不必要的GMI间经验移动所带来的显著吞吐量收益（约8倍）大大抵消。

表 5: 同步DRL训练中TDGEX和TCGEX的比较。

异步DRL训练：异步DRL训练（例如A3C【19, Volodymyr Mnih, et al. Asynchronous methods for deep reinforcement learning. International conference on machine learning (ICML), 2016.】）是另一种流行的DRL训练范式。与同步DRL训练相比，异步DRL并发执行DRL服务（经验收集）和策略训练（策略更新）。在这种情况下，模型更新将不再阻塞策略服务的执行。然而，这会带来策略参数陈旧性增加的问题，可能影响算法效率（如收敛速度）。异步DRL训练中的常见设计实践是最大化训练吞吐量，以便所有DRL Agent都能获得能反映最新收集经验的更新模型。因此，如何高效地在Agent GMI和Trainer GMI之间移动经验成为关键。我们引入了解耦的服务和训练方案进行GMI映射（图6(b)），其中服务GMI（用于经验收集）共置在一组GPU上以最大化服务吞吐量，而训练GMI则分组在另一组GPU上以最大化训练吞吐量。解决经验收集和Trainer同步性能的方案已在上述DRL服务部分和§4.1中详细说明。由于异步训练已被卸载到其他GPU上的GMI，现在的挑战变成了如何高效地在服务GMI和训练GMI之间移动经验（§4.2）。

工作负载感知的GMI选择

挑战与目标：除了探索高效的设计模板外，将GMI与DRL工作负载恰当匹配也至关重要，且需要非凡的努力。具体来说，我们需要 (i) 探究GMI资源预算与DRL工作负载性能之间的关系；(ii) 平衡本地GMI性能与全局系统性能。为此，我们引入了一种工作负载感知的GMI选择方法，以促进寻找DRL模拟引擎的正确配置和GMI的资源预算，从而实现我们的关键设计目标（最大化计算吞吐量）。

关键影响因素：我们考虑两个影响基于GMI的DRL性能和资源的关键因素。i) num_env：描述了每个GMI上并发运行的环境数量。它相当于传统DNN训练中的批量大小。考虑到算法（收敛）和系统（运行时执行）效率【18, Viktor Makoviychuk, et al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2108.10470, 2021.】，num_env的常见选择范围从128到16,384。通常，较大的num_env值会带来更高的计算吞吐量。然而，当GMI的所有计算资源被完全占用时，进一步增加num_env只会增加内存消耗，而计算资源成为主要瓶颈。这会导致吞吐量性能改善微乎其微甚至没有。ii) GMIperGPU：描述了分配给每个GMI的计算资源配额。计算资源的多少将决定GMI训练是否能成功启动和执行，更重要的是，决定了DRL计算的吞吐量性能。这是因为模拟和训练过程是计算密集型的。而内存资源，则只能决定DRL训练是否能被容纳在GPU上。

搜索算法：我们的方法将DRL环境规范（如名称和类型）和GPU数量作为输入，并返回num_env和GMIperGPU作为期望的运行时配置。如算法2所示，我们通过调整每个GPU的GMI数量来迭代GMI资源大小的设计空间（第2行）。例如，当GMIperGPU = 4时，每个GPU的资源将被平均分成4个GMI。另一个需要探索的设计维度是环境数量（num_env）。这将帮助我们在给定的GMI资源预算下，精确定位能够最大化吞吐量性能的设计。为了减少不必要的搜索，我们利用几个指标来剪枝设计点，并使用启发式规则进行提前停止。

算法 2: 工作负载感知的GMI选择
输入 :DRL环境规范EnvSpec，GPU数量num_GPUs
输出 :环境数量num_env，每个GPU的GMI数量GMIperGPU
1 best_perf, best_config = 0, (0, 0);
2 for GMIperGPU in [1, 2, 4, 8, ..] do
3    last_perf, last_mem = 0, 0;
4    last_num_env = 0;
5    for num_env in [128, 256, 512, ..] do
6        /*
7        * 性能剖析以获取当前配置下的性能和内存。
8        */
9        curr_perf, curr_mem = profile(GMIperGPU, num_env);
10       /* 无法运行，跳过后续步骤。 */
11       if curr_perf == 0 then
12           continue;
13       /* 计算性能和内存使用的变化。 */
14       d_perf = curr_perf - last_perf;
15       d_mem = curr_mem - last_mem;
16       /* 归一化吞吐量增益。 */
17       Sat = (d_perf/last_perf) / (d_mem/last_mem);
18       /* 提前停止的启发式规则。 */
19       if Sat < α then
20           break;
21       /* 估计的系统吞吐量。 */
22       est_perf = curr_perf * GMIperGPU * num_GPUs;
23       if est_perf > best_perf then
24           best_perf, best_config = est_perf, (num_env, GMIperGPU);
25       last_perf, last_mem = curr_perf, curr_mem;
26    end
27 end
28 num_env, GMIperGPU = best_config[0], best_config[1];
29 return num_env, GMIperGPU;

算法逻辑：profile函数将测试DRL基准测试在当前GMI资源和num_env配置下是否可运行。如果运行失败（例如挂起或崩溃），表明计算/内存资源不足，我们将跳过当前步骤的所有后续性能分析。否则，如果可运行，我们将比较其当前成功运行和上次成功运行之间的计算吞吐量变化和内存消耗变化（第13至14行）。为了确定GMI是否已达到其DRL模拟的最大能力，我们引入了一个饱和度（Sat）指标，通过将吞吐量增益相对于内存消耗惩罚进行归一化来计算（第15行）。我们的经验研究（§6.2关于并发环境的研究）表明，当我们增加num_env并观察到 (i) 内存消耗大幅增加 (ii) 计算吞吐量改善微乎其微时，GMI倾向于饱和其计算能力。我们利用这一见解，通过将Sat与一个阈值$\alpha$（通常$\alpha < 0.1$）进行比较来设置一个提前停止信号。如果当前GMI计算能力尚未饱和（$Sat \ge \alpha$），我们将估计整个系统的吞吐量，如果估计的吞吐量大于任何先前搜索到的设计，则记录最佳设置（第20至24行）。

A4 实验环境

A4 实验结果

总体性能


图 7: 端到端计算吞吐量比较。(a) DRL服务，(b) 使用NCCL的同步DRL训练，以及(c) 使用Horovod的同步DRL训练。计算吞吐量相对于单GPU上的NVIDIA Isaac Gym进行了归一化。

附加研究

A7 补充细节

相关工作

基于CPU的分布式DRL框架：随着深度强化学习的兴起，系统研究人员和工程师致力于提升其性能。一个研究方向是开发基于CPU的分布式框架用于DRL训练。Acme【11, Matt Hoffman, et al. Acme: A research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.00979, 2020.】是一个基于CPU的RL研究框架，用于运行不同规模的RL算法。Seed RL【8, Lasse Espeholt, et al. Seed rl: Scalable and efficient deep-rl with accelerated central inference. arXiv preprint arXiv:1910.06591, 2019.】是一个分布式RL框架，通过利用基于CPU的模拟器和基于TPU的Agent/Trainer实现批处理推理/训练。Ray【21, Philipp Moritz, et al. Ray: A distributed framework for emerging {AI} applications. In 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18), 2018.】是一个分布式计算框架，包含一个专门的RL包RLlib【16, Eric Liang, et al. Rllib: Abstractions for distributed reinforcement learning. arxiv e-prints, 2017.】。RLlib使用基于CPU的模拟器，并与Tensorflow/PyTorch协作作为Agent和Trainer的神经网络操作后端。Ray上的DRL性能在很大程度上受其CPU模拟器性能的制约，该模拟器依赖CPU线程/进程进行并行化。我们还探索了在Ray中使用基于GPU的模拟器【6, Steven Dalton, et al. Accelerating reinforcement learning through gpu atari emulation. 34th Conference on Neural Information Processing Systems (NeurIPS), 2020.】【18, Viktor Makoviychuk, et al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2108.10470, 2021.】的可能性，但由于对基于GPU的模拟任务调度不当，总是导致GPU模拟器崩溃。简而言之，这些设计在很大程度上受其缓慢的基于CPU的模拟器所限制。

基于GPU的DRL加速：另一个方向是利用GPU的强大能力来加速DRL。GA3C【2, Mohammad Babaeizadeh, et al. Ga3c: Gpu-based a3c for deep reinforcement learning. International Conference on Learning Representations (ICLR), 2016.】为A3C算法【19, Volodymyr Mnih, et al. Asynchronous methods for deep reinforcement learning. International conference on machine learning (ICML), 2016.】提出了一种CPU-GPU训练架构，通过在GPU上加速动作预测和策略训练，同时将环境模拟保留在CPU上。cuLE【6, Steven Dalton, et al. Accelerating reinforcement learning through gpu atari emulation. 34th Conference on Neural Information Processing Systems (NeurIPS), 2020.】指出了环境模拟是DRL系统的主要性能瓶颈，并实现了一个基于GPU的RL环境以进行快速模拟。NVIDIA Flex【17, Jacky Liang, et al. Gpuaccelerated robotic simulation for distributed reinforcement learning. In Conference on Robot Learning, 2018.】引入了一个基于GPU的模拟器，并通过直接数据并行实现分布式RL模拟和训练。Isaac Gym【18, Viktor Makoviychuk, et al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2108.10470, 2021.】代表了NVIDIA最新的基于单GPU的大规模并行RL环境模拟设计。Isaac Gym在DRL训练方面取得了比那些基于CPU的设计快200倍的SOTA性能。尽管环境模拟已经通过GPU加速，但环境模拟的性能仍然是关键的性能瓶颈（如图1所示）。这是因为单个基于GPU的模拟引擎（涉及复杂的物理计算）的吞吐量性能无法像常规的密集线性代数计算（基于GEMM的计算）那样得到很好的提升。

讨论

对未来方向的潜在贡献：我们的工作可能对几个未来方向有益。

对潜在用户：我们的工作为将GPU空间多路复用技术扩展到更广泛的应用领域提供了良好实践：1) 通过轻量级微基准测试来表征GPU任务的工作负载，关键是建模任务性能与GPU资源成本之间的关系；2) 通过在资源消耗和性能之间进行权衡，将任务与合适的GPU资源相匹配；3) 将通信密集型任务共置于GMI上，以最小化数据移动成本。

对GPU供应商：我们预见到更高效的进程内、GPU内通信技术（例如在GMI之间共享专用的设备内存空间）在进一步降低GMI间数据通信成本方面的巨大潜力。此外，动态可调的实例大小将是另一个关键特性，用于使GPU实例适应某些在运行时不同时间段具有不同计算和内存资源需求的特殊应用。

对DRL扩展：我们的工作可以作为扩展DRL训练（scale up和scale out）的一个良好起点。在单个多GPU节点上进行扩展（scale up）时，我们的自适应GMI管理策略可以指导任务到GMI的映射和GMI到GPU的布局设计，以最大化节点内吞吐量。在扩展到多个多GPU节点（scale out）时，我们的布局感知梯度归约技术可以扩展到支持高效的多节点模型同步，通过考虑节点内和节点间的GMI布局层次结构。我们的多通道通信可以针对不同类型的数据路径（如NVLink和InfiniBand）进行优化，从而最大化通信效率和带宽利用率。

对集群调度：我们的工作（专注于单个GPU作业）也可以提高现有集群调度器【36, Wencong Xiao, et al. Gandiva: Introspective cluster scheduling for deep learning. In 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18), 2018.】【37, Wencong Xiao, et al. Antman: Dynamic scaling on {GPU} clusters for deep learning. In 14th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 20), 2020.】在处理多个独立GPU作业时的效率。现有集群调度设计的一个关键挑战是不可调整的GPU计算布局（例如每个作业的GPU利用率低），这导致资源利用不足和次优的系统性能。GMI-DRL为未来的集群调度器开辟了新的设计维度，以优化其GPU计算布局（例如将零散的GPU作业压缩到更少的GPU中）和动态GPU计算调整（例如回收和重塑闲置的GPU资源以适应有GPU亲和性需求的作业）。这将减少GPU资源的闲置并提高集群调度的灵活性。

A5 结论

本文提出了GMI-DRL，一个通过基于实例的GPU空间多路复用技术来加速GPU上DRL的系统性设计。它包括一个资源高效的GPU多路复用实例（GMI）的新颖设计、一个基于进程的GMI编程方案，以及一个用于实现高GPU利用率和系统吞吐量的自适应GMI管理策略。我们还为跨GMI构建了高效的通信支持，以满足DRL训练的各种通信需求。全面的实验证明，GMI-DRL在单GPU设置上优于最先进的Isaac Gym，在多GPU设置上优于其使用NCCL/Horovod的变体。