CuPy 内存管理完全指南内存池、显存限制与流有序分配【免费下载链接】cupyNumPy SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupyCuPy 默认采用**内存池memory pool**机制进行 GPU 显存与固定内存pinned memory的分配与复用以显著降低内存分配开销和 CPU/GPU 同步带来的性能损耗。本文以官方用户指南 docs/source/user_guide/memory.rst 为骨架结合仓库源码cupy/cuda/memory.pyx、cupy/init.py、cupyx/_pinned_array.py深入剖析两大内存池的工作原理、统计接口、显存硬上限、自定义分配器、流有序内存分配器Stream Ordered Memory Allocator以及实验性的统一内存编程UMP支持。读完本文你将掌握如何监控、限制、替换乃至彻底禁用 CuPy 的内存池并能针对多 GPU、多进程共存等场景做出正确的内存策略选型。一、为什么需要内存池默认分配策略与性能收益CuPy 默认使用内存池来完成所有内存分配。内存池之所以能显著提升性能核心在于两点减少分配开销cudaMalloc这类底层分配调用本身成本高昂内存池通过缓存已释放的内存块使后续同尺寸分配直接命中复用绕过底层驱动调用。避免 CPU/GPU 同步从 MemoryPool 类定义 的源码注释可以看到当分配命中池中预分配的空闲块时不会调用cudaMalloc因此不会触发 CPU 与 GPU 之间的同步。这使得「内存分配 内核调用」的交错执行变得非常快。两大内存池CuPy 中存在两个相互独立的内存池分别服务不同的内存类型内存池管理的内存用途设备内存池Device Memory PoolGPU 显存所有 GPU 内存分配如cupy.ndarray的数据存储固定内存池Pinned Memory Pool不可交换的 CPU 内存页锁定内存CPU → GPU 数据传输过程中的临时缓冲固定内存pinned/pagelocked memory不会被操作系统换出到磁盘DMA 引擎可以直接访问因此是主机端到设备端数据传输的关键加速手段但其分配成本也更高所以同样值得池化复用。监控内存时的「内存不释放」是预期行为在使用nvidia-smi监控 GPU 显存或ps监控 CPU 内存观察内存占用时你可能会发现即使数组实例已经超出作用域内存依然没有被释放。这是完全正常的预期行为——默认内存池会把分配过的内存块「缓存」下来以便复用。池中的空闲块会在后续分配中再次被利用而不是立即归还给操作系统或 CUDA 驱动。这一点也意味着内存池会尽可能长时间地持有设备内存从 MemoryPool 源码注释 可以看出这可能导致与其他并行运行的 CUDA 程序产生显存竞争。二、内存池操作统计查询与块释放通过cupy.get_default_memory_pool与cupy.get_default_pinned_memory_pool可以拿到默认的两个内存池实例。这两个函数在 cupy/init.py 中实现而默认池本身在模块加载时就已创建并注册为 CuPy 的默认分配器# cupy/__init__.py简化 _default_memory_pool cuda.MemoryPool() _default_pinned_memory_pool cuda.PinnedMemoryPool() cuda.set_allocator(_default_memory_pool.malloc) cuda.set_pinned_memory_allocator(_default_pinned_memory_pool.malloc)MemoryPool和PinnedMemoryPool暴露了若干统计与操作接口核心方法定义在 cupy/cuda/memory.pyxused_bytes()当前被使用未释放的字节数free_bytes()已从底层分配但当前空闲的字节数total_bytes()池累计从底层分配的总字节数used_bytes free_bytesn_free_blocks()池中空闲内存块的数量free_all_blocks()释放池中所有空闲块归还给底层分配器set_limit(size..., fraction...)/get_limit()设置/查询当前设备的显存分配上限。完整示例观察内存池的完整生命周期以下示例完整展示了内存池的统计行为摘自官方指南并补充注释import cupy import numpy mempool cupy.get_default_memory_pool() pinned_mempool cupy.get_default_pinned_memory_pool() # 在 CPU 上创建一个数组。 # NumPy 在 CPU 上分配 400 字节不归 CuPy 内存池管理。 a_cpu numpy.ndarray(100, dtypenumpy.float32) print(a_cpu.nbytes) # 400 # 此时两个池都还没有任何分配。 print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 0 print(pinned_mempool.n_free_blocks()) # 0 # 将数组从 CPU 传输到 GPU。 # 这会从设备内存池分配 400 字节并从固定内存池分配另外 400 字节。 # 传输完成后固定内存会被立即释放回固定内存池。 # 注意出于性能考虑实际分配大小会被向上取整round up。 a cupy.array(a_cpu) print(a.nbytes) # 400 print(mempool.used_bytes()) # 512 print(mempool.total_bytes()) # 512 print(pinned_mempool.n_free_blocks()) # 1 # 当数组超出作用域时设备内存被释放并保留在池中供后续复用。 a None # 或使用 del a print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 512 print(pinned_mempool.n_free_blocks()) # 1 # 调用 free_all_blocks 清空内存池。 mempool.free_all_blocks() pinned_mempool.free_all_blocks() print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 0 print(pinned_mempool.n_free_blocks()) # 0这个示例中值得注意的细节分配大小取整请求 400 字节设备池实际分配了 512 字节这是 CuPy 底层内存分配对齐策略导致的向上取整行为固定内存的释放时机cupy.array(a_cpu)触发的 CPU→GPU 传输会临时使用固定内存作为中转传输完成后该固定内存块即被释放回池中体现为n_free_blocks() 1池缓存 vs 真正释放a None后used_bytes()归零但total_bytes()仍为 512说明内存只是归还给池子并未真正还给系统只有free_all_blocks()才会把内存归还给底层分配器。关于free_all_blocks还有一个细节从 源码注释 可以看出内存池为了空间利用率可能会拆分空闲块被拆分的块即使调用free_all_blocks也不会立即释放必须等到所有拆分部分重新合并为一个完整块后才会归还。此外free_all_blocks(streamNone)支持传入流参数以仅释放指定流 arena 中的块。三、固定内存高层 APIcupyx.empty_pinned 系列为了更方便地使用固定内存CuPy 在cupyx命名空间提供了四个高层工具函数它们在 cupyx/_pinned_array.py 中实现并在 cupyx/init.py 中导出cupyx.empty_pinned(shape, dtypefloat, orderC)cupyx.empty_like_pinned(a, dtypeNone, orderK, subokNone, shapeNone)cupyx.zeros_pinned(shape, dtypefloat, orderC)cupyx.zeros_like_pinned(a, dtypeNone, orderK, subokNone, shapeNone)这些函数返回的是由固定内存页锁定内存支撑的 NumPy 数组。从 empty_pinned 的实现 可以看到其本质先通过cuda.alloc_pinned_memory(nbytes)从 CuPy 的固定内存分配器默认即固定内存池取一块内存再用numpy.ndarray(..., buffermem)把这块内存包装成 NumPy 数组。因此只要 CuPy 的固定内存池处于启用状态这些 API 分配的固定内存就来自该池从而享受池化复用的性能收益。典型用法示例import cupy import cupyx # 直接创建固定内存 NumPy 数组 pinned_buf cupyx.empty_pinned((1024, 1024), dtypecupy.float32) pinned_zeros cupyx.zeros_pinned(4096) # 按现有数组的形状/布局创建 a_gpu cupy.random.random((64, 64)) pinned_like cupyx.empty_like_pinned(a_gpu) # 固定内存是 CPU→GPU 高效传输的天然载体 a_cpu cupyx.zeros_pinned(100) # 固定内存缓冲 gpu_arr cupy.asarray(a_cpu) # 由固定内存出发的传输更快empty_like_pinned和zeros_like_pinned支持orderK尽可能保持与a相同的内存布局其中subok参数目前尚未支持传入非None值会抛出TypeError见 源码第 75-76 行。四、限制 GPU 内存使用CUPY_GPU_MEMORY_LIMIT 与 set_limit在某些场景下如显存较小的 GPU、需要与其他进程共享 GPU你可能希望给 CuPy 设置一个显存分配的硬上限。方式一环境变量 CUPY_GPU_MEMORY_LIMIT设置CUPY_GPU_MEMORY_LIMIT环境变量即可为每个 GPU 设备设定可分配显存的硬上限默认值为0即不限制。该变量的完整说明见 docs/source/reference/environment.rst。它支持两种取值格式# 方式 A绝对字节数。设置硬上限为 1 GiB export CUPY_GPU_MEMORY_LIMIT1073741824 # 方式 B占 GPU 总显存的百分比。若 GPU 显存为 2 GiB # 则下面的配置与上面的绝对字节数配置等价。 export CUPY_GPU_MEMORY_LIMIT50%在 Python 中可以通过get_limit()验证限制是否生效import cupy print(cupy.get_default_memory_pool().get_limit()) # 1073741824方式二MemoryPool.set_limit按设备差异化限制环境变量为所有设备设置统一的默认值而cupy.cuda.MemoryPool.set_limit则可以为当前设备单独设置或覆盖环境变量指定的限制从而实现对每个 GPU 使用不同上限import cupy mempool cupy.get_default_memory_pool() with cupy.cuda.Device(0): mempool.set_limit(size1024**3) # 1 GiB with cupy.cuda.Device(1): mempool.set_limit(size2*1024**3) # 2 GiB从 set_limit 的源码实现 可以归纳出以下关键语义size与fraction不能同时指定两者都未指定或指定为0时表示取消限制fraction取值应在[0, 1]区间表示占「当前设备可用显存」的比例通过方法设置的限制优先级高于环境变量该方法只影响当前设备而环境变量为所有设备设置默认限制限制的修改不是线程安全的——在分配正在进行时修改限制其他线程可能读取到过期的值生产代码中需要注意。注意事项限制不覆盖池外内存官方指南明确提醒CUDA 还会在内存池之外分配一部分 GPU 内存如 CUDA context、cuBLAS/cuDNN 等库的句柄与工作区。这部分内存根据使用情况可能占用几十到几百 MiB且不计入上述限制。因此实际显存占用峰值会略高于你所设置的限制值。另外CUPY_GPU_MEMORY_LIMIT也会被 MemoryAsyncPool 在初始化时读取_parse_limit_string()即流有序内存池同样支持通过该环境变量设限。五、更换内存池与自定义分配器除了默认的内存池CuPy 允许你传入自定义的内存分配函数来替换默认分配策略相关接口为cupy.cuda.set_allocator(func)设置设备内存分配器cupy.cuda.set_pinned_memory_allocator(func)设置固定内存分配器。分配器函数需要满足统一签名接收一个参数请求的字节数返回cupy.cuda.MemoryPointer/cupy.cuda.PinnedMemoryPointer。CuPy 为此提供了几种开箱即用的分配器。5.1 托管内存Managed Memorymalloc_managed托管内存Unified Memory允许 GPU 显存超量订阅oversubscription适用于具备cudaDevAttrConcurrentManagedAccess属性通常为 Pascal 及更新的架构的 GPU。对应分配器为cupy.cuda.malloc_managed定义于 cupy/cuda/memory.pyx。推荐的做法是用它构造一个内存池而不是直接传入import cupy # 使用托管内存支撑的内存池 cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.malloc_managed).malloc)需要注意如果不构造MemoryPool直接把malloc_managed传给set_allocator那么内存释放时会立即归还给系统失去池化复用带来的性能优势——这种直接传递的方式在某些场景下可能是有意为之但要明确其代价。5.2 流有序内存分配器Stream Ordered Memory AllocatorMemoryAsyncPoolCUDA 11.2 起引入了流有序内存分配器Stream Ordered Memory AllocatorCuPy 通过cupy.cuda.MemoryAsyncPool提供了实验性接口。与 CuPy 自研内存池类似它也是以「流有序」的方式异步地从一个内存池中分配/释放内存关键区别在于它由NVIDIA CUDA 驱动内置实现因此同一进程内的其他 CUDA 应用程序也可以轻松地从同一个池中分配内存天然支持跨库共享。启用方式import cupy # 使用异步流有序内存 cupy.cuda.set_allocator(cupy.cuda.MemoryAsyncPool().malloc) # 创建自定义流 s cupy.cuda.Stream() # 在流 s 上异步分配内存 with s: a cupy.empty((100,), dtypecupy.float64)MemoryAsyncPool的构造参数pool_handles与MemoryPool不同用于指定使用哪个 CUDA mempooldefault表示设备默认 mempoolcurrent表示当前 mempool默认值也可以传入外部创建的cudaMemPool_t整数句柄还可以传入与可见设备数量等长的列表为每个设备独立指定详见 MemoryAsyncPool 定义。如果直接把malloc_async传给set_allocator而不构造MemoryAsyncPool则会使用设备当前的 mempool。使用流有序内存时有几点必须注意流语义由用户自己维护CuPy 不会自作聪明地替你管理流。你需要自行使用cupy.cuda.Stream与cupy.cuda.Event保证正确的流同步相关 API 见 docs/source/reference/cuda.rst 的流与事件章节。释放内存时会优先在分配它的流上异步释放其次在当前 CuPy 流上释放分配内存的流在其上所有内存被释放前被销毁也是允许的。与 cudaMalloc 的互操作问题内部使用cudaMallocCUDA 默认的同步分配器的应用与流有序分配器可能产生意外交互——归还到池中的内存对cudaMalloc而言可能不是立即可见的从而引发意外的 OOM内存不足错误。此时可以调用MemoryAsyncPool.free_all_blocks()或手动做一次事件/流/设备同步后重试。实验性状态与驱动要求MemoryAsyncPool的 API 与MemoryPool大体一致但部分方法需要足够新的驱动以及受支持的硬件、CUDA 版本与平台。从 源码 可以看到统计接口的可用性取决于驱动版本driverGetVersion() 11030。另外该池目前无法与内存钩子memory hooks配合使用且底层池初始大小为 0首次分配会按 32 MiB 的整数倍扩容该内部行为可能随驱动版本变化API 本身不依赖这些内部细节。5.3 彻底禁用内存池如果出于调试或与外部内存管理机制协同的目的你可以完全禁用默认内存池。必须在任何其他 CuPy 操作之前执行import cupy # 禁用设备内存GPU的内存池 cupy.cuda.set_allocator(None) # 禁用固定内存CPU的内存池 cupy.cuda.set_pinned_memory_allocator(None)禁用后每次分配都会直接调用底层cudaMalloc/cudaHostAlloc释放时立即归还给系统这通常会导致明显的性能下降仅建议在特殊场景下使用。六、统一内存编程UMP支持实验性对于启用了异构内存管理HMMHeterogeneous Memory Management或地址转换服务ATSAddress Translation Services的系统——例如 NVIDIA Grace Hopper 超级芯片——CuPy 提供了一种实验性的**统一内存编程Unified Memory ProgrammingUMP**能力让 NumPy 与 CuPy共享同一份系统内存CPU 与 GPU 之间不再有显式拷贝。启用步骤启用 UMP 需要完成以下 4 个步骤第 1 步安装numpy_allocator包一个提供 NumPy 对齐分配器的第三方扩展包。第 2 步设置环境变量CUPY_ENABLE_UMP1。第 3 步让 CuPy 从系统内存malloc_system分配。cupy.cuda.memory.malloc_system定义于 cupy/cuda/memory.pyx用于在 HMM/ATS 系统上分配可被 CPU 和 GPU 共同访问的系统内存import cupy as cp cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.memory.malloc_system).malloc)第 4 步让 NumPy 切换到对齐分配器使其也从系统内存分配此处直接复用 CuPy 编译产物中导出的对齐分配符号import cupy._core.numpy_allocator as ac import numpy_allocator import ctypes lib ctypes.CDLL(ac.__file__) class my_allocator(metaclassnumpy_allocator.type): _calloc_ ctypes.addressof(lib._calloc) _malloc_ ctypes.addressof(lib._malloc) _realloc_ ctypes.addressof(lib._realloc) _free_ ctypes.addressof(lib._free) my_allocator.__enter__() # 全局切换分配器启用后的效果完成上述配置后所有数据移动 API——如cupy.ndarray.get()、cupy.asnumpy、cupy.asarray——都变成no-op零拷贝因为 CPU 和 GPU 访问的是同一份物理内存。官方指南给出的加速示例# a, b, c 是 np.ndarrayd 是 cp.ndarray a np.random.random(100) b np.random.random(100) c np.add(a, b) # CPU 上执行 d cp.matmul(cp.asarray(a), cp.asarray(c)) # GPU 上执行但无数据拷贝本质上在这种模式下CPU/GPU「内存空间」的区分消失了np与cp退化为纯粹表征「执行空间」代码跑在 CPU 还是 GPU 上的标签。除了 NumPy/CuPy 的这组配置外用户的业务代码完全不需要任何改动。需要强调的是UMP 属于实验性特性依赖 HMM/ATS 硬件与系统支持如 Grace Hopper 平台在普通 x86 独立 GPU 的常见环境中并不适用。七、补充FFT 计划缓存与设备内存官方指南还特别提醒CuPy v8 及以上CuPy 提供了 FFT 计划缓存FFT plan cache当使用 FFT 及相关函数时该缓存可能占用一部分设备内存。如果希望回收这部分内存可以通过缩小或禁用缓存来实现。对于同时使用 FFT 与内存池监控的场景排查显存占用时不要忽略这一来源。八、总结与参考CuPy 的内存管理设计围绕「池化复用」这一核心展开设备内存池与固定内存池各司其职通过get_default_memory_pool/get_default_pinned_memory_pool暴露统计接口通过CUPY_GPU_MEMORY_LIMIT与set_limit控制显存上限通过set_allocator/set_pinned_memory_allocator支持托管内存、流有序内存乃至 UMP 系统内存等多种分配后端。理解这些机制可以帮助你在多进程共存、显存受限、跨库共享内存等真实场景下做出正确的配置决策。官方内存管理 API 细节docs/source/reference/cuda.rst环境变量完整清单docs/source/reference/environment.rst内存池与分配器核心实现cupy/cuda/memory.pyx默认池初始化与访问入口cupy/init.py固定内存高层 API 实现cupyx/_pinned_array.py【免费下载链接】cupyNumPy SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考