我们可以将计算操作与内存操作重叠，而不将内存固定在 CPU 上吗？

Question

我试图将计算和内存操作与 HuggingFace SwitchTransformer 重叠。

\n

这里\xe2\x80\x99有详细的解释。

\n

\n
• 内存操作用于将数据从CPU移动到GPU，其大小为每块4MB。
\n
• 块的数量是可变的（通常总共从 2 到 6 个）。
\n
• 计算操作包括几个非常小的计算操作，例如 GEMM，每个操作需要 10 到 100 微秒。
\n
• 我正在尝试使用 CudaStream，因此我创建了两个不同的 Cuda 流，并将内存操作和计算操作推送到每个流。
\n
• 但它并没有重叠。
\n

\n

            s_0 = torch.cuda.Stream()  # Create a new stream.\n            s_1 = torch.cuda.Stream()  # Create a new stream.\n\n            with torch.cuda.stream(s_0):\n                this_gate_info = router_mask, router_probs, router_logits\n                router_mask = router_mask.bool()\n                idx_mask = router_mask.transpose(1,2)\n                idx_mask = torch.cat(torch.split(idx_mask, 1, dim=0), dim=2)\n                idx_mask = idx_mask.sum(dim=2)\n                idx_mask = idx_mask.squeeze()\n                    \n                if next_blk is not None:\n                    active_idx = torch.nonzero(idx_mask, as_tuple=True)\n                    for idx in active_idx[0]:\n                        tmp = getattr(next_blk.layer[-1].mlp.experts, "expert_{}".format(idx))\n                        tmp.prefetching()    ## THIS IS MEMORY OPERATION COLORED GREEN IN THE FIGURE\n\n            with torch.cuda.stream(s_1):\n                delayed_router_mask, delayed_router_probs, delayed_router_logits = delayed_gate_info\n                delayed_expert_index = torch.argmax(delayed_router_mask, dim=-1)\n                    \n                delayed_router_mask = delayed_router_mask.bool()\n                delayed_idx_mask = delayed_router_mask.transpose(1,2)\n                delayed_idx_mask = torch.cat(torch.split(delayed_idx_mask, 1, dim=0), dim=2)\n                delayed_idx_mask = delayed_idx_mask.sum(dim=2)\n                delayed_idx_mask = delayed_idx_mask.squeeze()\n\n                for idx, expert in enumerate(self.experts.values()):\n                    if delayed_idx_mask[idx] != 0:\n                        expert_counter = expert_counter + 1\n                        next_states[delayed_router_mask[:, :, idx]] = expert(hidden_states[delayed_router_mask[:, :, idx]], None, None, None)\n

\n

\n

这是我的问题。

\n

\n

1. 首先，我了解到为了重叠内存操作（CPU->GPU）和计算操作，应该固定CPU中的内存。但就我而言，如图所示，它是可分页内存，而不是固定的。这是不能重叠的原因吗？

   \n
\n

2. 其次，我进行了一个实验，用一个简单的示例（将 GEMM 与 CPU->GPU 内存操作重叠）来证明这一点，这是输出。

   \n
\n

\n

