how to use your code in other models?

[zxgx]

Hi,

Your code seems pretty neat. I want to try your solution in other transformer models. Can I directly use another model name defined by huggingface as the ModelArguments?

[DachengLi1]

hi @zxgx thank you very much! The attention part should be generic to other autoregressive models: https://github.com/RulinShao/LightSeq/blob/main/lightseq/lightseq_async_attn.py#L436. You will only need to do a monkey patch to the attention. We didn't test it in other models, but let us know if you encounter any issues!

[RulinShao]

hi @zxgx thank you very much! The attention part should be generic to other autoregressive models: https://github.com/RulinShao/LightSeq/blob/main/lightseq/lightseq_async_attn.py#L436. You will only need to do a monkey patch to the attention. We didn't test it in other models, but let us know if you encounter any issues!

PS: we wrote the lightseq attention api in a similar way as FlashAttention, so one straightforward way is to find where the codes patch FlashAttention and replace FA with LightSeq attention that @DachengLi1 linked to.

[zxgx]

Great, thank you for your reply! 😄