NVIDIA / NVIDIA/TransformerEngine

Using system CUDA libraries

Open
#2,147 3 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug installation
Dominant language
Python
Stars
3.5k
Forks
831
Avg merge
3d 11h
Merged PRs (30d)
65

Description

Describe the bug

transformer-engine is currently searching for system CUDA binaries ( https://github.com/NVIDIA/TransformerEngine/blob/67fcc15255248a26be124de3854a47f84102f285/transformer_engine/common/__init__.py#L237). This is in conflict with Pytorch, which uses the CUDA Python packages (https://pypi.org/project/nvidia-cudnn-cu12/).

Steps/Code to reproduce bug

Tried using transformer-engine in a Docker container that did not have a system CUDA installed.

Expected behavior

transformer-engine should find the CUDA libraries inside the CUDA Python packages. Example:

import nvidia.cudnn
nvidia.cudnn.__file__
lib_path = os.path.join(nvidia.cudnn.__file__, "lib")

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in transformer_engine/common/init.py at the linked line around 237 and inspect how CUDA libraries are currently located. Compare that behavior with the nvidia.cudnn package path shown in the issue, then reproduce the Docker scenario without system CUDA and verify that transformer-engine finds the packaged libraries.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
56/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.