备注
点击 此处 下载完整示例代码
关于 PyTorch 中 non_blocking
和 pin_memory()
的良好使用指南¶
Created On: Jul 31, 2024 | Last Updated: Mar 18, 2025 | Last Verified: Nov 05, 2024
介绍¶
从 CPU 到 GPU 转移数据是在许多 PyTorch 应用中都会用到的一项基础操作。用户需要理解在设备间移动数据的最有效工具和选项。本教程将探讨 PyTorch 中设备间数据转移的两种关键方法:pin_memory()
和 to()
配合 non_blocking=True
选项。
您将学习的内容¶
通过异步传输和内存固定可以优化从 CPU 到 GPU 的张量传输。然而,需关注以下几点重要考虑事项:
使用
tensor.pin_memory().to(device, non_blocking=True)
可能比直接使用tensor.to(device)
慢至两倍。通常来说,
tensor.to(device, non_blocking=True)
是提高传输速度的有效选择。尽管
cpu_tensor.to("cuda", non_blocking=True).mean()
会正确执行,但尝试cuda_tensor.to("cpu", non_blocking=True).mean()
会导致错误的结果。
序言¶
本教程中报告的性能表现取决于用于构建教程的系统。尽管结论适用范围广泛,但具体观察结果可能会因硬件而略有差异,特别是在旧硬件上。本教程的主要目标是提供一种理论框架,用于理解 CPU 到 GPU 的数据传输。然而,任何设计决策都应根据具体情况进行调整,并以基准吞吐量测量以及任务需求为指导。
import torch
assert torch.cuda.is_available(), "A cuda device is required to run this tutorial"
本教程需要安装 tensordict。如果您的环境中还没有 tensordict,请在单独的代码单元中运行以下命令安装它:
# Install tensordict with the following command
!pip3 install tensordict
我们首先概述与这些概念相关的理论内容,然后转向这些功能的具体测试示例。
背景¶
内存管理基础¶
当用户在 PyTorch 中创建一个 CPU 张量时,这个张量的内容需要被放置在内存中。我们这里所讨论的内存概念相对复杂,值得仔细研究。我们区分由内存管理单元处理的两种类型的内存:RAM(为了简化说明)以及磁盘上的交换空间(可能是硬盘,也可能不是)。可用磁盘和 RAM(即物理内存)上的总空间汇总构成了虚拟内存,这是总资源一种抽象表示。简而言之,虚拟内存让可用空间看起来比单独的 RAM 空间要大,从而制造了主内存实际上大于其真实大小的假象。
正常情况下,一个普通的 CPU 张量是分页的,这意味着它被划分为称为页面的块,这些页面可以存在于虚拟内存中的任何地方(包括 RAM 或磁盘)。正如前面所提到的,这样的好处是内存看起来比主内存实际大小更大。
通常情况下,当程序访问某个不在 RAM 中的页面时,会发生“页面错误”,然后操作系统(OS)会将该页面带回 RAM(“换入”或“页面载入”)。反过来操作系统可能需要交换出(或“页面卸载”)另一个页面以给新页面腾出空间。
与分页内存对比,固定(或页面锁定或非分页)内存是一种无法交换到磁盘上的内存类型。它允许更快且时间更可预测的访问,但缺点是它比分页内存(即主内存)更加有限。

