Diminishing returns to JIT Speedup

I am working on speeding up a model with JIT compilation, as outlined here. I noticed in my own use case that there were diminishing returns to the JIT speedup as the problem size grew. I repeated the experiment with the simple Gaussian model at the link above, changing the number of observations from 100 to 1000 to 10000 and so on. At low problem sizes, the 2x speedup in the docs is replicated. At high problem sizes, though, the JIT-compiled SVI is no faster than the normal SVI.

Intuitively, the opposite should occur - compilation overhead should be amortized at high problem sizes. Has anyone observed this before?

The primary way the PyTorch JIT speeds up models is by compiling away Python overhead from tensor operations; it does very little optimization beyond that. As the sizes of the underlying tensor operations in a model increase, the amount of overhead relative to the amount of time actually spent performing the low level numerical operation decreases, and the performance increase from the JIT shrinks accordingly.

1 Like