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