torch.cuda.stream()未实现异步执行的问题排查与咨询

我目前在DualPipe框架里尝试实现两个模型分片的前向传播与反向传播的重叠执行，计划先通过torch.cuda.stream来实现异步逻辑，写了如下示例代码，但实际运行后发现前向和反向完全没有重叠，反向耗时和不使用stream时几乎一致：

from typing import List, Optional, Callable, tuple
import torch
from megatron.core.utils import get_attr_wrapped_model

def overlapped_forward_backward(
    module0: torch.nn.Module,
    inputs0: List[torch.Tensor],
    labels0: Optional[List[torch.Tensor]],
    loss_masks0: Optional[List[torch.Tensor]],
    loss1: Optional[torch.Tensor],
    outputs1: Optional[List[torch.Tensor]],
    output_grads1: Optional[List[torch.Tensor]],
    forward_step_func: Callable,
    is_last_stage0: bool,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    """ You should implement custom forward-backward overlap strategy. The code below is just an example. """
    device = inputs0[0].device
    if not hasattr(overlapped_forward_backward, 'backward_streams'):
        overlapped_forward_backward.backward_streams = {}
    if device not in overlapped_forward_backward.backward_streams:
        overlapped_forward_backward.backward_streams[device] = torch.cuda.Stream(device=device)
    backward_stream = overlapped_forward_backward.backward_streams[device]

    with torch.cuda.stream(backward_stream):
        if loss1 is not None:
            loss1.backward()
            loss1.detach_()
        else:
            run_backward(outputs1, output_grads1)

    if len(inputs0) == 1:
        set_input_tensor = get_attr_wrapped_model(module0, "set_input_tensor")
        set_input_tensor(inputs0)

    if is_last_stage0:
        inputs0_with_labels_loss_masks = list(inputs0)
        inputs0_with_labels_loss_masks.append(labels0)
        inputs0_with_labels_loss_masks.append(loss_masks0)
        outputs0, loss_func = forward_step_func(inputs0_with_labels_loss_masks, module0)
    else:
        outputs0, loss_func = forward_step_func(inputs0, module0)
    outputs0 = [outputs0] if isinstance(outputs0, torch.Tensor) else outputs0

    if is_last_stage0:
        loss0 = loss_func(outputs0[0])[0]
    else:
        loss0 = None

    torch.cuda.current_stream().wait_stream(backward_stream)
    return outputs0, loss0

为了排查问题，我单独写了一个小实验测试torch.cuda.stream的异步执行效果，结果更出乎意料：加了stream的代码耗时甚至比不加还长，看起来完全是阻塞执行的，和我理解的「with torch.cuda.stream()代码块内操作应该异步非阻塞、耗时接近0」完全不符：

import torch
import time

# GPU warmup
a = torch.randn(10000, 10000, device='cuda')
b = torch.randn(10000, 10000, device='cuda')
c = torch.mm(a, b)

# Count calculation time
calc_start = time.time()
a = torch.randn(10000, 10000, device='cuda')
b = torch.randn(10000, 10000, device='cuda')
for i in range(100):
    c = torch.mm(a, b)
calc_end = time.time()
print(f"calc time: {calc_end - calc_start}")

# Count stream time
calc_stream = torch.cuda.Stream()
torch.cuda.synchronize()
stream_start = time.time()
with torch.cuda.stream(calc_stream):
    a = torch.randn(10000, 10000, device='cuda')
    b = torch.randn(10000, 10000, device='cuda')
    for i in range(100):
        c = torch.mm(a, b)
stream_end = time.time()
print(f"stream time: {stream_end - stream_start}")
torch.cuda.synchronize()

我现在有几个核心疑问，希望能得到大家的建议：

• 是不是我的代码写法有误，导致stream没有真正触发异步执行？
• 有没有其他潜在因素（比如CUDA版本、硬件限制、框架特性）会导致这个现象？
• 针对我要实现「前反向传播重叠」的场景，有没有调试或修改的方向可以让stream真正发挥异步作用？

内容来源于stack exchange