【翻译】GMLake_ASPLOS 2024

粗略翻译一篇 GPU 内存碎片优化领域比较有影响力的论文。翻译过程使用了 DeepSeek，如有翻译不准确的地方还请读者直接阅读英文原文。

• 论文地址：[2401.08156] GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching
• 代码仓库：glake/GMLake at main · antgroup/glake
• 已有翻译：[2401.08156] GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching
• 解读博客：【分布式训练技术分享十】高效显存管理技术GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN

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

GMLake：基于虚拟内存拼接的高效透明 GPU 内存碎片整理方法，用于大规模深度神经网络训练

0 Abstract

大规模深度神经网络（DNNs），如大语言模型（LLMs），已彻底改变人工智能（AI）领域，并变得越来越流行。然而，训练或微调此类模型需要大量计算资源，其中，单个加速设备（如 GPU）的内存容量是最关键的瓶颈之一。由于 GPU 原生内存分配器的开销极其巨大（相较于现有缓存分配器约 10 倍的开销），PyTorch 和 TensorFlow 等 DNN 框架采用了一种缓存分配器，该分配器通过维护具有分割机制的内存池来实现快速的内存分配与释放（deallocation）。然而，在 recomputation（重计算）、offloading（卸载）、distributed training（分布式训练）和 low-rank adaptation（低秩适配）等常见的内存优化技术下，缓存分配器的效率会迅速下降。主要原因是这些内存优化技术会引入频繁且不规则的内存分配与释放请求，导致基于分割机制的缓存分配器出现严重的内存碎片化问题。

为缓解这一碎片化问题，我们提出了一种基于 GPU 低级虚拟内存管理的新型内存分配框架——GPU memory lake（GMLake）。GMLake 采用了一种新颖的虚拟内存拼接（virtual memory stitching，VMS）机制，该机制可以通过虚拟内存地址映射融合或组合非连续的内存块。GMLake 在 A100 GPU（80 GB 内存）上针对八个 LLM 进行测试，平均可减少 9.2 GB（最高可达 25 GB）的 GPU 内存使用，并降低 15%（最高可达 33%）的内存碎片率。GMLake 对 DNN 模型和内存优化技术完全透明，能够确保资源密集型深度学习任务的无缝执行。

关键词：内存碎片整理（Memory Defragmentation）、GPU（Graphics Processing Unit，图形处理器）、深度学习（Deep Learning）、虚拟内存拼接（Virtual Memory Stitching）

1 Introduction

大规模深度神经网络（DNN）模型，特别是大语言模型（LLMs），已经彻底改变了自然语言处理（NLP）和人工智能（AI）研究。LLMs（如 GPT-3 架构）是一类复杂的 DNN 模型，在理解、生成和处理人类语言方面展现出卓越的能力。这些模型利用海量文本数据，并采用基于 Transformer 的架构，该架构以注意力机制为核心，从而在各种 NLP 任务上实现最先进（state-of-the-art，SOTA）的性能。 然而，LLMs 的广泛应用伴随着巨大的计算挑战，因为训练或微调此类模型需要大量的计算资源。例如，具有 1750 亿参数的 OPT 模型在 1024 张 A100 GPU 上需要 34 天完成训练，而参数规模为 650 亿的 LLaMA 模型在 2048 张 A100 GPU 上处理 1.4 万亿个 token，训练耗时约 21 天。

因此，深度学习（DL）框架（如 PyTorch 和 TensorFlow）因其灵活性和计算效率，已成为 DNN 模型的基础基础设施。这些 DL 框架使得越来越大、越来越复杂的神经网络模型得以训练。同时，GPU 架构已成为支持 DNN 模型高性能执行的最广泛使用的硬件。另一方面，DNN 模型规模和复杂度的不断增长也给 GPU 内存管理带来了新的挑战。例如，使用 CUDA 的原生内存分配 API（如 cudaMalloc 和 cudaFree）会引入较大的开销。为了提高 GPU 内存分配的效率，DL 框架通常采用缓存分配器（caching allocator），该分配器通过最佳适配合并（best fit with coalescing，BFC）算法维护一个内存池。我们的实验表明，相较于原生内存分配器，缓存分配器的性能提升近 10 倍。

另一方面，大规模 DNN 模型对内存需求的快速增长推动了系统级和算法级的多种方法来缓解内存压力。其中包括重计算（recomputation）、卸载（offloading）、分布式训练（distributed training）和低秩适配（low-rank adaptation）等技术。 尽管这些优化方法能够有效减少训练或微调大规模 DNN 模型的内存占用，但它们可能会导致较低的内存利用率。其原因在于，这些方法会引入大量规律性和动态性的内存分配请求，从而导致 GPU 内存碎片率最高可达 30%。

图 1. 典型的内存分配问题示例。原始的分割方法能够提高 GPU 内存利用率，但会导致碎片化问题。而我们提出的虚拟内存拼接方法可以弥补并优化内存碎片化问题。

如图 1 左侧所示，DL 框架在内存池内管理内存分配，并采用“分割”（splitting）方法将内存池拆分，以适应 DNN 张量的任意大小，从而提高内存池的利用率。然而，这种方法在进行新内存分配时可能会导致严重的内存碎片化问题。例如，框架将第三行的内存块拆分，以存储新分配的 Block 4。但由于 Block 6 的大小大于 Block 5，而 Block 5 无法被利用，因此该部分内存变成碎片，导致内存池无法容纳 Block 6。最终，框架会报告内存不足（out-of-memory，OOM）错误，这是 DNN 模型训练过程中最常见的问题之一。前述的内存优化技术（如重计算和卸载）虽然能够缓解 OOM 问题，但它们会引入更加频繁且不规则的内存分配与释放请求，从而进一步加剧碎片化问题。

