【知识】pytorch中的pinned memory和pageable memory

原创

小锋学长生活大爆炸

发布于 2024-08-17 14:13:46

940

发布于 2024-08-17 14:13:46

文章被收录于专栏：学习之旅

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn] 如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

概念简介

默认情况下，主机（CPU）数据分配是可分页的。GPU 无法直接从可分页主机内存访问数据，因此当调用从可分页主机内存到设备内存的数据传输时，CUDA 驱动程序必须首先分配一个临时的页面锁定或“固定”主机数组，将主机数据复制到固定数组，然后将数据从固定阵列传输到设备内存。

如图所示，固定内存用作从设备到主机的传输暂存区域。通过直接在固定内存中分配主机阵列，我们可以避免在可分页主机阵列和固定主机阵列之间传输的成本。使用 cudaMallocHost（）或 cudaHostAlloc（）在 CUDA C/C++ 中分配固定主机内存，并使用 cudaFreeHost（）解除分配。固定内存分配可能会失败，因此应始终检查错误。

数据传输速率可能取决于主机系统的类型（主板、CPU 和芯片组）以及 GPU。通过运行BandwidthTest会产生以下结果。可见，固定传输的速度是可分页传输的两倍多。（我的测试发现，基本上能跑满PCIe的带宽。）

Device: NVS 4200M
Transfer size (MB): 16

Pageable transfers
  Host to Device bandwidth (GB/s): 2.308439
  Device to Host bandwidth (GB/s): 2.316220

Pinned transfers
  Host to Device bandwidth (GB/s): 5.774224
  Device to Host bandwidth (GB/s): 5.958834

不过，不应过度分配固定内存。这样做会降低整体系统性能，因为它会减少操作系统和其他程序可用的物理内存量。多少是太多是很难提前判断出来的，因此与所有优化一样，测试你的应用程序及其运行的系统以获得最佳性能参数。

用法示例

由于pinned memory后，可以使用DMA传输而不占用CPU，因此通常需要搭配non_blocking使用。

# tensor.pin_memory() 就行
pinned_tensor = torch.randn(data_size, dtype=torch.float32).pin_memory()

device = torch.device("cuda")
pinned_tensor.to(device, non_blocking=True)

速度测试

import torch
import time
import torch.multiprocessing as mp

# 数据大小
data_size = 10**7  # 例如，10M数据

def test_pinned_memory(rank, normal_tensor, pinned_tensor, device):
    # 测试普通内存到GPU传输时间
    start_time = time.perf_counter()
    normal_tensor_gpu = normal_tensor.to(device, non_blocking=True)
    torch.cuda.synchronize()  # 等待数据传输完成
    normal_memory_time = time.perf_counter() - start_time
    print(f"[进程 {rank}] 普通内存到GPU传输时间: {normal_memory_time:.6f} 秒")

    # 测试固定内存到GPU传输时间
    start_time = time.perf_counter()
    pinned_tensor_gpu = pinned_tensor.to(device, non_blocking=True)
    torch.cuda.synchronize()  # 等待数据传输完成
    pinned_memory_time = time.perf_counter() - start_time
    print(f"[进程 {rank}] 固定内存到GPU传输时间: {pinned_memory_time:.6f} 秒")

    # 比较结果
    speedup = normal_memory_time / pinned_memory_time
    print(f"[进程 {rank}] 固定内存的传输速度是普通内存的 {speedup:.2f} 倍")

if __name__ == '__main__':
    # 分配普通内存中的张量
    normal_tensor = torch.randn(data_size, dtype=torch.float32)

    # 分配固定内存中的张量
    pinned_tensor = torch.randn(data_size, dtype=torch.float32).pin_memory()

    # 目标设备
    device = torch.device("cuda")

    # 使用mp.spawn启动多进程测试
    mp.spawn(test_pinned_memory, args=(normal_tensor, pinned_tensor, device), nprocs=2, join=True)

输出：

