FlexGen icon indicating copy to clipboard operation
FlexGen copied to clipboard

Pytorch-Lightning strategy

Open k-sparrow opened this issue 2 years ago • 2 comments

Awesome work!

Any plans on having this as a strategy plugin for pytorch-lightning? (like DDP/DeepSpeed/ColossalAI) (https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html)

k-sparrow avatar Feb 22 '23 14:02 k-sparrow

Hey we'll be happy to help if anyone wants to undertake this, keeping in mind that this is intended to optimize inference.

In that sense it might also be a good strategy for Lightning Fabric, so you can just load the model, set it up with the strategy and call forward /cc @awaelchli

lantiga avatar Feb 24 '23 10:02 lantiga

Thanks for your interest! This looks cool but currently, this repo is specialized for some specific transformer models. It cannot be a general solution that works out of the box for any model. After we develop a more general interface and backend, we can think about the integration more seriously. I am happy to help with this as well.

Ying1123 avatar Feb 26 '23 03:02 Ying1123