import torch\nimport torch.nn as nn\nimport torch.cuda.nvtx as nvtx_cuda\n\ntorch.cuda.cudart().cudaProfilerStart()\n\ncuda = torch.device(\'cuda\')\nnvtx_cuda.range_push("STREAM INIT")\ns_0 = torch.cuda.Stream()  # Create a new stream.\ns_1 = torch.cuda.Stream()  # Create a new stream.\nnvtx_cuda.range_pop()\n\nA = torch.rand(size=(1024*4, 1024*4), device="cuda")\nB = torch.rand(size=(1024*4, 1024*4), device="cuda")\nC = torch.rand(size=(1024*4, 1024*4), device="cuda")\nD = torch.rand(size=(1024*4, 1024*4), device="cuda")\nE = torch.rand(size=(1024*4, 1024*4), device="cuda")\nF = torch.rand(size=(1024*4, 1024*4), device="cuda")\n\na = torch.rand(size=(1024*4, 1024*4), pin_memory=False)\nb = torch.rand(size=(1024*4, 1024*4), device="cuda")\n\niter = 10\n\nfor i in range(iter):\n\n    with torch.cuda.stream(s_0):\n        nvtx_cuda.range_push("S0")\n        C = A.matmul(B)\n        F = D.matmul(E)\n        nvtx_cuda.range_pop()\n        \n    with torch.cuda.stream(s_1):\n        nvtx_cuda.range_push("S1")\n        nvtx_cuda.range_pop()\n        b = a.to(cuda)\n\ntorch.cuda.cudart().cudaProfilerStop()\n

\n

\n这是可分页内存。\n\n这是固定内存。

\n

看来可分页内存也可以重叠。\n那么，我的应用程序不重叠的原因是什么？

\n

Answer 1

一般来说，对于任意大小和情况，为了将 D->H 或 H->D 复制操作与内核执行重叠，有必要：

1. 使用cudaMemcpyAsync()
2. 使用创建的流，而不是默认流
3. 使用固定指针/缓冲区进行主机分配

所讨论的内核启动还需要启动到不同的非空流中。（是的，修改流默认行为可能会影响其中一些。我假设默认为空流行为。）

关于最后一项（3），这是一个一般性的陈述。对此的支持来自文档中的多个位置，包括此处和此处。

然而，来自非固定缓冲区的 D->H 或 H->D 复制分阶段进行。CUDA 运行时创建其自己的固定缓冲区，用于所有可分页复制操作，并且对于适合 CUDA 运行时维护的缓冲区（大小未在任何地方指定）的足够小的传输，传输操作可以是异步的。在这种情况下，有可能见证可分页缓冲区的重叠。由于这没有正式指定，并且未指定大小和合理使用所需的内容，因此实际上人们通常不依赖于这种行为，通常的建议是使用固定内存来实现重叠。

然而，这并没有描述您的第二种情况，在这种情况下，分析器似乎表明可能存在“大”传输的重叠。然而，分阶段转移也是理解这一点的关键。

当传输操作满足第 1 项和第 2 项（上面），但不满足第 3 项，并且大小足够大而无法完全放入暂存缓冲区时，CUDA 运行时会将传输操作分成适合暂存缓冲区的块。然后，它将主机上的数据从可分页缓冲区传输到固定缓冲区。然后将固定的缓冲区内容传输到设备。重复此操作，直到传输完所有数据。操作cudaMemcpyAsync()本身不会返回，即不会解锁 CPU 线程，直到最终块传输到（固定的）暂存缓冲区。

因此，考虑到这一点，如果您启动内核，然后启动可分页传输（正是您的测试用例），您确实可能会在内核仍在执行时见证传输活动（即重叠）。但是，正如您的跟踪中所表明的，cudaMemcpyAsync()在传输操作完成或接近完成之前，该操作不会返回（解锁 CPU 线程）。

这是一个问题。这种行为（CPU 阻塞）对于尝试向 GPU 发出精心编排的并发/异步工作来说是灾难性的。因此，虽然您可能会看到精心构造的测试用例有一些重叠，但在一般情况下，使用可分页缓冲区使得启动不打算在将来某个时候进行的工作变得非常困难。这使得向 GPU 发出大量异步工作基本上是不可能的。

作为一个简单的示例，您的特定测试用例可分页传输正在重叠，因为该传输是在相关内核启动后发出的。内核启动对于 CPU 线程始终是非阻塞的，因此 CPU 线程可以开始可分页传输，从而实现重叠。然而，如果我们颠倒了执行顺序，那么在大多数情况下，特定的传输不会与特定的内核重叠，因为 CPU 线程在传输期间被阻塞，并且无法继续启动内核。这似乎是您原始案例中发生的情况：

（是的，我知道您的测试用例有一个循环，因此它很可能与其他一些内核启动重叠。）

因此，当向需要重叠的 GPU 发出异步工作时，一般建议使用固定缓冲区。很难依靠其他任何东西来高效执行工作。