CUDA 和(非)分页内存¶
为了理解 CUDA 如何从 CPU 到 CUDA 复制一个张量,让我们考虑上述两种情况:
如果内存是页面锁定的,设备可以直接从主内存访问数据。内存地址是定义良好的,因此需要读取这些数据的函数可以显著加速。
如果内存是分页的,则所有页面需要首先被载入主内存,才能被发送到 GPU。这种操作可能耗时,并且比在页面锁定张量上执行时不可预测。
更确切地说,当 CUDA 从 CPU 向 GPU 发送分页数据时,它首先需要创建该数据的页面锁定副本,然后再进行传输。
使用 non_blocking=True
的异步与同步操作(CUDA cudaMemcpyAsync
)¶
当从主机(例如 CPU)复制到设备(例如 GPU)时,CUDA 工具包提供了同步或异步操作主机执行的方式。
实际上,当调用 to()
时,PyTorch 总是会调用 cudaMemcpyAsync。如果 non_blocking=False``(默认),会在每次 ``cudaMemcpyAsync
后调用 cudaStreamSynchronize
,从而使对 to()
的调用在主线程上阻塞。如果 non_blocking=True
,不会触发同步,主机上的主线程不会阻塞。因此,从主机的角度来看,多个张量可以同时发送到设备,因为线程不需要等待一次传输完成后才能发起另一传输。
备注
总体而言,传输在设备端是阻塞的(即使主机端并非如此):设备上的复制操作不能与其他操作同时发生。然而,在某些高级场景中,设备端可以同时进行复制和内核执行。如以下示例所示,需要满足三个条件才能启用此功能:
设备必须至少有一个空闲的 DMA(直接内存访问)引擎。现代 GPU 架构如 Volterra、Tesla 或 H100 设备具有多个 DMA 引擎。
传输必须发生在一个单独的非默认 cuda 流上。在 PyTorch 中,cuda 流可以用
Stream
进行处理。源数据必须位于固定内存中。
我们通过运行以下脚本进行配置以证明这一点。
import contextlib
from torch.cuda import Stream
s = Stream()
torch.manual_seed(42)
t1_cpu_pinned = torch.randn(1024**2 * 5, pin_memory=True)
t2_cpu_paged = torch.randn(1024**2 * 5, pin_memory=False)
t3_cuda = torch.randn(1024**2 * 5, device="cuda:0")
assert torch.cuda.is_available()
device = torch.device("cuda", torch.cuda.current_device())
# The function we want to profile
def inner(pinned: bool, streamed: bool):
with torch.cuda.stream(s) if streamed else contextlib.nullcontext():
if pinned:
t1_cuda = t1_cpu_pinned.to(device, non_blocking=True)
else:
t2_cuda = t2_cpu_paged.to(device, non_blocking=True)
t_star_cuda_h2d_event = s.record_event()
# This operation can be executed during the CPU to GPU copy if and only if the tensor is pinned and the copy is
# done in the other stream
t3_cuda_mul = t3_cuda * t3_cuda * t3_cuda
t3_cuda_h2d_event = torch.cuda.current_stream().record_event()
t_star_cuda_h2d_event.synchronize()
t3_cuda_h2d_event.synchronize()
# Our profiler: profiles the `inner` function and stores the results in a .json file
def benchmark_with_profiler(
pinned,
streamed,
) -> None:
torch._C._profiler._set_cuda_sync_enabled_val(True)
wait, warmup, active = 1, 1, 2
num_steps = wait + warmup + active
rank = 0
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
schedule=torch.profiler.schedule(
wait=wait, warmup=warmup, active=active, repeat=1, skip_first=1
),
) as prof:
for step_idx in range(1, num_steps + 1):
inner(streamed=streamed, pinned=pinned)
if rank is None or rank == 0:
prof.step()
prof.export_chrome_trace(f"trace_streamed{int(streamed)}_pinned{int(pinned)}.json")
在 chrome 中加载这些跟踪文件(chrome://tracing
)显示以下结果:首先,让我们看看在主流中,将页面张量发送到 GPU 后对 t3_cuda
执行算术操作时发生了什么:
benchmark_with_profiler(streamed=False, pinned=False)

使用固定张量并未显著改变跟踪结果,同一操作仍按照顺序执行:
benchmark_with_profiler(streamed=False, pinned=True)

将分页张量发送到 GPU 在单独流上也是一个阻塞操作:
benchmark_with_profiler(streamed=True, pinned=False)

只有将固定张量复制到 GPU 并在单独流中执行重叠操作时,另一个 cuda 内核才会在主流上执行:
benchmark_with_profiler(streamed=True, pinned=True)

