[Python] Why is "with torch.cuda.stream()" not asynchronous?

Cause


I'm trying to achieve the overlap between the forward and backward passes of two model chunks in dualpipe. I plan to simply implement it using torch.cuda.stream first, as follows:

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:
from megatron.core.utils import get_attr_wrapped_model
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


I found that there is no overlap between the forward and backward passes. The running time of the backward code is the same as that without adding the stream.

Experiment


So, I did this experiment below:


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()


I tested the time with and without stream and found that the times were similar. Surprisingly, the time with stream was even longer:

enter image description here

Question


According to my understanding, "with torch.cuda.stream()" should be asynchronous and non - blocking, and the running time of the code block within it should approach zero. Only in this way can it perform parallel computing or communication with subsequent programs. However, the experimental results seem to indicate that it is blocking, and the program will only continue to execute after the entire calculation is completed. Is there something wrong with my code, or are there other possible reasons? Please give me some advice!

Continue reading...