Compute Flop Utilization in PyTorch

mfu_compute.py

import torch
from torch.utils.flop_counter import FlopCounterMode
from triton.testing import do_bench

def get_flops_achieved(f):
    flop_counter = FlopCounterMode(display=False)
    with flop_counter:
        f()
    total_flops = flop_counter.get_total_flops()
    ms_per_iter = do_bench(f)
    iters_per_second = 1e3/ms_per_iter
    print(f"{iters_per_second * total_flops / 1e12} TF/s")

from torchvision.models import resnet18
model = resnet18().cuda().half()
inp = torch.randn(128, 3, 224, 224, device='cuda', dtype=torch.half)
get_flops_achieved(lambda: model(inp).sum().backward())

compiled_model = torch.compile(model)
get_flops_achieved(lambda: compiled_model(inp).sum().backward())

I have just tried this for the first time and got the same error on the torch.compile segment.

pt=2.2.1+cu121, cuda=12.1,

@Chillee

Actually two issues here in some sense.

1. the flop counter actually isn't really intended to work when running a torch.compile'd model :P This is mostly "fine" because the flops of a compiled model is generally identical to the flops of an uncompiled model. But the interposition points we use (i.e. the dispatcher) are not guaranteed to trigger when running under the compiler (for example, if you're using cudagraphs).
2. The actual error that you're running into, which is fixed by pytorch/pytorch#123768

EDIT: To clarify, the intended usage of FlopCounterMode with a compiled model is to run it on an uncompiled version of the model.

super! Thank you, Horace

Do you want me to open an Issue to document FlopCounterMode? as its docs are missing so we users don't know what is kosher to use it for and what not

as its docs are missing so we users don't know what is kosher to use it for and what not

I was gonna do it for PyTorch 2.3 release but I didn't end up getting around to it 😭

Yeah please do open an issue :)