Is it possible to figure out how much device memory a pjit-compiled function will use

[hr0nix]

I want to experiment with different sharding strategies for my model and figure out which one results in less memory consumption. I think the easiest way to go about it is to get memory usage information directly from a pjit-compiled function, but I'm not sure if it's possible.

Any pointers or alternative ideas of how to go about such investigation will be much appreciated.

[jakevdp]

You can use the ahead-of-time cost analysis functionality, something like this:

import jax

def f(x):
  return jax.numpy.linalg.svd(x)[1]

x = jax.ShapeDtypeStruct((10000, 30000), 'float32')

print(jax.jit(f).lower(x).cost_analysis())

{'utilization operand 2 {}': 2.0,
 'utilization operand 3 {}': 1.0,
 'utilization operand 1 {}': 3.0,
 'bytes accessed output {}': 90003.0,
 'bytes accessed operand 1 {}': 40003.0,
 'flops': 10000.0,
 'utilization operand 5 {}': 1.0,
 'bytes accessed operand 4 {}': -1.0,
 'utilization operand 4 {}': 1.0,
 'bytes accessed operand 2 {}': 39999.0,
 'bytes accessed operand 3 {}': -1.0,
 'bytes accessed operand 6 {}': -1.0,
 'utilization operand 6 {}': 1.0,
 'utilization operand 0 {}': 15.0,
 'bytes accessed operand 0 {}': 10011.0,
 'optimal_seconds': -1.0,
 'bytes accessed': 180019.0,
 'bytes accessed operand 5 {}': -1.0}

Read more at https://jax.readthedocs.io/en/latest/aot.html#debug-information-and-analyses-when-available

[hr0nix]

Thanks for the answer @jakevdp !

As far as I understand, the particular stats that cost_analysis returns aren't documented anywhere and can change in future. Is there a way to better understand their current semantics? In particular, what does "bytes accessed" represent?

[hansbuehler]

Much appreciated @jakevdp.
Could I ask: even with the current undocumented structure

1. How do I read the result? At the moment I am taking the max of all "bytes accessed" integers. Is that correct?
2. Does this analysis include the memory required for the input and output data?

[jakevdp]

I don't know of any documentation of these outputs that address those questions.

[schance995]

Is the output of cost_analysis documented anywhere in XLA?

[jakevdp]

I don't know of any documentation of cost_analysis.