3 ms·

Related to this question, can someone explain the design goal of torch.jit to me? Is it supposed to boost performance or just give a means to export models? I f

Related to this question, can someone explain the design goal of torch.jit to me? Is it supposed to boost performance or just give a means to export models? I found my jitted code ran slower than interpreted pytorch, and the latter despite its asynchronous nature spent most of its time waiting for the next gpu kernel to start. Having got a working torch model on cpu, what's the best path to actually making it run as fast as I feel it has potential to?

bertr4nd 5y ago
It’s both. torch.jit started life as an optimizer. I think fusion of pointwise kernels on GPU - which we finally extended to CPU in this release - was one of the early wins via jit. But at some point it became a model export format for production environments that can’t use CPython for performance reasons. I’m surprised that you’re seeing worse performance with jit. It sometimes takes 20-ish iterations for the jit to “settle down” but I’d expect roughly equal performance at worst. If you can share a repro, I’d be happy to take a quick look if you file an issue on GitHub. (I’m @bertmaher there)
- tubby12345 5y ago
  you know what i still don't understand? what's taking so long to warm up? i see that there are graph passes that run to do various things at the TS IR level, but I don't see any stats being collected (on shapes) or something like that that then inform further optimization.
  - bertr4nd 5y ago
    There’s a “profiling graph executor” that records shapes and then hands them off to a fusion compiler. The profiling executor re-specializes on every new shape it sees, but stops at 20 re-specialization. We’re working on eliminating the dependence on shape specialization right now, since it’s kind of an unfortunate limitation for some workloads.
bertr4nd 5y ago
Oh also to answer your “as fast as possible” question: usually you’ll get the best performance by exporting your model to a perf-tuned runtime. We’ve seen really good results with TensorRT and (for transformers) FasterTransformer. I’ve also seen good results with ONNX runtime. Staying within pytorch, we recently added torch.jit.optimize_for_inference (I think it’s in 1.10, though not entirely sure) that can apply a bunch of standard optimizations to a model and often provides some nice wins.
NegatioN 5y ago
Torch.jit shouldn't impact your performance positively or negatively in my experience. Although I've only used it on cpu. It's as far as I know just used for model exports. The nice thing about it though, is that you can embed native python code (that's compiled to c++) into the model artifact. It's allowed us to write almost all of the serving logic of our models very closely to the model code itself, giving a better overview than having the server logic written in a separate repo. The server we use on top of this can be pretty "dumb", and just funnel all inputs to the model, which the Python code determines what to do with. As for model speedups, maybe you should look into quantization? I also find that there's usually lots of low hanging fruit if you go over code and rewrite to quicker ops which are mathematically equivalent, but allocate less memory, or do less ops.
m_ke 5y ago
It makes it possible to lift your model out of python while handling programming constructs like loops and if statements (see torch.jit.script). It also makes it possible to sidestep the GIL and remove the overhead of launching kernels from python, which only really makes a noticeable difference with models that queue up a lot of small operations on the GPU. (LSTMs are an example of where this would make a difference https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/ https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscrip...)

bertr4nd 5y ago

It’s both. torch.jit started life as an optimizer. I think fusion of pointwise kernels on GPU - which we finally extended to CPU in this release - was one of the early wins via jit. But at some point it became a model export format for production environments that can’t use CPython for performance reasons. I’m surprised that you’re seeing worse performance with jit. It sometimes takes 20-ish iterations for the jit to “settle down” but I’d expect roughly equal performance at worst. If you can share a repro, I’d be happy to take a quick look if you file an issue on GitHub. (I’m @bertmaher there)

tubby12345 5y ago
you know what i still don't understand? what's taking so long to warm up? i see that there are graph passes that run to do various things at the TS IR level, but I don't see any stats being collected (on shapes) or something like that that then inform further optimization.
- bertr4nd 5y ago
  There’s a “profiling graph executor” that records shapes and then hands them off to a fusion compiler. The profiling executor re-specializes on every new shape it sees, but stops at 20 re-specialization. We’re working on eliminating the dependence on shape specialization right now, since it’s kind of an unfortunate limitation for some workloads.

bertr4nd 5y ago

Oh also to answer your “as fast as possible” question: usually you’ll get the best performance by exporting your model to a perf-tuned runtime. We’ve seen really good results with TensorRT and (for transformers) FasterTransformer. I’ve also seen good results with ONNX runtime. Staying within pytorch, we recently added torch.jit.optimize_for_inference (I think it’s in 1.10, though not entirely sure) that can apply a bunch of standard optimizations to a model and often provides some nice wins.

NegatioN 5y ago

Torch.jit shouldn't impact your performance positively or negatively in my experience. Although I've only used it on cpu. It's as far as I know just used for model exports. The nice thing about it though, is that you can embed native python code (that's compiled to c++) into the model artifact. It's allowed us to write almost all of the serving logic of our models very closely to the model code itself, giving a better overview than having the server logic written in a separate repo. The server we use on top of this can be pretty "dumb", and just funnel all inputs to the model, which the Python code determines what to do with. As for model speedups, maybe you should look into quantization? I also find that there's usually lots of low hanging fruit if you go over code and rewrite to quicker ops which are mathematically equivalent, but allocate less memory, or do less ops.

m_ke 5y ago

It makes it possible to lift your model out of python while handling programming constructs like loops and if statements (see torch.jit.script). It also makes it possible to sidestep the GIL and remove the overhead of launching kernels from python, which only really makes a noticeable difference with models that queue up a lot of small operations on the GPU. (LSTMs are an example of where this would make a difference https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/ https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscrip...)