CUDA统一内存优化DeepUM: Tensor Migration and Prefetching in Unified Memory

ASPLOS'23

Abstract

利用CUDA Unified Memory，训练超出GPU内存容量的DNN。其实Unified Memory允许通过缺页异常来训练超出GPU内存容量的DNN，但是页面的迁移带来了很大的开销。DeepUM使用一种新的关联预取技术来隐藏页面传输的开销。作者和vDNN、SuperNeurons、capuchin等之前的sota比较，性能很好。DeepUM的关注点在memory swapping。

1. Introduction

现在流行的DNN都很大，尤其是那些大模型。它们的训练需要很大的GPU内存，但是好在，现在有这样一种方式：在较大的内存空间上进行预训练，然后在稍小的内存空间上进行微调。不过，尽管微调比训练需要的内存少，但是这个量仍然还是比较大的，很难在一般的单张GPU上进行。

DeepUM的关注点在于memory swapping。之前memory swapping的工作主要分为两类：

• 利用CUDA Unified Memory，使用页面预取。

• 使用纯粹的GPU内存（非统一的），swap in/swap out内存对象。比如vDNN、SuperNeurons、Capuchin都是这一类

Unified Memory在CPU和GPU之间构建了统一的内存空间，利用GPU的缺页异常机制来按需进行页面传输。

memory swapping工作很少基于UM来做，因为地址转换和缺页异常的开销很大：

• 由于UM是虚拟内存空间，因此对于GPU上的每个内存请求都要进行地址转换，这可能严重影响性能

• 处理缺页异常会在CPU和GPU之间迁移页面，这会带来很大的I/O开销。

但是UM也有很多好处：

• 使用纯粹的GPU内存，无法运行内存占用高于GPU内存容量的CUDA核函数，而UM可以利用缺页异常按需迁移页面，因此可以训练更大的DNN。

• 使用纯粹的GPU内存可能会有比较严重的内存碎片问题，DNN在训练过程中有频繁的内存分配和释放操作，尽管tensorflow和pytorch等主流DL框架都有自己的GPU内存管理机制，但是仍然存在一些内存碎片问题。而UM是虚拟内存空间，所有的内存对象都是以4KB的页为基础单元的，这会减小内存碎片问题，更有可能训练更大的DNN。

DeepUM实现了一些优化技术来减小UM带来的开销：

DeepUM修改了一项最初为cache line预取而开发的相关预取技术来预取GPU页面。DeepUM在许多预取技术中选择了相关预取（correlation prefetching）。

page fault handler监控异常访问，因此DeepUM可以从page fault handler这里获得fault pages之间的关系。

DeepUM利用了一个事实：在DNN训练时，核函数执行模式和内存访问模式大多是固定、重复的。因此可以记忆这些重复的模式，然后进行相关预取。

DeepUM的相关表记录了在训练时核函数的执行历史以及页面访问，它通过相关表预测接下来会执行哪个核函数，然后预期页面。传统的相关预取使用单个表来存储历史并记录CPU cache line之间的关系，DeepUM相关预取使用了两个不同的表（核函数的执行历史、页面访问模式）。

采取了两种优化技术来减少页面异常处理的时间：page pre-eviction，page invalidation。

2. Background

2.1 GPUs and CUDA Programming Model

略

2.2 CUDA Unified Memory

UM构建了一块统一的内存空间，CPU和GPU都能访问。从Pascal架构开始，NVIDIA GPU就有了页面迁移引擎（page migration engine），通过利用缺页异常机制（page fault mechenism），GPU可以拥有一个虚拟内存系统。当GPU访问一个不在GPU内存中的页面时，就会触发缺页中断，然后NVIDIA设备驱动将页面迁移到GPU内存里。因此，UM允许在GPU上执行超出GPU内存的程序。

​

​

如上图，左侧是CPU的统一内存空间（物理位置位于CPU内存，即主存）、右侧是GPU的统一内存空间（物理位置位于GPU内存），在图a中，page1和page2都位于CPU的统一内存空间，后来，如图b，GPU访问page1，触发缺页中断，NVIDIA设备驱动将page1从主存迁移到GPU内存上，然后重放对page1的访问。

