Slow fp16 GEMM on 4090

goodgame: <p>I opened a <a href="https://github.com/pytorch/pytorch/issues/123157" rel="noopener nofollow ugc">issue</a> recently in GitHub. The documented speed of fp16 GEMM on 4090 is 330 TFLOPS(with fp16 accumulation) and 165FLOPS(with fp32 accumulation). I run the following benchmark code on my 4090 machine.</p>
<pre><code class="lang-auto">import torch
print(torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction)
x = torch.randn(4096, 4096, device="cuda:0", dtype=torch.float16)
y = torch.randn(4096, 4096, device="cuda:0", dtype=torch.float16)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
# warmup 
for _ in range(1000):
    out = x.mm(y)
start.record()
for _ in range(1000):
    out = x.mm(y)
end.record()
end.synchronize()
time = start.elapsed_time(end)
print(f"TFLOPs {4096*4096*4096*2/1e9/(time/1000)}")
</code></pre>
<p>The output is:</p>
<pre><code class="lang-auto">True
TFLOPs 168.63091825340712
</code></pre>
<p>The default value of <code>allow_fp16_reduced_precision_reduction</code> is true. But Pytorch still uses fp32 as the accumulation mode for fp16 GEMM. Changing the accumulation mode to fp16 will double the performance on some GPU arch like 4090.</p>
<p>I want to submit a PR to change this behavior, making the option to actually use low-precision accumulation mode on arches that can benefit from it. Will this feature surprise people because of the change in precision? Maybe we should set the default value of this option to be false?</p>

ptrblck: <p>I don’t fully understand the idea as the reduced precision reduction in e.g. split-k kernels is already allowed (assuming a split-k kernel is used). A “pure” FP16 kernel using FP16 as the compute and accumulation type is not a good idea as convergence issues were shown in the past.</p>

goodgame: <p>Thanks for your reply! I have the same concerns with the precision problem. Mixed precision is important for NN training, but this feature is aimed at improving inference speed.</p>
<p>I’m currently optimizing the LLM inference speed on 4090. When changing the compute mode to FP16, there is very little perplexity increasing(0.009) in LLAMA2-7B, but 40% end2end speed up.</p>
<p>Can we add an option to support this kind of “pure” FP16 GEMM operation? When people turn on this option, they can have performance benefits at their own risk of precision issues.</p>

ptrblck: <aside class="quote no-group" data-username="goodgame" data-post="3" data-topic="200232">
<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/goodgame/48/68974_2.png" class="avatar"> goodgame:</div>
<blockquote>
<p>there is very little perplexity increasing(0.009) in LLAMA2-7B</p>
</blockquote>
</aside>
<p>I’m not too familiar with these metrics, but I assume this value is not considered large and would still yield the same outputs/predictions?</p>
<p>Could you create a feature request for the inference use case on GitHub so we could discuss it in more detail there?</p>

goodgame: <p>Sure, the new issue is <a href="https://github.com/pytorch/pytorch/issues/123558" class="inline-onebox" rel="noopener nofollow ugc">Support FP16 accumulation for faster LLM inference on 4090 like GPUs · Issue #123558 · pytorch/pytorch · GitHub</a>. Hope to hear voices from the community!</p>