小编Rya*_*yan的帖子

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

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

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