尽管UM有这些好处，但是UM的缺点很明显，那就是缺页异常处理的开销太大。SM中有TLB，在处理异常时，TLB无法运行，等到异常处理完毕后，TLB才能继续；缺页异常还会引入很大的I/O开销（CPU和GPU之间页面迁移）。因此可以使用CUDA API如cudaMemPrefetchAsync()、cudaMemAdvise()来减少缺页异常。

2.3 NVIDIA Page Fault Handler

fault buffer : fault buffer是GPU上的一个循环队列，存储了异常访问信息（faulted access information）。GPU可以同时产生多个异常。fault buffer里面可以存储对于同一个页面的多次异常访问信息。

UM block : 1个UM block最多由512个连续的页面组成，也就是4KB*512=2MB。每个UM block都包含了所有页面的信息，比如这些页面存在CPU还是GPU、读写权限等。NVIDIA以UM block的粒度来管理页面，如果一个UM block包含了faulted page，那么这个UM block被称为faulted UM block。

如下图Fig. 3展示了GPU页面异常处理（fault page handler）的流程。

一段翻译：

这个流程图展示了缺页异常处理的过程。首先，NVIDIA驱动程序从GPU中的异常缓冲区获取了页面地址和异常访问的访问类型（1）。然后，NVIDIA驱动程序对这些异常进行预处理（2）。它移除重复的地址并根据它们的UM块进行分组。接下来，NVIDIA驱动程序检查每个异常的UM块在GPU中的可用内存空间（3）。如果没有可用的GPU内存空间用于该异常的UM块，它会从GPU中逐出一些页面到CPU（4）。然后，它将异常的页面填充到GPU中（5）（即为异常的页面分配GPU内存空间），并将页面传输到GPU（6）。当传输完成后，UM块的异常页面被映射到了GPU（7）。这个过程重复进行，直到所有的异常UM块都被处理完（8）。最后，NVIDIA驱动程序向GPU发送一个重放信号，异常处理程序完成（9）。

3. Overall Structure of DeepUM

3.1 Structure of DeepUM

如下图Fig. 4展示了DeepUM的整体架构。DeepUM由两部分组成：DeepUM Runtime和DeepUM Driver。

DeepUM runtime

DeepUM runtime为CUDA内存分配函数提供封装函数，来把所有GPU内存分配请求转换为UM分配请求。通过在UM空间里面分配GPU内存对象，实现对GPU内存的超额分配。此外，DeepUM runtime为CUDA核函数启动命令和其他CUDA库函数（比如cuDNN cuBLAS）提供了封装函数。

DeepUM runtime管理着execution ID table，这个表记录了核函数启动历史，以及核函数名称、参数的哈希值。当一个新的核函数启动命令进入到deepUM runtime时，它会计算该核函数名称、参数对应的哈希值，然后在表里查询有没有同样的哈希值，如果有的话，就把对应的execution ID赋给这个核函数，如果没有的话，就把新的哈希值记录到表中，并赋给这个核函数一个新的execution ID。最后，DeepUM runtime在执行这个新的核函数启动命令之前，会把一个CUDA callback函数排队到CUDA runtime，这个callback函数会把接下来的核函数启动命令对应的execution ID传递给DeepUM driver，DeepUM driver在相关预取中使用这个execution ID。❓❓❓❓❓❓核函数启动命令需要经由DeepUM runtime才能得到execution ID，这里怎么直接传给DeepUM driver了？是经过了DeepUM runtime，但是论文中没有说明吗？

DeepUM driver

DeepUM driver处理GPU缺页异常，并且把页面预取到GPU。

我们观察到核函数执行模式和内存访问模式是固定的、重复的，因此可以记忆下来，进行相关预取。DeepUM driver管理correlation table，记录了在DNN训练过程中的核函数执行历史和页面访问。通过相关表中的信息预测接下来会执行哪个核函数，进行相关预取。

DeepUM driver中有四个线程：fault handling thread、correlator thread、prefetching thread、migration thread