为了缓解 GPU 内存碎片化并提高内存利用效率，本研究深入分析了 GPU 内存碎片化的成因，并提出了一种基于底层 GPU 虚拟内存管理的全新内存分配框架——GPU 内存池（GPU memory lake，GMLake）。该框架能够以较低的开销优化 GPU 内存管理。如图 1 右侧所示，GMLake 采用了一种新颖的虚拟内存拼接（virtual memory stitching，VMS）机制，其行为在某种程度上与传统的“分割”操作相反。与原始框架相比，VMS 能够通过虚拟内存地址映射融合或组合非连续的内存块。例如，VMS 可以将 Block 6 映射到由 Block 2 和 Block 5 组成的拼接块，并将 Block 6 存储在 Block 2 和 Block 5 的物理内存块中。显然，虚拟内存拼接能够有效减少内存碎片并提高内存利用率。我们在 DL 框架的底层实现了 GMLake，并替换了 DNN 训练的原始内存分配 API。GMLake 对 DNN 模型及其他内存优化方法（如重计算和卸载）完全透明，从而确保资源密集型深度学习任务的无缝执行。

总体而言，本研究做出了以下贡献：

• 我们进行了表征研究，表明现有 DL 框架中使用的缓存分配器在运行大规模 DNN 模型时，采用重计算、卸载、分布式训练和低秩适配等内存优化技术会导致最多 30% 的内存碎片化。
• 我们设计并实现了 GMLake，一种新型内存分配器，有效减少内存碎片化并提高内存利用率。GMLake 对上述现有内存优化技术完全透明，利用底层 CUDA 虚拟内存管理机制集成了虚拟内存拼接机制。
• 我们在多个知名的 LLM 优化平台上对 GMLake 进行了评估，使用一组代表性的开源 LLMs 来验证其有效性、效率和鲁棒性。在最佳情况下，我们可以将 GPU 内存使用量减少 33%，相比于 PyTorch 中的原生缓存分配器，在 A100 GPU（80 GB HBM 内存）上节省 25 GB 内存。

2 Background and Motivation

在本节中，我们提供了关于 DL 框架内存管理的必要背景，并通过实验观察概述了开展本研究的动机。首先，我们简要概述了大规模 DNN（如 LLM）增长的趋势。随后，我们进行了各种内存管理方法的比较分析。在比较之后，我们揭示了在分布式训练和内存高效优化策略的背景下，LLM 模型面临的显著碎片化挑战。因此，我们迫切需要开发一种新型且高效的内存分配器，以有效解决这些挑战。

2.1 Large-scale DNNs

LLM（如 OpenAI 的 GPT 系列）代表了大规模 DNN 的成功，并在多个语言处理任务中取得了显著进展。GPT-2 相较于其前身具有十倍的模型规模和复杂度，随后是拥有 1750 亿参数的 GPT-3，以及针对对话进行了微调的 ChatGPT。

然而，这些模型的规模和复杂度在训练和部署中带来了巨大的挑战。对庞大计算资源、海量数据和长时间（例如，OPT-175B 在 1024 A100 GPU 上训练需要 34 天）的需求，促使人们更加关注高效的训练优化。因此，本研究强调了在 LLM 训练中进行高效内存管理的重要性。

2.2 Memory Management of DL Framework

像 PyTorch 和 TensorFlow 这样的框架在 DNN 模型的训练和推理中起着至关重要的作用。同时，GPU 已成为高性能模型执行的重要硬件。本研究聚焦于这些流行框架在 GPU 上的内存管理优化。我们比较了三种类型的内存管理方法：GPU 原生分配器、缓存分配器和虚拟内存（VM）分配器，并通过多次实验展示了每种分配器的效率和相关开销。

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

原生分配器（Native Allocator）。如图 2(a) 所示，原生分配器通过 GPU 厂商提供的 API 提供，即 cudaMalloc 和 cudaFree，这些函数需要设备同步。原生分配器的设计简单，缺乏灵活性。这使得它不适用于需要动态和可调整大小内存结构或复杂内存管理的应用，尤其是在深度学习（DL）中。如果 DL 框架实现原生 GPU 分配器而没有适当的同步优化，它可能会在训练 DNN 模型时导致不可接受的开销。

我们的实验结果量化了原生分配器的开销。我们禁用 PyTorch 的缓存分配器（将在下文介绍）并在四个 A100-80G GPU 上训练 OPT-1.3B 模型。PyTorch 中的原生分配器为用户提供了相同的编程模型，用户可以通过设置环境变量来更改分配器。GPU 原生分配器的吞吐量比原始 PyTorch 分配器低 9.7 倍。因此，高效的内存管理设计应当是 DL 框架中最关键的组件之一。

缓存分配器（Caching Allocator）。深度学习框架通常使用带有内存池的缓存分配器来实现快速的内存分配和释放，而无需设备同步。图 2(b) 展示了 PyTorch 和 TensorFlow 中使用的缓存分配器的 BFC 算法 [76]。PyTorch 和 TensorFlow 中的 BFC 实现几乎相同，唯一的区别在于它们的数据结构。

BFC 算法主要有四个操作步骤：

1. 寻找最佳匹配块：算法首先寻找最合适的已分配但未激活的内存块，称为“最佳匹配” （best fit）。如果没有合适的已分配内存块候选，缓存分配器会调用原生 GPU 分配器 API 来分配新的内存块。
2. 拆分内存块：如果请求的内存小于最佳匹配块，算法将拆分该块为两个块以提高内存利用率。其中一个拆分块分配给内存请求，而另一个保持在内存池中以供将来重新分配。为了有效管理内存，这两个块通过双向链表连接，每个块都会监控相邻块的可用状态。
3. 释放操作：对于释放（去分配）操作，算法不调用原生 GPU API 的 cuMemFree，而仅释放块的指针并将该块设置为非活动状态。
4. 合并操作：最后，缓存分配器会检查左侧或右侧的相邻块是否也为非活动状态。如果是，缓存分配器将这些相邻的非活动块合并为一个块。

