• 教學 >
  • PyTorch 中 non_blockingpin_memory() 的良好使用指南
捷徑

關於 PyTorch 中 non_blockingpin_memory() 良好使用的指南

建立於:2024 年 7 月 31 日 | 最後更新:2024 年 8 月 01 日 | 最後驗證:2024 年 11 月 05 日

作者Vincent Moens

簡介

將資料從 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,則不會觸發同步,並且主機上的主執行緒不會被阻塞。因此,從主機的角度來看,可以同時將多個張量發送到裝置,因為執行緒不需要等待一個傳輸完成才能啟動另一個傳輸。

注意

一般而言,傳輸在裝置端是阻塞的(即使在主機端不是):當另一個操作正在執行時,無法在裝置上進行複製。但是,在一些進階情況下,可以在 GPU 端同時完成複製和核心執行。如下面的範例所示,必須滿足三個要求才能啟用此功能

  1. 該裝置必須至少有一個空閒的 DMA(直接記憶體存取)引擎。現代 GPU 架構,例如 Volterra、Tesla 或 H100 裝置,具有多個 DMA 引擎。

  2. 傳輸必須在單獨的、非預設的 cuda stream 上完成。在 PyTorch 中,可以使用 Stream 來處理 cuda stream。

  3. 來源資料必須在鎖定記憶體中。

我們透過在以下腳本上執行分析來演示此操作。

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)會顯示以下結果:首先,讓我們看看如果 t3_cuda 上的算術運算在可分頁張量在主 stream 中發送到 GPU 之後執行會發生什麼情況

benchmark_with_profiler(streamed=False, pinned=False)

使用鎖定張量不會改變追蹤太多,兩個操作仍然是連續執行的

benchmark_with_profiler(streamed=False, pinned=True)

在單獨的 stream 上將可分頁張量發送到 GPU 也是一個阻塞操作

benchmark_with_profiler(streamed=True, pinned=False)

只有鎖定張量在單獨的 stream 上複製到 GPU 會與在主 stream 上執行的另一個 cuda kernel 重疊

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()
Device casting runtime (pin-memory)
pageable_tensor.to('cuda:0'):  0.4688 ms
pinned_tensor.to('cuda:0'):  0.3729 ms
pageable_tensor.pin_memory():  0.3623 ms
pageable_tensor.pin_memory().to('cuda:0'):  0.7267 ms

我們可以觀察到,將鎖定記憶體張量轉換為 GPU 的速度確實比可分頁張量快得多,因為在底層,可分頁張量必須先複製到鎖定記憶體,然後才能發送到 GPU。

但是,與某種常見的觀點相反,在將可分頁張量轉換為 GPU 之前,呼叫 pin_memory() 不應帶來任何顯著的速度提升,相反,此呼叫通常比僅執行傳輸慢。這是合理的,因為我們實際上是在要求 Python 執行 CUDA 在將資料從主機複製到裝置之前無論如何都會執行的操作。

注意

pin_memory 的 PyTorch 實作依賴於透過 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()
Device casting runtime (non-blocking)
copy_to_device(*tensors):  26.9607 ms
copy_to_device_nonblocking(*tensors):  19.6101 ms

為了更好地了解這裡發生的情況,讓我們分析這兩個函式

from torch.profiler import profile, ProfilerActivity


def profile_mem(cmd):
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        exec(cmd)
    print(cmd)
    print(prof.key_averages().table(row_limit=10))

讓我們首先看看使用常規 to(device) 的呼叫堆疊

print("Call to `to(device)`", profile_mem("copy_to_device(*tensors)"))
copy_to_device(*tensors)
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                 aten::to         3.74%       1.203ms       100.00%      32.181ms      32.181us          1000
           aten::_to_copy        13.68%       4.402ms        96.26%      30.978ms      30.978us          1000
      aten::empty_strided        24.77%       7.970ms        24.77%       7.970ms       7.970us          1000
              aten::copy_        18.44%       5.934ms        57.82%      18.605ms      18.605us          1000
          cudaMemcpyAsync        17.62%       5.669ms        17.62%       5.669ms       5.669us          1000
    cudaStreamSynchronize        21.76%       7.003ms        21.76%       7.003ms       7.003us          1000
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 32.181ms

