Why are batched operations slower than a for loop?

username: howardh
cooked: <p>I’m trying to speed up some computations converting them to a single matrix operation. I thought this would be faster because it would allow the computations to be run in parallel, whereas the for loop would perform each computation sequentially, but that was not the case when I tested it. So I’m here to get a better understanding of how this works. Why is it that the batched version is so much slower than the loop?</p>
<pre><code class="lang-python">import timeit
import torch

def test_thing():
    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    num_iterations = 1_000

    num_modules = 12
    num_inputs = 13
    batch_size = 16
    embed_dims = 512

    # Initialize inputs and weights/biases
    x = torch.rand([num_inputs, batch_size, embed_dims], device=device)
    weight = torch.rand([num_modules, embed_dims, embed_dims], device=device)                                                                    
    bias = torch.rand([num_modules, embed_dims], device=device)                                                                                  
                                                                                                                                                 
    x2 = x.unsqueeze(0)                                                                                                                          
    weight2 = weight.unsqueeze(1).transpose(-2,-1)                                                                                               
    bias2 = bias.unsqueeze(1).unsqueeze(2)                                                                                                       
                                                                                                                                                 
    # Functions to evaluate                                                                                                                      
    def run_batched():                                                                                                                           
        output = x.unsqueeze(0) @ weight.unsqueeze(1).transpose(-2,-1) + bias.unsqueeze(1).unsqueeze(2)                                          
        torch.cuda.synchronize()  
        return output                                                                                                                    
                                                                                                                                         
    def run_batched2():                                                                                                                  
        output = x2 @ weight2 + bias2                                                                                                    
        torch.cuda.synchronize()
        return output                                                                                                                    
                                                                                                                                         
    def run_non_batched():                                                                                                               
        output = [                                                                                                                       
            x @ w.transpose(-2,-1) + b                       
            for w,b in zip(weight, bias)                     
        ]                                                    
        torch.cuda.synchronize()
        return torch.stack(output)                           
                                                             
    # Ensure that they are all computing the same thing      
    assert torch.allclose(run_batched(), run_non_batched())  
    assert torch.allclose(run_batched2(), run_non_batched()) 
                                                             
    # Measure run time                                                          
    total_time_batched = timeit.Timer(run_batched).timeit(number=num_iterations)  
    total_time_batched2 = timeit.Timer(run_batched2).timeit(number=num_iterations)     
    total_time_nonbatched = timeit.Timer(run_non_batched).timeit(number=num_iterations)
                                                             
    print(f"Batch: {total_time_batched}")                    
    print(f"Batch2: {total_time_batched2}")                  
    print(f"NonBatch: {total_time_nonbatched}")              
                                                             
    assert total_time_batched2 &lt; total_time_nonbatched       
    assert total_time_batched &lt; total_time_nonbatched
</code></pre>
<p>Console output:</p>
<pre><code class="lang-auto">Batch: 3.460433868924156
Batch2: 3.4635026860050857
NonBatch: 0.6445531530771405
</code></pre>

username: KFrank
cooked: <p>Hi Howard!</p>
<aside class="quote no-group" data-username="howardh" data-post="1" data-topic="152835" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/howardh/48/49742_2.png" class="avatar"> howardh:</div>
<blockquote>
<p>Why is it that the batched version is so much slower than the loop?</p>
</blockquote>
</aside>
<p>Good question – I don’t know.</p>
<p>I can reproduce your result. including on a recent nightly build.</p>
<p>It’s interesting to note that a non-loop <code>einsum()</code> version is competitive<br>
with your “non-batched” loop version, and, at least on my system, turns<br>
out to be somewhat faster.</p>
<p>(I see slower, but similar relative timings using just the cpu, with <code>einsum()</code><br>
again being the fastest version.)</p>
<p>Running a tweaked version of your test code, with <code>einsum()</code> added as<br>
“Batch3,” I get:</p>
<pre><code class="lang-plaintext">1.10.0
10.2
GeForce GTX 1050 Ti
check run_batched: True
check run_batched2: True
check run_batched3: True
Batch: 6.828574229999504
Batch2: 11.512633088001166
Batch3: 0.9353041959984694
NonBatch: 1.3501966660005564
</code></pre>
<p>Here is my version of your test script:</p>
<pre data-code-wrap="python"><code class="lang-python">import timeit
import torch

def test_thing():
    print (torch.__version__)
    if torch.cuda.is_available():
        print (torch.version.cuda)
        print (torch.cuda.get_device_name())
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')
    
    num_iterations = 1_000
    
    num_modules = 12
    num_inputs = 13
    batch_size = 16
    embed_dims = 512
    
    # Initialize inputs and weights/biases
    x = torch.rand([num_inputs, batch_size, embed_dims], device=device)
    weight = torch.rand([num_modules, embed_dims, embed_dims], device=device)
    bias = torch.rand([num_modules, embed_dims], device=device)
    
    x2 = x.unsqueeze(0)
    weight2 = weight.unsqueeze(1).transpose(-2,-1)
    bias2 = bias.unsqueeze(1).unsqueeze(2)
    
    # Functions to evaluate
    def run_batched():
        output = x.unsqueeze(0) @ weight.unsqueeze(1).transpose(-2,-1) + bias.unsqueeze(1).unsqueeze(2)
        return output
    
    def run_batched2():
        output = x2 @ weight2 + bias2
        torch.cuda.synchronize()
        return output
    
    def run_batched3():
        output = torch.einsum ('ijk,lmk -&gt; lijm', x, weight) + bias2
        torch.cuda.synchronize()
        return output
    
    def run_non_batched():
        output = [
            x @ w.transpose(-2,-1) + b                       
            for w,b in zip(weight, bias)                     
        ]                                                    
        torch.cuda.synchronize()
        return torch.stack(output)                           
    
    # Ensure that they are all computing the same thing
    print ('check run_batched:', torch.allclose(run_batched(), run_non_batched()))
    print ('check run_batched2:', torch.allclose(run_batched2(), run_non_batched()))
    print ('check run_batched3:', torch.allclose(run_batched3(), run_non_batched()))
    
    # Measure run time                                                          
    total_time_batched = timeit.Timer(run_batched).timeit(number=num_iterations)  
    total_time_batched2 = timeit.Timer(run_batched2).timeit(number=num_iterations)     
    total_time_batched3 = timeit.Timer(run_batched3).timeit(number=num_iterations)     
    total_time_nonbatched = timeit.Timer(run_non_batched).timeit(number=num_iterations)
    
    print(f"Batch: {total_time_batched}")
    print(f"Batch2: {total_time_batched2}")
    print(f"Batch3: {total_time_batched3}")
    print(f"NonBatch: {total_time_nonbatched}")

test_thing()
</code></pre>
<p>For completeness, here are my nightly-build results:</p>
<pre><code class="lang-plaintext">1.12.0.dev20220410
11.3
GeForce GTX 1050 Ti
check run_batched: True
check run_batched2: True
check run_batched3: True
Batch: 6.188123900999926
Batch2: 10.426687854000193
Batch3: 0.92185746800169
NonBatch: 1.3100559570011683
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: howardh
cooked: <p>Thanks KFrank, that’s good to know. I get a significant speed up using einsum on my machine as well.</p>