Error with the env setup

[Jingyu6]

Hi,

I followed the setup procedure as described in the readme and no error showed up after completing all. However, when I try the following:

>>> from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
>>> from mad.model.model import LanguageModel, AutoEncoder
No ROCm runtime is found, using ROCM_HOME='/opt/rocm'
No module named 'causal_attention_cuda'
Segmentation fault

The same error appears when I run:

python -m train --task in-context-recall --layers mh-attention swiglu mh-attention swiglu

I'm running the script on a single A100. Any idea on how to get it to work?

Best,
Jingyu