[进程 0] 普通内存到GPU传输时间: 1.054590 秒 [进程 0] 固定内存到GPU传输时间: 0.012945 秒 [进程 0] 固定内存的传输速度是普通内存的 81.47 倍 [进程 1] 普通内存到GPU传输时间: 1.169124 秒 [进程 1] 固定内存到GPU传输时间: 0.013019 秒 [进程 1] 固定内存的传输速度是普通内存的 89.80 倍

可以看到速度还是非常快的。

反直觉情况

我再瞎试的过程中发现，如果将pinned memory放在一个class中，那么多进程时候，pinned memory的移动很慢。暂不清楚为什么。

示例代码（反例，仅供观看，请勿使用）：

import torch
import torch.multiprocessing as mp
class PinnedMemoryManager:
    def __init__(self, data_size):
        self.data_size = data_size
        self.normal_tensor = None
        self.pinned_tensor = None

    def allocate_normal_memory(self):
        # 分配普通内存
        self.normal_tensor = torch.randn(self.data_size, dtype=torch.float32)

    def allocate_pinned_memory(self):
        # 分配固定内存
        self.pinned_tensor = torch.randn(self.data_size, dtype=torch.float32).pin_memory()

    def transfer_to_device(self, device, use_pinned_memory=False):
        # 选择使用普通内存或固定内存
        tensor = self.pinned_tensor if use_pinned_memory else self.normal_tensor
        if tensor is None:
            raise ValueError("Tensor not allocated. Call allocate_memory first.")

        # 数据传输
        start_time = torch.cuda.Event(enable_timing=True)
        end_time = torch.cuda.Event(enable_timing=True)

        start_time.record()
        tensor_gpu = tensor.to(device, non_blocking=True)
        end_time.record()

        # 同步并计算传输时间
        torch.cuda.synchronize()
        transfer_time = start_time.elapsed_time(end_time) / 1000.0  # 转换为秒
        return tensor_gpu, transfer_time

    def free_memory(self):
        # 释放内存
        del self.normal_tensor
        del self.pinned_tensor
        self.normal_tensor = None
        self.pinned_tensor = None


def test_pinned_memory(rank, manager, device):
    # 测试普通内存到GPU传输时间
    normal_gpu, normal_memory_time = manager.transfer_to_device(device, use_pinned_memory=False)
    print(f"[进程 {rank}] 普通内存到GPU传输时间: {normal_memory_time:.6f} 秒")

    # 测试固定内存到GPU传输时间
    pinned_gpu, pinned_memory_time = manager.transfer_to_device(device, use_pinned_memory=True)
    print(f"[进程 {rank}] 固定内存到GPU传输时间: {pinned_memory_time:.6f} 秒")

    # 比较结果
    speedup = normal_memory_time / pinned_memory_time
    print(f"[进程 {rank}] 固定内存的传输速度是普通内存的 {speedup:.2f} 倍")


if __name__ == '__main__':
    # 数据大小
    data_size = 10**7  # 例如，10M数据

    # 初始化固定内存管理器
    manager = PinnedMemoryManager(data_size)
    manager.allocate_normal_memory()
    manager.allocate_pinned_memory()

    # 目标设备
    device = torch.device("cuda")

    # 使用mp.spawn启动多进程测试
    mp.spawn(test_pinned_memory, args=(manager, device), nprocs=2, join=True)

    # 释放内存
    manager.free_memory()

输出：

[进程 1] 普通内存到GPU传输时间: 0.013695 秒 [进程 1] 固定内存到GPU传输时间: 0.013505 秒 [进程 1] 固定内存的传输速度是普通内存的 1.01 倍 [进程 0] 普通内存到GPU传输时间: 0.013752 秒 [进程 0] 固定内存到GPU传输时间: 0.013593 秒 [进程 0] 固定内存的传输速度是普通内存的 1.01 倍

可以看到基本上没啥改进。

暂不清楚原因，只能先无脑避免这种用法了。

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

pytorch

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

pytorch

登录后参与评论

0 条评论

热度