AttributeError: module 'transformer_engine' has no attribute 'pytorch' when importing megatron.core

[xdobetter]

Issue Type

• Bug Report
• Feature Request
• Question
• Documentation
• Security
• Other

Description

Hello! First of all, thank you for your amazing work on this project.
While setting up the environment and running megatron-core / transformer-engine, I encountered the issue described below.
I’ve tried several possible fixes (such as changing Python versions, upgrading/downgrading PyTorch and Transformer Engine, etc.), but the problem still persists.
I’m not entirely sure whether this is caused by a version compatibility issue or an environment configuration problem, so I would really appreciate any advice or suggestions from the maintainers.
Thanks a lot for your time and support!

Traceback (most recent call last):
File "/root/miniconda3/envs/ct/megatron/lib/python3.10/site-packages/megatron/core/extensions/transformer_engine.py", line 94, in
class TELinear(te.pytorch.Linear):
AttributeError: module 'transformer_engine' has no attribute 'pytorch'

Environment

Python version: 3.10.15 (Anaconda)

CUDA version: 12.6

PyTorch version: 2.7.0 (cu121)

Transformer Engine version: 1.2.0

Megatron-Core version: 0.8.0

OS: Ubuntu 20.04 (Linux, x86_64)

Additional Context

Deliverable