error when installing jax.

[lynshwoo2022]

when i run pip install --no-build-isolation transformer_engine[jax] -v
i get error info as:

Looking in indexes: http://aptps.cloud.hihonor.com/repository/pypi-proxy/simple
  Requirement already satisfied: pybind11 in /opt/nas/p/conda/envs/zls-vlm/lib/python3.10/site-packages (2.13.6)
  WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
  running bdist_wheel
  running build
  running build_ext
  /usr/bin/gcc -Wno-unused-result -Wsign-compare -DNDEBUG -fwrapv -O2 -Wall -fPIC -O2 -isystem /opt/nas/p/conda/envs/zls-vlm/include -fPIC -O2 -isystem /opt/nas/p/conda/envs/zls-vlm/include -fPIC -I/opt/nas/p/conda/envs/zls-vlm/include/python3.10 -c flagcheck.cpp -o flagcheck.o -std=c++17
  building 'transformer_engine_jax' extension
  error: unknown file type '.cu' (from 'csrc/utils.cu')
  error: subprocess-exited-with-error
  
  × python setup.py bdist_wheel did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  full command: /opt/nas/p/conda/envs/zls-vlm/bin/python -u -c '
  exec(compile('"'"''"'"''"'"'
  # This is <pip-setuptools-caller> -- a caller that pip uses to run setup.py
  #
  # - It imports setuptools before invoking setup.py, to enable projects that directly
  #   import from `distutils.core` to work with newer packaging standards.
  # - It provides a clear error message when setuptools is not installed.
  # - It sets `sys.argv[0]` to the underlying `setup.py`, when invoking `setup.py` so
  #   setuptools doesn'"'"'t think the script is `-c`. This avoids the following warning:
  #     manifest_maker: standard file '"'"'-c'"'"' not found".
  # - It generates a shim setup.py, for handling setup.cfg-only projects.
  import os, sys, tokenize, traceback
  
  try:
      import setuptools
  except ImportError:
      print(
          "ERROR: Can not execute `setup.py` since setuptools failed to import in "
          "the build environment with exception:",
          file=sys.stderr,
      )
      traceback.print_exc()
      sys.exit(1)
  
  __file__ = %r
  sys.argv[0] = __file__
  
  if os.path.exists(__file__):
      filename = __file__
      with tokenize.open(__file__) as f:
          setup_py_code = f.read()
  else:
      filename = "<auto-generated setuptools caller>"
      setup_py_code = "from setuptools import setup; setup()"
  
  exec(compile(setup_py_code, filename, "exec"))
  '"'"''"'"''"'"' % ('"'"'/tmp/pip-install-iboyvwr3/transformer-engine-jax_eda60e55d6bb43ef9d343953f3e271f1/setup.py'"'"',), "<pip-setuptools-caller>", "exec"))' bdist_wheel -d /tmp/pip-wheel-z6lvdt6w
  cwd: /tmp/pip-install-iboyvwr3/transformer-engine-jax_eda60e55d6bb43ef9d343953f3e271f1/
  Building wheel for transformer_engine_jax (setup.py) ... error

[ksivaman]

It looks like gcc (instead of nvcc) is being used to compile the cuda files in the JAX extensions. Do you have CUDA toolkit installed?

[ksivaman]

Note: In release_v2.4 of TE, the framework extensions will be purely in c++, and thus won't require a cuda toolkit or nvcc for compilation.

[lynshwoo2022]

It looks like gcc (instead of nvcc) is being used to compile the cuda files in the JAX extensions. Do you have CUDA toolkit installed?

yes i already have cudatoolkit and gcc/g++.
problem solved by installing system cudnn

[lynshwoo2022]

It looks like gcc (instead of nvcc) is being used to compile the cuda files in the JAX extensions. Do you have CUDA toolkit installed?

yes i already have cudatoolkit and gcc/g++. problem solved by installing system cudnn

further more, if this problem remains, i'll get same errors like #355 (comment).
adding cudnn path of pytorch dosen't work, system cudnn must be installed in my case

[ksivaman]

Yes, for current release versions, system cuDNN is a requirement