Abstract
一种明显减轻基于GPU的AI计算中内存困难的方法是通过CPU卸载，即数据在GPU和CPU RAM之间移动。虽然CPU卸载很有用，但由于CPU RAM与GPU RAM之间的相对传输速率，它可能会大大减慢计算速度。为了解决这个问题，重叠内存传输和计算是必要的，但这种异步性引入了不确定性，因此无法事先知道操作的最佳顺序。作者描述了Turnip系统，它是一个使用CPU卸载运行AI计算的系统，旨在处理这种不确定性。Turnip的关键创新在于将AI计算编译成一个依赖图，这使得Turnip运行时可以自由地以许多不同的顺序运行操作，如GPU Kernel 调用；在运行时，Turnip会根据实时事件动态选择最佳顺序。作者发现，Turnip的性能显著优于支持受限GPU RAM的标准PyTorch系统，并且在内存严重受限的环境中避免了内存不足错误。
1 Introduction
现代AI计算的内存管理是困难的。例如，在LLaMA大型语言模型（Touvron等人，2023年）中，对长度为的序列的注意力计算会产生一个包含个浮点数的中间结果。因此，对于一个长达100,000个标记的长输入序列，注意力计算将导致1.2万亿个数字，以半精度存储则需要2.4太字节。正是由于这些原因，AI程序员常常受到内存不足错误的困扰（常见问题解答）。
CPU卸载——即数据被移动到CPU RAM中存储——可以有所帮助。由于CPU RAM比GPU RAM便宜得多，且在GPU服务器中安装数太字节的CPU RAM几乎是“免费的”，因此利用CPU RAM临时存储数据是合理的。这个想法在几个系统中得到了探索，如pofo(Beaumont等人，2021年)、AutoTM (Hildebrand等人，2020年)、SwapAdvisor (Huang等人，2020年)、Capuchin (Peng等人，2020年)和POET (Patil等人，2022年)。这些系统将GPU计算视为数据流图，并计划如何通过使用CPU RAM卸载来将计算放入GPU RAM中。
虽然CPU卸载是一个明显的想法，但由于CPU RAM和GPU RAM之间的相对较慢的传输速率，它可能会大大减慢计算速度。因此，任何用于CPU卸载的系统都必须确保在进行此类传输时，没有计算因为等待传输完成而被阻塞。
在本文中，作者提出了Turnip（“非确定性GPU运行时带CPU卸载”的简称），这是为多GPU服务器设计的一种运行时，旨在系统地支持CPU RAM卸载。Turnip的关键创新在于它将一个预计算的内存访问计划（称为memgraph）与一个“非确定性”的事件驱动系统运行时结合起来。memgraph是一个依赖图，其中顶点代表任务（如在大型语言模型的一层中执行GPU Kernel 以完成注意计算的较小部分），边代表数据或内存依赖。任何尊重memgraph中依赖关系的执行顺序都是有效的，任务可以在任何时候其依赖关系得到满足且适当的资源空闲时被分派。因此，两次相同的memgraph执行可能导致在GPU上执行的不同操作序列，或不同的张量序列被分页到CPU RAM——这就是非确定性的由来。然而，memgraph中的依赖关系使得最终的输出始终是“正确”的，无论执行顺序如何。Turnip的事件驱动、完全异步的运行时是独特的。因为操作可以在依赖关系满足时随时分派，不受特定顺序的限制，这降低了任何GPU因等待内存传输完成而停滞的可能性。如果一个任务由于memgraph中的未满足依赖关系而无法运行，那么可能存在另一个可以运行的任务。
关键技术挑战是如何有效地构建尽可能少的依赖关系的memgraph，以便运行时尽可能多地自由分派操作，使其不会因为等待内存传输完成而阻塞。Turnip通过模拟计算的一次执行来构建memgraph，将张量映射到GPU内存位置，并在必要时添加代表内存依赖关系的边，以及卸载和重新加载操作。
2 Why Is Non-Determinacy of Execution Order Crucial?
Turnip的设计基于一个简单的假设：当运行一个利用CPU RAM的基于GPU的计算时，像卸载（offload）和重新加载（reload）这样的异步操作将具有看似非确定性的运行时间，这是很难预先计划的。系统运行时必须适应由此产生的非确定性，否则性能可能会受到影响。
考虑图1，它描绘了一个具有卸载（从GPU RAM到CPU RAM的数据移动）和重新加载（从CPU RAM到GPU RAM）的单GPU系统的memgraph。顶点代表产生数据的操作（例如，GPU Kernel 调用）。黑色边表示数据或消费依赖；红色边表示内存依赖（详细描述memgraphs的部分将在第4节中介绍）。注意从offload到的内存依赖。这是因为的输出将被写入输出的位置，因此直到卸载完成，才能执行。从reload到的数据依赖存在是因为与顶点相关的 Kernel 将消耗重新加载的张量。
设想一个系统已经执行了与顶点和相关的GPU Kernel 。它当前正在执行卸载和重新加载。在这一点上，不可能知道哪个 Kernel 应该接下来运行_（或），因为这取决于哪个内存传输先完成。理想情况下，_这个决定将在运行时做出_。如果系统在编译时决定确定性运行然后在之前，并且重新加载先完成，GPU将闲置，等待卸载完成。这就是为什么需要特殊的、“非确定性”运行时：以适当处理通过添加内存操作引起的非确定性。## 3 相关工作
处理有限GPU内存的系统采取了两种方法。一些系统，如Turnip，接受一种通用的GPU计算的抽象版本。其他系统则更具体地针对某些类型的模型、优化算法，或是特定的任务，如训练或推理。与Turnip不同，这些现有系统都没有考虑卸载和重新加载操作的非确定性对系统性能的影响，也没有关注系统运行时。
采用第一种更通用方法的系统接受一个通用的数据流图，并像Turnip一样在有限内存中规划执行：pofo(Beaumont等人，2021)，AutoTM(Hildebrand等人，2020)，SwapAdvisor(Huang等人，2020)，Checkmate(Jain等人，2020)，Capuchin(Peng等人，2020)，以及POET(Patil等人，2022)都假设了一个机器学习计算的输入数据流图，然后在有限内存中规划执行。Checkmate只考虑张量重新物化，而POET、pofo和Capuchin考虑重新物化和卸载；AutoTM和SwapAdvisor只考虑卸载。
更针对性的方法由DeepSpeed项目（Deepspeed）和不同的ZeRO优化采用。对于 Transformer 和其他类似模型，DeepSpeed推理（包括ZeRO-Inference）(Aminabadi等人，2022)有两个关键思想。首先，DeepSpeed推理“在未使用时将一些激活从GPU卸载到CPU内存。”其次，DeepSpeed推理“将模型权重固定在DRAM（如果足够大）或NVMe中，并在需要时将每一层流式传输到GPU内存进行计算。”FlexGen(Sheng等人，2023)试图在有限硬件条件下，采用多种方法加快 Transformer 推理，包括将模型权重卸载到CPU，量化（Yao等人，2022；Frantar等人，2022）和稀疏注意力（Child等人，2019）。后两个想法与本文中的想法正交。对于CPU卸载，FlexGen优化了一种“之”字形块调度，通过 Transformer 层和批量中的序列，卸载和重新加载KV缓存（Pope等人，2023）和模型权重。PagedAttention(Kwon等人，2023)处理 Transformer 中低内存利用率的问题，为KV缓存开发了一个分页系统。
ZeRO-Offload（Ren等人，2021年）是一种针对有限内存训练的全面解决方案，主要可以看作是使用CPU RAM来运行ADAM优化器，按照精心控制的计划将权重移动到GPU RAM。ZeRO-Offload是对ZeRO（Rajbhandari等人，2020年）的增强，后者旨在提高内存效率，通过在多个GPU之间分割优化器和数据。ZeRO-Infinity（Rajbhandari等人，2021年）与之相似，并包含了一个CPU卸载引擎，以及操作符分块以利用多个GPU的RAM。
4 Taskgraphs and Memgraphs in Turnip
Turnip接受的输入是一个任务图（taskgraph）。任务图是一个数据流图（有向无环图），描述了如何执行多GPU计算。在任务图中，边代表数据流，顶点代表对张量的操作。没有输入的顶点（称为_input vertex_）与输入张量相关联。与非输入顶点相关的操作可能是在特定GPU上执行的 Kernel 调用，或者是GPU到GPU的数据传输。
Turnip对任务图是如何创建的持不可知态度；它可能是使用诸如FlexFlow（Jia等人，2019）或Alpa（Zheng等人，2022）之类的框架创建的。考虑一个矩阵乘法，假设作者希望在这三个GPU上执行这个矩阵乘法。为了生成一个任务图，像FlexFlow这样的框架可能会选择如图2所示分解这个矩阵乘法，可能对应于图3中的任务图。
给定这样一个任务图，Turnip首先将其编译成一个memgraph，最终执行这个memgraph。与任务图一样，memgraph也是一个有向无环图。原始任务图中的每个顶点都会在相应的memgraph中出现。此外，编译过程可能会添加额外的卸载和重新加载操作，这些操作将内存从GPU RAM移动到CPU RAM，反之亦然。在编译过程中，将每个顶点在memgraph中关联的输出映射到一个内存位置。与输入任务图不同，memgraph不是数据流图；它是一个依赖图。如果从到有一条边，这意味着依赖于，并且可能要等到执行后才能执行。在memgraph中，有两种依赖关系。一种是数据依赖，它从任务图继承而来（或者通过添加卸载或重新加载而创建；见下文）。第二种是内存依赖，是为了确保图中没有竞态条件而添加的。_竞态条件_发生时，有一些顶点使得图的两个有效执行可能会产生不同的输出。当两个顶点写入同一个内存位置时，这种情况就可能发生，根据执行顺序，第三个顶点可能读取任一输出。
让作者说明一下如何将图3中的任务图编译成memgraph。假设作者的三个GPU每个都有五个内存位置，为了简化，每个张量大小相同，并且恰好占据一个内存位置。在编译过程中，任务图中每个操作关联的张量被分配到一个内存位置，如图4所示。GPU 1必须处理总共七个张量（两个输入张量和五个通过某些操作创建的额外张量），鉴于作者的五个位置，作者无法将这七个张量都装入内存。因此，操作和输出的张量都被映射到GPU1-Loc1，操作和输出的张量都被映射到GPU1-Loc3。
图5展示了一个相应的记忆图。请注意，已经添加了两条表示记忆依赖的新边。这些边保证了图中不存在竞态条件。具体来说，如果一个图中，每当顶点和的输出都被映射到同一个内存位置时，要么安全地覆盖的结果，要么安全地覆盖的结果，那么这个图将不存在竞态条件。作者说“安全地覆盖的结果”，当且仅当对于每个消耗输出的，存在从（或的某个后代）到（或的某个祖先）的记忆依赖。为什么？如果要安全地覆盖的结果，作者需要确保在所有的消费者完成执行之前不能执行——这样的记忆依赖确保了这一点。
图6：张量到GPU RAM的可能映射。
例如，从图4中作者看到顶点的输出与顶点的输出映射到相同的位置。在图3的相关记忆图中，为了确保安全地覆盖的结果，作者从（的唯一消费者）添加到的记忆依赖。从图4作者还看到和的结果映射到相同的位置。为了确保安全地覆盖的结果，作者从（的唯一消费者）添加到的记忆依赖。请注意，这个记忆依赖用虚线表示；这表示它是多余的，因为已经有一个从3到8的数据依赖，所以这个记忆依赖不是正确性所必需的。
当内存受到更多限制时，情况可能会变得更加复杂。考虑这样一个场景：每个GPU上作者只有四个内存位置，而作者希望编译同一个任务图。图3中的任务图顶点到GPU 1上内存位置的 一种可能映射如图6所示；相关的任务图如图7所示。特别要注意的是增加了一个卸载-重新加载（offload-reload）对。卸载和重新加载都是编译期间添加到memgraph中的新操作，以便在内存受限的情况下辅助执行。作者总是可以将任务图中的编译成memgraph中的。在offload之后，的结果不再占用GPU内存，但它不能被任务图映射到GPU的内存位置。唯一消费A )到4的原因。或者，考虑被映射到GPU1-Loc4的reload和2 。为了确保reload能够安全地覆盖2 的结果，从2 的唯一消费者（顶点4 ）到reload存在一个内存依赖。
5 The Turnip Execution Engine
一旦生成了记忆图（memgraph），Turnip引擎就会使用一种非确定性的、基于事件的框架来执行它。一旦GPU处于未使用状态，或者一个张量准备好被卸载到RAM，Turnip运行时可以立即将任何可用的任务分配给GPU，或者开始传输，而无需考虑计算的整体状态。还要注意，在记忆图执行过程中没有调用诸如cudaMalloc或cudaFree这样的内存管理例程，因为内存管理不再是动态的。张量的放置在执行前是预先确定的，如果尊重依赖关系，就不会由于竞态条件导致内存损坏。
为了执行记忆图，Turnip运行一个中心事件处理循环，该循环反复处理回调函数，这些回调函数是在与记忆图顶点相关的工作完成后被调用的（完成GPU到GPU的传输、GPU Kernel 的完成，或者卸载或重新加载的完成）。当一个顶点完成并且调用了回调时，事件循环检查是否有任何其他顶点可以执行。也就是说，它搜索一个顶点，其中（a）在记忆图中所有以边指向的顶点也都已完成；（b）如果是一个 Kernel 调用，那么分配给的GPU当前是空闲的。当事件循环找到这样的顶点时，它就启动它，并搜索另一个这样的顶点。当它找不到可执行的顶点时，它就会进入休眠状态，直到被另一个回调唤醒。
6 Building a memgraph
本文作者要解决的关键技术问题是： 如何从taskgraph构建memgraph？ 编译过程的主要要求是正确性。
正确性要求：
(a) taskgraph中存在的每个数据依赖在memgraph中也存在，或者用一系列的卸载-重载操作来替换；
(b) memgraph中没有竞态条件；
(c) memgraph中没有循环。此外，作者还希望memgraph具有高性能。
如果内存依赖严重限制了顶点的执行顺序，那么memgraph将不会具有高性能。这种限制可能会减少并行性和GPU利用率。
即使在非确定性执行下，内存依赖也可能导致GPU一段时间内未被使用，因为GPU被内存传输或与两个张量映射到GPU RAM中相同位置的相关依赖阻塞。
通过仔细实现各种malloc变体，可以减少memgraph中添加的依赖导致此类停顿的可能性。
例如，当simMallocOffld搜索受害者时，作者寻找下一次使用在将来最远的受害者——即，在的末尾（在张量大小不同且可能有不止一个受害者的通用情况下，作者力求最大化任何被驱逐张量的最小年龄）。当simMalloc找到一个张量的空闲槽位时，它应该选择最后一次使用在最早过去的可用槽位（在的开始处最近）。
参考
[1].Turnip: A "Nondeterministic" GPU Runtime with CPU RAM Offload.