GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching

作者/机构: Cong Guo∗ (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Rui Zhang∗ (Ant Group), Jiale Xu (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Jingwen Leng† (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Zihan Liu (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Ziyu Huang (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Minyi Guo† (Shanghai Jiao Tong University, Shanghai Qi Zhi Institute), Hao Wu (Ant Group), Shouren Zhao (Ant Group), Junping Zhao† (Ant Group), Ke Zhang (Ant Group)

A1 主要贡献

本文研究的核心问题是，在训练大规模深度神经网络（DNN），特别是大语言模型（LLM）时，现有的内存缩减技术（如重计算、卸载、分布式训练和低秩自适应）会引入频繁且不规则的内存分配/释放请求，这导致深度学习框架（如PyTorch和TensorFlow）中基于“分裂”机制的缓存分配器产生严重的内存碎片问题，碎片率高达30%。这种碎片化降低了GPU内存的利用率，并可能导致过早的内存不足（OOM）错误。

为了解决这一问题，本文的研究目标是设计并实现一种新颖的内存分配框架，以有效减少GPU内存碎片，提高内存利用率，同时对现有的DNN模型和内存优化技术保持完全透明。

本文提出的主要创新点是 GMLake (GPU Memory Lake)，一个基于低层GPU虚拟内存管理的内存分配框架。其核心机制是 虚拟内存拼接 (Virtual Memory Stitching, VMS)，该机制能够通过虚拟内存地址映射，将物理上不连续的内存块融合成一个逻辑上连续的内存块，从而有效利用碎片化的内存空间。如图1所示，传统的“分裂”方法虽然能提高内存池利用率，但在处理新的、较大的内存请求时会因无法找到足够大的连续空间而失败（OOM）。而GMLake的VMS机制可以将不连续的空闲块（如Block 2和5）拼接起来，以满足新的内存分配请求（如Block 6），从而消除内存碎片。

图1. 内存分配问题的代表性示例。原始的分裂方法可以提高GPU内存利用率，但会导致碎片化。我们提出的虚拟内存拼接可以补充和优化内存碎片问题。

本文的主要贡献如下：
* 我们进行了一项特性研究，表明在运行带有各种内存缩减技术（如重计算、卸载、分布式训练和低秩自adaptation）的大规模DNN模型时，现有DL框架中使用的缓存分配器会遭受高达30%的内存碎片。
* 我们设计并实现了GMLake，一种新颖的内存分配器，它有效减少了内存碎片并提高了内存利用率。GMLake对上述现有的内存缩减技术是透明的。它通过使用低层CUDA虚拟内存管理接口集成了虚拟内存拼接机制。
* 我们在多个著名的LLM优化平台上，使用一系列代表性的开源LLM对GMLake进行了评估，以证明其有效性、效率和鲁棒性。在最佳情况下，与PyTorch在A100 GPU（80GB HBM内存）上的原生缓存分配器相比，我们可以减少33%的GPU内存使用，相当于节省25GB的内存。

A3 背景知识与关键观察

2.1 大规模DNNs

LLM，如OpenAI的GPT系列【索引33，Employing large language models in survey research，2023，Natural Language Processing Journal；索引57，Recent advances in natural language processing via large pre-trained language models: A survey，2021，ACM Computing Surveys；索引88，A survey of large language models，2023，arXiv】，代表了大规模DNN的成功，并在各种语言处理任务中取得了显著进展。从GPT-2【索引66，Language models are unsupervised multitask learners，2019，OpenAI blog】到GPT-3【索引7，Language models are few-shot learners，2020，NeurIPS】，模型规模和复杂性急剧增加。然而，这些模型的庞大规模和复杂性给训练和部署带来了巨大挑战，例如OPT-175B【索引87，Opt: Open pre-trained transformer language models，2022，arXiv】需要在1024个A100 GPU上花费34天进行训练。这凸显了针对LLM训练进行高效内存管理的重要性。

2.2 DL框架的内存管理

像PyTorch【索引61，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】和TensorFlow【索引1，Tensorflow: Large-scale machine learning on heterogeneous distributed systems，2016，arXiv】这样的框架在DNN模型训练和推理中起着至关重要的作用。本文重点关注在GPU上对这些流行框架的内存管理优化，并比较了三种内存管理方式：GPU原生分配器、缓存分配器和虚拟内存（VM）分配器。

原生分配器。如图2(a)所示，原生分配器通过GPU厂商提供的API（即cudaMalloc和cudaFree）提供，这些API需要设备同步。其设计简单，缺乏灵活性，不适合需要动态和可调整大小内存结构的应用，尤其是在深度学习领域。实验表明，在PyTorch中禁用缓存分配器并使用原生分配器训练OPT-1.3B模型【索引87，Opt: Open pre-trained transformer language models，2022，arXiv】，其吞吐量比原始PyTorch分配器低9.7倍。因此，高效的内存管理设计是DL框架最关键的组件之一。

缓存分配器。DL框架通常使用带有内存池的缓存分配器，以实现无需设备同步的快速内存分配和释放。图2(b)描述了PyTorch和TensorFlow中采用的BFC（Best Fit with Coalescing）算法【索引76，Bfc allocator，2022，TensorFlow】。BFC算法主要有四个操作：1）“最佳适配”：搜索最合适的已分配但非活动的内存块。若无合适块，则调用原生GPU分配器API分配新内存块。2）“分裂”：若请求的内存小于最佳适配块，则将该块分裂为两块以提高内存利用率。一个块分配给请求，另一个保留在内存池中。3）“释放”：不调用原生cuMemFree，仅释放块的指针并将其状态设为非活动。4）“合并”：检查相邻块是否也为非活动状态，如果是，则将它们合并成一个大块。缓存分配器显著减少了原生API的调用次数，效率远高于原生分配器。然而，当内存分配请求不规则且大小差异巨大时，“分裂”机制会引发严重的内存碎片问题。对于以往的正则模型（如Transformer【索引78，Attention is all you need，2023，arXiv】），由于层级结构相同，张量大小一致，碎片问题不明显。但随着LLM规模的增长，分布式训练和复杂的内存策略加剧了碎片问题。

图2. 三种内存管理策略。

2.3 内存高效优化

为缓解大规模DNN模型日益增长的内存需求，学术界和工业界开发了多种系统级和算法级的优化方法，如重计算（Recomputation）【索引40，A graph theoretic framework of recomputation algorithms for memory-efficient backpropagation，2019，NeurIPS】、卸载（Offloading）【索引69，ZeRO-Offload: Democratizing Billion-Scale model training，2021，USENIX ATC】和低秩自适应（Low-Rank Adaptation, LoRA）【索引29，Lora: Low-rank adaptation of large language models，2021，arXiv】。重计算通过在反向传播期间重新计算特定层的输出来节省内存；卸载（如ZeRO-Offload）将优化器内存和计算从GPU转移到CPU；LoRA通过引入秩分解矩阵来减少可训练参数和GPU内存需求。然而，这些优化技术虽然有效减少了GPU内存占用，但其固有的动态和不规则内存分配请求却导致了内存利用率下降和碎片化加剧。

如图3所示，在四卡A100-80G上训练OPT-1.3B模型时，仅使用PyTorch（P）的内存利用率很高，而结合使用LoRA（L）、重计算（R）或卸载（O）等技术后，内存利用率显著下降。为了探究这种不规则性的来源，我们分析了GPT-NeoX-20B模型的内存足迹（图5）。与原始PyTorch相比，结合了LR（LoRA和重计算）优化的内存足迹显示出更多的不规则性。统计数据显示，原始PyTorch进行了4.6万次分配，平均大小为93MB；而优化后的版本有7.6万次分配，平均大小为85MB，这表明复杂策略导致了更频繁、更小规模的分配，从而引发碎片化。

图3. 五种方法组合下的内存利用率。
图4. 不同GPU数量下的内存利用率。

图5. GPT-Neox-20B训练的内存足迹。

这些结果促使我们解决内存碎片问题，以实现更高效、可扩展的大规模DNN模型训练。

观察 1：使用的内存优化策略越复杂、越不规则，产生的碎片就越多。

2.4 分布式训练

随着DNN模型复杂性的增加，分布式训练变得至关重要。数据并行（如PyTorch DDP【索引43，Pytorch distributed: Experiences on accelerating data parallel training，2020，arXiv】）和模型并行（包括流水线并行【索引42，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv；索引72，Megatron-lm: Training multibillion parameter language models using model parallelism，2019，arXiv；索引83，An efficient 2d method for training superlarge deep learning models，2023，IPDPS】和张量并行【索引28，Pipedream: Fast and Efficient Pipeline Parallel DNN Training，2018，arXiv；索引30，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS；索引45，Chimera: efficiently training largescale neural networks with bidirectional pipelines，2021，SC】）是常用的策略。然而，使用更多的GPU会因其不规则的内存分配和释放而导致更多的碎片。如图4所示，在PyTorch上对OPT-13B进行测试，当GPU数量从1增加到16时，内存利用率从>90%下降到76%。这种碎片化浪费了内存资源，并限制了LLM训练的批处理大小。

观察 2：随着GPU数量的扩展，内存碎片问题可能会变得更加突出。

2.5 低层虚拟内存（VM）管理

为了满足应用对快速高效内存管理日益增长的需求，CUDA引入了一项名为低层虚拟内存管理【索引62，Introducing low-level gpu virtual memory management，2020，NVIDIA Developer Blog】的新特性，类似于Windows的VirtualAlloc【索引56，Virtualalloc function (memoryapi.h)，2022，Microsoft Docs】和Linux的mmap【索引48，mmap(2) — Linux manual page，Linux man-pages project】。该特性提供了reserve和map等原始操作来操纵虚拟地址空间。我们称之为虚拟内存分配器，并展示其可用于减少大规模DNN训练中的内存碎片。

如图2(c)所示，其基本思想是：cuMemAddressReserve为新内存分配保留一个虚拟内存地址，cuMemCreate在GPU上分配物理内存块（不保证物理连续性），cuMemMap将物理句柄映射到保留的虚拟地址，从而桥接物理和虚拟内存。这种低层VM API的优势在于它可以分配和映射非连续的物理块，从而解决GPU内存碎片问题。然而，虚拟内存分配器的开销远大于原生GPU分配器。

图6. 原生分配器（第一个）和虚拟内存分配器的分配延迟。

为了验证VM分配器的开销，我们测试了三种不同总大小（512MB, 1GB, 2GB）的内存分配。图6显示了原生内存分配器和虚拟内存分配器的延迟对比。x轴表示构成内存块的内部物理块的大小。例如，一个1GB的内存块需要映射512个2MB大小的物理块。结果显示，虚拟内存的延迟非常高。特别是当虚拟内存块被划分为2MB的物理块时，其速度比原生分配器慢100倍以上，这是完全不可接受的。

表1. VMM API执行时间分解（以cuMalloc执行时间为基准进行归一化）。

为了进一步探究VMM API的瓶颈，我们分析了其执行时间分解。表1显示了分配2GB GPU内存时VMM API的延迟分解，所有延迟都相对于cuMalloc进行了归一化。GMLake中的每次分配只需要一个cuMemAddressReserve，但需要为每个物理块调用多次cuMemCreate、cuMemMap和cuMemSetAccess（一个使映射可用的特殊函数）。可以看出，使用2MB的小块来分配2GB内存比原生cuMalloc慢115倍。

观察 3：尽管虚拟内存可以减少内存碎片，但GPU上原始的虚拟内存分配器仍存在许多挑战，需要进一步优化。

A2 方法细节

GMLake

在这项工作中，我们介绍GMLake，一种专为GPU内存设计的高效内存分配方法（即GPU内存湖）。GMLake利用CUDA的低层VM管理特性来加速内存分配和释放过程。图7概述了GMLake，它向现有的缓存分配器提供了相同的内存（去）分配接口，但内部集成了虚拟内存拼接（VMS）机制。这种集成是通过精确利用CUDA的低层VM管理API实现的。

图7. 缓存分配器和GMLake的概述。

GMLake的组成。原始DL框架中的缓存分配器采用BFC算法。为了避免同步并提高内存管理效率，这些框架内部管理内存池以处理（去）分配，而不是直接使用原生API。遵循这种方法，我们的GMLake也包含以下三个组件：
* 虚拟内存API：指用于指示GPU使用虚拟内存地址分配和释放内存的低层API，如果未经充分优化，该过程会产生巨大开销。
* 虚拟内存池：作为基础数据结构，专为缓存虚拟内存而设计。其实现对于提高效率至关重要。
* GMLake分配器：包括管理VM池所需的所有功能、算法和策略。
本节将从基础到顶层详细描述这三个组成部分的设计和细节。

3.1 虚拟内存API

pBlock的设计与操作。如2.5节介绍和图2(c)所示，低层VMM API是GPU和应用程序之间的基础接口。如图8底部所示，我们利用VMM API构建了primitive block (pBlock)，这是GMLake分配器的关键数据结构。pBlock内的操作包括：
* AddrReserve: 首先，分配一个pBlock需要指定分配大小并保留相应的虚拟地址（VA）。
* Create: 接着，pBlock创建物理上存储数据的物理块。
* Map: 最后，pBlock将所有物理块映射到虚拟地址，使张量能够无缝访问。

优化策略。为了优化碎片整理，我们对所有物理块应用了统一的2MB大小。尽管这种2MB块大小的开销很大（见2.5节），但可以通过高效的数据结构和精心设计的拼接策略来缓解，从而实现最佳的碎片整理效果和与PyTorch代码库的最佳兼容性。同时，我们使用以下DNN特定的优化来减少拼接频率，使其端到端开销可以忽略不计：GMLake使用VMM处理大于2MB的分配。对于小于2MB的内存分配，我们使用原始PyTorch缓存分配器的分裂方法来处理其内部碎片问题。此外，在LLM训练中，小于2MB的分配很少见。Map操作是虚拟内存拼接的基础，它允许我们连接多个在物理内存中可能不连续的物理连续块。通过使用VMM API，我们可以将虚拟内存组织到虚拟内存池中，这是GMLake分配器的底层数据结构。

3.2 虚拟内存池

设计动机与结构。鉴于原始VMM API耗时较长，减少其使用次数对GMLake实现高效率至关重要。受缓存分配器的启发，我们设计了具有缓存功能的虚拟内存池（VMP），从而显著减少了物理内存（去）分配的次数。如图8所示，我们区分了两种类型的内存池：原始内存池（pPool）和拼接内存池（sPool）。

图8. 原始和拼接内存池的数据结构。

pPool和pBlock。pPool的数据结构使用一个有序集合来存储pBlocks。对于每个pBlock，pPool首先构建一个结构来记录指向pBlock的指针，包括pBlock的活动状态等基本属性。随后，新分配的pBlock被插入到集合中，所有pBlocks按块大小降序排列。pBlock作为原始块，代表了高层张量可访问的最小单元，并且它是一个可以被多个sBlocks拼接和指向的基础数据结构。

sPool和sBlock。sPool也组织成一个有序集合，类似于pPool。其元素包括拼接块结构，该结构集成了多个pBlocks。例如，如图8所示，sBlock 3包含了被拼接在一起的pBlock 3和pBlock 5，而pBlock 3也可能被sBlock 2指向和拼接。因此，sBlock的属性受pBlocks的影响，即只要有一个pBlock是活动的，所有对应的sBlocks都被标记为活动。在实践中，sBlock将虚拟内存重新映射到所指向的pBlocks的所有物理块上，使其可被高层张量访问。为简化流程，我们规定sBlock只能被分配给与sBlock大小相匹配的张量分配。

层级关系。比较物理块、pBlock和sBlock，揭示了它们在不同级别上的数据结构角色。物理块由低层API控制，对高层张量保持透明；而pBlock和sBlock分别位于pPool和sPool中，为高层张量的访问提供虚拟内存地址。此外，sBlock在更高级别上运作，由多个pBlocks组成。接下来，我们将描述GMLake如何使用这些数据结构来实现高效的内存管理。

3.3 分配器

分配器包含内存分配和释放所需的所有基本功能和算法。由于篇幅限制，我们仅简要解释分配和释放模块中最重要的功能。

3.3.1 分配模块

Alloc函数。该函数负责分配一个新的pBlock并将其插入pPool，如图8所示。它是分配新物理块和增加已分配GPU内存的唯一接口。

Split函数。该函数将一个pBlock（原始块）分割成两个较小的pBlocks，类似于图2中描述的“Split”操作，但底层实现完全不同。具体来说，GMLake中的Split函数基于pBlock结构操作，产生两个新的pBlocks，它们具有相应的虚拟内存地址和重新映射的物理块。之前的pBlock结构随后从pPool集合中移除。

Stitch函数。该函数是创建sBlock并将其插入sPool的唯一机制，如图8顶部所示。此函数是我们分配器的一个组成部分，可以将多个pBlocks拼接成一个sBlock。我们使用VMM API（即NVIDIA专门为虚拟内存管理提供的低层API）来“拼接”两个pBlocks，如图2(c)所示。假设我们有两个pBlocks，p1（1GB）和p2（2GB）。我们采用VMM API cuMemCreate创建相应的物理块，并通过cuMemAddressReserve保留虚拟地址（VA）。然后，VA使用cuMemMap映射物理地址（PA）。实际上，我们不需要取消p1和p2的原始VA-PA映射，因为VMM中的PA可以被多个VA指向。因此，我们只需使用cuMemAddressReserve为sBlock s1保留一个3GB的VA。所有sBlocks从不通过cuMemCreate创建新的物理块。我们使用cuMemMap API将s1的VA（3GB）映射到p1和p2的物理块。由于多个sBlocks可以包含相同的物理块，我们需要更多属性（如pBlock的活动状态）来确保每个物理块仅由单个张量使用。

BestFit函数。该函数识别最适合内存分配的pBlock或sBlock，返回状态和候选块以供后续处理。如算法1所详述，我们设计了四种状态，覆盖了GMLake可能遇到的所有场景。它假设sPool和pPool都按大小降序排列。
* 精确匹配 (第2-4行)：当候选块的大小与分配大小匹配时出现。该块可以是sPool中的sBlock或pPool中的pBlock。这是sBlock可以被分配用于新分配的唯一情况。所有其他状态仅涉及pBlock。
* 单个块 (第12行)：此时，BestFit识别出大于请求分配大小的最佳适配（最小）pBlock。
* 多个块 (第14行)：在所有pBlocks都小于所需分配大小，但其总大小满足分配要求的情况下，BestFit函数会贪婪地寻找多个候选pBlocks进行拼接。
* 块不足 (第16行)：当没有足够的pBlocks来满足请求的分配大小时发生，尽管BestFit仍会返回一个块列表。

GMLake替代了PyTorch缓存分配器模块的几个内部函数。“拼接”操作对用户完全透明，不增加用户修改代码的负担。我们所采用的CUDA API不仅包括VMM相关的API，还包括像cuMemAlloc这样的常规API。GMLake API的实现利用CUDA API实现细粒度的内存拼接和重用，因此它们属于不同层次并服务于相应的功能。

3.3.2 释放模块

设计原则。释放模块避免使用低层VMM API主动释放物理GPU内存，而是仅更新或恢复拼接的虚拟内存块。

Update函数。当收到一个高层张量的释放请求时，我们用Update函数替换原始的VMM释放函数。此函数更改活动pBlocks和sBlocks的状态，从而实现张量与块之间的链接和分配的移除。在整个程序运行期间，实际的物理内存仍然由相应的pBlock控制。

StitchFree函数。此函数的目的是释放sPool中持有的最近最少使用（LRU）的sBlocks。由于空间限制，我们在此省略了复杂的细节。我们已经实现了完整的算法和数据结构来支持基于LRU的StitchFree。值得注意的是，我们只从sPool中释放非活动的sBlock结构。

Defragmentation Strategy

本节中，我们介绍减少内存碎片问题的策略。我们首先提出一个基于GMLake分配器的复杂算法，该算法理论上可以消除所有碎片。然后，我们讨论并描述我们的优化措施，以保证其效率和鲁棒性。

4.1 内存分配策略

整体流程。图9展示了GMLake为新的张量分配请求分配或指派新内存块的策略。该策略基于BestFit模块提供的四种状态，我们为每种状态设计了后处理步骤。

图9. GMLake中的内存分配策略。

• 状态S1（精确匹配）: 对于分配请求1，如果找到大小完全匹配的现有pBlock或sBlock，则立即返回该块。
• 状态S2（单个块）: 如果在非活动的sPool和pPool中未找到精确匹配的块，我们进入状态S2。BestFit函数会提供一个比请求大的pBlock（pBlock 1）。该块通过Split操作被分割成两个不同的pBlocks，并都插入到pPool中。新创建的pBlock 1替换其前身并分配给请求2。同时，GMLake执行Stitch函数，将pBlock 1和pBlock 2合并成sBlock 1，并将其添加到sPool中。
• 状态S3（多个块）: GMLake通过Stitch操作合并多个pBlocks，为分配请求3创建一个整合的sBlock 3。如果需要，最后一个pBlock可以像S2中的过程一样被Split。此阶段结束时，新的pBlocks（特别是pBlock 4和pBlock 5）被添加到pPool中，而sBlock 2和sBlock 3被添加到sPool中。
• 状态S4（块不足）: 对于分配请求4，可能没有足够的可用块进行拼接和分配。此时，GMLake触发Alloc函数，使用低层API创建带有新物理块和相应虚拟内存地址的pBlock 7。这个新的pBlock被添加到pPool中。此外，我们将pBlock 5和pBlock 7通过Stitch合并成一个新的sBlock 4，返回给分配请求4并添加到sPool中。如果没有合格的候选块（即不存在pBlock 5），GMLake会直接分配pBlock 7，而不使用Stitch函数。
• 状态S5（内存溢出）: 如果Alloc函数调用失败，GMLake在状态S5中立即报告内存不足（OOM）错误。

4.2 策略分析

我们从有效性、效率和鲁棒性三个方面分析GMLake的内存分配策略。

有效性。GMLake分配策略有效地确保了GPU系统中几乎所有的碎片都被消除。
* 接口设计：我们将所有操作整合到三个主要函数中：Alloc、Split和Stitch。Alloc是唯一可以创建新pBlock的函数，只有Stitch可以生成新的sBlock，而Split不增加已分配的内存。
* 数据结构：pPool代表了GPU内存的严格一对一映射，每个pBlock都与其他pBlock不同，它是一个没有重复元素和重叠地址的已分配GPU内存集合。sPool被设计用于存储sBlocks，以类似于软链接机制的方式保留指向pBlocks的链接和指针。GMLake禁止分裂sBlocks，因为这可能会影响pBlocks。sPool被视为pPool的一个子集。
最终，每当程序达到GPU内存使用的新高峰时（例如，调用Alloc函数），pPool可能无法提供足够的块进行拼接和分配，从而导致内存被完全利用而没有碎片。这与原始的缓存分配器形成对比，后者可能会留下许多未使用的子块。

效率。我们采用了几种方法来实现高效率。
* 算法复杂度：S3中的算法问题是一个经典的NP-hard打包问题【索引55，An exact approach to the strip-packing problem，2003，INFORMS journal on Computing】。然而，通过应用Split和Stitch函数，可以生成一个精确匹配的块来适应分配，从而实现了线性复杂度。
* 利用DNN训练的规律性：由sPool和pPool组成的联合集合记录了每次张量分配的所有大小和相应块，就像一个记录DNN模型张量分配模式的磁带。幸运的是，DNN模型训练是高度规律化的，因为每次迭代都处理相同的模型参数和输入数据大小。因此，经过几次迭代后，GMLake将不再执行S2、S3和S4。GMLake将只使用S1“精确匹配”策略来完成剩余的训练，这与原始缓存分配器需要持续进行分裂和合并操作形成对比。

鲁棒性。在实践中，由于GPU的总容量限制，拼接和创建新的sBlocks不能无限进行。此外，过多的拼接操作在sPool上运行分配模块（如BestFit）时会损害GMLake分配器的效率。当总容量超过此限制或阈值时，GMLake会使用StitchFree来释放LRU的sBlocks并清除模式磁带，从而作为鲁棒性的后备机制。此外，当DNN训练呈现出极其不规则的模式时，可能会产生大量小块，导致频繁的分裂和拼接，从而提前达到限制。为避免不必要的性能损失，我们建立了一个最小碎片限制。如果一个块小于此限制，GMLake将避免对其进行拼接或分裂。因此，所有算法和模块都遵循碎片限制（例如，128MB），以确保DNN训练的高效率和鲁棒性。

A4 实验环境

• 硬件配置:

  • 单节点: 一台配备Intel Xeon Platinum 8369B CPU、1TB DRAM和8个通过NVLink互连的NVIDIA A100 GPU（每个80GB显存）的服务器。
  • 多节点: 两台与单节点配置相同的服务器，通过RDMA进行分布式训练。

• 软件配置:

  • 代码实现: GMLake用5000行C++代码实现，并集成到PyTorch的缓存分配器中。
  • 依赖库: 适配PyTorch-1.13.1和PyTorch-2.0版本，CUDA 11.4，cuDNN 8.5。
  • 训练平台: Deepspeed, FSDP, Colossal-AI。

• 模型与数据集:

  • 模型: 评估了一系列代表性的开源LLM，包括OPT-1.3B, GPT-2, GLM-10B, OPT-13B, Vicuna-13B, GPT-NeoX-20B。
  • 数据集: 使用了各模型官方示例中的默认数据集，因为微调期间的内存使用不受数据集质量影响。
  • 具体的模型、策略和DDP框架组合见下表。

表2. 基准测试规格。(L: LoRA; R: 重计算; O: 卸载; FSDP: 全分片数据并行.)

A4 实验结果

GMLake在8个不同模型的76个工作负载中，平均减少了15%（最高33%）的碎片率，平均减少了9.2GB（最高25GB）的预留GPU内存。

5.2 GMLake的可扩展性

• 对内存高效策略的可扩展性:
  • 实验内容: 在OPT-13B、Vicuna-13B和GPT-NeoX-20B模型上，使用Deepspeed Zero3和四块A100 GPU，评估了不同内存高效策略组合（无策略N、重计算R、重计算+LoRA LR、重计算+卸载RO、重计算+LoRA+卸载LRO）下的性能。
  • 实验结果: 如图10所示，与PyTorch相比，GMLake将利用率提高了5%至24%，预留内存减少了约10GB至17GB。PyTorch在使用复杂优化策略时碎片率超过20%，而GMLake能有效将碎片率控制在5%到10%。这表明GMLake能透明且有效地与这些日益复杂的优化策略协同工作。

图10. 在内存高效策略组合上的内存利用率比较。

• 对GPU横向扩展的可扩展性:
  • 实验内容: 使用LR策略和Deepspeed平台，将GPU数量从1个扩展到16个，评估GMLake的性能。
  • 实验结果: 如图11所示，GMLake始终表现出较低的碎片率（高利用率）和预留内存消耗。随着GPU数量增加，PyTorch的利用率逐渐下降，而GMLake能有效维持约90%的利用率。这是因为分布式策略（如DeepSpeed ZeRO-3）将张量切分得更小，加剧了碎片化，而GMLake能很好地应对这一挑战，同时保持与PyTorch相似的高吞吐量，显示了其低开销和优异的可扩展性。

图11. 在GPU横向扩展上的内存利用率比较。

• 对不同平台的可扩展性:
  • 实验内容: 在Deepspeed、FSDP和Colossal-AI等不同平台上，对OPT-13B、GLM-10B和GPT-2模型进行微调。
  • 实验结果: 如图12所示，GMLake在各种优化训练平台上均表现出高可扩展性，碎片减少了约9%至33%，预留内存减少了7GB至25GB。

图12. 在不同平台上的内存利用率比较。

5.3 GMLake的端到端有效性

• 实验内容: 通过改变批处理大小，对一系列LLM（1.3B至20B参数）进行端到端微调，比较GMLake和PyTorch缓存分配器的性能。
• 实验结果: 如图13所示，GMLake在各种模型和批处理大小下都显著降低了峰值内存消耗。随着模型规模增大，内存效率超过95%，而基线PyTorch难以达到80%。更重要的是，GMLake使得在PyTorch会因OOM错误而失败的大批量场景下能够成功训练。吞吐量方面，GMLake与基线相当，在极大批处理大小下甚至可能因其高效的内存管理而超过基线。

图13. 在不同批处理大小下，端到端有效性的内存利用率和吞吐量比较。

5.4 内存轨迹分析

• 实验内容: 在4个GPU上使用LoRA和重计算策略，追踪GPT-NeoX-20B模型的内存分配行为。
• 实验结果与分析: 如图14所示：
  1. PyTorch在大约200秒时因OOM异常终止，而GMLake在此批处理大小下正常运行。
  2. GMLake和PyTorch的活动内存水平相当，但GMLake的预留内存远低于PyTorch，这直接证明了GMLake在解决碎片问题上的有效性。
  3. 内存轨迹显示了LLM微调阶段（特别是前向和后向传播）的周期性内存模式。
  4. GMLake在四次迭代后达到稳定状态，并实现了与PyTorch相同的吞吐量。这表明GMLake的分配策略能迅速适应训练过程中的内存波动，找到最佳的缓存和拼接策略，并通过重用拼接的sBlocks来摊销其成本。

图14. 在批处理大小为72时，GPT-NeoX的内存轨迹。

A7 补充细节

相关工作与讨论

内存碎片整理。内存碎片整理在不同领域已有广泛研究，包括使用细粒度固定大小块的方法【索引73，Eliminating external fragmentation in a non-moving garbage collector for java，2000，CASES】，基于数据移动的压缩策略【索引59，The z garbage collector，2023，openjdk；索引79，Parallel memory defragmentation on a GPU，2012，MSPC】，以及复制式垃圾回收系统【索引31，Sapphire: copying GC without stopping the world，2001，Java Grande；索引53，Parallel generational-copying garbage collection with a blockstructured heap，2008，ISMM；索引74，Improving locality with parallel hierarchical copying GC，2006，ISMM】。GMLake采用的基于拼接的技术在概念上与将小块合并为大块相似，但它最小化了频繁数据移动和复制的需求，从而显著提高了内存效率。

高效LLM。针对基于Transformer的LLM，内存已成为计算系统中最宝贵的资源。研究人员提出了多种算法优化来抑制内存消耗，如量化、剪枝和KV-cache压缩。系统级内存优化方面，vLLM【索引82，vllm: Easy, fast, and cheap llm serving with pagedattention，2023，https://vllm.ai/】利用基于页面的虚拟内存管理技术提升资源效率；FlashAttention【索引13 ，Flashattention-2: Faster attention with better parallelism and work partitioning，2023，CoRR；索引14，Flashattention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS】通过分块技术优化注意力计算；FlexGen【索引71，High-throughput generative inference of large language models with a single GPU，2023，CoRR】则优化内存-计算的安排。在这一背景下，GMLake作为一个用户透明的内存管理系统，通过智能高效的内存重用，为LLM训练提供了新的优化维度。

与其他工作的新颖性对比。GMLake在一个独特的内存作用域（DNN训练）上工作，这与vLLM和vMalloc/CUDA VMM不同。

• 与vLLM对比: vLLM是针对自注意力操作的算法级解决方案，作用于张量内部，通过查找表移除填充的token。而GMLake是一个更通用的内存池管理器，面向整个DNN训练过程，应用场景更广。
• 与vMalloc/CUDA VMM对比: 这些是低层的系统工具，用于碎片整理。然而，它们无法感知DL框架中普遍使用的内存池设计。没有内存池设计，性能会显著下降。因此，这些系统级内存工具不能直接应用于DL框架。GMLake的创新在于将VMM的能力与缓存分配器的设计相结合，解决了特定于DL训练的碎片问题。

表3. 与相关工作的技术差异。

A5 结论

本研究介绍了GMLake，一种高效、低碎片的新型内存分配器。它基于低层CUDA虚拟内存管理接口，将多个非连续的物理内存块整合成一个单一、连续的实体，从而缓解了碎片问题。为了无缝集成到现有的DL框架中，我们制定了一套虚拟内存管理API。此外，我们提出了一种多阶段的碎片整理策略，以保证分配的效率和鲁棒性。我们的评估表明，GMLake将内存碎片降低到5%至10%，同时保持了与基线相当的吞吐量。