search

NVIDIA / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch
BSD 3-Clause "New" or "Revised" License
8.28k stars 1.38k forks source link

No module named 'fast_self_multihead_attn' #958

Open hecmay opened 3 years ago

hecmay commented 3 years ago

Hi, Iwas trying run the fast-MHA example here: https://github.com/NVIDIA/apex/tree/master/apex/contrib/multihead_attn

I got the following errors when running the performance:

Traceback (most recent call last):
  File "perf_test_multihead_attn.py", line 5, in <module>
    from apex.contrib.multihead_attn import SelfMultiheadAttn
  File "/home/apex/apex/contrib/multihead_attn/__init__.py", line 1, in <module>
    from .self_multihead_attn import SelfMultiheadAttn
  File "/home/apex/apex/contrib/multihead_attn/self_multihead_attn.py", line 9, in <module>
    from .fast_self_multihead_attn_func          import fast_self_attn_func
  File "/home/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py", line 2, in <module>
    import fast_self_multihead_attn
ModuleNotFoundError: No module named 'fast_self_multihead_attn'

The Apex packaged is installed with CPP and CUDA extension turned on. And I was not able to find the definition of this missing function anywhere.

jehoon1ee commented 3 years ago

Hi,I'm facing the same issue. Have you found the solution?

lendrick commented 3 years ago

Removing pyprof from my code fixed this for me.

sanjaychari commented 3 years ago

I am facing the same issue. I am not using pyprof in my code.

chebbyChefNEQ commented 3 years ago

Looks like apex doesn't build this extension by default.

https://github.com/NVIDIA/apex/blob/6b7e77b0c34ca2be4243d495423e17d49090b49e/setup.py#L371

https://github.com/NVIDIA/apex/blob/6b7e77b0c34ca2be4243d495423e17d49090b49e/setup.py#L301

adding --global-option="--fast_multihead_attn" when installing apex fixed the problem for me. (I'm running pyprof, and I get this exception during pyprof.init())