search

OpenNLPLab / lightning-attention

Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models
MIT License
182 stars 15 forks source link

The code stuck when running example_lightning_attn.py #9

Closed yang-yk closed 6 months ago

yang-yk commented 6 months ago

Thanks for your excellent work. I have created the running environment as you described in readme.md. However, when I run the ./examples/ops/ example_lightning_attn.py, the code will get stuck in the lightning_attn_func function. Further debugging revealed that the problem lies in the _fwd_kernerl[grid] function. Could you please help me solve this problem? The Pytorch and Triton versions are 2.0.1 and 2.0.0, respectively and the GPU is Nvidia 3090Ti with Cuda 11.7. Could you please help me solve this problem? Looking forward to hearing from you. Thank you very much.

Doraemonzzz commented 6 months ago

Hello, have you try this?

pip install triton-nightly==2.1.0.dev20230728172942 --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
yang-yk commented 6 months ago

Hello, have you try this?

pip install triton-nightly==2.1.0.dev20230728172942 --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/

Yeah. It did work. Something went wrong when I installed the Triton package. The problem was solved when I reinstalled Triton. Many thanks for your help.