Skip to content

FLOPS counter for all your GPU benchmarking needs

Notifications You must be signed in to change notification settings

mitkotak/fast_flops

Repository files navigation

Fast FLOPS

... and where to find them 🐇

Ever needed to report FLOPs for your Torch/JAX code. We got you covered ! (More importantly, here's a survey on why bunnies flop)

Borrowed from https://gitlab.com/NERSC/roofline-on-nvidia-gpus/-/tree/roofline-hackathon-2020

Nsight Compute installation needed. Make sure the GPU counters are enabled

bunny

(Image Credit: Tuong Phung)

Workflow (Warning: Extremely clunky right now)

  • Define your function that you wanna profile

    • JAX

      SIZE = 500
      x = jax.random.normal(jax.random.PRNGKey(0), (SIZE, SIZE))
      y = jax.random.normal(jax.random.PRNGKey(1), (SIZE, SIZE))
      
      def func(x, y):
          return jnp.einsum("ij,jk->ik", x, y)
    • Torch

      SIZE = 500
      x = torch.randn(SIZE, SIZE).to(device='cuda')
      y = torch.randn(SIZE, SIZE).to(device='cuda')
      
      def func(x, y):
          return torch.einsum("ij,jk->ik", x, y)
  • Wrap the function in the flops_counter decorator

    from fast_flops import func_flops
    
    @flops_counter
    def func_flops(func, x, y):
        return func(x, y)
  • Let it run through JIT (Don't worry we have warmups to cleanse the JIT overhead) and execute !

    • JAX

      func = jax.jit(func)
      func_flops(func, x, y)
    • Torch

      func = torch.compile(func, fullgraph=True, mode='max-autotune')
      func_flops(func, x, y)
  • The pipeline can be executed using

    bash launch_profiler.sh examples/matmul/test_matmul_torch.py

    with the output looking something like

    Measured Time: 3.196823494576654e-05
    Measured GFLOP/s: 2353.1630879907675
    Measured FLOPS: 67108864.0

TODOs

  • Only turn on ncu during the hot loop in addition to nvtx
  • Axe run_profiler.py and postprocess.py. Ideally should work with python decorated_file.py
  • Plotting utilieis and csv plumbing

About

FLOPS counter for all your GPU benchmarking needs

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published