Backword MFU Calculation Too Low Using Pytorch Flash AttentionV2

Uwwal: <p>Hello everyone,</p>
<p>I am currently working on calculating the Model-Flo-Utilization (MFU) using Flash AttentionV2 and have encountered a couple of issues that I hope to get some insights on:</p>
<ol>
<li><strong>Discrepancy with Paper-Reported MFU</strong>: Despite referring to many docs including the use of sdpa, the calculation of flops, we are unable to achieve the same MFU values reported in the paper. Could there be specific implementation details or considerations that I might be overlooking?</li>
<li><strong>Higher MFU values for large model training</strong>: In our experiments, we have found that the MFU values for large model training are higher than the values we calculated. Are there common pitfalls or problems in MFU calculations that could lead to this discrepancy? Are there specific aspects that might explain why we observed lower than expected MFU values?</li>
</ol>
<p>The following figure shows our test results with data type bfloat16, batchsize = 32, num head = 32, running on the A800, pytorch=2.3.1, cuda=11.8 (Image 1)</p>
<p>We can see that the running time of the test results at seqlen=1024 and dims=128 is close to the results in <a href="https://github.com/pytorch/pytorch/pull/105602#issuecomment-1755453126" rel="noopener nofollow ugc">flash attention v2 pr</a>. (Image 2)</p>
<p>But converting mfu to flops we are far from the results in the <a href="https://tridao.me/publications/flash2/flash2.pdf" rel="noopener nofollow ugc">v2 paper</a>, which is very confusing to me. (Image 3)</p>
<p>(I’m sorry because new users can only upload one image so I’ve combined all three in one)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/0/101def7fcc3102009318a623a90d4b94f3b776e9.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/101def7fcc3102009318a623a90d4b94f3b776e9" title="Image 1, Image 2, Image 3"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/101def7fcc3102009318a623a90d4b94f3b776e9_2_356x499.jpeg" alt="Image 1, Image 2, Image 3" data-base62-sha1="2izM7ss4TyX77yrcgVhnSbBVIcF" width="356" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/101def7fcc3102009318a623a90d4b94f3b776e9_2_356x499.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/101def7fcc3102009318a623a90d4b94f3b776e9_2_534x748.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/0/101def7fcc3102009318a623a90d4b94f3b776e9_2_712x998.jpeg 2x" data-dominant-color="B5B6B2"><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">Image 1, Image 2, Image 3</span><span class="informations">1920×2695 295 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>test code:</p>
<pre data-code-wrap="Python"><code class="lang-Python">import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend
from torch.cuda import amp

TFLOPS = 1000000000000
A800_PEAK_FLOPS = 312 * TFLOPS

LOOP_TIME = 10

batch_size = 32
num_heads = 32

class SDPA(torch.nn.Module):
    def __init__(self):
        super(SDPA, self).__init__()

    def forward(self, Q, K, V):
        with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            output = F.scaled_dot_product_attention(Q, K, V, is_causal=False)

        return output


def cal_flops_attn_fwd(seq_len, dims):
    flops = 0
    # scores = torch.matmul(Q, K.transpose(-2, -1))
    # [s h] * [h s] = [s s] -&gt; 2 * s * s * h
    flops += 2 * seq_len * seq_len * dims

    # output = torch.matmul(attention_weights, V)
    # [s s] * [s h] = [s h] -&gt; 2 * s * h * s
    flops += 2 * seq_len * dims * seq_len
    return flops * batch_size * num_heads


def cal_flops_attn_bwd(seq_len, dims):
    return 2 * cal_flops_attn_fwd(seq_len, dims)


if __name__ == "__main__":
    def measure_time(func):
        total_fwd_time = 0
        total_bwd_time = 0

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)

        with amp.autocast(enabled=True, dtype=torch.bfloat16):
            for i in range(LOOP_TIME):
                start.record(stream=torch.cuda.current_stream())
                r = func()
                end.record(stream=torch.cuda.current_stream())
                end.synchronize()

                fwd_time_ = start.elapsed_time(end)

                start.record(stream=torch.cuda.current_stream())
                a = r.sum()
                a.backward()
                end.record(stream=torch.cuda.current_stream())

                end.synchronize()

                bwd_time_ = start.elapsed_time(end)

                if i == 0:
                    continue

                total_fwd_time += fwd_time_
                total_bwd_time += bwd_time_

        torch.cuda.empty_cache()
        return total_fwd_time + total_bwd_time, total_fwd_time, total_bwd_time


    for seq_len in (1024, 1536, 2048):
        for dimension in (64, 128):
            dtype = torch.bfloat16
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

            query = torch.rand(batch_size, num_heads, seq_len, dimension, dtype=dtype, device=device,
                               requires_grad=True)
            key = torch.rand(batch_size, num_heads, seq_len, dimension, dtype=dtype, device=device, requires_grad=True)
            value = torch.rand(batch_size, num_heads, seq_len, dimension, dtype=dtype, device=device,
                               requires_grad=True)

            warmup_iter = 3

            # warm_up
            for _ in range(warmup_iter):
                res = torch.matmul(query, key.transpose(-2, -1))

            torch.cuda.synchronize()
            torch.cuda.empty_cache()

            sdpa = SDPA()

            with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
                total_time, fwd_time, bwd_time = measure_time(
                    lambda: sdpa(query, key, value)
                )

            total_fwd_flops = cal_flops_attn_fwd(seq_len, dimension) * (LOOP_TIME - 1)
            total_bwd_flops = cal_flops_attn_bwd(seq_len, dimension) * (LOOP_TIME - 1)

            total_flops = total_fwd_flops + total_bwd_flops

            fwd_mfu = (total_fwd_flops / A800_PEAK_FLOPS) / fwd_time * 1000
            bwd_mfu = (total_bwd_flops / A800_PEAK_FLOPS) / bwd_time * 1000
            total_mfu = (total_flops / A800_PEAK_FLOPS) / total_time * 1000

            print(f"sdpa sum time: fwd time: {fwd_time}, bwd time: {bwd_time}, total_time: {total_time}")

            print(f"sdpa avg time: fwd time: {fwd_time / (LOOP_TIME - 1)}, "
                  f"bwd avg time: {bwd_time / (LOOP_TIME - 1)}, "
                  f"total avg time: {total_time / (LOOP_TIME - 1)}")

            print(f"sdpa mfu({seq_len}, {dimension}): fwd {fwd_mfu}, bwd {bwd_mfu}, {total_mfu}")

</code></pre>
<p>Any guidance or advice on these issues would be greatly appreciated. Thank you!</p>
<p>Best regards,<br>
Uwwal</p>

xgbj: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, could you take a look at this issue?</p>