• fault handling thread : 使用NVIDIA驱动中的函数（比如访问fault buffer、给GPU发送重放信号）处理缺页异常。DeepUM driver拦截发往NVIDIA driver的缺页异常中断信号，然后fault handing thread读取fault buffer，将读取到的异常访问信息传递给其他三个线程。fault queue存储异常页面对应的UM block的地址，fault queue拥有最高优先级，以保证GPU尽快重放异常访问。

• correlator thread : 管理相关表，它根据来自fault handing thread的异常访问信息更新相关表。详见4.2

• prefetching thread : prefetching thread查找相关表并根据出现异常的块地址计算要预取的UM block地址。然后，它将预取命令加入预取队列，这是一个单生产者/单消费者队列。预取命令包括要预取的UM block地址和预测将使用的UM block的execution ID。❓❓❓❓❓❓啥意思？

• migration thread : 在CPU和GPU之间迁移UM block。fault queue的优先级高于prefetch queue，因此当fault queue空时，才会处理prefetch queue。

4. Correlation Prefetching for GPU Pages

4.1 Pair-Based Correlation Prefetching

pair-based correlation prefetching在相关表中记录了miss addrssed的序列。当miss时，会查询相关表，预取所有和miss address相关的address。

如下图Fig. 5(a)所示，相关表中每一行记录的是一个不同的miss address。miss address序列是abca，当前的miss address是c，因此Last指向c，SecondLast指向b。

4.2 Correlation Prefetching in DeepUM

DeepUM中的相关预取技术旨在减少缺页异常。和4.1中原始的基于对的相关预取不同，DeepUM中的相关预取维护两个表：execution ID和UM block。这两个表的NumLevels都为1。注意：DeepUM中的相关预取是以UM block为单位的！而不是cache line或者page！

Execution ID correlation table (execution table)

该表记录了CUDA核函数的execution ID，execution ID来自DeepUM runtime。（详见 3.1中对DeepUM runtime的描述）

如下图Fig. 6所示，每个execution ID可能对应一个或多个集合（可能被调用了一次或多次）。比如Execution ID 0对应的集合（7，9，92，75），前三个ID 7、9、92表示执行0之前所执行的三个核函数，第四个ID 75表示预测下一次执行的核函数对应的execution ID是75。

UM block correlation table (block table)

UM block correlation table在UM block的粒度上记录历史，而不是记录异常页面的地址。两点原因：

• 对于大规模的DNN来说，异常页面地址太多

• NVIDIA驱动管理页面的粒度是UM block，UM block correlation table选择以UM block的粒度来记录历史，可以与NVIDIA驱动保持一致，更有效

如下图Fig. 7所示，每个execution ID都有一个对应的块表，记录了相应核函数的UM block访问历史。此外，块表记录了相应核函数第一个异常UM block（start）的地址、最后一个预取的UM block（end）的地址，这两个指针用来实现chaining。（start UM block是在当前核函数执行时第一个缺页异常的页面所在的UM block；end UM block是在当前核函数执行结束之前，最后一个缺页异常的页面所在的UM block）

Prefetching mechanisms and chaining

当缺页异常发生时，DeepUM driver查询当前正在执行的核函数的block table，预取所有和异常UM block相关的UM block。

当prefetching thread遇到block table中的end UM block时，就结束对当前执行的核函数的UM block预取，然后通过execution table预测下一个要执行的核函数，然后开始预取被预测要执行的核函数对应的block table里面的start UM block。

比如上图Fig. 7，假设正在执行ID 0、正在预取b，然后预取b的successor block e、q，而q是end，因此结束，然后根据execution table预测接下来要执行的核函数，然后开始预取它的start......

这个过程被称作chaining，因为在缺页异常发生之后，会持续预取UM block。当出现新的缺页异常中断信号，或者prefetching thread无法预测下一个要执行的核函数时，chainging结束。

5. Optimizations for GPU Page Fault Handling

5.1 Page Pre-eviction

