An module error when running cm_train.py

[RICKand-MORTY]

OS: Ubuntu20.04 x86_64
Python: 3.11.5
flash-attn: 2.3.6/2.1.2.post3

I use pip to install flash-attn module and the version is 2.3.6.When I run python cm_train to test, I find python throw this error.I change my flash-attn version to 2.1.2.post3,but the error is still exist.
The following is my output:
raceback (most recent call last): File "/root/autodl-tmp/conmodel/scripts/cm_train.py", line 178, in <module> main() File "/root/autodl-tmp/conmodel/scripts/cm_train.py", line 54, in main model, diffusion = create_model_and_diffusion(**model_and_diffusion_kwargs) # 创建model和diffusion ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/autodl-tmp/conmodel/scripts/../cm/script_util.py", line 76, in create_model_and_diffusion model = create_model( ^^^^^^^^^^^^^ File "/root/autodl-tmp/conmodel/scripts/../cm/script_util.py", line 140, in create_model return UNetModel( ^^^^^^^^^^ File "/root/autodl-tmp/conmodel/scripts/../cm/unet.py", line 612, in __init__ AttentionBlock( File "/root/autodl-tmp/conmodel/scripts/../cm/unet.py", line 293, in __init__ self.attention = QKVFlashAttention(channels, self.num_heads) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/autodl-tmp/conmodel/scripts/../cm/unet.py", line 344, in __init__ from flash_attn.flash_attention import FlashAttention ModuleNotFoundError: No module named 'flash_attn.flash_attention'

I guess 'flash_attention' may change to elsewhere during update.

[RICKand-MORTY]

suitable flash-attn version: 1.0.2 or 0.2.8
I finally solve the problem.I changed my flash-attn version to 1.0.2 and 0.2.8 to test and both of these version is suitable.Besides, when I trained the model, I found that mpiexec -n 8 would cause deadlock,which I haven't find the reason yet.
Therefore, I delete mpiexec -n -8 from command to train.