显然，缓存分配器显著减少了原生 GPU 内存分配器 API 的调用次数。在最好的情况下，所有内存块只通过原生分配器分配和释放一次。因此，缓存分配器比原生 GPU 分配器更高效，且广泛应用于现有的深度学习框架中。

然而，“拆分”机制会导致当内存分配请求不规则且其大小差异较大时，出现大量潜在的内存碎片问题。此前由于模型通常比较规则且不够大，这个碎片化问题并不显著。例如，基于 Transformer 的模型 [78] 是由多个相同层堆叠而成，每层的张量大小相同，从而导致的内存碎片问题较小。

然而，随着 LLM 的规模增长，碎片化问题随着分布式训练和复杂训练内存策略的引入而显著恶化，导致批次限制、内存管理低效及训练不畅。在接下来的子节中，我们将观察这些复杂训练场景中，碎片化问题如何变得更加具有挑战性。

2.3 Memory-efficient Optimization

大规模深度神经网络（DNN）模型的内存需求迅速增长，这推动了在系统级和算法级开发方法以缓解内存需求。此类方法的例子包括重计算、卸载和低秩适应。

• 重计算（Recomputation），也称为检查点（checkpointing）法，在反向传播过程中重新计算特定层的输出，而不是存储它们，从而节省内存。
• 卸载（Offloading），也称为交换（swap），如 ZeRO-Offload，将优化器的内存和计算从 GPU 转移到 CPU，使得在单个 GPU 上训练大型模型，并能够在多个 GPU 之间扩展。
• 除了这些系统级方法，算法方法也能有效减少内存需求。例如，专为大规模模型设计的 LoRA 引入了秩分解矩阵，以最小化可训练参数和 GPU 内存需求，在减少成本和时间的同时，达到类似于全模型微调的准确性。

然而，尽管这些优化方法可以有效减少 GPU 内存占用，但有时可能会导致内存利用率较低。如图 3 所示，我们在四个 A100-80G GPU上使用不同的优化方法组合训练 OPT-1.3B 模型。仅使用 PyTorch（P）时，内存利用率较高，而采用如 LoRA（L）、重计算（R）或卸载（O）等技术时，内存利用率显著降低。根据我们的调查，组合这些技术会导致较高的内存碎片化。原因在于，这些内存优化技术本身会引发动态和不规则的内存分配请求。

图 3. 五种方法组合的内存利用率。

为了探讨这种不规则性的来源，我们展示了在 GPT-NeoX20B 模型训练过程中的内存占用情况。如图 5 所示，左图显示了原始 PyTorch 的内存占用情况，右图则是使用 LR（LoRA 和重计算）优化后的 PyTorch 收集的数据。显然，由于像重计算这样的使用策略，右图显示的内存占用更加不规则。统计数据显示，左图进行了 46 千次分配，平均大小为 93 MB，而右图进行了 76 千次分配，平均大小为 85 MB，表明复杂策略导致了更频繁且更小的内存分配，从而引发了碎片化。

图 5. GPT-Neox-20B 训练的内存占用。

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

观察1：使用的内存优化策略越复杂和不规则，碎片化问题就越严重。

2.4 Distributed Training

分布式训练不在笔者关注的范围，感兴趣还请直接阅读英文原文。

观察 2：随着 GPU 数量的增加，内存碎片化问题可能会变得更加严重。

2.5 Low-level Virtual Memory (VM) Management

鉴于应用程序对快速高效的内存管理需求日益增长，CUDA 引入了一种名为低级虚拟内存管理（low-level virtual memory management）的新特性，这类似于 Windows 的 VirtualAlloc 和 Linux 的 mmap。该特性打破了 malloc 类似的抽象，并提供了诸如保留（reserve）和映射（map）等原始操作，以操控虚拟地址空间。在我们的工作中，我们展示了这种低级虚拟内存管理可以用于减少内存碎片化并提高大规模 DNN（深度神经网络）训练的内存利用率，我们将其称为虚拟内存分配器（virtual memory allocator）。

图 2© 说明了使用这种低级虚拟内存管理的基本思想。

• cuMemAddressReserve 函数为新的内存分配保留一个虚拟内存地址，
• 而 cuMemCreate 在 GPU 上分配物理内存块。底层系统如何将内存转换为物理地址空间尚未公开。此外，物理块之间没有保证是连续的。
• cuMemMap 函数将物理内存和虚拟内存连接起来，将物理句柄映射到保留的地址。
• CUDA 还提供了一套内存释放函数，如 cuMemUnmap、cuMemAddressFree 和 cuMemRelease。

显然，低级虚拟内存 API 的优点在于，它可以分配和映射不连续的物理块，从而解决 GPU 内存碎片化问题。然而，虚拟内存分配器的开销远高于原生 GPU 分配器。

为了验证虚拟内存分配器的开销，我们进行了三个不同大小内存分配的实验：512 MB、1 GB 和 2 GB，这些是总分配块的大小。图 6 展示了原生内存分配器与虚拟内存分配器之间的比较结果。纵轴表示分配延迟，采用对数尺度。横轴上，2 MB、4 MB、… 和 1024 MB 表示构成分配块的内部物理块的大小。例如，1 GB 的分配块需要映射 512 个 2 MB 的块。最终，图 6 中的结果显示，虚拟内存的延迟非常高。具体而言，如果虚拟内存块被划分为 2 MB 的小块，其延迟将比原生分配器慢超过 100 倍，这是完全不可接受的。

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

为了进一步探讨 VMM（虚拟内存管理）API 的瓶颈，我们提供了分配的执行时间细分。表 1 显示了 2 GB GPU 内存分配的 VMM API 延迟细分。所有延迟都已归一化到 cuMalloc。GMLake 中的每个分配只需要一个 cuMemAddressReserve，但每个物理块需要多个 cuMemCreate、cuMemMap 和 cuMemSetAccess。cuMemSetAccess 是一个特殊函数，用于使 VMM 提供的映射可用。我们可以看到，使用 2 MB 小块来分配 2 GB 内存的速度比原生 cuMalloc 慢了 115 倍。

