Using torchaudio alongside TransformerEngine

[tylerweitzman]

Hi,
I do not know how to use torchaudio alongide TransformerEngine. The NGC docker nvcr.io/nvidia/pytorch:23.04-py3 doesn't come installed with torchaudio. If I do a normal torch install I lose CUDA 12.1. So, I'm using torch-nightly instead with cuda 12.1: pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121

But then when I import import transformer_engine.pytorch as te I get this error:
ImportError: /usr/local/lib/python3.8/dist-packages/transformer_engine_extensions.cpython-38-x86_64-linux-gnu.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE

By the way, does TransformerEngine support Python 3.11 or do I have to use Python 3.8?

Thank you

[timmoon10]

Are you rebuilding Transformer Engine after reinstalling PyTorch? If not, any C++ differences in PyTorch at build time and run time could lead to your linker errors.

Alternatively, it seems simpler to run pip3 install torchaudio within the NGC container. If that is doing something drastic like reinstalling PyTorch, you could try something like:

pip3 install kaldi_io SoundFile  # https://github.com/pytorch/audio/blob/main/requirements.txt
pip3 install --no-deps torchaudio

I don't think there is anything preventing you from running TE with Python 3.11, but it just depends on what environment you've installed PyTorch.

[hiddefolkertsma]

I've had a similar issue, using nvcr.io/nvidia/pytorch:23.06-py3. Installing torchaudio nightly with torch nightly also installed a lower CublasLT version (120100 vs 120103 that ships with the container), which meant I couldn't use FP8 on Ada:

TransformerEngine/transformer_engine/paddle/fp8.py

Line 27 in 7a30ba4

return False, "CublasLt version 12.1.3.x or higher required for FP8 execution on Ada."

I solved it by installing torchaudio from source like so:

git clone https://github.com/pytorch/audio
cd audio
pip install -v -e . --no-use-pep517

tested with

torch                   2.1.0a0+4136153
torchaudio              2.1.0a0+6324462
transformer-engine      0.12.0.dev0+7a30ba4