利用CuPy在多GPU上统一分配大容量内存教程

import cupy as cp
import numpy as np

# 配置CuPy使用统一内存池
pool = cp.cuda.MemoryPool(cp.cuda.malloc_managed)
cp.cuda.set_allocator(pool.malloc)

# 设定期望分配的内存大小（例如，大于32 GB）
desired_memory_gb = 42
element_size_bytes = np.dtype(np.float64).itemsize
desired_memory_bytes = desired_memory_gb * (1024**3)
num_elements = desired_memory_bytes // element_size_bytes

# 尝试创建一个大型数组
try:
    array = cp.full(num_elements, 1.1, dtype=np.float64)
    print(f"Array allocated on unified memory with shape {array.shape}")
except cp.cuda.memory.OutOfMemoryError as e:
    print(f"Failed to allocate array: {e}")
    print("This might indicate that even with unified memory, a single GPU + CPU combination is insufficient or not optimally utilized.")

# 注意：此方法可能导致内存集中在一个GPU上，剩余部分溢出到CPU。
# 并没有强制其在所有可用GPU上进行分配。

import cupy as cp
import numpy as np

# 设定期望分配的总内存大小（例如，大于32 GB）
desired_total_memory_gb = 42
element_size_bytes = np.dtype(np.float64).itemsize
desired_total_memory_bytes = desired_total_memory_gb * (1024**3)

# 获取所有可用的GPU设备列表
available_gpus = cp.cuda.Device.available_gpus()
num_gpus = len(available_gpus)

if num_gpus == 0:
    print("No NVIDIA GPUs found. Please ensure CUDA is installed and GPUs are available.")
    # 可以选择在此处回退到CPU NumPy数组
    # arrays = [np.full(num_elements_per_gpu, 1.1, dtype=np.float64)]
else:
    # 计算每个GPU需要分配的元素数量
    # 简单起见，这里假设均匀分配，实际应用中可能需要更复杂的策略
    num_elements_per_gpu = desired_total_memory_bytes // element_size_bytes
    # 如果希望总数不变，而每个GPU分配一部分，则需要调整 num_elements_per_gpu
    # num_elements_per_gpu = (desired_total_memory_bytes // element_size_bytes) // num_gpus

    arrays = []
    print(f"Attempting to allocate {desired_total_memory_gb} GB across {num_gpus} GPUs.")

    for i, device_id in enumerate(available_gpus):
        with cp.cuda.Device(device_id): # 激活当前GPU设备
            # 确保每个GPU使用统一内存池，或者直接在当前设备上分配
            # 如果希望使用统一内存且跨GPU，可以在此设置全局分配器
            # 但更直接的方式是让CuPy根据当前设备上下文在GPU内存中分配
            # cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.malloc_managed).malloc) # 可选，如果希望每个GPU都使用其自己的统一内存池

            try:
                # 在当前激活的GPU上创建数组
                # 注意：cp.full 默认在当前活动的GPU上分配内存
                # 如果希望明确指定设备，可以使用 device=device_id 参数，但with cp.cuda.Device(device_id)已足够
                arr = cp.full(num_elements_per_gpu, 1.1, dtype=np.float64)
                arrays.append(arr)
                print(f"Array allocated on GPU-{device_id} with shape {arr.shape} and {arr.nbytes / (1024**3):.2f} GB.")
            except cp.cuda.memory.OutOfMemoryError as e:
                print(f"Error: Out of memory on GPU-{device_id}. Could not allocate {num_elements_per_gpu} elements. {e}")
                # 可以在这里添加错误处理逻辑，例如减少分配量或跳过此GPU

    if arrays:
        print(f"\nSuccessfully allocated parts of the array across {len(arrays)} GPUs.")
        # 示例：访问第一个GPU上的数组
        # print(arrays[0][:5])
    else:
        print("No arrays were successfully allocated on any GPU.")