表 1. 按 cuMalloc 执行时间标准化的 VMM API 执行时间分解。

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

3 GMLake

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

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

深度学习框架中的原始缓存分配器采用 BFC 算法，如第 2.2 节所述，并在图 2(b) 中展示。为了避免同步并提高内存管理效率，这些框架在内部管理内存池来处理（释放）分配，而不是直接使用原生 API。沿用这一方法，我们的 GMLake 也包含以下三个组成部分：

• 虚拟内存 API：指的是用于指示 GPU 使用虚拟内存地址分配和释放内存的低级 API，如果没有充分优化，该过程通常会产生显著的开销。
• 虚拟内存池：作为基础数据结构（Efficient Stitching Data Structure），设计用于缓存虚拟内存。其实现对于提高效率至关重要。
• GMLake 分配器：包含所有管理虚拟内存池所需的函数（Extended High Level APIs）、算法和策略。

在本节中，我们将描述这三个重要组成部分的设计和细节，并从基础层到最上层构建 GMLake 框架。

3.1 Virtual Memory API

如第 2.5 节所述，并在图 2© 中展示，低级虚拟内存管理（VMM）API 作为 GPU 与应用程序之间的基本接口。如图 8 底部所示，我们利用 VMM API 构建了原始块（primitive block，pBlock），这是 GMLake 分配器的关键数据结构。pBlock 中的操作包括：

• AddrReserve：最初，分配原始块需要指定分配大小并预留相应的虚拟地址（VA）。
• Create：随后，原始块创建物理块（physical chunks），在其中存储数据。
• Map：最后，原始块将所有物理块映射到虚拟地址，从而实现张量的无缝访问。

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

为了优化碎片整理，我们对所有块（chunk）应用了统一的 2 MB 块大小。尽管这种 2 MB 块大小的开销相当可观（见第 2.5 节），但可以通过高效的数据结构和精心设计的拼接策略来缓解，从而实现最佳的碎片整理效果，并与 PyTorch 代码库兼容。同时，我们使用以下专门针对深度神经网络（DNN）的优化来减少拼接的频率，以使其端到端的开销可以忽略不计。GMLake 使用 VMM 来处理大于 2MB 的分配。对于小于 2MB 的内存分配，我们采用原始 PyTorch 缓存分配器的拆分方法来解决其内部碎片问题。此外，在大规模语言模型（LLM）训练中，分配 < 2MB 的情况较为罕见。

map 操作作为虚拟内存拼接的基本操作，允许我们连接多个物理上不连续的块（chunks），这些块在物理内存中可能并不连续。通过使用 VMM API，我们可以将虚拟内存编排到虚拟内存池中，虚拟内存池是 GMLake 分配器的底层数据结构。虚拟内存池的详细信息如下所述。

3.2 Virtual Memory Pool

考虑到原始的虚拟内存管理（VMM）API 需要消耗大量时间，为了实现 GMLake 的高效性，减少其使用是至关重要的。受到缓存分配器的启发，我们设计了具有缓存能力的虚拟内存池（VMP），从而显著减少物理内存（分配/释放）实例的次数。如图 8 所示，我们区分了两种类型的内存池：原始内存池（primitive memory pool，pPool）和拼接内存池（stitched memory pool，sPool）。

pPool 和 pBlock。

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

sPool 和 sBlock。

• sPool 也组织成一个排序集合，类似于 pPool。其元素包括拼接块结构，该结构整合了多个 pBlocks。例如，如图 8 所示，sBlock 3 包含拼接在一起的 pBlocks 3 和 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 Allocator

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

3.3.1 Allocation Module

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 专为虚拟内存管理（VMM）提供的低级 API 来“拼接”两个 pBlocks，如图 2© 所示。假设我们有两个 pBlocks，𝑝1（1GB）和 𝑝2（2GB）。

• 我们使用 VMM API cuMemCreate 来创建相应的物理块，并通过 cuMemAddressReserve 保留虚拟地址（VA）。
• 然后，使用 cuMemMap 将 VA 映射到物理地址（PA）。实际上，我们不需要取消映射 𝑝1 和 𝑝2 的原始 VA-PA 映射，因为 VMM 中的 PA 可以被多个 VA 指向。
• 因此，我们只需要使用 cuMemAddressReserve 为 sBlock 𝑠1 保留一个 3GB 的 VA。对于所有 sBlocks，它们从不创建 cuMemCreate 新的物理块。
• 我们使用 API cuMemMap 将 𝑠1（3GB）的 VA 映射到 𝑝1 和 𝑝2 的物理块上。由于多个 sBlocks 可以包含相同的物理块，我们需要更多的属性来确保每个物理块仅被单个张量使用，例如 pBlock 的活动状态。

BestFit 函数用于识别最适合内存分配的 pBlock 或 sBlock，并返回其状态和候选块，以供后续处理。如算法 1 所示，我们设计了四种状态，涵盖了 GMLake 可能面临的所有情况。它假设 sPool 和 pPool 都是按块大小降序排列的。

• 精确匹配（Exact match，第 2-4 行）：当候选块的大小与分配大小匹配时，会出现这种状态。该块可以是来自 sPool 的 sBlock，也可以是来自 pPool 的 pBlock。这是唯一可以为新分配分配 sBlock 的情况。所有其他状态仅涉及 pBlock。
• 单个块（Single block，第 12 行）：在此情况下，BestFit 会识别出一个大于请求分配大小的最佳匹配（最小）pBlock。
• 多个块（Multiple blocks，第 14 行）：在所有 pBlocks 都小于所需分配大小，但它们的总大小满足分配要求的情况下，BestFit 函数贪心地寻找多个候选 pBlocks 进行拼接。
• 块不足（Insufficient blocks，第 16 行）：当没有足够的 pBlocks 来满足请求的分配大小时，仍然会返回一个块列表。

