Register-level FLOPS counting in JAX programs ?

[mitkotak]

Hey I am on the hunt for a FLOPS counter that can count up the ADD/MUL happening at the register level since JAX's FLOP counter can be unreliable. I was wondering whether there's any way to do this using nsys-jax.

If yes, then is there a way to set up nsys-jax up outside a docker container i.e are there a list of instructions that I can send to my sysadmin to help them set it up ?

Thanks !

[mitkotak]

Just wanted to share an update. What I have been doing so far is running ncu and then counting up the FLOPS using instructions similar to here but as you can imagine, its extremely tedious to run.
Are there any alternatives out there that would be useful particularly for ML pipelines ? I am also open to some form of XLA/HLO level tracer.

Thanks for your input !

[olupton]

At the moment nsys-jax does not fill this gap, unfortunately. We understand that there is a need for improving the profiling tools here, and hope to support it later.

re: openxla/xla#10479, are there any particular patterns that are not cost-modelled correctly that are affecting you?