page eviction发生在驱动无法为faulted page分配GPU内存空间时（已满）。而此时，需要等待驱逐完成，才能继续下一步，因此这会增加缺页异常的处理时间。因此提出page pre-eviction

DeepUM对这样的页面执行pre-eviction：

• least recently migrated

• 在当前正在执行的核函数以及接下来N个被预测会执行的核函数内都不会被用到

由于NVIDIA driver跟踪GPU的空闲空间，DeepUM从NVIDIA driver中获取可用空间信息。它从execution ID correlation table中获取下一个执行的核函数的信息。

5.2 Invalidating UM Blocks of Inactive PyTorch Blocks

翻译：

PyTorch针对CPU和GPU有不同的内存分配器。PyTorch的GPU内存分配器管理设备内存池，以最小化内存分配/释放时间并减少内存碎片化。GPU内存分配器管理两种类型的内存池：大内存池和小内存池。在PyTorch中，内存对象被称为块。在本文中，我们将其称为PT块，以区别于UM块。大内存池由大于1MB的PT块组成，小内存池由小于或等于1MB的PT块组成。当内存分配请求到来时，如果请求的大小大于1MB，则内存分配器从大内存池中找到一个PT块。否则，它从小内存池中找到一个PT块。当内存池中有多个PT块与请求的大小匹配时，分配器返回最小可用的PT块。此外，当PT块的大小远大于请求的大小时，PT块会被拆分。所选的PT块被从内存池中移除并标记为活动状态。然而，当内存池中没有可用的PT块时，GPU内存分配器通过向CUDA运行时请求设备内存空间来分配新的PT块。在DNN模型使用完PT块并将其返回给GPU内存分配器后，分配器会将PT块插入到适当的内存池中并将其标记为非活动状态。非活动的PT块，即内存池中的PT块，仅在池中没有可用内存空间时才被释放以产生新的内存空间。当我们将PyTorch内存分配器与UM一起使用时，问题就出现了。当GPU内存中的非活动PT块被驱逐到CPU内存时，会产生不必要的大量数据传输。此外，它们会占用CPU内存空间。当非活动PT块被标记为活动并再次被DNN模型使用时，问题会变得更加严重。由于PT块中的页面已被驱逐到CPU内存，它们应再次迁移到GPU，导致大量数据传输。为了解决这个问题，我们向PyTorch内存分配器中添加了几行代码，用于告知DeepUM驱动程序何时将PT块标记为非活动状态。如果受害页面属于非活动的PT块，DeepUM驱动程序会简单地使GPU内存中对应的UM块失效。

6. Evaluation

6.1 Methodology

LMS (Large Model Support)是IBM开发的一个开源项目，它支持PyTorch框架和自动GPU内存交换，我们通过实际执行LMS来直接比较 DeepUM和LMS的性能。由于其他五种方法基于TensorFlow或其他框架，并且大多是闭源的，我们通过间接比较它们与DeepUM的性能。Ren等人[50]实现了这些方法，并测量了它们在训练吞吐量上相对于NVIDIA UM的加速比。因此，我们从Ren等人的研究中获取这些方法的数据，并将其与DeepUM相对UM的加速比进行比较。我们使用相同的模型、数据集和GPU进行公平比较。

系统配置、模型、数据集

6.2 Comparison with Naïve UM and IBM LMS

使用一块单独的 V100 32GB GPU

UM：使用 NVIDIA UM

LMS：原始 LMS

LMS-mod：对 LMS 进行修改以定期释放 PyTorch 内存池中缓存的 PT 块。通过定期清理缓存的内存对象，可以减少因内存碎片化导致的内存溢出（OOM）的发生。

图a图b上，DLRM效果不明显的原因

和其他内存预取策略一样，性能取决于应用程序的内存访问模式以及内存传输时间和计算时间之间的比率。例如，对于 LMS 和 DeepUM，在DLRM上几乎没有显示出速度提升。DLRM 是 Facebook 提出的一个推荐模型，其输入数据包括用户的偏好、购买清单等。该模型针对每个输入数据项查找嵌入表，并使用适当的嵌入向量转换输入数据项。在 DLRM 中，大部分内存空间用于存储嵌入表。此外，其内存访问模式是不规则的，因为嵌入表的查找高度依赖于输入数据。这就是为什么 LMS 和 DeepUM 的预取策略效果不佳的原因。