PyTorch 的视角¶
pin_memory()
¶
PyTorch 提供了通过 pin_memory()
方法和构造函数参数创建和发送张量到页面锁定内存的可能性。在启用了 CUDA 的机器上,CPU 张量可以通过 pin_memory()
方法转换为固定内存。重要的是,pin_memory
在主机主线程上是阻塞的:它会等待张量被复制到页面锁定内存后再继续执行下一操作。新张量可以通过像 zeros()
、ones()
和其他构造函数直接创建于固定内存中。
让我们来检查固定内存和发送张量到 CUDA 的速度:
import torch
import gc
from torch.utils.benchmark import Timer
import matplotlib.pyplot as plt
def timer(cmd):
median = (
Timer(cmd, globals=globals())
.adaptive_autorange(min_run_time=1.0, max_run_time=20.0)
.median
* 1000
)
print(f"{cmd}: {median: 4.4f} ms")
return median
# A tensor in pageable memory
pageable_tensor = torch.randn(1_000_000)
# A tensor in page-locked (pinned) memory
pinned_tensor = torch.randn(1_000_000, pin_memory=True)
# Runtimes:
pageable_to_device = timer("pageable_tensor.to('cuda:0')")
pinned_to_device = timer("pinned_tensor.to('cuda:0')")
pin_mem = timer("pageable_tensor.pin_memory()")
pin_mem_to_device = timer("pageable_tensor.pin_memory().to('cuda:0')")
# Ratios:
r1 = pinned_to_device / pageable_to_device
r2 = pin_mem_to_device / pageable_to_device
# Create a figure with the results
fig, ax = plt.subplots()
xlabels = [0, 1, 2]
bar_labels = [
"pageable_tensor.to(device) (1x)",
f"pinned_tensor.to(device) ({r1:4.2f}x)",
f"pageable_tensor.pin_memory().to(device) ({r2:4.2f}x)"
f"\npin_memory()={100*pin_mem/pin_mem_to_device:.2f}% of runtime.",
]
values = [pageable_to_device, pinned_to_device, pin_mem_to_device]
colors = ["tab:blue", "tab:red", "tab:orange"]
ax.bar(xlabels, values, label=bar_labels, color=colors)
ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime (pin-memory)")
ax.set_xticks([])
ax.legend()
plt.show()
# Clear tensors
del pageable_tensor, pinned_tensor
_ = gc.collect()
我们可以观察到,将固定内存张量转换到 GPU 的速度确实比分页张量快,因为在底层,分页张量必须先复制到固定内存,然后再发送到 GPU。
然而,与一些常见的印象相反,在将分页张量转换到 GPU 之前调用 :meth:`~torch.Tensor.pin_memory()`通常不会带来任何显著的速度提升,相反,此调用通常比直接执行传输还要慢。这很合理,因为我们实际上是请求 Python 执行一个 CUDA 无论如何都会在主机到设备复制数据之前完成的操作。
备注
PyTorch 的 pin_memory 实现依赖于通过 cudaHostAlloc 创建一个全新的页面锁定存储,在少数情况下可能比如 cudaMemcpy
转换数据的方式更快。同样,这种观察可能会因可用硬件、目标张量大小或可用 RAM 的多少而有所不同。
non_blocking=True
¶
如前所述,许多 PyTorch 操作可以通过 non_blocking
参数异步地与主机执行。
为了准确评估使用 non_blocking
的好处,我们将设计一个稍复杂的实验,因为我们希望评估使用和不使用 non_blocking
时发送多个张量到 GPU 的速度。
# A simple loop that copies all tensors to cuda
def copy_to_device(*tensors):
result = []
for tensor in tensors:
result.append(tensor.to("cuda:0"))
return result
# A loop that copies all tensors to cuda asynchronously
def copy_to_device_nonblocking(*tensors):
result = []
for tensor in tensors:
result.append(tensor.to("cuda:0", non_blocking=True))
# We need to synchronize
torch.cuda.synchronize()
return result
# Create a list of tensors
tensors = [torch.randn(1000) for _ in range(1000)]
to_device = timer("copy_to_device(*tensors)")
to_device_nonblocking = timer("copy_to_device_nonblocking(*tensors)")
# Ratio
r1 = to_device_nonblocking / to_device
# Plot the results
fig, ax = plt.subplots()
xlabels = [0, 1]
bar_labels = [f"to(device) (1x)", f"to(device, non_blocking=True) ({r1:4.2f}x)"]
colors = ["tab:blue", "tab:red"]
values = [to_device, to_device_nonblocking]
ax.bar(xlabels, values, label=bar_labels, color=colors)
ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime (non-blocking)")
ax.set_xticks([])
ax.legend()
plt.show()
为了更好地理解这里发生的情况,让我们对这两个函数进行性能剖析:
首先看看调用常规 to(device)
时的调用栈:
print("Call to `to(device)`", profile_mem("copy_to_device(*tensors)"))
然后是 non_blocking
版本:
print(
"Call to `to(device, non_blocking=True)`",
profile_mem("copy_to_device_nonblocking(*tensors)"),
)
毫无疑问,使用 non_blocking=True
时的结果更为优秀,因为所有传输都在主机端同时启动,并且只进行一次同步操作。
收益将因张量的数量和大小以及所使用的硬件而异。
备注
有趣的是,阻塞的 to("cuda")
实际上执行与设置 non_blocking=True
后一样的异步设备转换操作(cudaMemcpyAsync
),但在每次复制后都有一个同步点。
协同效应¶
现在我们已经提出,将已经处于固定内存中的张量数据传输到 GPU 比从分页内存中传输更快,并且异步传输比同步传输更快,我们可以对这些方法的组合进行基准测试。首先,让我们编写几个新函数,这些函数将在每个张量上调用 pin_memory
和 to(device)
:
def pin_copy_to_device(*tensors):
result = []
for tensor in tensors:
result.append(tensor.pin_memory().to("cuda:0"))
return result
def pin_copy_to_device_nonblocking(*tensors):
result = []
for tensor in tensors:
result.append(tensor.pin_memory().to("cuda:0", non_blocking=True))
# We need to synchronize
torch.cuda.synchronize()
return result
使用 pin_memory()
的优点在大批量大型张量时尤为显著:
tensors = [torch.randn(1_000_000) for _ in range(1000)]
page_copy = timer("copy_to_device(*tensors)")
page_copy_nb = timer("copy_to_device_nonblocking(*tensors)")
tensors_pinned = [torch.randn(1_000_000, pin_memory=True) for _ in range(1000)]
pinned_copy = timer("copy_to_device(*tensors_pinned)")
pinned_copy_nb = timer("copy_to_device_nonblocking(*tensors_pinned)")
pin_and_copy = timer("pin_copy_to_device(*tensors)")
pin_and_copy_nb = timer("pin_copy_to_device_nonblocking(*tensors)")
# Plot
strategies = ("pageable copy", "pinned copy", "pin and copy")
blocking = {
"blocking": [page_copy, pinned_copy, pin_and_copy],
"non-blocking": [page_copy_nb, pinned_copy_nb, pin_and_copy_nb],
}
x = torch.arange(3)
width = 0.25
multiplier = 0
fig, ax = plt.subplots(layout="constrained")
for attribute, runtimes in blocking.items():
offset = width * multiplier
rects = ax.bar(x + offset, runtimes, width, label=attribute)
ax.bar_label(rects, padding=3, fmt="%.2f")
multiplier += 1
# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_ylabel("Runtime (ms)")
ax.set_title("Runtime (pin-mem and non-blocking)")
ax.set_xticks([0, 1, 2])
ax.set_xticklabels(strategies)
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
ax.legend(loc="upper left", ncols=3)
plt.show()
del tensors, tensors_pinned
_ = gc.collect()
其他方向的复制操作(GPU -> CPU,CPU -> MPS)¶
到目前为止,我们假设从 CPU 到 GPU 的异步复制是安全的。这通常是正确的,因为 CUDA 会自动处理同步,以确保在读取时访问的数据是有效的,__只要张量处于分页内存中__。
然而,在其他情况下,我们不能做出同样的假设:当一个张量被放置在固定内存中时,在调用主机到设备传输之后修改原始副本可能会导致 GPU 上接收到的数据损坏。同样,当传输方向为从 GPU 到 CPU,或从任何非 CPU 或 GPU 的设备到非 CUDA 管理的 GPU(例如,MPS)时,如果没有显式同步,就无法保证 GPU 上读取的数据是有效的。
在这些场景中,这些传输无法确保在访问数据时复制已经完成。因此,主机上的数据可能不完整或不正确,从而导致数据无效。
让我们首先用一个固定内存中的张量演示这一点:
DELAY = 100000000
try:
i = -1
for i in range(100):
# Create a tensor in pin-memory
cpu_tensor = torch.ones(1024, 1024, pin_memory=True)
torch.cuda.synchronize()
# Send the tensor to CUDA
cuda_tensor = cpu_tensor.to("cuda", non_blocking=True)
torch.cuda._sleep(DELAY)
# Corrupt the original tensor
cpu_tensor.zero_()
assert (cuda_tensor == 1).all()
print("No test failed with non_blocking and pinned tensor")
except AssertionError:
print(f"{i}th test failed with non_blocking and pinned tensor. Skipping remaining tests")
使用分页张量总是有效的:
i = -1
for i in range(100):
# Create a tensor in pageable memory
cpu_tensor = torch.ones(1024, 1024)
torch.cuda.synchronize()
# Send the tensor to CUDA
cuda_tensor = cpu_tensor.to("cuda", non_blocking=True)
torch.cuda._sleep(DELAY)
# Corrupt the original tensor
cpu_tensor.zero_()
assert (cuda_tensor == 1).all()
print("No test failed with non_blocking and pageable tensor")
现在让我们演示从 CUDA 到 CPU 的传输在没有同步时也无法产生可靠的输出:
tensor = (
torch.arange(1, 1_000_000, dtype=torch.double, device="cuda")
.expand(100, 999999)
.clone()
)
torch.testing.assert_close(
tensor.mean(), torch.tensor(500_000, dtype=torch.double, device="cuda")
), tensor.mean()
try:
i = -1
for i in range(100):
cpu_tensor = tensor.to("cpu", non_blocking=True)
torch.testing.assert_close(
cpu_tensor.mean(), torch.tensor(500_000, dtype=torch.double)
)
print("No test failed with non_blocking")
except AssertionError:
print(f"{i}th test failed with non_blocking. Skipping remaining tests")
try:
i = -1
for i in range(100):
cpu_tensor = tensor.to("cpu", non_blocking=True)
torch.cuda.synchronize()
torch.testing.assert_close(
cpu_tensor.mean(), torch.tensor(500_000, dtype=torch.double)
)
print("No test failed with synchronize")
except AssertionError:
print(f"One test failed with synchronize: {i}th assertion!")
通常情况下,只有当目标是支持 CUDA 的设备且原始张量在分页内存中时,异步复制到设备才在不显式同步的情况下是安全的。
总之,当使用 non_blocking=True
时,从 CPU 到 GPU 的数据复制是安全的,但对于任何其他方向,虽然可以使用 non_blocking=True
,但用户必须确保在访问数据之前执行设备同步。
实践建议¶
根据我们的观察,我们现在可以总结一些早期的建议:
一般来说,使用 non_blocking=True
会提供良好的吞吐量,无论原始张量是否位于固定内存中。如果张量已经在固定内存中,传输可以加速,但从 Python 主线程手动发送到固定内存是一种阻塞操作,从而抵消了使用 non_blocking=True
的大部分优势(因为 CUDA 已经自动执行 pin_memory 转移)。
现在可能有人会问,pin_memory()
方法有什么用途。在接下来的部分中,我们将进一步探讨如何使用它使数据传输更快。
额外考虑¶
PyTorch 提供了一个著名的 DataLoader
类,其构造函数接受一个 pin_memory
参数。考虑到我们之前的讨论,你可能会想知道,如果内存固定本质上是阻塞的,DataLoader
如何管理加速数据传输。
关键在于 DataLoader 使用一个单独的线程来处理从分页内存到固定内存的数据传输,从而防止主线程受阻。
为了说明这一点,我们将使用同名库中的 TensorDict 原语。当调用 to()
时,默认行为是将张量异步发送到设备,然后进行一次 torch.device.synchronize()
调用。
此外, TensorDict.to()
包括一个 non_blocking_pin
选项,该选项启动多个线程来执行 pin_memory()
,然后进行 to(device)
操作。此方法可以进一步加速数据传输,如以下示例所示。
from tensordict import TensorDict
import torch
from torch.utils.benchmark import Timer
import matplotlib.pyplot as plt
# Create the dataset
td = TensorDict({str(i): torch.randn(1_000_000) for i in range(1000)})
# Runtimes
copy_blocking = timer("td.to('cuda:0', non_blocking=False)")
copy_non_blocking = timer("td.to('cuda:0')")
copy_pin_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=0)")
copy_pin_multithread_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=4)")
# Rations
r1 = copy_non_blocking / copy_blocking
r2 = copy_pin_nb / copy_blocking
r3 = copy_pin_multithread_nb / copy_blocking
# Figure
fig, ax = plt.subplots()
xlabels = [0, 1, 2, 3]
bar_labels = [
"Blocking copy (1x)",
f"Non-blocking copy ({r1:4.2f}x)",
f"Blocking pin, non-blocking copy ({r2:4.2f}x)",
f"Non-blocking pin, non-blocking copy ({r3:4.2f}x)",
]
values = [copy_blocking, copy_non_blocking, copy_pin_nb, copy_pin_multithread_nb]
colors = ["tab:blue", "tab:red", "tab:orange", "tab:green"]
ax.bar(xlabels, values, label=bar_labels, color=colors)
ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime")
ax.set_xticks([])
ax.legend()
plt.show()
在此示例中,我们将许多大张量从 CPU 传输到 GPU。对于使用多线程 pin_memory()
的场景非常理想,这可以显著提高性能。然而,如果张量很小,与多线程相关的开销可能超过其收益。同样,如果张量数量有限,在单独线程中固定张量的优势也会受到限制。
此外,尽管在固定内存中创建永久缓冲区以从分页内存传输张量到 GPU 看起来很有优势,但这种策略不一定能加快计算速度。将数据复制到固定内存的固有限制仍然是一个瓶颈因素。
此外,将磁盘上的数据(无论是共享内存还是文件)传输到 GPU 通常需要一个中间步骤,即将数据复制到固定内存(位于 RAM 中)。在此上下文中,对大数据传输使用 non_blocking
可以显著增加 RAM 消耗,从而可能导致不利影响。
实际上,并没有一种通用的解决方案。结合使用多线程 pin_memory
和 non_blocking
的传输效果取决于各种因素,包括特定系统、操作系统、硬件以及任务的性质。以下是一些在尝试加速 CPU 和 GPU 之间的数据传输或比较不同情境下的吞吐量时需要检查的因素:
可用核心数量
有多少 CPU 核心可用?系统是否与其他用户或可能竞争资源的进程共享?
核心利用率
CPU 核心是否被其他进程严重利用?应用程序是否在数据传输时并发执行其他 CPU 密集型任务?
内存利用率
当前使用了多少分页内存和页面锁定内存?是否有足够的可用内存来分配额外的固定内存而不影响系统性能?请记住,一切都是有代价的,例如
pin_memory
会消耗 RAM 并可能影响其他任务。CUDA 设备能力
GPU 是否支持多 DMA 引擎以进行并发数据传输?使用的 CUDA 设备的具体能力和限制是什么?
需要发送的张量数量
在典型操作中会传输多少张量?
需要发送的张量大小
传输的张量大小是多少?少量大张量或大量小张量可能无法从相同的传输方案中受益。
系统架构
系统架构如何影响数据传输速度(例如,总线速度、网络延迟)?
此外,在固定内存中分配大量张量或大尺寸张量可能会占用 RAM 的相当一部分。这会减少可用内存用于其他关键操作(例如分页),从而对算法的整体性能产生负面影响。
结论¶
在整个教程中,我们探讨了从主机向设备发送张量时影响传输速度和内存管理的几个关键因素。我们了解到使用 non_blocking=True
通常可以加速数据传输,并且如果正确实施,pin_memory()
也可以提高性能。然而,这些技术需要精心设计和校准才能有效。
请记住,分析你的代码并关注内存消耗对于优化资源使用和实现最佳性能至关重要。