GMLake 替代了 PyTorch 缓存分配器模块的多个内部函数。“拼接”操作对用户完全透明，无需用户修改代码。CUDA VMM API 是 NVIDIA 专门为虚拟内存管理（VMM）提供的低级 API。我们采用的 CUDA API 不仅包括与 VMM 相关的 API，还包含常规 API，如 cuMemAlloc。GMLake API 的实现利用 CUDA API 来实现精细的内存拼接和重用，因此它们属于不同的层次，承担相应的功能。

3.3.2 Deallocation Module

释放模块。释放模块避免通过低级 VMM API 主动释放物理 GPU 内存。相反，它仅更新或恢复拼接的虚拟内存块。

更新。当收到高层张量的释放请求时，我们用更新（Update）函数替代原始的 VMM 释放函数。该函数交替切换活动 pBlock 和 sBlock 的状态，从而移除张量与块之间的链接和分配。在程序运行期间，实际的物理内存仍由相应的 pBlock 控制。

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

我们已经介绍了 GMLake 中各个组件的详细信息，包括低级 API、数据结构和高级函数。借助这些精心设计的特性，我们能够在 GMLake 中执行高效的分配策略，有效应对大规模深度神经网络训练中的内存碎片问题。

4 Defragmentation Strategy

在本节中，我们提出了解决内存碎片问题的策略。我们首先提出了一种复杂的算法，该算法基于 GMLake 分配器理论上消除了所有碎片问题。接着，我们讨论并描述了我们的优化措施，以保证其效率和鲁棒性。

4.1 Memory Allocation Strategy

图 9 显示了 GMLake 分配策略，用于将新内存块分配或分配给新的张量分配请求。该策略基于 BestFit 模块提供的四种状态，并为每种状态设计了一个后处理步骤。

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

在状态 S1 中，当出现分配请求 1 时，会立即返回现有的 pBlock 或 sBlock。

如果在非活动的 sPool 和 pPool 中未找到完全匹配的块，我们将进入状态 S2，由 BestFit 函数生成一个单块。这需要将较大的 pBlock 1 通过 Split 操作分割成两个不同的 pBlock，并将它们插入 pPool。新创建的 pBlock 1 替换其前身，并分配给分配请求 2。同时，GMLake 执行 Stitch 函数，将 pBlock 1 和 pBlock 2 合并成 sBlock 1，然后将其添加到 sPool。

在状态 S3 中，GMLake 将多个 pBlock 合并（Stitch）成合并后的 sBlock 3，以处理分配请求 3。如果需要，最终的 pBlock 可以像在 S2 中那样进行划分（Split）。该阶段结束时，新 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 函数。

如果 Alloc 函数调用失败，GMLake 会立即在状态 S5 中报告内存不足（Out-of-Memory，OOM）错误。

4.2 Strategy Analysis

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

4.2.1 Effectiveness

GMLake 分配策略有效地确保几乎完全消除了 GPU 系统中的内存碎片。

接口。该策略的有效性得益于我们的接口设计，将所有操作整合为三个主要功能：Alloc、Split 和 Stitch。

• Alloc 是唯一可以创建新 pBlock 的函数，
• 只有 Stitch 可以生成新的 sBlock，
• 而 Split 并不会增加分配的内存。

数据结构。

• pPool 表示 GPU 内存的严格一对一映射，其中每个 pBlock 都与其他 pBlock 区别开来。它是一个分配的 GPU 内存集合，不包含重复元素和重叠地址。
• sPool 旨在存储 sBlock，并保留指向 pBlock 的链接和指针，类似于**软链接（soft link）**机制。
• GMLake 禁止分割 sBlock，因为这可能影响 pBlock。
• sPool 被视为 pPool 的一个子集。

最终，每当程序达到 GPU 内存使用的新峰值时，例如在调用 Alloc 函数时，pPool 可能无法提供足够的块进行拼接和分配，从而导致内存完全利用而没有碎片化。这与原始缓存分配器形成对比，后者可能会留下许多未使用的子块。

4.2.2 Efficiency

我们采用了多种方法来实现高效率。最初，S3 中的算法问题表示经典的 NP 难度的打包问题（packing problem）。然而，通过应用 Split 和 Stitch 函数，可以生成一个完全匹配的块来适配分配，从而实现线性复杂度。

其次，包含 sPool 和 pPool 的并集（union set）记录了每个张量分配的所有大小和对应的块，类似于录音带（tape）记录了 DNN 模型的张量分配模式。幸运的是，DNN 模型训练具有高度规律性，因为每次迭代处理相同的模型参数和输入数据大小。因此，在经过几次迭代后，GMLake 将不再执行 S2、S3 和 S4。GMLake 只会在剩余的训练过程中使用 S1“精确匹配”策略，这与原始缓存分配器形成对比，后者持续需要执行分割和合并操作。

4.2.3 Robustness

在实践中，由于 GPU 的总容量限制，拼接和创建新 sBlock 不可能无限进行。此外，过度的拼接操作可能会降低 GMLake 分配器的效率，特别是在运行如 BestFit 等分配模块时。当总容量超过这一限制或阈值时，GMLake 会使用 StitchFree 函数释放 LRU sBlock 并清除模式记录（pattern tape），从而作为一种回退（fallback）机制以提高鲁棒性。

此外，当 DNN 训练呈现极不规则的模式时，可能会生成大量小块，导致频繁的分割和拼接，从而提前达到限制。为避免不必要的性能损失，设定了最小碎片化限制。如果块的大小小于该限制，GMLake 将避免对其进行拼接或分割。因此，所有算法和模块都遵循这一碎片化限制（例如 128 MB），以确保 DNN 训练中的高效率和鲁棒性。

