NVIDIA / NVIDIA/TransformerEngine

CUDA extension loading should be modular

Open
#2,405 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

  1. Modular loading: Core extensions should load independently of optional extensions with feature flags (e.g. HAS_FA3).
  2. 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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.