如何让CuPy在多GPU集群上分配跨GPU的统一内存？

AI 小北

2026-7-3

问题

我用Python和CuPy访问NVIDIA V100多GPU集群（共4张GPU，单卡32GB），需要分配无法放入单张GPU的大型数组，因此尝试通过CuPy的ManagedMemory使用统一内存，期望在所有可用GPU上分配内存。

我的代码：

import cupy as cp
import numpy as np

# allocate unified memory
pool = cp.cuda.MemoryPool(cp.cuda.malloc_managed)
cp.cuda.set_allocator(pool.malloc)

# Desired memory in GB
desired_memory_gb = 42 # SET this value to greater than 32 GB

# Calculate the number of elements required to achieve desired memory
element_size_bytes = np.dtype(np.float64).itemsize
desired_memory_bytes = desired_memory_gb * (1024**3)  # Convert GB to bytes
num_elements = desired_memory_bytes // element_size_bytes

# Create the array with the calculated number of elements
array = cp.full(num_elements, 1.1, dtype=np.float64)

# GPU
print("Array allocated on unified memory...")

遇到的问题：
上述代码虽能分配超过32GB的统一内存数组，但统一内存并未分配到所有GPU，仅单张GPU占用32GB，剩余部分分配到CPU，而非GPU-2、3、4。如何强制程序使用所有GPU而非CPU来分配统一内存？

解决方案

CuPy默认的统一内存分配逻辑不会自动跨GPU分布内存，要让统一内存优先使用GPU显存而非CPU内存，可通过以下方式调整：

设置CUDA环境变量强制设备优先分配
在启动Python脚本前，设置CUDA_MANAGED_FORCE_DEVICE_ALLOC环境变量为1，让统一内存优先分配到所有GPU的显存池中，仅当总显存耗尽时才会使用CPU内存：

export CUDA_MANAGED_FORCE_DEVICE_ALLOC=1

也可以在Python代码开头通过os模块直接设置：

import os
os.environ["CUDA_MANAGED_FORCE_DEVICE_ALLOC"] = "1"

显式触发布局到所有GPU
统一内存的页迁移是按需触发的，默认只会在当前活跃GPU上分配内存页。你需要遍历所有GPU并执行简单操作，触发内存页迁移到对应GPU：

# 创建数组后，遍历所有GPU触发内存分布
for dev_id in range(cp.cuda.runtime.getDeviceCount()):
    with cp.cuda.Device(dev_id):
        # 执行轻量操作触发内存页迁移到当前GPU
        cp.sum(array)

这样每个GPU都会加载一部分数组的内存页，避免剩余内存落到CPU上。

验证内存分配情况
可以通过nvidia-smi命令查看各GPU显存占用，或者在代码中打印每个GPU的内存使用详情：

for dev_id in range(cp.cuda.runtime.getDeviceCount()):
    with cp.cuda.Device(dev_id):
        mem_info = cp.cuda.memory_info()
        print(f"GPU {dev_id}: 已用 {mem_info.used/(1024**3):.2f} GB, 剩余 {mem_info.free/(1024**3):.2f} GB")

注：NVIDIA V100完全支持CUDA统一内存的设备优先分配特性，设置环境变量后，CUDA会尝试将所有统一内存分配到所有可用GPU的显存总和中，仅当总显存不足时才会溢出到CPU内存。配合跨GPU的显式操作，就能实现内存在所有GPU上的分布。

内容的提问来源于stack exchange，提问作者skm