5 Evaluation

我们用 5000 行 C++ 代码实现了 GMLake，并将其集成到 PyTorch 的缓存分配器中。我们已经将 GMLake 适配到不同版本的 PyTorch，例如 PyTorch-1.13.1 和 PyTorch-2.0。

我们评估了 GMLake 在微调多个流行的大型语言模型（LLM）上的性能。我们在不同条件下比较了 GMLake 虚拟内存分配器与 PyTorch 缓存分配器的表现，考虑了不同的训练框架、GPU 扩展性以及优化策略的组合。通过这一分析，我们展示了 GMLake 在复杂环境中的可扩展性，并有效解决了内存碎片化问题。总体而言，GMLake 实现了平均 15% 的碎片化比率降低，最多可达到 33%，同时预留的 GPU 内存平均减少了 9.2GB，最多减少了 25GB，这些结果来自于 76 个工作负载和 8 个不同模型。

5.1 Evaluation Methodology

测试平台。GMLake 的评估使用了两种不同的配置：单节点多 GPU 和多节点多 GPU 实验。单节点评估在一台配备有 Intel Xeon Platinum 8369B CPU、1 TB DRAM 和 8 张 NVIDIA A100 GPU（每张 80 GB 内存）的服务器上进行。这些 GPU 通过 NVLink 相互连接，运行 CUDA 11.4 和 cuDNN 8.5。相应地，多节点评估包含两台服务器，每台服务器的配置与单节点服务器相同。在微调阶段，这两台节点通过 RDMA 进行分布式训练。

训练场景。我们在多个知名的 LLM 优化平台上评估了 GMLake，包括 Deepspeed 和 FSDP，涵盖了如 LORA、梯度检查点（即重计算）和卸载等不同的优化策略。我们的重点是微调场景。

模型和数据集。我们在一组来自官方示例的代表性开源 LLM 上进行评估。微调阶段评估的模型、数据集、分布式数据并行（DDP）框架以及优化策略的完整列表见表 2。值得注意的是，选择开源 LLM 应用程序的默认数据集是基于在微调过程中数据集质量不会影响内存使用的考虑。

表 2. 基准规格。（L: LoRA；R: 重计算；O: 离载；FSDP: 完全分片数据并行。）

基准对比。我们将 GMLake 与 PyTorch 2.0 进行比较，涉及不同的 LLM 训练，例如 Deepspeed 和 ColossalAI。据我们所知，PyTorch 缓存分配器代表了一种应用无关的最先进内存分配器，具有广泛的兼容性，适用于各种深度学习应用。值得注意的是，GMLake 与 PyTorch 原始分配器之间的转换非常方便，只需切换某些配置。此外，GMLake 配备了一组特定的超参数，这些参数经过经验配置，以通过最佳实践实现最佳性能结果。

指标。

• 与用于虚拟内存页面的固定长度碎片化指标 FMFI 不同，GMLake 中使用的块可以按任意大小拆分。因此，我们经验性地定义了一个特定的碎片化比率，该比率等于 $(1-利用率)$。
• 为了衡量内存碎片化，我们首先计算内存利用率，该值等于峰值活动内存（peak active memory）除以峰值预留内存（peak reserved memory）。
  • 术语**“活动内存”指所有活动块占用的累计内存，这些块当前由高层张量分配并用于 DNN 计算**。
  • 另一方面，“预留内存”指 PyTorch 和 GMLake 分配的总内存。
  • 这些指标在它们的峰值时被记录。
• 为了描述 GMLake 应用前后的内存减少关系，我们计算多个工作负载的内存减少比率的算术平均结果。计算算术平均数的公式是，

$$\text{MemReductionRatio}=\frac{\sum \text{Reserved}-\sum \text{GMLakeReserved}}{\sum \text{Reserved}}$$

其中“Reserved”和“GMLakeReserved”分别指 PyTorch 和 GMLake 的预留内存。虽然预留内存通常在训练过程中趋于稳定，但由于模型执行期间张量的分配和释放，活动内存的波动较大。鉴于缓存机制，峰值活动内存对于利用率或碎片化分析至关重要。我们还使用吞吐量（以每秒样本数表示） 来量化 DNN 训练的速度。

5.2 Scalability of GMLake

略，感兴趣还请直接阅读英文原文。

5.3 End-to-End Effectiveness of GMLake

略，感兴趣还请直接阅读英文原文。

5.4 Memory Trace Analysis

最后，我们跟踪了在 4 个 GPU 上使用 LoRA 和重计算策略的 PyTorch 和 GMLake 对 GPTNeoX-20B 的内存分配行为，如图 14 所示。

图 14. GPT-NeoX 上批量大小为 72 的内存追踪。

有三个值得注意的要点突出了 GMLake 的优势。

• 首先，PyTorch 因 OOM（内存溢出）异常在大约 200 秒时终止，而 GMLake 在此批次大小下能够正常运行。
• 其次，尽管 GMLake 和 PyTorch 的活动内存处于相同水平，但它们的预留内存差异很大，这表明 PyTorch 存在较大的碎片化问题。
• 第三，在 100 秒到 400 秒之间，PyTorch 和 GMLake 的活动内存定期波动，显示了 LLM 微调阶段的内存模式，特别是在前向传播和反向传播过程中。活动内存的上升表示前向传播内存分配，而活动内存的下降表示反向传播。两个相似模式之间的间隔反映了单次迭代的时间成本。
• 最后，经过四次迭代，GMLake 达到稳定状态，并实现了与 PyTorch 相同的吞吐量。这表明，GMLake 中使用的内存分配策略能够快速适应前向/反向训练过程中的内存波动，并找到最佳的缓存和拼接策略。

我们利用了 DNN（深度神经网络）训练的周期性特性。图 14 显示，DNN 训练中存在一个稳定的周期，在这个周期中，类似的虚拟内存管理（VMM）分配请求会重复出现。这种周期性为重用拼接的块（sBlock）提供了机会，从而分摊它们的成本。为了实现这一点，我们在 GMLake 中设计了拼接内存池（sPool）来覆盖 sBlock。

