3 ms·

All diffusion models are quite inefficent due to running in PyTorch eager mode (with torch.compile being kinda janky in practice on 2.0/2.1). I would be more i

All diffusion models are quite inefficent due to running in PyTorch eager mode (with torch.compile being kinda janky in practice on 2.0/2.1). 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.

brucethemoose2 3y ago
(And for reference, AITemplate roughly doubles SD 1.5's speed. Not sure if thats good or if Paella would have even more room for auto optimization).
londons_explore 3y ago
A framework could make the programmer think they are in eager mode, while actually making graphs behind the scenes. 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.
- brucethemoose2 3y ago
  ... Kinda like torch.compile? 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.
saiojd 3y ago
I find you post intriguing. What would you say are the major janks with torch.compile, and what issues are addressed by TVM/AITemplate but not by torch.compile? EDIT: If I understand correctly these libraries target deployment performance, while torch.compile is also/mostly for training performance?
- brucethemoose2 3y ago
  - The gain in stable diffusion is modest (15%-25% last I checked?) - 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#issuecomment-1608738743 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.
  - saiojd 3y ago
    Thanks for the info. I didn't know about the TomeSD stuff, really interesting. Why do you think that AITemplate is so much faster?

brucethemoose2 3y ago

(And for reference, AITemplate roughly doubles SD 1.5's speed. Not sure if thats good or if Paella would have even more room for auto optimization).

londons_explore 3y ago

A framework could make the programmer think they are in eager mode, while actually making graphs behind the scenes. 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.

brucethemoose2 3y ago
... Kinda like torch.compile? 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.

saiojd 3y ago

I find you post intriguing. What would you say are the major janks with torch.compile, and what issues are addressed by TVM/AITemplate but not by torch.compile? EDIT: If I understand correctly these libraries target deployment performance, while torch.compile is also/mostly for training performance?

brucethemoose2 3y ago
- The gain in stable diffusion is modest (15%-25% last I checked?) - 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#issuecomment-1608738743 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.
- saiojd 3y ago
  Thanks for the info. I didn't know about the TomeSD stuff, really interesting. Why do you think that AITemplate is so much faster?