Speed up on GPUs

[nhartman94]

Hello! Following some very helpful feedback on pyhf GPU support in
#896
and
#1244
I tested out some limits w/ different #s of channels (and combination of histosys and shapesys NPs) to see the timing comparisons.

The following table shows the root finding time (in min), and I was a bit confused why the gpu wasn't speeding thing up, or what would be nice x-checks to do.
(The "zeros" entries on the table timed out after 15 hrs)

Ty so much!!

[kratsg]

One thing to note is that there is some (small) time required to transfer data over to the GPU. You should try benchmarking jax using something like a 5-point mu-scan just to understand the speed up. You might not have a significantly complex-enough workspace to benefit from GPU in some cases. Also, jax will jit the PDF (in both cases) which is going to contribute some time as well when comparing to other backends, but again, it really depends.

I have a (ATLAS-only) workspace that takes at least an hour to converge using the CPU, but about 40 seconds w/ a GPU. There's certainly a lot of factors involved to make this happen.

[nhartman94]

Hi Giordan,

Ty so much for the suggestion. I tried a 5 point mu scan (with pyhf.infer.intervals.upperlimit) with JAX backend + w/ and w/o the GPU, but I'm seeing a similar result that I'm not seeing a speed up with the gpu.

Does this mean my workspaces might not be complex-enough to get the GPU speed-up? Or is there anything else I can check?

Thx so much!

[matthewfeickert]

The following table shows the root finding time (in min), and I was a bit confused why the gpu wasn't speeding thing up, or what would be nice x-checks to do.

@nhartman94 Can you share with us the code that you're running? I assume that this is ATLAS Internal code, so can you create a private GitLab or GitHub repository and give the pyhf dev team permissions to it?

As @kratsg said, there's a lot of components and factors that go into fit time performance, so the more information we have the more we can help.

[nhartman94]

Thanks so much @matthewfeickert !

I put a (close to minimal) example of one of these workspaces at:
https://gitlab.cern.ch/hartman/pyhf-4b-timing . Please let me know what other information from me would be useful.

[nhartman94]

Hi - just to clarify, the numbers in my table are for table I included above is the time to find the the CLs = 0.05 mu value, so the #s in this table actually include (approximately) 40 hypothesis tests per number that I cite.