Cannot reach theoretical performance with `torch.matmul`

username: ha2ka
cooked: <p>Hi, I’m working with following script for benchmarking my RTX 3080 GPU. In case of <code>torch.float16</code>, <code>torch.matmul</code> is performed with 29 [TFLOPS]. But, in case of <code>torch.float32</code>, <code>torch.matmul</code> is performed with only 9.8 [TFLOPS]. Because the theoretical performance of RTX 3080 is 29.77 [TFLOPS] and the GPU has no half units, the results may be something wrong.</p>
<p>Do you have any information on this?</p>
<pre><code class="lang-python">import torch
import time

torch.backends.cudnn.benchmark = True

size = 8192*2
repeat = 100

with torch.no_grad():

    x = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float32)
    w = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float32)
    torch.cuda.synchronize()
    start = time.time()
    for i in range(repeat):
        y = torch.matmul(x, w)
    torch.cuda.synchronize()
    print(time.time() - start, " secs")
    print(size**3 / ((time.time() - start) / repeat) / 1e9, " GFLOPS")
</code></pre>
<pre><code class="lang-auto">// for torch.float16
15.150275945663452  secs
29029.38026442699  GFLOPS

//for torch.float32
44.60238480567932  secs
9860.550400878896  GFLOPS
</code></pre>

username: eqy
cooked: <p>I believe that stated number is for TF32 TFLOPs, and you might be disabling TF32 by default depending on your setup.</p>
<p>Could you try running <code>torch.set_float32_matmul_precision('high')</code> before your benchmark and check if that changes the results?</p>

username: ha2ka
cooked: <p>Thank you for your reply.</p>
<p>I inserted the function you told me, then results are changed as follows:</p>
<pre><code class="lang-auto">// for torch.float16
14.611926078796387  secs
30098.913633454387  GFLOPS

// for torch.float32
28.9222469329834  secs
15206.416530498134  GFLOPS
</code></pre>
<p>In addition, I changed the <code>repeat</code> variable to <code>10</code> from <code>100</code>, then results are changed as follows:</p>
<pre><code class="lang-auto">// for torch.float16 (with repeat = 10)
1.8297147750854492  secs
24036.135977214086  GFLOPS

// for torch.float32 (with repeat = 10)
3.2416188716888428  secs
13567.226088349447  GFLOPS
</code></pre>
<p>The modified script is as follows:</p>
<pre><code class="lang-python">import torch
import time

torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('high')

size = 8192*2
repeat = 10

with torch.no_grad():

    x = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float32)
    w = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float32)
    torch.cuda.synchronize()
    start = time.time()
    for i in range(repeat):
        y = torch.matmul(x, w)
        torch.cuda.synchronize()
    print(time.time() - start, " secs")
    print(size**3 / ((time.time() - start) / repeat) / 1e9, " GFLOPS")
</code></pre>

username: eqy
cooked: <p>I would also recommend moving the second <code>torch.cuda.synchronize()</code> out of the for loop (so you synchronize only once after running all iterations) and see if that helps somewhat.</p>

username: ha2ka
cooked: <p>I changed the script is as follows:</p>
<pre><code class="lang-python">for i in range(repeat):
        torch.cuda.synchronize()
        start = time.time()
        y = torch.matmul(x, w)
        torch.cuda.synchronize()
        print(size**3 / ((time.time() - start)) / 1e12, " TFLOPS")
</code></pre>
<p>So the results changed like this:</p>
<pre><code class="lang-auto">// torch.float16
21.586235349878155  TFLOPS
28.278445566394296  TFLOPS
29.119723232835405  TFLOPS
29.780876695030035  TFLOPS
29.610455979112597  TFLOPS
29.58355453208967  TFLOPS
29.539698391619098  TFLOPS
29.567379058328427  TFLOPS
29.4742516322403  TFLOPS
29.494398091095096  TFLOPS
</code></pre>
<p>I think that the first one is always low performance because of memory transferring or something.<br>
And this settings, results of <code>float32</code> case is as follows:</p>
<pre><code class="lang-auto">6.2574480417038565  TFLOPS
14.587065018163548  TFLOPS
15.773094195853083  TFLOPS
15.74772563022044  TFLOPS
15.447554939534173  TFLOPS
15.692985973957423  TFLOPS
15.671055399278538  TFLOPS
15.581757487920573  TFLOPS
15.683139624736594  TFLOPS
15.634119138360994  TFLOPS
</code></pre>
<p>As you can see from the results, the suggestion to put <code> torch.cuda.synchronize()</code> outside the loop did not work. (In fact, when I checked the results with your suggestion before modifying the script, results was not changed.)</p>

username: eqy
cooked: <p>I am able to get ~16TFLOPs with a slightly modified script, but that is still only ~50% of “speed-of-light”</p>
<pre><code class="lang-auto">import torch
import time

torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('high')

size = 8192*2
repeat = 20

with torch.no_grad():
    x = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float)
    w = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float)
    y = torch.matmul(x, w)
    torch.cuda.synchronize()
    start = time.perf_counter()
    for i in range(repeat):
        y = torch.matmul(x, w)
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, " secs")
    print(size**3 / ((end - start) / repeat) / 1e9, " GFLOPS")

print("now using CUDA Graphs...")
with torch.no_grad():
   x = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float)
   w = torch.randn(size, size, device=torch.device("cuda"), dtype=torch.float)
   s = torch.cuda.Stream()
   s.wait_stream(torch.cuda.current_stream())
   with torch.cuda.stream(s):
       for i in range(3):
           y = torch.matmul(x, w)
   torch.cuda.current_stream().wait_stream(s)
   g = torch.cuda.CUDAGraph()
   with torch.cuda.graph(g):
       y = torch.matmul(x, w)
   torch.cuda.synchronize()
   start = time.perf_counter()
   for i in range(repeat):
       g.replay()
   torch.cuda.synchronize()
   end = time.perf_counter()
   print(end - start, " secs")
   print(size**3 / ((end - start) / repeat) / 1e9, " GFLOPS")
</code></pre>
<pre><code class="lang-auto">5.425337762571871  secs
16212.986927542423  GFLOPS
now using CUDA Graphs...
5.438446355983615  secs
16173.90785243318  GFLOPS
</code></pre>
<p>I will follow up to see if there is an explanation for the less than expected performance.</p>

username: eqy
cooked: <p>Ah, I think we might be undercounting the number of operations e.g., <a href="https://math.stackexchange.com/questions/3512976/proof-of-of-flops-in-matrix-multiplication" class="inline-onebox" rel="noopener nofollow ugc">linear algebra - Proof of # of FLOPs in Matrix Multiplication - Mathematics Stack Exchange</a></p>
<p>which would suggest the correct number is something closer to 2<em>size</em>size*(size-1). Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> for the pointer</p>

username: ha2ka
cooked: <p>That was a trivial mistake… I wasn’t aware of it at all lol.</p>
<p>Thank you for your patient discussions.</p>