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

霞舞

发布时间：2025-11-23 13:30:19

994人浏览过

来源于php中文网

原创

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

本教程旨在指导用户如何使用CuPy库在多GPU集群中高效地分配和管理超出单个GPU容量的大型数组。文章将详细解释CuPy统一内存（Managed Memory）的基本概念，并针对其在多GPU环境下可能遇到的分配不均问题，提供一种明确的解决方案，即通过遍历可用GPU并为每个设备独立创建数组，确保内存资源得到充分利用，避免数据溢出到CPU。

引言：多GPU环境下的内存分配挑战

在高性能计算领域，尤其是在处理大规模数据集时，单个GPU的内存容量往往成为瓶颈。NVIDIA的统一内存（Unified Memory）技术旨在简化GPU编程，允许CPU和GPU访问同一块内存空间，由系统自动管理数据迁移。CuPy作为Python中用于GPU计算的库，也提供了对统一内存的支持。然而，当我们需要分配一个远超单个GPU容量的数组，并希望它能跨多个GPU进行分配时，直接使用CuPy的统一内存分配器可能会遇到一个常见问题：内存可能只分配在一个GPU上，剩余部分则溢出到CPU，而非均匀分布到集群中的其他GPU。

CuPy统一内存的初步尝试与局限性

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上进行分配。

尽管上述代码成功地使用了统一内存并可能分配了超过单个GPU容量的数组，但它并未强制内存均匀分布到所有可用的GPU上。实际上，CuPy的默认行为可能是在当前活动的GPU上尝试分配，如果不足，则由统一内存机制将其余部分映射到主机（CPU）内存。这与我们希望将数据分布到多个GPU以充分利用集群资源的目标不符。

解决方案：显式地在每个GPU上分配内存

要强制程序在所有可用的GPU上分配统一内存（或至少是GPU内存），我们需要显式地遍历每个GPU设备，并在每个设备上创建相应的数组。CuPy提供了方便的接口来识别和切换GPU设备。

微厦在线学习培训系统2.0

《微厦在线学习考试系统》将“ 视频学习、试题练习、在线考试 ”紧密相联，打造成为集 “学、练、考” 于一体的在线学习系统。“点播/直播”、“刷题/测试”、“组卷/考试”，根据学习内容的不同权重汇总综合成绩，生成学习证明。支持在线支付（微信支付、支付宝支付）；利用充值卡、学习卡配合线下营销；Web端、APP、小程序，多终端方便学习。适用：大学、中职中专、培训机构、企事业单位行业：企业内训、成人教育、

下载

以下是实现这一目标的优化代码：

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.")

代码解释：

cp.cuda.Device.available_gpus(): 这个函数返回一个列表，其中包含所有可用的GPU设备的ID。这使得程序能够动态地适应不同配置的GPU集群。
with cp.cuda.Device(device_id):: 这是一个上下文管理器，用于临时切换到指定的GPU设备。在其作用域内执行的所有CuPy操作（如 cp.full）都将在该设备上进行。
cp.full(num_elements_per_gpu, 1.1, dtype=np.float64): 在当前激活的GPU上创建一个新的CuPy数组。通过将总元素数量除以GPU数量（或根据具体策略计算），确保每个GPU分配一部分数据。
arrays.append(arr): 将在每个GPU上创建的数组存储在一个列表中，方便后续管理和访问。

注意事项与最佳实践

数据访问模式： 这种在多个GPU上分配独立数组的方法，意味着数据是分片的。在进行计算时，需要确保计算逻辑能够处理这种分片数据。例如，可以使用CuPy的cupy.concatenate将所有分片合并（如果总内存允许），或者编写并行计算逻辑，让每个GPU处理其本地数据。
内存池管理： 虽然上述解决方案没有显式地为每个GPU设置统一内存池，但cp.full在with cp.cuda.Device(device_id):上下文内会尝试在当前设备的GPU内存中分配。如果需要更精细的统一内存管理（例如，希望主机和设备都能访问这些分片），可以在每个with cp.cuda.Device(device_id):块内调用cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.malloc_managed).malloc)来确保该设备上的后续分配使用统一内存。然而，对于简单的多GPU分配，显式切换设备通常就足够了。
负载均衡： 示例代码简单地将总元素数量平均分配给每个GPU。在实际应用中，如果GPU的型号或内存容量不同，可能需要更复杂的负载均衡策略。
错误处理： 在分配过程中，如果某个GPU内存不足，cp.cuda.memory.OutOfMemoryError异常会被抛出。在生产环境中，应捕获此异常并进行适当处理，例如跳过该GPU或调整分配策略。
分布式计算框架： 对于更复杂的分布式计算任务，例如需要跨节点通信或动态任务调度，可以考虑使用更高级的框架，如Dask-CUDA。这些框架提供了更强大的抽象和工具来管理多GPU和多节点资源。

总结

通过显式地遍历并激活每个GPU设备，我们能够有效地将大型数组的内存分配任务分散到整个GPU集群中，从而克服单个GPU内存容量的限制，并避免不必要的内存溢出到CPU。这种方法为在CuPy中进行大规模GPU计算提供了坚实的基础，使得开发者能够充分利用多GPU集群的计算能力。在实际应用中，结合对数据访问模式的理解和适当的错误处理，可以构建出高效且鲁棒的GPU加速应用程序。

Python 单元测试设计与 pytest 实战

Python itertools 常用迭代组合技巧

Python 模块拆分与依赖控制技巧

Python 对象 ID 与内存地址的对应关系

Python 可调用对象的类型判断