PyTorch 源码学习：GPU 内存管理之初步探索 expandable_segments

zhulu5062025-03-04 16:06

本文是笔者 PyTorch 源码学习 系列的收尾之作，围绕 PyTorch 2.1 推出的 expandable_segments 机制进行分析讨论：以深度学习框架面临的 GPU 内存碎片化问题为背景，首先介绍了 CUDA 10.2 推出的 VMM API，接着分析了 CUDA VMM API 在 TensorFlow 和 PyTorch 两个主流深度学习框架的应用。遗憾的是，TensorFlow 目前并没有继续维护其 GpuVirtualMemAllocator 相关的代码，但因其简洁的实现和易读的代码而仍然值得学习。最后，笔者简单分析了 PyTorch expandable_segments 机制的实现原理和核心代码，期待未来有更好的博客深入分析 PyTorch expandable_segments 机制的代码实现。
更多内容请参考：

文章目录

  • [1 背景：内存碎片化问题](#1 背景：内存碎片化问题)
  • [2 关于 CUDA VMM](#2 关于 CUDA VMM)
  • [3 在 TensorFlow 的实现](#3 在 TensorFlow 的实现)
    • [3.1 阴差阳错](#3.1 阴差阳错)
    • [3.2 追根溯源](#3.2 追根溯源)
    • [3.3 代码实现](#3.3 代码实现)
      • [3.3.1 gpu_driver.h & cuda_driver.cc](#3.3.1 gpu_driver.h & cuda_driver.cc)
      • [3.3.2 gpu_virtual_mem_allocator.h](#3.3.2 gpu_virtual_mem_allocator.h)
      • [3.3.3 gpu_virtual_mem_allocator.cc](#3.3.3 gpu_virtual_mem_allocator.cc)
  • [4 在 PyTorch 的实现](#4 在 PyTorch 的实现)
    • [4.1 初见端倪](#4.1 初见端倪)
    • [4.2 探本穷源](#4.2 探本穷源)
    • [4.3 实现原理](#4.3 实现原理)
    • [4.4 核心代码](#4.4 核心代码)
  • 参考资料

1 背景：内存碎片化问题

以下内容来自：一文读懂 PyTorch 显存管理机制

如上图所示，假设当前想分配 800MB 显存，虽然空闲的总显存有 1000MB，但是上方图的空闲显存由地址不连续的两个 500MB 的块组成，不够分配这 800MB 显存；而下方的图中，如果两个 500MB 的空闲块地址连续，就可以通过显存碎片的整理组成一个 1000MB 的整块，足够分配 800MB。上方图的这种情况就被称为显存碎片化。

哪些 Block 的地址是连续的？见原文下面的内容👇

被 split 的操作很简单，当前的 Block 会被拆分成两个 Block，第一个大小正好为请求分配的 size，第二个则大小为 remaining，被挂到当前 Block 的 next 指针上。这样一来，这++两个 Block 的地址自然而然成为连续的了++ 。随着程序运行，较大的 Block（只要仍小于阈值 max_split_size_mb）会不断被分成小的 Block。值得注意的是，由于新 Block 的产生途径只有一条，即通过步骤三中的alloc_block函数经由cudaMalloc申请，++无法保证新 Block 与其他 Block 地址连续++ ，因此所有被维护在双向链表内的有连续地址空间的 Block 都是由一个最初申请来的 Block 拆分而来的。

以下内容来自：PyTorch显存管理介绍与源码解析（一）

如下所示，是一种典型的碎片问题。用户申请了一个显存块A，尽管框架托管的显存剩余总数满足用户需求（a+c > A），但是无法返回一个满足需求的显存地址给用户。

2 关于 CUDA VMM

关于什么是 CUDA VMM，最权威的说法还请直接参考 NVIDIA 官方文档和博客。

以下内容来自：Nvidia GPU Virtual Memory Management | by Bruce-Lee-LY | Medium 及中文版，内容主要还是参考了【翻译】Introducing Low-Level GPU Virtual Memory Management。

CUDA Driver API 与 CUDA Runtime API 的区别：

Nvidia 在用户态 CUDA Driver 提供了一套 API 用于显存申请和释放，返回的结果与 CUDART API 没有区别，只是在使用层面与 CUDART API 有区别，比如调用cuMemAlloc之前需要编程人员手动使用 CUDA Driver API 进行初始化（cuInit）和创建 Context（cuCtxCreate），而对于调用cudaMalloc来说，这些都是隐式 完成，对编程人员是透明的。

Virtual Memory Management 特性：

就常用显存管理 API 来说，由于编程人员只能获取到显存的虚拟地址，如果有动态调整显存大小的需求（比如 GPU 上 vector 扩容），用户必须显式地申请更大的一块显存，并从原始显存中复制数据到新显存，再释放原始显存，然后继续跟踪新分配的显存地址，这样的操作通常会导致应用程序的性能降低 和较高的显存带宽峰值利用率。

在 CUDA 10.2 中引入 VMM API 为应用程序提供了一种直接管理统一虚拟地址空间的方法，可以将显存的虚拟地址和物理地址解耦，允许编程人员分别处理它们 。VMM API 允许编程人员在合适的时候将显存的虚拟地址与物理地址进行映射和解映射。借助 VMM API 可以更好地解决动态调整显存大小 的需求，只需要申请额外的物理地址，再与原始虚拟地址扩展的空间进行映射，既不需要更换追踪的显存地址，也不需要将数据从原始显存拷贝到新显存。因此，VMM API 能够帮助编程人员构建更高效的动态数据结构 ，并更好地控制应用程序中的显存使用。

以下内容来自：CUDA虚拟地址管理-解决显存碎片的屠龙宝刀

CUDA 在 10.2 引入了一个重要的 feature，virtual memory management API（VMM），给用户暴露了操作 CUDA 虚拟地址的接口，用户可以使用类似 mmap 一样的接口为给定的虚拟地址映射物理内存。其中有四个比较重要的接口：

  • cuMemCreate---创建一个代表物理内存的 handle
  • cuMemAddressReserve---保留一段虚拟地址.
  • cuMemMap---将一个物理内存的 handle 映射到虚拟地址空间上.
  • cuMemSetAccess---给分配的内存管理访问权限

使用 VMM API 的流程图：

文中后面的分析主要还是参考了【翻译】Introducing Low-Level GPU Virtual Memory Management

试想这样一个场景，我分配一个 1GB 大小连续的 tensor，我这个 tensor 因为某种需求其 size 会增长，这个时候问题就来了，如果我使用cudaMalloc，我必须要先调用cudaMalloc分配一个更大的 tensor，然后再调用cudaMemcpy把原来的数据拷贝到新的 tensor 中，再调用cudaFree把旧的 tensor 释放掉。这会引入两个非常昂贵的开销，一个是cudaMalloc/Free的 API 调用，一个是数据拷贝，并且会在新 tensor 创建时造成一个显存峰值，此时的显存总大小是 tensor1+tensor2，更容易造成 OOM。

而如果使用 VMM 接口，我们不需要释放并拷贝原有的 tensor 对应物理内存 block1，我只需要再申请新增的显存大小的物理内存 block2，并重新申请一段 VA 映射到 block1 和 block2 这两段不连续的物理内存即可。如果我们提前 buffer 一些大小不一的物理内存 block，这一开销会相比于前者大大降低。我们仍然用图对比下二者的不同：

作者举的例子参考了【翻译】Introducing Low-Level GPU Virtual Memory Management，而且在"VMM 接口"中描述的思想和[2401.08156] GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching论文中的虚拟内存拼接机制类似。

以下内容来自：PyTorch 内存管理之 expandable_segments

内存申请差异：

  • 开启前：使用cudaMalloc、cudaFree接口申请释放内存，按照业务需要申请指定大小的内存块，申请到的虚拟地址已经映射好物理地址。
  • 开启后：使用 cuda 提供的虚拟内存管理接口申请，申请流程：申请虚拟地址 -> 申请物理地址 -> 地址映射；释放流程相反。

以下内容来自：PyTorch显存管理介绍与源码解析（三）

为什么需要引入虚拟地址？使用物理显存时经常遇到一个场景：当物理地址上面有多个离散的显存片段时，我们无法将其作为一个整体使用。举个例子：物理空间上有多个片段容量小于 1GB 离散空闲空间，且这些片段的空间总和超过了 1GB，但是当需要申请一个连续的 1GB 显存时会报 OOM 错误。

有了虚拟地址管理，可通过 VMM 把离散的物理显存映射到一个新的虚拟地址段，并把这个地址给到应用，这样我们就能够获得足够大的空间。

而且当发现虚拟地址空间不足时，我们能够继续把离散的物理地址片段映射到新的虚拟地址空间，从而增大总的可用显存空间。

3 在 TensorFlow 的实现

笔者偶然间发现 PyTorch expandable_segments 的思想其实很早就出现在 TensorFlow 的 GpuVirtualMemAllocator 中，但很遗憾 TensorFlow 目前已经把 GpuVirtualMemAllocator 相关的内容废弃了，但其代码实现仍然有很高的研究价值。

3.1 阴差阳错

笔者在 GPU 底层机制分析：显存分配开销 这篇博客发现 TensorFlow 其实有与 PyTorch expandable_segments 相似的实现。文中有以下内容：

VirtualMem 分配器

GpuVirtualMemAllocator 和 BFCAllocator 一样，都是 SubAllocator 的具体实现。 这个分配器会在构造时申请一大块连续的虚拟地址空间作为预留空间，当真正申请时再和物理地址空间进行映射关联 。这块涉及到 Driver API 虚拟地址到物理地址的映射和绑定......

3.2 追根溯源

在 GPU 底层机制分析：显存分配开销 中，作者给了 GpuVirtualMemAllocator 的源码地址，是 2021 年 3 月 6 日提交的代码版本，当时是 TensorFlow 2.5.0。

从这版代码来看，与 GpuVirtualMemAllocator 相关的核心文件包括：

但有意思的是，TensorFlow 在 2024 年 3 月 14 日已经把 GpuVirtualMemAllocator 移除了，具体见#63659。

我们在 TensorFlow 仓库中搜索包含GpuVirtualMemAllocator关键字的 Commits 得到如下 Commits search results：

2020 年 12 月 18 日

Add GpuVirtualMemAllocator, a suballocator which reduces fragmentation. This SubAllocator users the virtual memory capabilities introduced in CUDA 10.2 to suballocate contiguous memory regions for the BFC allocator without requiring ahead-of-time allocation of all physical GPU memory.

新增了 GpuVirtualMemAllocator，这是一个用于减少内存碎片化的子分配器。该子分配器利用了 CUDA 10.2 引入的虚拟内存功能，能够为 BFC 分配器（Best-Fit with Coalescing）分配连续的内存区域，而无需预先分配全部的物理 GPU 内存。

With this suballocator, allow_growth should have no fragmentation impact.

通过使用该子分配器，allow_growth 参数将不再对内存碎片化产生影响。

关于 allow_growth 可以参考：限制 GPU 内存增长

2021 年 1 月 22 日

Use device handle instead of gpu ordinal in GpuVirtualMemAllocator for configuring peer access.

在 GpuVirtualMemAllocator 中，使用设备句柄（device handle）代替 GPU 序号（GPU ordinal）来配置对等访问（peer access）。

Check that peers support virtual address management and have peer access at GpuVirtualMemAllocator creation time.

在 GpuVirtualMemAllocator 创建时，检查对等设备是否支持虚拟地址管理以及是否具备对等访问权限。

2021 年 1 月 30 日

Use GpuVirtualMemAllocator in GPUBFCAllocator.

在 GPUBFCAllocator 中使用 GpuVirtualMemAllocator。

The GpuVirtualMemAllocator is a suballocator which guarantees contiguous suballocations, allowing the BFC allocator to grow without fragmentation. This should reduce the risk of ooming when setting the allow_growth flag.

GpuVirtualMemAllocator 是一个能够保证连续子分配的子分配器，使得 BFC 分配器能够在无碎片化的情况下扩展内存。这将降低在启用 allow_growth 标志时发生内存不足（OOM）的风险。

2021 年 2 月 3 日

Disable use of the new GpuVirtualMemAllocator in BfcAllocator. We've obsered some OOMing with the new allocator; will re-enable with a follow-up patch that resolves these issues.

在 BfcAllocator 中禁用新的 GpuVirtualMemAllocator。我们观察到在使用新分配器时出现了一些内存不足（OOM）问题，将在后续补丁中解决这些问题后重新启用。

2024 年 3 月 14 日

Remove unused GpuVirtualMemAllocator.

移除未使用的 GpuVirtualMemAllocator。

综上可以看出，TensorFlow 早在 2020 年 12 月 18 日就已经实现了 GpuVirtualMemAllocator，但在 2021 年 2 月 3 日因为可能带来的 OOM 问题而禁用了 GpuVirtualMemAllocator，后来也没有解决该问题，最终在 2024 年 3 月 14 日移除了 GpuVirtualMemAllocator 的相关代码。

3.3 代码实现

从 3.2 节可知，TensorFlow 目前已经不支持 GpuVirtualMemAllocator，但其算法思想和 PyTorch 的 expandable_segments 机制基本一致，且 GpuVirtualMemAllocator 代码的可读性更高，总之，值得学习。

3.3.1 gpu_driver.h & cuda_driver.cc

GpuVirtualMemAllocator 的实现离不开 CUDA VMM API 的支持，gpu_driver.h 中的class GpuDriver就提供了一些可供 GpuVirtualMemAllocator 调用的接口，这些接口最终在 cuda_driver.cc 实现。

这里提一下 gpu_driver.h 和 cuda_driver.cc 的关系：

cpp 复制代码 
// tensorflow/stream_executor/cuda/cuda_driver.h 文件里 👇
#include "tensorflow/stream_executor/gpu/gpu_driver.h"
cpp 复制代码 
// tensorflow/stream_executor/cuda/cuda_driver.cc 文件里 👇
#include "tensorflow/stream_executor/cuda/cuda_driver.h"

所以，可以这样理解：gpu_driver.h → cuda_driver.h → cuda_driver.cc

1）GpuDriver里关于 VMM 的代码实现：

cpp 复制代码 
// GpuDriver contains wrappers for calls to the userspace library driver. It's
// useful to isolate these calls and put basic wrappers around them to separate
// userspace library driver behaviors from the rest of the program.
// GpuDriver 包含了对用户空间库驱动程序的调用封装。将这些调用隔离并封装起来有助于将用户空间库驱动程序的行为与程序的其他部分分离。
//
// At the moment it's simply used as a namespace.
// 目前它仅用作一个命名空间。
//
// The calls log any specific errors internally and return whether the operation
// was successful to the caller.
// 这些调用会在内部记录任何特定的错误，并返回操作是否成功给调用者。
//
// The order of parameters is generally kept symmetric with the underlying CUDA
// driver API.
// 参数的顺序通常与底层的 CUDA 驱动程序 API 保持对称。
//
// Links on functions are to specific documentation under
// http://docs.nvidia.com/cuda/cuda-driver-api/
// 函数链接指向 http://docs.nvidia.com/cuda/cuda-driver-api/ 下的具体文档。
//
// Thread safety: these functions should not be used from signal handlers.
// 线程安全：这些函数不应从信号处理程序中使用。
class GpuDriver {
 public:

  // Virtual memory support was added to CUDA in 10.2
  // 虚拟内存支持在 CUDA 10.2 中添加
#if CUDA_VERSION >= 10020

  // Reserves a range of virtual device memory addresses via
  // cuMemAddressReserve. bytes must be a multiple of the host page size.
  // Returns nullptr base address in VmemSpan if the reservation fails.
  // 通过 cuMemAddressReserve 保留一段虚拟设备内存地址范围。bytes 必须是主机页面大小的倍数。
  // 如果保留失败，VmemSpan 中的基地址将返回 nullptr。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1ge489256c107df2a07ddf96d80c86cd9b
  struct VmemSpan {
    GpuDevicePtr base;
    // Size in bytes.
    uint64_t size_bytes;
  };
  static port::StatusOr<VmemSpan> ReserveVirtualMemory(GpuContext* context,
                                                       uint64_t bytes);

  // Frees a range of virtual addresses that were previously reserved through
  // ReserveVirtualMemory via cuMemAddressFree.
  // 通过 cuMemAddressFree 释放之前通过 ReserveVirtualMemory 保留的虚拟地址范围。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1g6993ecea2ea03e1b802b8255edc2da5b
  static void FreeVirtualMemory(GpuContext* context, VmemSpan reservation);

  // Calculates the minimum alignment for memory allocations done through
  // cuMemCreate via cuMemGetAllocationGranularity.
  // 通过 cuMemGetAllocationGranularity 计算通过 cuMemCreate 进行内存分配的最小对齐方式。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1g30ee906c2cf66a0347b3dfec3d7eb31a
  static port::StatusOr<uint64_t> GetMinAllocationGranularity(
      GpuDeviceHandle device);

  // Allocates physical memory and returns a handle that can be mapped to
  // virtual addresses via cuMemCreate. bytes must be a multiple of the
  // granularity returned by GetMinAllocationGranularity.
  // 通过 cuMemCreate 分配物理内存并返回一个可以映射到虚拟地址的句柄。
  // bytes 必须是 GetMinAllocationGranularity 返回的粒度的倍数。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1g899d69a862bba36449789c64b430dc7c
  struct GenericMemoryHandle {
    uint64_t handle;
    uint64_t bytes;
  };
  static port::StatusOr<GenericMemoryHandle> CreateMemoryHandle(
      GpuContext* context, uint64_t bytes);

  // Frees memory represented by the provided MemoryHandle via cuMemRelease.
  // 通过 cuMemRelease 释放由提供的 MemoryHandle 表示的内存。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1g3014f0759f43a8d82db951b8e4b91d68
  static void ReleaseMemoryHandle(GpuContext* context,
                                  GenericMemoryHandle handle);
  
  // Maps a memory allocation handle to a reserved virtual address range via
  // cuMemMap and sets the appropriate access settings via cuMemSetAccess.
  // 通过 cuMemMap 将内存分配句柄映射到保留的虚拟地址范围，并通过 cuMemSetAccess 设置适当的访问权限。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1gff1d395423af5c5c75375516959dae56
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1g1b6b12b10e8324bf462ecab4e7ef30e1
  static port::Status MapMemory(
      GpuContext* context, GpuDevicePtr va, const GenericMemoryHandle& handle,
      const std::vector<GpuDeviceHandle>& device_handles);

  // Unmaps the backing memory from the given virtual address range. This range
  // must fully unmap a memory handle that was mapped using MapMemory; partial
  // unmapping is not supported.
  // 从给定的虚拟地址范围取消映射后备内存。此范围必须完全取消映射使用 MapMemory 映射的内存句柄；不支持部分取消映射。
  // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html#group__CUDA__VA_1gfb50aac00c848fd7087e858f59bf7e2a
  static void UnmapMemory(GpuContext* context, GpuDevicePtr va, uint64_t bytes);

#endif  // CUDA_VERSION >= 10200

};

2）cuda_driver.cc 里关于 VMM 的代码实现：

cpp 复制代码 
namespace stream_executor {
namespace gpu {

#if CUDA_VERSION >= 10020

// ReserveVirtualMemory 保留一段虚拟设备内存地址范围
/* static */ port::StatusOr<GpuDriver::VmemSpan>
GpuDriver::ReserveVirtualMemory(GpuContext* context, uint64_t bytes) {
  ScopedActivateContext activation(context);
  CUdeviceptr base;

  // 1) 调用 cuMemAddressReserve 保留一段虚拟设备内存地址范围
  CUresult res = cuMemAddressReserve(&base, bytes, /*alignment=*/0,
                                     /*addr=*/0, /*flags=*/0);
  if (res != CUDA_SUCCESS) {
    return port::InternalError(
        absl::StrFormat("error reserving %d bytes of virtual GPU memory: %s",
                        bytes, ToString(res)));
  }

  // 2) 返回保留的虚拟内存地址范围和大小
  return {{base, bytes}};
}

// FreeVirtualMemory 释放之前保留的虚拟地址范围
/* static */ void GpuDriver::FreeVirtualMemory(
    GpuContext* context, GpuDriver::VmemSpan reservation) {
  ScopedActivateContext activation(context);

  // 调用 cuMemAddressFree 释放之前保留的虚拟地址范围
  CUresult res = cuMemAddressFree(reservation.base, reservation.size_bytes);
  if (res != CUDA_SUCCESS) {
    LOG(ERROR) << "error freeing vmem reservation of size "
               << reservation.size_bytes << " at address " << reservation.base;
  }
}

// GetMinAllocationGranularity 获取内存分配的最小粒度
/* static */ port::StatusOr<uint64_t> GpuDriver::GetMinAllocationGranularity(
    GpuDeviceHandle device) {

  // 1) 设置内存分配属性
  CUmemAllocationProp props = {};
  props.type = CU_MEM_ALLOCATION_TYPE_PINNED;
  props.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
  props.location.id = device;
  
  // 2) 调用 cuMemGetAllocationGranularity 获取最小分配粒度
  size_t granularity;
  CUresult res = cuMemGetAllocationGranularity(
      &granularity, &props, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
  if (res != CUDA_SUCCESS) {
    return port::InternalError(absl::StrCat(
        "failed to get min allocation granularity: ", ToString(res)));
  }

  // 3) 返回最小分配粒度
  return granularity;
}

// CreateMemoryHandle 创建物理内存句柄
/* static */ port::StatusOr<GpuDriver::GenericMemoryHandle>
GpuDriver::CreateMemoryHandle(GpuContext* context, uint64_t bytes) {
  ScopedActivateContext activation(context);

  // 1) 从上下文中获取设备句柄
  auto device = DeviceFromContext(context);
  if (!device.ok()) {
    LOG(ERROR) << "Failed to get device from context" << device.status();
    return device.status();
  }
  
  // 2) 设置内存分配属性
  CUmemAllocationProp props = {};
  props.type = CU_MEM_ALLOCATION_TYPE_PINNED;
  props.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
  props.location.id = device.ValueOrDie();
  
  // 3) 调用 cuMemCreate 创建物理内存句柄
  CUmemGenericAllocationHandle mem_handle;
  CUresult res = cuMemCreate(&mem_handle, bytes, &props, 0);
  if (res != CUDA_SUCCESS) {
    return port::InternalError(
        absl::StrFormat("failed to create memory allocation of size %d: %s",
                        bytes, ToString(res)));
  }

  // 4) 返回创建的内存句柄及其大小
  return GpuDriver::GenericMemoryHandle{mem_handle, bytes};
}

// ReleaseMemoryHandle 释放内存句柄
/* static */ void GpuDriver::ReleaseMemoryHandle(
    GpuContext* context, GpuDriver::GenericMemoryHandle handle) {
  ScopedActivateContext activation(context);
  
  // 调用 cuMemRelease 释放内存句柄
  CUresult res = cuMemRelease(handle.handle);
  if (res != CUDA_SUCCESS) {
    LOG(ERROR) << "Failed to release memory handle " << handle.handle
               << " of size " << handle.bytes << ": " << ToString(res);
  }
}

// MapMemory 将内存句柄映射到虚拟地址范围并设置访问权限
/* static */ port::Status GpuDriver::MapMemory(
    GpuContext* context, CUdeviceptr va,
    const GpuDriver::GenericMemoryHandle& handle,
    const std::vector<GpuDeviceHandle>& device_handles) {
  ScopedActivateContext activation(context);
  
  // 1) 从上下文中获取设备句柄
  auto device = DeviceFromContext(context);
  if (!device.ok()) {
    return device.status();
  }
  
  // 2) 调用 cuMemMap 将内存句柄映射到虚拟地址范围
  // NB: Zero is the only valid value for both flags and offset.
  // 注意：flags 和 offset 参数的唯一有效值是 0。
  CUresult res =
      cuMemMap(va, handle.bytes, /*offset=*/0, handle.handle, /*flags=*/0);
  if (res != CUDA_SUCCESS) {
    return port::InternalError(absl::StrFormat(
        "Failed to map %d bytes at %d: %s", handle.bytes, va, ToString(res)));
  }
  
  // 3) 设置访问描述符，指定哪些设备可以访问该内存
  std::vector<CUmemAccessDesc> access_descriptors(device_handles.size());
  for (int i = 0; i < access_descriptors.size(); ++i) {
    access_descriptors[i].location.id = device_handles[i];
    access_descriptors[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
    access_descriptors[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
  }
  
  // 4) 调用 cuMemSetAccess 设置内存访问权限
  res = cuMemSetAccess(va, handle.bytes, access_descriptors.data(),
                       access_descriptors.size());
  if (res != CUDA_SUCCESS) {
    // Unmap the memory that we failed to set access for.
    // 在设置访问权限失败的情况下，需要清理之前映射的内存。
    if (cuMemUnmap(va, handle.bytes) != CUDA_SUCCESS) {
      LOG(ERROR)
          << "Failed to unmap memory in GpuDriver::MapMemory error path.";
    }
    return port::InternalError(absl::StrFormat(
        "Failed to set read/write access on memory mapped at %d: %s", va,
        ToString(res)));
  }
  // 5) 返回成功状态
  return port::Status::OK();
}

// UnmapMemory 取消虚拟地址范围的内存映射
/* static */ void GpuDriver::UnmapMemory(GpuContext* context, CUdeviceptr va,
                                         uint64_t bytes) {
  ScopedActivateContext activation(context);
  
  // 调用 cuMemUnmap 取消虚拟地址范围的内存映射
  CUresult res = cuMemUnmap(va, bytes);
  if (res != CUDA_SUCCESS) {
    LOG(ERROR) << "Failed to unmap memory at " << va << " of size " << bytes
               << ": " << ToString(res);
  }
}

#endif

}  // namespace gpu
}  // namespace stream_executor

3.3.2 gpu_virtual_mem_allocator.h

首先，让我们结合注释简单分析 gpu_virtual_mem_allocator.h 的代码实现：

cpp 复制代码 
// GpuVirtualMemAllocator is a SubAllocator for use with BFCAllocator which
// provides contiguous allocations with each call to Alloc. This is done by
// reserving a large chunk of virtual addresses at construction and then mapping
// physical memory pages to this virtual address range as requested.
//
// GpuVirtualMemAllocator 是一个用于 BFCAllocator 的子分配器，
// 它在每次调用 Alloc 时提供连续的分配。这是通过在构造时保留一大块虚拟地址，
// 然后根据需要将物理内存页映射到此虚拟地址范围来实现的。
// 
// This class is not thread-safe.
class GpuVirtualMemAllocator : public SubAllocator {
 public:
  // 创建一个 GpuVirtualMemAllocator 实例，并初始化相关的参数。
  static stream_executor::port::StatusOr<
      std::unique_ptr<GpuVirtualMemAllocator>>
  Create(const std::vector<Visitor>& alloc_visitors,
         const std::vector<Visitor>& free_visitors,
         stream_executor::gpu::GpuContext& gpu_context, PlatformDeviceId gpu_id,
         size_t virtual_address_space_size,
         const std::vector<PlatformDeviceId>& peer_gpu_ids);
  // 析构函数。
  ~GpuVirtualMemAllocator() override;

  // Allocates memory at least as large as requested by num_bytes. Will be
  // aligned to the min allocation granularity (typically 2MiB).
  // alignment is ignored by this allocator.
  // 分配至少为 num_bytes 请求大小的内存。将按照最小分配粒度（通常为 2MiB）对齐。
  // alignment 参数在此分配器中被忽略。
  void* Alloc(size_t alignment, size_t num_bytes,
              size_t* bytes_received) override;

  // Frees should only happen at the end of the contiguous memory allocations or
  // else we introduce pointless fragmentation...But, this is supported. If the
  // allocation happens at the end, then the next_alloc_offset_ is moved back,
  // otherwise a hole is created.
  // 
  // 释放操作应该只发生在连续内存分配的末尾，否则会引入不必要的碎片...但是，这是被支持的。
  // 如果分配发生在末尾，则 next_alloc_offset_ 会向后移动，否则会创建一个空洞。
  //
  // Holes are not re-used, all allocations continue to come at the end of the
  // next_alloc_offset_. To accommodate this, the virtual_address_space_size
  // should be much larger than the max physical size of the allocator.
  //
  // 空洞不会被重用，所有分配将继续在 next_alloc_offset_ 的末尾进行。
  // 为了适应这一点，virtual_address_space_size 应该比分配器的最大物理大小大得多。
  //
  // In practice, since the BFC allocator coalesces adjacent AllocationRegions,
  // this free function should never be invoked.
  //
  // 实际上，由于 BFC 分配器会合并相邻的 AllocationRegions，
  // 因此这个 free 函数应该永远不会被调用。
  void Free(void* ptr, size_t num_bytes) override;
  
  // 如果分配器支持合并，则返回 true。
  bool SupportsCoalescing() const override { return true; }

 private:
  // 构造函数。
  GpuVirtualMemAllocator(
      const std::vector<Visitor>& alloc_visitors,
      const std::vector<Visitor>& free_visitors,
      stream_executor::gpu::GpuContext& gpu_context, PlatformDeviceId gpu_id,
      std::vector<stream_executor::gpu::GpuDeviceHandle> access_device_handles,
      stream_executor::gpu::GpuDriver::VmemSpan vmem, size_t granularity);
  
  // 与此分配器关联的 GPU 上下文。
  stream_executor::gpu::GpuContext& gpu_context_;
  // GPU 设备的 ID。
  PlatformDeviceId gpu_id_;

  // Peer access is configured at mmap time so the allocator must be aware of
  // all gpus that may want to read the memory. This list also includes the
  // above gpu_id_ to facilitate the invocation of the GpuDriver::MapMemory
  // function.
  // 
  // 对等访问在 mmap 时配置，因此分配器必须知道可能想要读取内存的所有 GPU。
  // 此列表还包括上面的 gpu_id_，以便于调用 GpuDriver::MapMemory 函数。
  const std::vector<stream_executor::gpu::GpuDeviceHandle> access_gpu_handles_;

  // The virtual memory span held by this allocator.
  // // 此分配器持有的虚拟内存范围。
  stream_executor::gpu::GpuDriver::VmemSpan vmem_;

  // The next offset from the vmem base address that will be allocated. This
  // corresponds to the size of physically pinned memory if holes haven't been
  // created with "free".
  //
  // 从 vmem 基地址开始的下一个分配偏移量。如果没有通过 "free" 创建空洞，
  // 这对应于物理固定内存的大小。
  size_t next_alloc_offset_ = 0;

  // Smallest allocation as determined by CUDA.
  // 由 CUDA 确定的最小分配粒度。
  const size_t granularity_;
  
  // 表示虚拟地址和物理内存句柄之间的映射。
  struct Mapping {
    stream_executor::gpu::GpuDevicePtr va;
    stream_executor::gpu::GpuDriver::GenericMemoryHandle physical;
  };

  // 映射列表，按虚拟地址排序。
  // List of mappings, sorted by va.
  std::vector<Mapping> mappings_;
  
  // 禁止复制和赋值。
  TF_DISALLOW_COPY_AND_ASSIGN(GpuVirtualMemAllocator);
};

3.3.3 gpu_virtual_mem_allocator.cc

gpu_virtual_mem_allocator.h 定义了 GpuVirtualMemAllocator 类的接口和成员变量，而 gpu_virtual_mem_allocator.cc 则实现了这些接口的具体逻辑。让我们结合注释简单分析 gpu_virtual_mem_allocator.cc 的代码实现：

1）GpuVirtualMemAllocator::Create创建 GpuVirtualMemAllocator 实例：

cpp 复制代码 
/* static */ stream_executor::port::StatusOr<
    std::unique_ptr<GpuVirtualMemAllocator>>
GpuVirtualMemAllocator::Create(
    const std::vector<Visitor>& alloc_visitors,
    const std::vector<Visitor>& free_visitors, GpuContext& gpu_context,
    PlatformDeviceId gpu_id, size_t virtual_address_space_size,
    const std::vector<PlatformDeviceId>& peer_gpu_ids) {

  // 1) 初始化 access_gpu_handles，包括主 GPU 和所有支持虚拟地址管理的对等 GPU。
  // - 首先获取主 GPU 的句柄，并检查其是否支持虚拟地址管理。
  // - 遍历 peer_gpu_ids，获取每个对等 GPU 的句柄，并检查是否支持虚拟地址管理和对等访问。
  // - 如果支持，则将对等 GPU 的句柄添加到 access_gpu_handles 中。
  std::vector<GpuDeviceHandle> access_gpu_handles;
  access_gpu_handles.reserve(peer_gpu_ids.size() + 1);

  GpuDeviceHandle gpu_handle;
  TF_RETURN_IF_ERROR(GpuDriver::GetDevice(gpu_id.value(), &gpu_handle));
  TF_RETURN_IF_ERROR(CheckVirtualAddressManagementSupport(gpu_handle, gpu_id));

  access_gpu_handles.push_back(gpu_handle);
  for (const auto& peer_id : peer_gpu_ids) {
    GpuDeviceHandle peer_handle;
    TF_RETURN_IF_ERROR(GpuDriver::GetDevice(peer_id.value(), &peer_handle));
    TF_ASSIGN_OR_RETURN(bool supports_virtual_address_management,
                        SupportsVirtualAddressManagement(peer_handle));
    if (GpuDriver::CanEnablePeerAccess(gpu_handle, peer_handle) &&
        supports_virtual_address_management) {
      access_gpu_handles.push_back(peer_handle);
    }
  }

  // 2) 遍历 access_gpu_handles，获取每个设备的最小分配粒度，并取最大值。

  // Find the min granularity for all devices that have access to this memory;
  // that is, the maximum min granularity among all devices.
  // 找到所有可以访问此内存的设备的最小粒度，即所有设备中最大的最小粒度。
  size_t max_granularity = 1;
  for (const auto device_handle : access_gpu_handles) {
    TF_ASSIGN_OR_RETURN(size_t granularity,
                        GpuDriver::GetMinAllocationGranularity(device_handle));
    max_granularity = std::max(max_granularity, granularity);
  }

  // 3) 调用 GpuDriver::ReserveVirtualMemory 保留虚拟内存，并确保其大小对齐到 max_granularity。

  // Create the virtual memory reservation. Must be aligned to system page size,
  // and larger than the CUDA min granularity. Empirically, the granularity
  // check is sufficient as the granularity is some multiple of the page size.
  // TODO(imintz): Create OS agnostic page size utility for completeness.
  // 创建虚拟内存保留区域。必须与系统页大小对齐，并且大于 CUDA 的最小粒度。
  // 经验表明，粒度检查是足够的，因为粒度是页大小的某个倍数。
  TF_ASSIGN_OR_RETURN(
      GpuDriver::VmemSpan vmem,
      GpuDriver::ReserveVirtualMemory(
          &gpu_context, AlignUp(virtual_address_space_size, max_granularity)));
  VLOG(1) << "Reserved GPU virtual memory at " << vmem.base << " of size "
          << strings::HumanReadableNumBytes(vmem.size_bytes) << " bytes";
  
  // 4) 使用构造函数创建 GpuVirtualMemAllocator 实例，并将其包装在 std::unique_ptr 中返回。
  return std::unique_ptr<GpuVirtualMemAllocator>(new GpuVirtualMemAllocator(
      alloc_visitors, free_visitors, gpu_context, gpu_id,
      std::move(access_gpu_handles), vmem, max_granularity));
}

2）GpuVirtualMemAllocator 的构造函数 和析构函数：

cpp 复制代码 
// 1) 构造函数，初始化成员变量。
GpuVirtualMemAllocator::GpuVirtualMemAllocator(
    const std::vector<Visitor>& alloc_visitors,
    const std::vector<Visitor>& free_visitors, GpuContext& gpu_context,
    PlatformDeviceId gpu_id,
    const std::vector<GpuDeviceHandle> access_gpu_handles,
    GpuDriver::VmemSpan vmem, size_t granularity)
    : SubAllocator(alloc_visitors, free_visitors),
      gpu_context_(gpu_context),
      gpu_id_(gpu_id),
      access_gpu_handles_(access_gpu_handles),
      vmem_(vmem),
      granularity_(granularity) {}

// 2) 析构函数，遍历 mappings_，取消映射并释放物理内存，最后释放虚拟内存。
GpuVirtualMemAllocator::~GpuVirtualMemAllocator() {
  for (const auto mapping : mappings_) {
    GpuDriver::UnmapMemory(&gpu_context_, mapping.va, mapping.physical.bytes);
    GpuDriver::ReleaseMemoryHandle(&gpu_context_, std::move(mapping.physical));
  }
  GpuDriver::FreeVirtualMemory(&gpu_context_, vmem_);
}

3）GpuVirtualMemAllocator::Alloc根据请求的大小和对齐要求分配内存：

cpp 复制代码 
void* GpuVirtualMemAllocator::Alloc(size_t alignment, size_t num_bytes,
                                    size_t* bytes_received) {

  // 1) 如果请求的内存大小为 0，直接返回 nullptr。                                    
  if (num_bytes == 0) return nullptr;

  // 2) 计算对齐后的内存大小。granularity_ 是最小分配粒度，确保分配的内存大小是 granularity_ 的倍数。
  size_t padded_bytes = (num_bytes + granularity_ - 1) & ~(granularity_ - 1);
  
  // 3) 计算下一个可用的虚拟地址。
  GpuDevicePtr next_va = vmem_.base + next_alloc_offset_;
  
  // 4) 检查是否有足够的虚拟内存空间来满足请求。
  // TODO(imintz): Attempt to extend the vmem allocation by reserving additional
  // virtual memory at the specific address at the end of the initial vmem
  // reservation.
  // TODO(imintz): 尝试通过在初始虚拟内存保留区域的末尾保留额外的虚拟内存来扩展虚拟内存分配。
  if (next_va + padded_bytes > vmem_.base + vmem_.size_bytes) {
    LOG(ERROR) << "OOM in GPU virtual memory allocator when attempting to "
                  "allocate {request: "
               << strings::HumanReadableNumBytes(num_bytes)
               << ", aligned: " << padded_bytes << "} bytes.";
    return nullptr;  // 如果空间不足，返回 nullptr 并记录错误日志。
  }

  // Create physical memory backing allocation.
  // 5) 创建物理内存句柄，大小为对齐后的内存大小。
  auto maybe_handle =
      GpuDriver::CreateMemoryHandle(&gpu_context_, padded_bytes);
  if (!maybe_handle.ok()) {
    LOG(ERROR) << maybe_handle.status();
    return nullptr;  // 如果创建失败，记录错误日志并返回 nullptr。
  }
  GpuDriver::GenericMemoryHandle handle = std::move(maybe_handle).ValueOrDie();

  // Map VAs for this physical memory.
  // 6) 将物理内存映射到虚拟地址。
  auto status =
      GpuDriver::MapMemory(&gpu_context_, next_va, handle, access_gpu_handles_);
  if (!status.ok()) {
    LOG(ERROR) << status;
    GpuDriver::ReleaseMemoryHandle(&gpu_context_, std::move(handle));
    return nullptr;  // 如果映射失败，记录错误日志并释放物理内存句柄。
  }

  // 7) 更新下一个分配偏移量，并将映射信息保存到 mappings_ 中。
  next_alloc_offset_ += handle.bytes;
  mappings_.push_back({next_va, std::move(handle)});

  // 8) 调用 VisitAlloc 通知分配事件。
  VisitAlloc(reinterpret_cast<void*>(next_va), gpu_id_.value(), padded_bytes);

  // 9) 返回分配的内存地址，并通过 bytes_received 返回实际分配的大小。
  *bytes_received = padded_bytes;
  return reinterpret_cast<void*>(next_va);
}

逻辑总结：

  1. 对齐内存大小 ：根据最小分配粒度 granularity_ 对齐请求的内存大小。
  2. 检查虚拟内存空间：确保有足够的虚拟内存空间来满足请求。
  3. 创建物理内存句柄 ：调用 GpuDriver::CreateMemoryHandle 创建物理内存句柄。
  4. 映射虚拟地址：将物理内存映射到虚拟地址。
  5. 更新状态 ：更新 next_alloc_offset_ 和 mappings_，并通知分配事件。

4）GpuVirtualMemAllocator::Free释放之前分配的内存：

cpp 复制代码 
void GpuVirtualMemAllocator::Free(void* ptr, size_t num_bytes) {

  // 1) 如果指针为 nullptr，直接返回。
  if (ptr == nullptr) return;
  
  // 2) 使用二分查找在 mappings_ 中找到与 ptr 对应的映射。
  auto mapping_it =
      std::lower_bound(mappings_.begin(), mappings_.end(), ptr,
                       [](const Mapping& mapping, const void* ptr) {
                         return reinterpret_cast<const void*>(mapping.va) < ptr;
                       });
  if (mapping_it == mappings_.end() ||
      (reinterpret_cast<void*>(mapping_it->va) != ptr)) {
    LOG(ERROR) << "Could not find GPU vmem mapping for address at "
               << reinterpret_cast<uintptr_t>(ptr);
    return;  // 如果找不到对应的映射，记录错误日志并返回。
  }
  
  // 3) 计算需要释放的映射数量和总字节数。
  int num_mappings_to_free = 0;
  int total_bytes = 0;
  for (auto it = mapping_it; it != mappings_.end() && total_bytes < num_bytes;
       ++it) {
    ++num_mappings_to_free;
    total_bytes += it->physical.bytes;
  }
  if (total_bytes != num_bytes) {
    LOG(ERROR) << "Invalid size requested for freeing GPU vmem mapping. Got "
               << strings::HumanReadableNumBytes(num_bytes) << " but expected "
               << strings::HumanReadableNumBytes(mapping_it->physical.bytes);
    return;  // 如果请求释放的字节数与实际分配的字节数不匹配，记录错误日志并返回。
  }
  
  // 记录日志，显示释放的映射数量和总字节数。
  VLOG(1) << "Freeing " << num_mappings_to_free << " mappings for a total of "
          << total_bytes << " bytes";

  // 4) 遍历需要释放的映射，取消映射并释放物理内存句柄。
  for (auto it = mapping_it; it < mapping_it + num_mappings_to_free; ++it) {
    GpuDriver::UnmapMemory(&gpu_context_, it->va, it->physical.bytes);
    GpuDriver::ReleaseMemoryHandle(&gpu_context_, std::move(it->physical));
  }

  // Move back the next_alloc_offset_ if this free was at the end.
  // 5) 如果释放的映射位于 mappings_ 的末尾，则更新 next_alloc_offset_。
  if (mapping_it + num_mappings_to_free == mappings_.end()) {
    next_alloc_offset_ = mapping_it->va - vmem_.base;
  }
  
  // 6) 从 mappings_ 中移除已释放的映射。
  mappings_.erase(mapping_it, mapping_it + num_mappings_to_free);

  // 7) 调用 VisitFree 通知释放事件。
  VisitFree(ptr, gpu_id_.value(), num_bytes);
}

逻辑总结：

  1. 查找映射 ：使用二分查找在 mappings_ 中找到与 ptr 对应的映射。
  2. 检查映射有效性 ：确保找到的映射与 ptr 匹配。
  3. 计算释放范围：确定需要释放的映射数量和总字节数。
  4. 取消映射并释放内存：遍历需要释放的映射，取消映射并释放物理内存句柄。
  5. 更新状态 ：如果释放的映射位于末尾，则更新 next_alloc_offset_，并从 mappings_ 中移除映射及通知释放事件。

4 在 PyTorch 的实现

4.1 初见端倪

第一次知道 PyTorch 的 expandable_segments 机制还是在最初研究 GMLake 的时候，当时 GMLake 的作者正在知乎上和网友们激烈讨论 GMLake 和 expandable_segments 的区别和优劣，这也引起了我的兴趣。

PyTorch 的 expandable_segments 机制在 2023 年 3 月 17 日由 zdevito 提交了第一个 PR。

zdevito 在提交这个 PR 时的评论：

我们通常建议通过预先分配一个大内存块来缓解内存碎片问题，这块内存可以在后续使用中逐步分割。对于++张量大小动态变化的程序++，这种方法尤其有效，因为它可以避免在首次遇到新的最大输入时发生内存不足（OOM）问题，否则程序将不得不分配新的内存段。

然而，预先分配内存块的一个主要问题是很难准确估计所需的内存块大小。如果分配的内存块过小，内存块的空间会很快耗尽，分配器将不得不额外分配独立的内存块；如果分配的内存块过大，其他非 PyTorch 库可能会因为无法分配到足够内存而无法正常运行。

本补丁提供了与预分配内存块相同的好处，但++无需提前确定内存块的大小++ 。通过使用 cuMemMap 风格的 API，它增加了++在需要更多内存时扩展段中最后一个内存块++的能力。

与全面使用 cudaMallocAsync 来避免内存碎片相比，本补丁能够在保留现有分配器行为的同时解决这一常见的碎片问题。此外，这一行为可以动态启用或禁用。例如，用户可以将长期存在的参数和状态分配到独立的缓冲区中，而将临时状态放入可扩展的大内存块中，从而进一步减少内存碎片。

有关实现细节及其限制，请参阅代码中的内联注释。

我们可以在该 PR 的 Files changed 中看到与 expandable_segments 机制有关的代码文件主要是：

另外，我们用与之前相似的方法，在 GitHub 中搜索repo:pytorch/pytorch "expandable segments" OR "expandable blocks" author:zdevito，可以查看与 expandable_segments 机制相关的所有 PR 及 Commits。

关注几个比较有信息量的 Commits：

2023 年 5 月 3 日

Support expandable_segments:True in fbcode for caching allocator

在缓存分配器的 fbcode 中支持 expandable_segments:True

Now that expandable_segments has been merged from OSS, we can enable it in the internal build. It still defaults to off, so this should not change any behavior changes in the allocator unless the flag is explicitly set.

现在 expandable_segments 已从开源版本（OSS）合并，我们可以在内部构建中启用它。++该选项默认仍为关闭状态，因此除非显式设置该标志，否则不会改变分配器的任何行为++。

2023 年 11 月 19 日

Don't recommmend max_split_size_mb first

推荐不要首先 max_split_size_mb。

I've run into a couple of cases now where max_split_size_mb has been set in projects as a workaround for fragmentation but it ends up causing problems later, such as degraded performance from freeing empty segments. While it is a useful setting to have, expandable_segments is probably a better first resort for fixing fragmentation since when it works it is less likely to need synchronous GPU operations to continue running.

我最近遇到了几个案例，项目中设置了 max_split_size_mb 作为解决内存碎片问题的临时方案，但最终却导致了其他问题，例如因释放空内存段而导致的性能下降。虽然 max_split_size_mb 是一个有用的设置，但 expandable_segments 可能是解决碎片化问题的更好首选方案，因为当它生效时，不太可能需要同步 GPU 操作来继续运行。

2024 年 8 月 30 日

expandable_segments <-> other allocator options

expandable_segments 与其他分配器选项

Previously setting garbage_collection_threshold or max_split_size_mb along with expandable_segments:True could cause the allocator to hit assert failures when running nearly out of memory. This PR ensures garbage_collection and max_split freeing do not accidentally try to release expandable segments.

此前，如果将 garbage_collection_threshold 或 max_split_size_mb 与 expandable_segments:True 一起设置，可能会导致++分配器在内存几乎耗尽时触发断言失败++ 。本次 PR 确保垃圾回收（garbage_collection）和 max_split 释放操作++不会意外尝试释放可扩展段（expandable segments）++。

综上可以看出，PyTorch 在 2023 年 3 月 17 日首次提出 expandable_segments 机制后，一直到 2024 年 8 月 30 日，中间也迭代了几个小版本。

4.2 探本穷源

PyTorch 在 2.1 的官方文档第一次提到 expandable_segments

expandable_segments (experimental, default: False) If set to True, this setting instructs the allocator to create CUDA allocations that can later be expanded to better handle cases where a job changing allocation sizes frequently, such as having a changing batch size. Normally for large (>2MB) allocations, the allocator calls cudaMalloc to get allocations that are the same size as what the user requests. In the future, parts of these allocations can be reused for other requests if they are free. This works well when the program makes many requests of exactly the same size or of sizes that even multiples of that size. Many deep learning models follow this behavior. However, one common exception is when the batch size changes slightly from one iteration to the next, e.g. in batched inference. When the program runs initially with batch size N, it will make allocations appropriate for that size. If in the future, it runs at size N - 1, the existing allocations will still be big enough. However, if it runs at size N + 1, then it will have to make new allocations that are slightly larger. Not all the tensors are the same size. Some might be (N + 1)A and others (N + 1)A B where A and B are some non-batch dimensions in the model. Because the allocator reuses existing allocations when they are big enough, some number of (N + 1)A allocations will actually fit in the already existing NB A segments, though not perfectly. As the model runs it will partially fill up all of these segments leaving unusable free slices of memory at the end of these segments. The allocator at some point will need to cudaMalloc a new (N + 1)AB segment. If there is not enough memory, there is now no way to recover the slices of memory that are free at the end of existing segments. With models 50+ layers deep, this pattern might repeat 50+ times creating many slivers.

expandable_segments（实验性功能，默认值：False）如果设置为True，该设置指示分配器创建可以扩展的 CUDA 分配 ，以更好地处理那些在作业中频繁更改分配大小的情况 ，例如++批量大小变化的情况++。

  • 通常，对于大于 2MB 的分配，分配器会调用cudaMalloc来获取与用户请求的大小相同的分配。++未来，如果这些分配中的部分空间是空闲的，可以重新用于其他请求。当程序发出许多相同大小或大小为该尺寸的整数倍的请求时，这种方式效果较好++。许多深度学习模型遵循这种行为。
  • 然而，一个常见的例外是当批量大小在每次迭代之间稍微变化时 ，例如++批量推理++ 。当程序最初使用批量大小 N N N运行时，它会进行适合该大小的分配。如果将来程序运行时批量大小变为 N − 1 N-1 N−1，现有的分配仍然足够大。然而，如果批量大小变为 N + 1 N+1 N+1，程序将不得不进行稍微更大的新分配。并不是所有的张量大小都相同。有些可能是 ( N + 1 ) ∗ A (N+1) * A (N+1)∗A，其他则是 ( N + 1 ) ∗ A ∗ B (N+1) * A * B (N+1)∗A∗B，其中 A A A和 B B B是模型中的++非批量维度++ 。由于分配器会在现有分配足够大的情况下重用它们，因此一些 ( N + 1 ) ∗ A (N+1) * A (N+1)∗A的分配实际上可以适应已经存在的 N ∗ B ∗ A N * B * A N∗B∗A段，尽管不是完全匹配。随着模型的运行，这些段会部分填充，导致这些段的末尾留下无法使用的空闲内存片段。分配器在某个时刻将需要进行新的cudaMalloc来分配 ( N + 1 ) ∗ A ∗ B (N+1) * A * B (N+1)∗A∗B大小的段。如果内存不足，那么无法回收现有段末尾的空闲内存片段。对于 50 层以上的深度模型，这种模式可能会重复 50 多次，产生许多无法使用的小内存片段。

expandable_segments allows the allocator to create a segment initially and then expand its size later when more memory is needed. Instead of making one segment per allocation, it tries to make one segment (per stream) that grows as necessary. Now when the N + 1 case runs, the allocations will tile nicely into the one large segment until it fills up. Then more memory is requested and appended to the end of the segment. This process does not create as many slivers of unusable memory, so it is more likely to succeed at finding this memory.

expandable_segments 允许分配器最初创建一个段，并在需要更多内存时扩展其大小 。它不会为每个分配创建一个段，而是尝试为每个流创建一个段，该段会根据需要增长 。现在，当 N + 1 N+1 N+1的情况发生时，分配将很好地拼接到一个大段中，直到它被填满。然后会请求更多的内存，并将其附加到该段的末尾。这个过程不会产生那么多无法使用的内存片段，因此更有可能成功地找到合适的内存。

以下内容来自：PyTorch高性能编程（持续更新）

对于输入数据 size 频繁变化的场景，使用 Expandable Segments。

bash 复制代码 
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

跟cudaMalloc直接分配 Kernel 可访问的内存地址不同，该机制操作的是虚拟内存空间（++对应的物理内存地址不具备访问权限++ ），可以++通过驱动 map 更多的物理内存在已分配的 block 的后面，从而使得 segments 可向上扩展++ ，一定程度上提高了 cache match 的效率，减少内存碎片。

以下内容来自：使用PyTorch Profiler进行模型性能分析，改善并加速PyTorch训练 - overfit.cn

设置 PyTorch 相对较新的分配器模式：

bash 复制代码 
PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

这告诉 PyTorch 分配器分配可以在将来扩展的块 。但是，++如果大小变化太大，它仍然可能无法解决问题++。

所以我们只能手动来进行优化，那就是是++使数据形状一致++ 。这样分配器就更容易++找到合适的数据块进行重用++。比如最简单的将数据填充到相同的大小。或者可以通过运行具有最大输入大小的模型来预热分配器。

以下内容来自：PyTorch 内存管理之 expandable_segments

训练和推理中优化场景的应用：

  • 单个 step 内临时内存的申请释放，先申请小块内存，后申请大块内存；
  • step 间输入 shape 变化，输入 shape 变大，导致前一 step 申请的内存块小于本 step 需要的大小，从而需要调用 cuda 接口新申请，造成内存碎片增加。

如，先申请 3 块 80M 内存，使用完释放，再申请 3 块 100M 内存，最终剩余内存碎片 240M。

以下内容来自：PyTorch显存管理介绍与源码解析（三）

以一个 Transformer 用例来对比测试，开启 expandable 功能前后的显存块申请差异。示例运行后装载数据进行可视化，跳转到 "allocator state hisotry" 页面可看到：

expandable_segments:False：

expandable_segments:True：

两个方式下 segment 形状有了明显的区别：

  • expandable_segments 模式下创建的块更少；
  • expandable_segments 模式下的块利用得更加的紧凑，留出的空白面积（未使用显存）更大。

这样的显存管理带来的好处是碎片量更少，用户能够申请更大的整块显存，相比之下性能更优。

另外，相比cudaMalloc，VMM 其调用链路多了几个流程，操作时间会更长，而虚拟地址的映射操作调用频率会更高，而频繁的动作可能会让整个 GPU 利用率变差。

简而言之 expandable_segments 特点：降低了碎片、提升了有效的可用空间；但初始化时间更久，当前有些功能不支持。

以下内容来自：[2405.04437] vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

In a recent work, GMLake (gmlakeasplos24, ) showed that using CUDA virtual memory support can mitigate fragmentation in DNN training jobs, increasing training batch size. In particular, GMLake uses CUDA support to coalesce multiple smaller physical memory pages into a single virtually contiguous object that can prevent out-of-memory errors for large object allocations.

在最近的一项研究中，GMLake 展示了使用 CUDA 虚拟内存支持可以缓解 DNN 训练任务中的碎片化问题，从而增加训练批量大小。特别地，GMLake 利用 CUDA 支持将多个较小的物理内存页合并成一个虚拟连续的对象，从而防止大对象分配时出现内存不足错误。

Similar to our approach, PyTorch also recently added support for expandable segments (expandable-segments, ) that dynamically attach physical memory pages to a pre-reserved virtual memory buffer, if enabled. However, PyTorch uses expandable segments opportunistically, using synchronous memory allocation, and allocates physical memory pages only in multiples of 2MB granularity.

与我们的方法类似，PyTorch 最近也添加了对可扩展段（expandable-segments）的支持，在启用时，动态地将物理内存页附加到预先保留的虚拟内存缓冲区中。然而，PyTorch 的可扩展段使用是机会性的，它采用同步内存分配，并且只按 2MB 粒度分配物理内存页。

以下内容来自：[2407.20018] Efficient Training of Large Language Models on Distributed Infrastructures: A Survey

6.3.2 VMM-based Defragmentation

GMLake [308] and PyTorch expandable segments [309] propose to mitigate fragmentation by utilizing the virtual memory management (VMM) functions of the low-level CUDA driver application programming interface. This low-level API provides developers with direct control over the GPU's virtual memory operations, such as reserving, mapping, and managing virtual memory addresses. Building on this, GMLake [308] introduces a virtual memory stitching mechanism that consolidates non-contiguous memory blocks into larger ones through virtual memory address mapping, minimizing data movement and copying. Similarly, PyTorch's expandable segments [309] enable allocated memory segments to be expanded to larger sizes for reuse. Both approaches are transparent to different models and memoryefficient training techniques and can be seamlessly integrated into existing deep learning frameworks. Furthermore, GMLake demonstrates excellent scalability on multiGPUs with minimal overhead and does not require modification to user code. PyTorch-v2.1 has also integrated expandable segments.

6.3.2 基于虚拟内存管理（VMM）的碎片整理

GMLake [308] 和 PyTorch 可扩展段 [309] 提出了通过利用低级 CUDA 驱动程序应用编程接口（API）的虚拟内存管理（VMM）功能来缓解碎片化问题。这个低级 API 为开发人员提供了对 GPU 虚拟内存操作的直接控制，例如预留、映射和管理虚拟内存地址。在此基础上，

  • GMLake [308] 引入了一种虚拟内存拼接机制，通过虚拟内存地址映射将不连续的内存块合并成更大的内存块，从而最小化数据的移动和复制。
  • 类似地，PyTorch 的可扩展段 [309] 使已分配的内存段能够扩展为更大的尺寸以供重用。

这两种方法对不同的模型和内存高效训练技术透明，且能够无缝地集成到现有的深度学习框架中。此外，

  • GMLake 在多GPU环境下表现出出色的可扩展性，且开销最小，不需要修改用户代码。
  • PyTorch v2.1 也已集成了可扩展段功能。

4.3 实现原理

PyTorch expandable_segments 机制的核心代码改动都在 CUDACachingAllocator.cpp，让我们先通过代码注释了解一下 PyTorch expandable_segments 机制的实现原理：

Implementation

实现

The expandable_segments:True option is used to enable/disable this behavior. We
use cuda's low-level memory APIs, which are similar to mmap, to extend the
memory segments. These APIs separate the allocation of physical memory
(cuMemCreate) from the allocation of virtual address space (cuMemAddressReserve)
and the associate between them cuMemMap/cuMemSetAccess.

expandable_segments:True 选项用于启用/禁用此行为。我们使用 CUDA 的低级内存 API（类似于 mmap）来扩展内存段。++这些 API 将物理内存的分配（cuMemCreate）与虚拟地址空间的分配（cuMemAddressReserve）以及它们之间的关联（cuMemMap/cuMemSetAccess）分开++。

When we allocate a new segment, we allocate enough address space to map
basically the entire physical memory of the GPU (there is 256TiB of address
space), but we only map enough physical memory to handle the current amount of
memory needed by the program. As more is requested, we add more physical memory
to the segment. This can work at the granularity of GPU pages which are 2MiB
currently.

当我们分配一个新段时，我们分配足够的地址空间来映射基本上整个 GPU 的物理内存（有 256TiB 的地址空间），但我们只映射足够的物理内存来处理程序当前所需的内存量。随着更多内存的请求，我们向段中添加更多物理内存。这可以在 GPU 页面的粒度上工作，目前为 ++2MiB++。

If we end up out of memory, we can unmap all the memory in our segment
corresponding to empty physical pages, and return it to CUDA for use at another
address in the segment or in a segment for a different stream.

如果我们最终内存不足，我们可以取消映射段中对应于空物理页面的所有内存，并将其返回给 CUDA，以便在段中的另一个地址或不同流的段中使用。

A current limitation of CUDA's API is that physical memory
(CUmemGenericAllocationHandle) cannot be split up after it is mapped even if the
handle holds multiple GPU pages. The cost to map/unmap memory is proportional to
the number of physical memory chunks that were allocated (mapping 10 separately
allocated 2MiB pages takes 10x time compared to mapping one 20MiB physical
allocation of 10 pages). Changing memory mappings also appears to involve at
least some synchronous actions with the GPU and so should be considered an
expensive operation. To limit overhead, we use 2MiB pages for our small pool and
20MiB pages for our large pool. Initially allocation using expandable_blocks
will be slower than cudaMalloc, though still in the milliseconds range for
mapping the entire memory.

CUDA API 的一个当前限制是，物理内存（CUmemGenericAllocationHandle）在映射后无法拆分，即使句柄包含多个 GPU 页面。++映射/取消映射内存的成本与分配的物理内存块的数量成正比++ （映射 10 个单独分配的 2MiB 页面所需的时间是映射一个 20MiB 物理分配的 10 倍）。++更改内存映射似乎还涉及至少一些与 GPU 的同步操作，因此应被视为昂贵的操作++ 。为了限制开销，我们对小池使用 2MiB 页面，对大池使用 20MiB 页面。++最初使用 expandable_blocks 进行分配将比cudaMalloc慢++，尽管映射整个内存的时间仍在毫秒范围内。

When mapping new memory to expand the segment, we look for the lowest address at
which we can fit a new allocation by adding new pages. Normally this will be at
the end of the block. But if have previously unmapped blocks earlier in the
segment during an OOM, it will first try to fill in those gaps to keep the
segment as a single block. By allocating at the lowest address we encourage
the split up parts of the block to merge into a single block again, reducing
fragmentation potential.

当映射新内存以扩展段时，我们通过添加新页面来寻找适合新分配的++最低地址++ 。通常这将在++块的末尾++ 。但如果之前在 OOM 期间取消了段中较早的块的映射，它将首先尝试填补这些空白，以保持段为单个块。通过以最低地址分配，我们++鼓励块的拆分部分再次合并为单个块，减少碎片化的可能性++。

Allocation of blocks in the segment uses the same best-fit heuristics of the
rest of the allocator.

段中的块分配使用与分配器其余部分相同的++最佳适应启发式方法++。

Expandable blocks can be enabled/disabled throughout the run of a program. When
disabled, the allocator will not put new allocations in an expandable block.

可扩展块可以在程序运行期间启用/禁用。禁用时，分配器不会将新分配放入可扩展块中。

以下内容来自：PyTorch 内存管理之 expandable_segments

PyTorch expandable_segments 内存管理流程：

1）初始化

业务首次内存申请会触发初始化，申请大块虚拟地址大小为 gpu 显存的 1.125 倍，例 GPU 显存 32G，则申请 40G 的虚拟地址。

2）内存申请

PyTorch 内存管理会按照申请内存的大小 (1M) 分为两个内存池，以大块内存池为例。

为了方便管理物理内存块，物理内存按照 20M 申请。假如业务第一次内存申请为 89M，向上取整 90M，则先申请 5 个 20M 的物理内存块，从虚拟地址起始指针依次 map。多次申请的物理内存块按顺序 map 到虚拟地址，从虚拟地址的角度来看是一整块连续的虚拟内存。

3）内存释放

PyTorch 内存池通常不会释放内存池中的内存给 GPU，除非用户显示调用torch.cuda.empty_cache()，会将未在使用的内存释放。另外在内存使用极限场景，PyTorch 内存池调用 cuda 接口申请内存失败时，程序不会直接 OOM 退出，而是先释放内存池中未使用的内存碎片，再尝试申请内存，如果依然申请不到才 OOM 退出。

判断申请的内存块是否在用，未使用则释放，并重新申请，重新申请的内存块在虚拟地址尾端 map。

如图，假设内存极限场景 32G 物理内存已经全被内存池申请，内存池中有两块不连续空闲 20M 内存；业务申请 40M 内存使用，由于物理内存申请完，会先触发 OOM，umap 后释放物理地址（灰色）；重新申请 2*20M 释放的物理地址，并 map 在末尾（ptr+32G），返回虚拟地址供业务使用。

以下内容来自：PyTorch显存管理介绍与源码解析（三）

如何去解决碎片问题？这里给一个常见场景，如下所示，虚拟地址上面有 6 个在使用的块（segment_size=2M），这些子块的大小是一致，当前总共使用了 12M 的空间。假设其中第 2 号、4 号块在某个时刻被应用释放，这时就有了两个离散的空闲块，且大小都为 2M。当物理内存没有新的可用空间时，如何申请到 4M 的 segment 块？ 直接在物理地址申请会报错 OOM 错误。

在 VMM 中，由于虚拟地址可以重映射，所以该场景下显存变得可申请到。具体的操作是：先解映射（unmap），让 2 号块和 4 号块的虚拟地址与物理块解除绑定，然后把物理块映射到虚拟地址的尾部，就能够获得一个连续的更大的块：6 号 + 7 号。

这个操作过程中块 2 与块 4 的物理地址并不需要真正的释放，只是做了一下重映射操作，所以过程不会产生 OOM 问题。

4.4 核心代码

因笔者精力有限，且相比于 TensorFlow 的 GpuVirtualMemAllocator，PyTorch expandable_segments 机制的代码实现更为复杂，所以笔者目前只求初步理解其核心代码，并不准备深入分析其完整的代码实现。之后如果发现有质量不错的源码分析博客，笔者会在这里引用学习~😊

ExpandableSegment是 PyTorch 实现 expandable_segments 机制的核心数据结构，其代码实现如下：

cpp 复制代码 
// 用于表示一段内存的范围
struct SegmentRange {
  char* ptr;    // 指向内存段起始地址的指针
  size_t size;  // 内存段的大小
  SegmentRange(void* p, size_t s) : ptr(static_cast<char*>(p)), size(s) {}
};

struct ExpandableSegment {

  // 构造函数
  ExpandableSegment(
      int device,             // 当前 GPU 设备 ID
      cudaStream_t stream,    // CUDA 流，用于同步操作
      size_t size,            // 每个内存段的大小
      std::vector<int> peers) // 需要共享该内存的其他 GPU 设备列表
      : device_(device),      
        stream_(stream),      
        max_handles_(0),
        // 2MB for small pool, 20MB for large pool
        segment_size_(size),  
        peers_(std::move(peers)) {
    // 获取 GPU 设备属性
    cudaDeviceProp prop{};
    C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device_));

    // we allocate enough address space for 1 1/8 the total memory on the GPU.
    // This allows for some cases where we have to unmap pages earlier in the
    // segment to put them at the end.
    // 计算并预留足够的虚拟地址空间，通常是 GPU 总内存的 1.125 倍，
    // 这是为了应对某些情况下需要取消映射并重新映射内存页的需求。
    max_handles_ = numSegments(prop.totalGlobalMem + prop.totalGlobalMem / 8);

    // cuMemAddressReserve 预留虚拟地址空间
    C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressReserve_( 
        &ptr_, segment_size_ * max_handles_, 0ULL, 0, 0ULL));
  }

  // 将物理内存映射到虚拟地址空间

  // begin must be aligned to segment_size_.
  // returns the actual range mapped, which may be
  // greater than requested if size is not aligned to segment_size_.
  // return size of 0 indicates OOM
  // begin 必须与 segment_size_ 对齐。
  // 返回实际映射的内存范围，该范围可能比请求的范围更大（如果请求的大小未与 segment_size_ 对齐）。
  // 返回的大小为 0 表示内存不足（OOM）。
  SegmentRange map(SegmentRange range) {
    // 1. 计算需要映射的内存段的起始和结束索引
    auto begin = segmentLeft(range.ptr);
    auto end = segmentRight(range.ptr + range.size);
    // 2. 断言检查：确保 range.ptr 是某个内存段的起始地址
    TORCH_INTERNAL_ASSERT(ptr() + begin * segment_size_ == range.ptr);
    // 3. 如果 begin == end，说明没有需要映射的内存段，直接返回空范围
    if (begin == end) {
      return rangeFromHandles(begin, end);
    }
    // 4. 扩展 handles_ 向量，确保其大小足够容纳新的内存段
    while (end > handles_.size()) {
      handles_.emplace_back(c10::nullopt);
    }
    // 5. 遍历需要映射的内存段，为每个内存段分配物理内存
    for (auto i : c10::irange(begin, end)) {
      // 5.1 断言检查：确保当前内存段尚未分配
      TORCH_INTERNAL_ASSERT(!handles_.at(i));
      // 5.2 创建 CUDA 内存分配句柄
      CUmemGenericAllocationHandle handle = 0;
      CUmemAllocationProp prop = {};
      prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;        // 内存类型为固定内存
      prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // 内存位于 GPU 设备
      prop.location.id = device_;                       // 指定 GPU 设备 ID
      // 5.3 调用 cuMemCreate_ 分配物理内存
      auto status =
          DriverAPI::get()->cuMemCreate_(&handle, segment_size_, &prop, 0);
      // 5.4 如果内存不足（OOM），释放已分配的内存并返回空范围
      if (status == CUDA_ERROR_OUT_OF_MEMORY) {
        for (auto j : c10::irange(begin, i)) {
          auto h = handles_.at(j).value();  // 获取已分配的内存句柄
          handles_.at(j) = c10::nullopt;    // 将句柄置为空
          C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h));  // 释放内存
        }
        trimHandles();  // 移除 handles_ 末尾的空句柄
        return rangeFromHandles(begin, begin);  // 返回空范围
      }
      // 5.5 检查 cuMemCreate_ 的返回值，确保内存分配成功
      C10_CUDA_DRIVER_CHECK(status);
      // 5.6 将分配的内存句柄存储到 handles_ 中
      handles_.at(i) = handle;
    }
    // 6. 将虚拟地址映射到物理内存
    for (auto i : c10::irange(begin, end)) {
      C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemMap_(
          ptr_ + i * segment_size_, // 虚拟地址
          segment_size_,            // 内存段大小
          0,                        // 映射标志
          handles_.at(i).value(),   // 内存句柄
          0ULL));                   // 保留参数
    }
    // 7. 设置当前设备对内存段的访问权限
    setAccess(device_, begin, end);
    // 8. 设置其他设备（peers_）对内存段的访问权
    for (auto p : peers_) {
      setAccess(p, begin, end);
    }
    // 9. 返回实际映射的内存范围
    return rangeFromHandles(begin, end);
  }

  // 取消映射指定范围内的内存段。

  // unmaps all the completely empty segment_size_ segments between
  // [begin, begin + size), returns the offset where the range begin,
  // and the actual size unmapped (multiple of segment_size_)
  // 取消映射 [begin, begin + size) 范围内所有完全空闲的 segment_size_ 段，
  // 返回该范围的起始偏移量以及实际取消映射的大小（大小为 segment_size_ 的倍数）。
  SegmentRange unmap(SegmentRange range) {
    // 1. 计算需要取消映射的内存段的起始和结束索引
    auto begin = segmentRight(range.ptr);
    auto end = segmentLeft(range.ptr + range.size);
    // 2. 如果 begin >= end，说明没有需要取消映射的内存段，返回空范围
    if (begin >= end) {
      return SegmentRange{range.ptr, 0};  // 返回起始地址和大小为 0 的范围
    }
    // 3. 调用 unmapHandles 取消映射 [begin, end) 范围内的内存段
    unmapHandles(begin, end);
    // 4. 返回实际取消映射的内存范围
    return rangeFromHandles(begin, end);
  }

  // 返回虚拟地址空间的起始地址
  char* ptr() const {
    return (char*)ptr_;
  }

  // 返回虚拟地址空间的总大小
  size_t size() const {
    return max_handles_ * segment_size_;
  }
  
  // 添加一个新的 GPU 设备到 peers_ 列表中，并为其设置内存访问权限
  void addPeer(int device) {
    // 1. 将新的 GPU 设备 ID 添加到 peers_ 列表中
    peers_.push_back(device);
    // 2. 遍历所有已分配的内存段，为新设备设置访问权限
    forEachAllocatedRange(
        [&](size_t begin, size_t end) { setAccess(device, begin, end); });
  }

  // 析构函数，释放所有已分配的内存段并 cuMemAddressFree 释放虚拟地址空间
  ~ExpandableSegment() {
    // 1. 遍历所有已分配的内存段，取消映射并释放资源
    forEachAllocatedRange(
        [&](size_t begin, size_t end) { unmapHandles(begin, end); });
    // 2. 释放虚拟地址空间
    C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressFree_(
        ptr_, segment_size_ * max_handles_));
  }

 private:
  // cuMemSetAccess 设置指定设备对内存段的访问权限
  void setAccess(int device, size_t begin, size_t end) {
    // 1. 定义 CUmemAccessDesc 结构体，用于描述内存访问权限
    CUmemAccessDesc desc;
    desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // 内存位于 GPU 设备
    desc.location.id = device;                        // 指定 GPU 设备 ID
    desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;  // 设置读写权限
    // 2. 调用 cuMemSetAccess_ 设置内存访问权限
    C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemSetAccess_(
        ptr_ + begin * segment_size_, (end - begin) * segment_size_, &desc, 1));
  }

  // cuMemUnmap 取消映射指定范围内的内存段并 cuMemRelease 释放相关资源
  void unmapHandles(size_t begin, size_t end) {
    // note: unlike cudaFree, MemUnmap and MemRelease do
    // not appear to synchronize in all cases, so we have to wait for the
    // stream to finish before this memory is truly free.
    // 注意：与 cudaFree 不同，MemUnmap 和 MemRelease 在某些情况下似乎不会同步，
    // 因此我们必须等待流（stream）执行完成，才能确保内存真正被释放。

    // cannot call c10::cuda::stream_synchronize because
    // it might grab the GIL which can lead to a deadlock
    // Locking order must be GIL -> Allocator Lock
    // 不能调用 c10::cuda::stream_synchronize，因为
    // 它可能会获取 GIL（全局解释器锁），从而导致死锁。
    // 锁的顺序必须是：GIL -> 分配器锁（Allocator Lock）。

    // 1. 同步 CUDA 流，确保所有操作完成
    C10_CUDA_CHECK(cudaStreamSynchronize(stream_));
    // 2. 遍历 [begin, end) 范围内的内存段
    for (auto i : c10::irange(begin, end)) {
      // 3. 获取当前内存段的句柄
      CUmemGenericAllocationHandle h = handles_.at(i).value();
      // 4. 将当前内存段的句柄置为空
      handles_.at(i) = c10::nullopt;
      // 5. 调用 cuMemUnmap_ 取消映射虚拟地址
      C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemUnmap_(
          ptr_ + segment_size_ * i, segment_size_));
      // 6. 调用 cuMemRelease_ 释放物理内存
      C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h));
    }
    // 7. 调用 trimHandles 移除 handles_ 末尾的空句柄
    trimHandles();
  }

  // 移除 handles_ 向量末尾的空句柄
  void trimHandles() {
    while (!handles_.empty() && !handles_.back()) {
      handles_.pop_back();
    }
  }

  // 遍历所有已分配的内存段，并对每个连续的内存段调用 fn 函数
  void forEachAllocatedRange(std::function<void(size_t, size_t)> fn) {
    auto start = 0;// 初始化起始索
    // 1. 遍历 handles_ 向量
    for (auto i : c10::irange(handles_.size())) {
      // 2. 如果当前内存段已分配，且前一个内存段未分配，则更新起始索引
      if (handles_.at(i) && (i == 0 || !handles_.at(i - 1))) {
        start = i;
      }
      // 3. 如果当前内存段已分配，且后一个内存段未分配，则调用 fn 处理连续的内存段
      if (handles_.at(i) && (i + 1 == handles_.size() || !handles_.at(i + 1))) {
        fn(start, i + 1);
      }
    }
  }

  // 计算给定大小需要多少个内存段
  size_t numSegments(size_t size) {
    return (size + segment_size_ - 1) / segment_size_;
  }

  // 计算指针 p 所在的内存段的起始索引
  size_t segmentLeft(char* p) {
    auto size = p - ptr();
    return size / segment_size_;
  }

  // 计算指针 p 所在的内存段的结束索引
  size_t segmentRight(char* p) {
    auto size = p - ptr();
    return numSegments(size);
  }

  // 根据内存段的起始和结束索引返回对应的 SegmentRange
  SegmentRange rangeFromHandles(size_t begin, size_t end) {
    return SegmentRange(
        ptr() + segment_size_ * begin, segment_size_ * (end - begin));
  }

  int device_;
  cudaStream_t stream_;
  CUdeviceptr ptr_{};
  size_t max_handles_;
  size_t segment_size_;
  std::vector<c10::optional<CUmemGenericAllocationHandle>> handles_;
  // devices on which this memory should be mapped in addition
  // to the device where the physical memory lives (device_).
  // 除了物理内存所在的设备（device_）之外，
  // 这些内存还应映射到的其他设备。
  std::vector<int> peers_;
};

关于 ExpandableSegment，还可以参考PyTorch 内存管理之 expandable_segments 及 PyTorch显存管理介绍与源码解析（三）这两篇博客的分析。

参考资料

文献

CUDA VMM

TensorFlow

PyTorch

相关推荐
谢宁华为战略管理研发管理10 分钟前
《向华为学习：BEM战略解码》课程大纲
学习·华为
windyrain22 分钟前
AI 学习之路（一）- 重新认识 AI
人工智能·机器学习·aigc
IT、木易38 分钟前
大白话TypeScript 第十章TypeScript 学习全阶段详细总结
前端·学习·typescript
tt55555555555539 分钟前
嵌入式学习笔记-卡尔曼滤波,PID,MicroPython
笔记·学习·嵌入式
柃歌1 小时前
【UCB CS 61B SP24】Lecture 19 & 20: Hashing & Hashing II 学习笔记
java·数据结构·笔记·学习·算法
北京青翼科技1 小时前
【PCIE737】基于全高PCIe x8总线的KU115 FPGA高性能硬件加速卡
图像处理·人工智能·信号处理·智能硬件
桥Dopey1 小时前
MAC 本地搭建部署 dify（含 github访问超时+Docker镜像源拉取超时解决方案）
人工智能·docker·github·ai编程
摩尔线程1 小时前
SEKI —— 基于大型语言模型的自进化与知识启发式神经架构搜索
人工智能·语言模型·架构
道法自然，人法天1 小时前
微服务学习（4）：RabbitMQ中交换机的应用
运维·学习·docker·微服务·rabbitmq
breaksoftware1 小时前
51单片机编程学习笔记——74HC245八路三态输出双向收发器
笔记·学习·51单片机
热门推荐
01DeepSeek各版本说明与优缺点分析02DeepSeek R1本地化部署 Ollama + Chatbox 打造最强 AI 工具03如何在WPS和Word/Excel中直接使用DeepSeek功能04从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑05本地部署DeepSeek教程（Mac版本）06DeepSeek本地部署详细指南07Anaconda +Jupyter Notebook安装（2025最新版）08保姆级教程：利用Ollama与Open-WebUI本地部署 DeedSeek-R1大模型09DeepSeek r1本地安装全指南10【docker】Windows10安装Docker Desktop - WSL update failed