Unable to import transformer_engine.pytorch using TE v1.9.0

[snarayan21]

Hey, when installing transformer engine as pip install transformer-engine==1.9.0 I'm unable to simply import transformer_engine.pytorch -- it causes this error:

>>> import transformer_engine.pytorch
Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/transformer_engine/pytorch/__init__.py", line 22, in _load_library
    so_path = next(so_dir.glob(f"*.{extension}"))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3/dist-packages/transformer_engine/pytorch/__init__.py", line 34, in <module>
    _load_library()
  File "/usr/lib/python3/dist-packages/transformer_engine/pytorch/__init__.py", line 25, in _load_library
    so_path = next(so_dir.glob(f"*.{extension}"))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopIteration

I tried modifying the library load code but to no avail. This also happens with the newer post1 version as well. Is there something that's broken for installing from pypi? I'm able to install using the command below:

NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install git+https://github.com/NVIDIA/TransformerEngine.git@fc98961

This installs TE from this commit, which was immediately prior to the build system refactor commit here.

Would be great to resolve this so we can directly start adding TE as a pip dependency instead of having to install it via docker images. Thanks!

[snarayan21]

Resolved by doing pip install transformer-engine[pytorch]!

Thank you @ksivaman for the help!