图a图b上，ResNet效果不明显的原因

ResNet 由多个称为残差块的构建块组成。一旦为残差块预取了内存对象，计算时间将主导残差块的处理过程。

图c：总能耗与 UM 相比的比率（UM为1）

表 3 显示了 IBM LMS 和 DeepUM 的最大可能批量大小。DeepUM 可以运行批量大小需要的峰值内存使用量几乎与总 CPU 内存大小相同的模型。表中的数字表明，在 DeepUM 中利用 UM 会遇到较少的内存碎片问题，并且更有可能运行大型 DNN 模型而不出现任何内存问题。

表 4 显示了用于存储相关表的内存空间大小。由于 DeepUM 在发现具有新的执行 ID 的核函数时动态分配 UM 块相关表，因此存储相关表的内存大小因每个 DNN 模型和批量大小而异。请注意，相关表存储在 CPU 端。

表 5 显示了每个模型和不同批量大小下每个训练迭代的平均页面错误数量。结果表明，DeepUM 预取页面的准确性相当高，并且可以显著减少页面错误。

对预取程度的敏感性。DeepUM 需要一些其他 CPU 线程来更新相关表和管理队列。这些任务不会产生显著的能源/功耗或性能开销，因为它们的大部分工作是表查找/更新和排队。然而，预取可能会导致更多的能源/功耗消耗。此外，过度的预取可能会损害性能，因为预测可能是错误的，并且可能会将不必要的页面迁移到 GPU。因此，这可能会浪费内存带宽并驱逐即将被访问的页面。DeepUM 预取了预测将被接下来的 𝑁 个核函数访问的页面。用户可以通过修改 𝑁 来静态控制预取的程度。为了验证过度预取的影响，我们在不同的预取程度（𝑁 ）下测量应用程序的执行时间和能源消耗。图 11 显示了不同 𝑁 值的结果。图 11(a) 显示了加速比，图 11(b) 显示了不同 𝑁 值的总能源消耗与 𝑁 = 8 的比率。它们表明加速比和能源消耗成反比关系。此外，存在一个最佳点（𝑁 = 32），在该点上加速比最高，而能源消耗最低。

6.4 Comparison with TensorFlow-Based Approaches

我们将DeepUM的性能与基于TensorFlow的方法进行了比较：vDNN[51]、AutoTM[21]、SwapAdvisor[24]、Capuchin[45]和Sentinel[50]。图13显示了每个DNN模型的加速比。加速比是在V100 16GB GPU上获得的。请注意，这些数字来自Ren等人的研究[50]，他们测量了训练吞吐量相对于未进行预取的NVIDIA UM的加速比。此外，我们以与图9相同的方式获得了加速比的上限（理想值）。表7显示了基于TensorFlow的方法和DeepUM的最大可能批量大小。为了测量最大可能批量大小，我们将DeepUM的总CPU内存使用限制为128GB，以匹配与基于TensorFlow的方法的系统配置。总的来说，DeepUM比IBM LMS和其他TensorFlow方法（除Sentinel外）更快。DeepUM显示出与Sentinel相当的性能。请注意，Sentinel的换页机制对用户不透明，而DeepUM的是完全自动的，如表8所示。此外，DeepUM允许比其他先前方法更大的批量大小。请注意，以前的方法在DNN层或张量级别管理数据。这意味着在一层或一个张量中访问的所有数据一起移动。性能差异来自于DeepUM更细粒度的数据移动，其中通过相关表进行准确预测，内存对象以更精细的方式进行预取和驱逐。

7. Related Work

已经进行了许多研究，以克服GPU内存容量问题，通过利用交换机制。

在交换方法中有两个类别。其中一个类别是将整个GPU内存对象交换到CPU内存或NVMe设备中[6, 21, 24, 33, 45, 49–51, 55]。

