NVIDIA / NVIDIA/TransformerEngine
CUDA extension loading should be modular
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
Describe the bug
TE's CUDA extension loading uses an all-or-nothing approach that causes the entire transformer_engine.pytorch module to become broken when any optional dep (e.g. flash-attn) fails.
Root cause: Version checking and optional dependency imports happen during core CUDA extension loading. When an optional dependency check fails (e.g., flash-attn 2.8.3 > maximum supported 2.8.1), the exception is caught silently, leaving te.pytorch partially initialized.
Steps/Code to reproduce bug
# 1. Install TransformerEngine 2.9.0 and flash-attn 2.8.3
pip install transformer-engine==2.9.0
pip install flash-attn==2.8.3 # Note: TE 2.9.0 only supports <= 2.8.1
# 2. Try to use TransformerEngine
python3 -c "
import transformer_engine as te
print(f'TransformerEngine imported: {te.__version__}')
print(f'Has pytorch attr: {hasattr(te, \"pytorch\")}')
# Try to use functionality that doesn't need flash-attn
try:
linear = te.pytorch.Linear(128, 256)
print('✓ te.pytorch.Linear works')
except AttributeError as e:
print(f'✗ Error: {e}')
Output:
See https://github.com/NVIDIA/Megatron-LM/issues/696#issuecomment-3558275306
Expected behavior
- Modular loading: Core extensions should load independently of optional extensions with feature flags (e.g.
HAS_FA3). - Warnings instead of silent failures and hard version pinning
# If flash-attn version incompatible
UserWarning: flash-attn 2.8.3 is newer than tested maximum 2.8.1. This may work, but report issues if encountered.
Environment overview (please complete the following information)
Python: 3.10.12
Transformer Engine: 2.9.0
flash-attn: 2.8.3
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the documented import with Transformer Engine 2.9.0 and flash-attn 2.8.3, then trace transformer_engine.pytorch initialization and its CUDA extension and optional-dependency checks. Done means core functionality such as te.pytorch.Linear remains available when flash-attn is incompatible, while optional features expose flags and emit warnings instead of failing silently.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100