• 当一个新的 sBlock 被创建时，我们将其添加到 sPool 中。
• 当 sBlock 被释放时，我们仍然保留它的存在。下次当相同的 sBlock 需要创建时，可以直接重用先前创建的 sBlock。
• 只要我们保持足够的 sPool 实例，所有的分配只需搜索最佳匹配的 sBlock，而无需创建新的 sBlock。我们称之为收敛，例如，在图 14 的四次迭代后。

6 Related Works and Discussion

我们在两个方面将 GMLake 与现有工作进行了比较：内存碎片整理和 LLM（大规模语言模型）的内存优化。

内存碎片整理（Memory Defragmentation）。内存碎片整理在各种背景下得到了广泛研究。在解决碎片相关问题方面，早期文献提出了一种简单的方案，涉及使用细粒度的固定大小块。虽然这种方法消除了数据移动开销，但引入了更高的访问开销和有限的灵活性。为了提高效率和灵活性，研究人员提出了基于压缩的策略，例如通过数据移动将多个小块合并为一个更大的连续单位。其他碎片整理技术，包括基于拷贝的垃圾回收系统，减少了数据移动逻辑的复杂性，但以临时内存浪费为代价。尽管这些方法与将小块合并为更大的块在概念上有一些相似之处，GMLake 采用了一种基于拼接的技术，这减少了频繁的数据移动和复制，从而显著提高了内存效率。除了传统的内存系统，近期的研究还探讨了持久性内存的碎片整理。

高效的 LLM。对于基于 Transformer 的 LLM，内存已经成为计算系统中的重要资源。注意力机制的平方复杂度导致 LLM 的内存消耗大幅上升，从而使得有效的内存管理变得尤为重要。研究人员提出了各种算法优化方案，旨在降低内存消耗，包括量化（quantization）技术、剪枝（pruning）策略和键值缓存压缩（KV-cache compression）方法、编译（compilation）和调度（scheduling）等。

有各种系统级的内存优化。vLLM 工作通过基于页面的虚拟内存管理技术，显著提高了资源效率和服务吞吐量。FlashAttention 使用平铺技术优化注意力计算，并显著减少了内存消耗。FlexGen 框架引入了一种优化策略，以确定最佳的内存计算安排，从而实现高效的管道执行。在这个领域中，GMLake 成为一个用户透明的内存管理系统，协调智能高效的内存重用。

与其他工作的创新性。GMLake 在 DNN 训练中采用了独特的内存范围，这与 vLLM 和 vMalloc/CUDA VMM 不同。它们都无法解决 GMLake 所解决的问题。

vLLM 是一个基于算法的解决方案，专门用于自注意力操作，并作用于张量内部。LLM 最初将所有序列用变长的标记填充到最大长度，从而造成了显著的冗余。vLLM 使用查找表来移除填充的标记。因此，vLLM 仅适用于自注意力操作，而 GMLake 则面向 DNN 训练，适用范围更广。CUDA VMM 类似于 vMalloc，是一个低级系统工具，用于碎片整理。然而，CUDA VMM/vMalloc 无法感知 DL 框架中常用的内存池。如果没有内存池设计，性能会显著下降。因此，这些基于系统的内存工具不能直接在 DL 框架中实现。

7 Conclusion

本研究介绍了 GMLake，一种高效且低碎片的内存分配器。

• 它建立在低级 CUDA 虚拟内存管理机制的基础上，将多个不连续的物理内存块合并为一个连续的实体（entity），从而缓解了碎片问题。
• 为了便于与现有深度学习（DL）框架的无缝集成，我们制定了一个虚拟内存管理 API。
• 此外，我们提出了多阶段碎片整理策略，以保证分配的效率和稳健性。
• 我们的评估结果表明，GMLake 将内存碎片化减少到 5% ∼ 10%，同时保持相同的吞吐量。

扩展阅读1：GMLake 和 GLake

以下内容来自：glake/docs/readme_cn.md at master · antgroup/glake

GLake 是个工作在底层（虚拟与物理显存管理）与系统层（包括多卡、多通道、多任务）的加速库以及相关工具集，尝试对显存+传输进行一体优化。

GLake 总体架构如下图所示，基于分层设计（当前主要是基于PyTorch和NVIDIA GPU，正在对接更多国产卡）：

• 显存池：提供全局、异构显存池，内置显存碎片优化、多stream-进程复用、安全等特性。
• 核心优化层：提供增值优化功能，包括全局分配、多路并发、分层、重复数据消重、KV-cache 优化等。
• 扩展层：结合框架和团队自研的 VGPU，提供参考集成方案或扩展，例如 PyTorch 等。

扩展阅读2：ASPLOS’24 - Lightning Talks

以下内容来自：ASPLOS’24 - Lightning Talks - Session 8B - GMLake: Efficient and Transparent GPU Memory Defragmentat

扩展阅读3：ASPLOS’24 会议上分享的 slides

以下内容来自：GMLake-asplos-after-1

扩展阅读4：AICon-2024 大会上的演讲材料

以下内容来自：

• glake/docs/蚂蚁-GLake显存与传输优化-AIConf-V1.0.pdf at main · antgroup/glake
• 蚂蚁集团：大模型推理显存优化的深度探索与实践

在 GPU 显存管理方面，我们有一些新的观察。传统上，开发者在 GPU 上进行显存分配时，通常使用 CUDA Malloc。然而，从 CUDA 10.2 版本开始，引入了一个新的 API——虚拟内存管理（Virtual Memory Management，简称 VMM）。这个 API 已经存在一段时间了，它的作用与传统的显存分配方式有所不同。