Call to `to(device)` None

現在是 non_blocking 版本

print(
    "Call to `to(device, non_blocking=True)`",
    profile_mem("copy_to_device_nonblocking(*tensors)"),
)
copy_to_device_nonblocking(*tensors)
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                 aten::to         4.68%       1.123ms        99.90%      23.996ms      23.996us          1000
           aten::_to_copy        17.47%       4.196ms        95.22%      22.873ms      22.873us          1000
      aten::empty_strided        32.60%       7.831ms        32.60%       7.831ms       7.831us          1000
              aten::copy_        22.05%       5.297ms        45.15%      10.846ms      10.846us          1000
          cudaMemcpyAsync        23.10%       5.549ms        23.10%       5.549ms       5.549us          1000
    cudaDeviceSynchronize         0.10%      24.911us         0.10%      24.911us      24.911us             1
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 24.021ms

Call to `to(device, non_blocking=True)` None

毫無疑問,使用 non_blocking=True 的結果更好,因為所有傳輸都在主機端同時啟動,並且只完成一次同步。

好處會因張量的數量和大小以及所使用的硬體而異。

注意

有趣的是,阻塞的 to("cuda") 實際上執行與 non_blocking=True 相同的非同步裝置轉換操作(cudaMemcpyAsync),並在每次複製後都有一個同步點。

協同作用

既然我們已經說明了將已固定在釘選記憶體 (pinned memory) 中的張量資料傳輸到 GPU 比從可分頁記憶體 (pageable memory) 傳輸更快,並且我們知道以非同步方式進行這些傳輸也比同步方式更快,我們可以針對這些方法的組合進行基準測試。首先,讓我們編寫幾個新的函式,在每個張量上調用 pin_memoryto(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()
Runtime (pin-mem and non-blocking)
copy_to_device(*tensors):  612.0971 ms
copy_to_device_nonblocking(*tensors):  537.7727 ms
copy_to_device(*tensors_pinned):  374.0692 ms
copy_to_device_nonblocking(*tensors_pinned):  347.7238 ms
pin_copy_to_device(*tensors):  968.0525 ms
pin_copy_to_device_nonblocking(*tensors):  647.8676 ms

其他複製方向 (GPU -> CPU, CPU -> MPS)

到目前為止,我們一直假設從 CPU 到 GPU 的非同步複製是安全的。這通常是真的,因為 CUDA 會自動處理同步,以確保在讀取時存取的資料是有效的。但是,這種保證不適用於相反方向,從 GPU 到 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!")
0th test failed with non_blocking. Skipping remaining tests
No test failed with synchronize

相同的考量適用於從 CPU 到非 CUDA 裝置 (例如 MPS) 的複製。通常,只有當目標是啟用 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 參數。考慮到我們之前關於 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()
Device casting runtime
td.to('cuda:0', non_blocking=False):  622.4414 ms
td.to('cuda:0'):  546.2743 ms
td.to('cuda:0', non_blocking_pin=True, num_threads=0):  664.7217 ms
td.to('cuda:0', non_blocking_pin=True, num_threads=4):  358.8236 ms

在此範例中,我們將許多大型張量從 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() 也可以提高效能。但是,這些技術需要仔細的設計和校準才能有效。

請記住,分析您的程式碼並密切關注記憶體消耗對於優化資源使用並實現最佳效能至關重要。

其他資源

如果您在使用 CUDA 裝置時遇到記憶體複製問題,或者想了解更多關於本教程中討論的內容,請查看以下參考文獻

腳本的總執行時間: ( 1 分鐘 19.063 秒)

由 Sphinx-Gallery 生成的圖片集

文件

取得 PyTorch 的完整開發人員文件

查看文件

教學

取得針對初學者和進階開發人員的深入教程

查看教程

資源

尋找開發資源並獲得您的問題解答

查看資源