How do I profile layer-by-layer Jax performance?

[abhinavgoel95]

I am looking to break down Jax performance, layer-by-layer. Is there a way to place a breakpoint during forward propagation?

[YouJiacheng]

Unfortunately, in general JAX might do cross layer fusion, exact performance is hard to measure. However, you can use host_callback to measure it roughly: check #10908 (comment)
i.e. you can print the time difference between input ready and output ready of a layer, but you cannot ensure there isn't any other operation irrelevant to the layer executed on GPU during this period.

[abhinavgoel95]

Thank you for your reply.

[hawkinsp]

The main tool we tend to use is the tracing/profiling infrastructure described here: https://jax.readthedocs.io/en/latest/profiling.html

As YouJiacheng says, we don't compile layers in isolation, so it's not always easy to point to a single thing on a trace which exactly corresponds to a layer.

One thing that will certainly work would be to jit-compile that layer in isolation and write a microbenchmark that exercises it. But I'd certainly start with the tracing tools.

[abhinavgoel95]

Thank you for your reply.

[sharadmv]

Also consider using jax.named_scope or jax.named_call to annotate your layers so they appear in the trace viewer.