search

NX-AI / xlstm

Official repository of the xLSTM.
GNU Affero General Public License v3.0
918 stars 66 forks source link

cuda version of mLSTM #8

Open bio-mlhui opened 3 weeks ago

bio-mlhui commented 3 weeks ago

will there be a cuda version of mLSTM? The mLSTM is O(n^2) complexity in both training and inference. I got OOM error in 24GB GPU, the context_lenght is set to 512, batch_size is 12, dim is 320, heads=10.

liujike commented 3 weeks ago

Did you have the erro when config the .yaml： RuntimeError: Error building extension 'slstm_HS64BS8NH1NS4DBfDRbDWbDGbDSbDAfNG4SA1GRCV0GRC0d0FCV0FC0d0

bio-mlhui commented 3 weeks ago

no, you can try delete this temporary folder and rebuild