Apex: No module named 'fast_self_multihead_attn'

[Hecmay]

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]

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

[lendrick]

Removing pyprof from my code fixed this for me.

[sanjaychari]

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

[robin5540]

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())