传统的 CUDA Malloc 是一次性分配显存，分配后程序员或框架无法进行更多的操作，只能直接使用。而 VMM 提供了一个分两步进行显存申请的机会。

• 首先，它允许我们申请一个虚拟地址，这个地址可以非常大，但实际上在没有创建 handle 之前，这个地址是廉价的，几乎不需要成本。这个虚拟地址是连续的，但有一个最小单位，称为 chunk，规定为两兆。
• 接下来，我们可以通过另一个 API 申请一个 handle，这个 handle 可以简单理解为代表一个物理显存。
• 然后，我们可以通过第三个 API 将虚拟地址和 handle 关联起来，进行映射。映射之后，程序才能正常访问显存。
• 这样，虚拟地址是连续的，而物理地址可以不连续。

这种机制为我们提供了新的创新机会。例如，在典型的场景中，我们可以解决显存碎片问题。如果只有一层地址空间，我们很难进行有效的管理。有了 VMM API，我们可以进行创新，解决训练中的显存碎片问题，以及推理场景中的 KV 缓存碎片问题。

我们的工作就是基于这个 API 进行的。当然，仅仅有一个简单的 malloc 操作是不够的，就像我们不能仅用一个 malloc 就写出一个数据库或 Redis 一样。在实际应用中，我们还需要做很多工作，比如**内存池（pooling）和缓存（caching）**等。

我们的显存优化工作从训练场景开始，最早可以追溯到 2023 年的一篇相关工作。该工作主要优化了训练场景中的显存碎片问题，核心技术是使用了 CUDA 的 VMM API 接口。

针对训练场景，我们有一个单独的工作，名为 GMLake。其核心思路是：

• 使用虚拟显存和物理显存的动态管理和映射。
• 通过动态映射技术，优化显存的使用效率，从而带来显著的收益。

扩展阅读5：某乎上关于 GMLake 的问答

以下内容来自：蚂蚁集团称自研 GMLake 已被 PyTorch 集成，企业自研产品有何用途？
下面有作者团队的回答，也有网友的回答，具体内容见仁见智。

作者 George 对网友 东尼大佬 的回应：

• 我们和 expandable segment 的工作算是处于同时研发，我们在 22 年底就已经构思出了方案，然后最早是在 pytorch 1.13.1 的版本上做的，在 23 年 2 月份就已经完成，开始进行了大量验证，在 3 月份就集成到 pytorch 2.0 上，彼时我们没有对外宣传主要是需要先在蚂蚁内部大模型训练用起来，期间也暴露出一些问题，我们也是进行了多次优化，包括我们未开源出来做的一些多 stream 之间显存复用的工作，我们在 4 月份开始做大量实验以及构思论文写作，在 8 月份投出去论文，然后就是在等到论文消息，我们对外仓库确实是 10 月份建立的，因为我们代码对外开源在内部也是经过审批之后才开始规划开源的事项，论文一月份挂也是因为我们刚确认了终稿，论文中了之后还是要根据 reviewer 的意见进行修改的，expandsble segment 的思路也非常好，但是这个 pr 并没有放到 pytorch 2.0 的正式 release 版中，所以我们没有及时关注到，毕竟一些实验性质的 feature 我们内部也不太可能会使用的，我们在 pytorch 2.1 正式 release 之后，也将 gmlake 适配到 2.1 上。
• 然后针对我们论文中测试的模型针对两个 feature 测试了一些 case，我们的效果还是要比 expandable segment 好一些，因为 es 是属于链式的块，链上存在 used 和 unused block 时，unused block 显存大小不合适的话是没法被用起来的，而 gmlake 是可以将多个不连续的 stitch 起来，但是 es 逐渐扩大 size 这个 idea 是很好的，我们当时也在想是否有必要做这个事情，但内部模型训练确实不存在这种 case，所以一直没有做，从 es 的注释中也可以看到，作者主要是解决推理场景下 batch 变化导致的显存问题，我们当前也在做进一步的优化。
  • 首先，因为 es 的作者只是在 pytorch 的注释中解释了一下为什么这么做，举的例子是推理场景 batch 变化，所以我们也没办法进一步得知作者起初做这个 feature 的真实想法。
  • es 训练效果也是很明显的，我们也测试过一些 case，因为在一个 step 内部，backward 阶段使用的显存要比 forward 阶段使用的显存多，而加上激活重计算后，forward 阶段释放的显存在 backward 阶段拼接新的一些块就可以使用起来。es 的链式方式是很不错的 idea，但我们测试，仍然是可以进一步优化的，我们也在考虑将 GMLake 和 es 的方式进行整合。

网友 ooosimdexe 的回答：

• GMLake 的做法是 override 了torch/c10/cuda/CUDACachingAllocator.cpp，往 PyTorch 原生的 DeviceAllocator 和 NativeCachingAllocator 这两个类里面糊了自己 Virtual Memory Stitching (VMS) 的一套东西，集成的方式也是 override 原生的.so。
• 技术角度个人觉得本质上就是把 CPU 上 virtual memory pool 的 techniques 搬到 GPU 显存管理上来。

作者 George 对网友 ooosimdexe 的回应：

• 我们其实没有使用 pytorch 提供的 plugin 的方式，是因为我们的逻辑是与 pytorch bfc 的分配方式强相关的，就是因为原有的方式在面向一些训练场景时加剧了显存碎片化问题，主要是因为内部的实现机制，所以我们最早就是直接集成进去，后面我们也在规划剥离出来，但因为人手有限，当前还没有腾出时间来做。
• 第三个问题，vmm 接口确实出现的很早，但一直没有利用起来，其实 vmm 接口还是有很多的限制，在我们的论文里是有一些分析的，我们也对 vmm 的方式做过大量实验，才最终确定了这种方案。

扩展阅读6：关于 CUDA VMM 和 PyTorch expandable_segments

具体内容可以参考笔者的博客：PyTorch 源码学习：GPU 内存管理之初步探索 expandable_segments