I would be more interested to see Paella vs SD running on a ML compiler framework, like TVM or AITemplate. Maybe one or the other is more amenable to optimization.
I would be more interested to see Paella vs SD running on a ML compiler framework, like TVM or AITemplate. Maybe one or the other is more amenable to optimization.
The trick is to simply not do any calculations till the last possible moment - ie. the time the program tries to convert the finished image to a jpeg. Only at that point do you compile the graph and run the actual computation on the GPU.
You then also cache the graph, so that the compilation step can be avoided if the program tries to do the same computation again with different data.
That approach is limited though. AITemplate and TVM take a looong time to compile and produce standalone executable files, hence the gains are much larger than torch triton.
EDIT: If I understand correctly these libraries target deployment performance, while torch.compile is also/mostly for training performance?
- Torch 2.0 only supports static inputs. In actual usage scenarios, this means frequent lengthy recompiles.
- Eventually, these recompiles will overload the compilation cache and torch.compile will stop functioning.
- Some common augmentations (like TomeSD) break compilation, force recompiles, make compilation take forever, or kill the performance gains.
- There are othdr miscellaneous bugs, like compilation freezing the Python thread and causing networking timeouts in web UIs, or errors with embeddings.
- Dynamic input in Torch 2.1 nightly fixes many of these issues, but was only maybe working a week ago? See https://github.com/pytorch/pytorch/issues/101228#issuecommen...
- TVM and AITemplate have massive performance gains. ~2x or more for AIT, not sure about an exact number for TVM.
- AIT supported dynamic input before torch.compile did, and requires no recompilation after the initial compile. Also, weights (models and LORAs) can be swapped out without a recompile.
- TVM supports very performant Vulkan inference, which would massively expand hardware compatibility.
Note that the popular SD Web UIs don't support any of this, with two exceptions I know of: VoltaML (with WIP AIT support) and the Windows DirectML fork of A1111 (which uses optimized ONNX models, I think). There is about 0% chance of ML compilation support in A1111, and the HF diffusers UIs are less bleeding edge and performance/compatibility focused.
And yes, triton torch.compile is aimed at training. There is an alternative backend (Hidet) that explicitly targets inference, but it does not work with Stable Diffusion yet.