Cuda stream not in parallel for element-wise/matrix extraction

qi1: <p>Dear all,</p>
<p>I’m trying to perform element-wise operations (torch.multiply) or extract weight matrices during torch.matmul using a different stream. However, after analyzing the CUDA kernel with Nsight System, I realized these operations can’t be parallelized; for example, I can’t asynchronously run torch.multiply while performing torch.matmul or asynchronously run 2 torch.multiply in 2 streams.</p>
<p>My guess is that CUDA requires the operation to be large enough to launch asynchronously, but I lack sufficient GPU DRAM or threads for this. Conversely, when I have adequate GPU resources, small operations don’t launch asynchronously. How can this issue be resolved? Thank you!</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/6/267cc4297c34399d000b4886d373b25cb76bc5da.png" data-download-href="https://discuss.pytorch.org/uploads/default/267cc4297c34399d000b4886d373b25cb76bc5da" title="2024-04-03 08-48-35 的屏幕截图"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/267cc4297c34399d000b4886d373b25cb76bc5da_2_690x75.png" alt="2024-04-03 08-48-35 的屏幕截图" data-base62-sha1="5utsawOkKPlqHeMOMDXfijkZPWa" width="690" height="75" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/267cc4297c34399d000b4886d373b25cb76bc5da_2_690x75.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/267cc4297c34399d000b4886d373b25cb76bc5da_2_1035x112.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/6/267cc4297c34399d000b4886d373b25cb76bc5da_2_1380x150.png 2x" data-dominant-color="D9D9DB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">2024-04-03 08-48-35 的屏幕截图</span><span class="informations">3057×333 44.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Here is my system info:</p>
<p>print(“PyTorch version:”, torch.<strong>version</strong>)<br>
print(“CUDA available:”, torch.cuda.is_available())<br>
if torch.cuda.is_available():<br>
print(“CUDA version:”, torch.version.cuda)<br>
print(torch.<strong>version</strong>)<br>
print(torch.version.cuda)<br>
print(torch.backends.cudnn.version())<br>
‘’’<br>
PyTorch version: 2.0.1+cu117<br>
CUDA available: True<br>
CUDA version: 11.7<br>
2.0.1+cu117<br>
11.7<br>
8906<br>
System: ubuntu 20.04<br>
GPU: NVIDIA GeForce RTX 4090<br>
‘’’</p>
<p>Here is my code:</p>
<pre><code class="lang-auto">import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
a = torch.rand(1000, 1000, device=device, dtype=torch.float16)
b = torch.rand(1000, 1000, device=device, dtype=torch.float16)
a_large = torch.rand(10000, 10000, device=device, dtype=torch.float16)
b_large = torch.rand(10000, 10000, device=device, dtype=torch.float16)

# Create two CUDA streams
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()

# ----test multiply async
with torch.cuda.stream(stream1):
    for i in range(100000):
        result1 = torch.multiply(a, a)  # a * a
       
with torch.cuda.stream(stream2):
    for i in range(100000):
        result2 = torch.multiply(b, b)  # b * b

torch.cuda.synchronize()

# ----test matmul, multiply async
layer1 = torch.nn.Linear(80, 80, device=device)
layer2 = torch.nn.Linear(80, 80, device=device)
input_tensor = torch.randn(128, 80, device=device)

with torch.cuda.stream(stream1):
    for i in range(100000):
        result1 = layer2(layer1(input_tensor))  # layer1(input_tensor)

with torch.cuda.stream(stream2):
    for i in range(100000):
        result2 = torch.multiply(b, b)  # b * b

torch.cuda.synchronize()

# ----test matmul, select indices async
layer1 = torch.nn.Linear(80, 80, device=device)
layer2 = torch.nn.Linear(80, 80, device=device)
input_tensor = torch.randn(128, 80, device=device)
select_indices = torch.arange(500, device=device)

with torch.cuda.stream(stream1):
    for i in range(100000):
        result1 = layer2(layer1(input_tensor))  # layer1(input_tensor)

with torch.cuda.stream(stream2):
    for i in range(100000):
        _ = b[:, select_indices]

torch.cuda.synchronize()

# ----test matmul, select indices async (large matrix)
layer1 = torch.nn.Linear(4000, 4000, device=device)
layer2 = torch.nn.Linear(4000, 4000, device=device)
input_tensor = torch.randn(128, 4000, device=device)
select_indices = torch.arange(5000, device=device)

with torch.cuda.stream(stream1):
    for i in range(100000):
        result1 = layer2(layer1(input_tensor))  # layer1(input_tensor)

with torch.cuda.stream(stream2):
    for i in range(100000):
        _ = b_large[:, select_indices]

torch.cuda.synchronize()
exit()
</code></pre>
<p>I ran it with nsys profile -w true --gpu-metrics-device=0 -x true --force-overwrite=true -o my_profile_simpletest python simpletest.py</p>

qi1: <p>Is there a way to make CUDA directly parallelize small operations in pytroch?</p>

ptrblck: <p><a href="https://discuss.pytorch.org/t/running-independent-nn-module-instances-in-nn-modulelist-truly-in-parallel-in-pytorch/188166/3">This post</a> might be helpful explaining the compute resources with a link to a great GTC talk explaining it in more details.</p>

qi1: <p>thanks, I will take a look. I’ve noticed that I cannot limit CUDA usage in PyTorch, for example, limit the maximum GPU thread usage by a stream. I want to double-check whether this is correct.</p>

ptrblck: <p>I assume you mean you cannot limit the sm usage? If so, then yes, you cannot limit it in PyTorch and all kernels will have the ability to use all available compute resources.</p>