Use AITemplate for faster inference?

[jacobbieker]

Detailed Description

https://github.com/facebookincubator/AITemplate

This is a new thing for PyTorch that compiles the models down to CUDA kernels for fast inference. There are still some caveats, but it might be worth it for larger models, like Power Perceiver, or possibly MetNet for making it quicker and cheaper inference.

Context

Possible Implementation

[peterdudfield]

Thanks @jacobbieker for this. How big is the MetNet model? How long is training taking?

[jacobbieker]

Metnet is quite large, haven't trained it yet, still doing the last debugging with the datapipe, but I'd expect it to be quite a bit faster than power perceiver and probably smaller?

[peterdudfield]

Metnet is quite large, haven't trained it yet, still doing the last debugging with the datapipe, but I'd expect it to be quite a bit faster than power perceiver and probably smaller?

When you start, do you mind saying how many parameters it is? Just to get some comparison to PP

[peterdudfield]

PP is 366 million parameters Metnet ~ 2.3 million, Metnet2 ~ 173 million

[jacobbieker]

Another thing for speeding up the loading of models can be HF's new safe tensors: https://github.com/huggingface/safetensors