vDNN[51]是第一个引入GPU内存交换用于深度神经网络（DNN）工作负载的方法。然而，DNN模型必须使用vDNN API函数，并且它仅支持卷积神经网络（CNNs）。

TFLMS[33]是IBM开发的开源项目，包含在IBM Watson Machine Learning软件包中。它通过修改TensorFlow的计算图来调度交换命令。它需要修改TensorFlow和TensorFlow用户脚本。IBM还发布了TFLMS的PyTorch版本。

Superneurons[55]和FlashNeuron[6]都将DNN模型作为输入，并导出最佳的张量卸载调度。Superneurons将GPU内存中的张量卸载到CPU内存，而FlashNeuron[6]则利用NVMe SSD来卸载张量。

AutoTM[21]和SwapAdvisor[24]也将DNN模型作为输入来调度其内存操作。AutoTM使用整数线性规划（ILP）来安排数据移动，同时减少GPU内存碎片。另一方面，SwapAdvisor[24]使用遗传算法来调度运算符、内存分配和交换决策。

Capuchin[45]在运行时识别张量访问模式，并调度张量驱逐、预取和重计算。

Sentinel[50]基于TensorFlow，并动态地从TensorFlow运行时和操作系统中收集张量访问信息。它类似于DeepUM，因为它利用页面故障机制来进行分析并获取内存访问模式。然而，Sentinel使用CPU的页面故障机制，而DeepUM使用GPU的页面故障机制。在Sentinel的分析阶段，张量被分配在CPU的固定内存中，而GPU对其进行访问。Sentinel和DeepUM之间的另一个区别是，Sentinel需要修改TensorFlow和TensorFlow用户脚本来插入Sentinel分析API函数调用。Sentinel的评估结果显示，它在训练吞吐量方面优于vDNN、SwapAdvisor、AutoTM和Capuchin。

DeepSpeed[49]是一个高度优化且广泛使用的多GPU深度学习框架。它通过钩取PyTorch API来跟踪DNN操作序列。然后，它提供了GPU内存与CPU主内存或NVMe SSD的交换机制。然而，DeepSpeed管理的是卸载模型参数、梯度和优化器状态，而不是激活。程序员应该通过激活检查点手动管理激活和临时缓冲区。这意味着当激活量超过GPU内存大小时，程序员必须修改他们的代码进行激活检查点。

交换方法的另一类是利用带预取的CUDA UM[5, 35]。OC-DNN[5]在每个DNN操作前手动插入预取命令。DRAGON[35]使用NVMe SSD作为UM的后备存储。它针对一般应用，并使用DNN工作负载作为展示。它还需要用户代码修改和设备驱动程序修改。

8. Conclusions

在本论文中，我们提出了DeepUM，通过利用CUDA统一内存和使用CPU内存作为后备存储，允许对DNN进行GPU内存超额订阅。虽然CUDA UM允许使用页面故障机制进行GPU内存超额订阅，但会引入巨大的开销。我们在UM块级别使用了新的相关预取技术来隐藏故障处理开销。我们还引入了两种优化技术，以最小化GPU故障处理时间。一种是页面预驱逐技术，其中一个新的页面预驱逐策略与相关预取技术相结合。另一种是在PyTorch中使不必要的页面失效，当它们被选为页面驱逐的受害者时。DeepUM驱动程序钩取了NVIDIA设备驱动程序的页面故障处理程序，以便监视页面故障并预取所需的页面。它还允许通过LD_PRELOAD环境变量加载DeepUM运行时。通过设置环境变量，可以轻松地打开和关闭DeepUM。虽然我们的方法需要在原始PyTorch源代码中进行少量代码修改来改变PyTorch内存分配器的行为，但不需要修改用户Python代码。我们使用来自多个来源的九个大规模DNN模型（包括MLPerf、PyTorch示例和Hugging Face）对DeepUM的性能进行了评估，并将其与六种最先进的GPU内存交换方法进行了比较。评估结果显示，DeepUM达到了一种不透明的交换机制的性能相当。DeepUM比其他五种方法表现出更好的性能。同时，DeepUM可以处理其他方法无法处理的更大模型。