Safely import TELinearAdapter

[marcromeyn]

Problem

The file src/megatron/hub/peft/lora.py was importing TELinearAdapter directly in the import statements:

from megatron.hub.peft.lora_layers import LinearAdapter, LoRALinear, TELinearAdapter, patch_linear_module

This causes an ImportError when transformer-engine is not installed, even though the code properly handles the optional dependency elsewhere in the file.

Solution

Move the TELinearAdapter import to the try-except block where transformer-engine is imported:

# Before
from megatron.hub.peft.lora_layers import LinearAdapter, LoRALinear, TELinearAdapter, patch_linear_module

# After
from megatron.hub.peft.lora_layers import LinearAdapter, LoRALinear, patch_linear_module

# Later in the file
try:
    import transformer_engine.pytorch as te
    from megatron.hub.peft.lora_layers import TELinearAdapter
    
    HAVE_TE = True
except ImportError:
    te = None
    TELinearAdapter = None
    HAVE_TE = False

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[marcromeyn]

/ok to test 1e499d9

[marcromeyn]

/ok to test 1e499d9

[ko3n1g]

What do you think about raising a user-friendly error inside of LoRA explaining that they need TE to proceed?

[ko3n1g]

Also, what do you think about using nemo's safe import method?

[marcromeyn]

/ok to test 3864904

[ko3n1g]

Thanks for refactoring this!

[pruprakash]

QA RCCA Analysis

1. Fix Reference

• PR: Safely import TELinearAdapter #103 - Safely import TELinearAdapter
• Issue: PR is the fix itself (no linked issue)

2. Root Cause

TELinearAdapter was imported directly in import statements, causing ImportError when transformer-engine is not installed. The code properly handled optional dependency elsewhere but not in the import block.

3. Trigger Configuration

• Using LoRA without transformer-engine installed
• Importing megatron.hub.peft.lora

4. Nature of the Bug

Classification: CODE BUG - ImportError when transformer-engine not installed

5. Existing Test Coverage

In Fix PR: NO - No test files

In NMFW Tests: Not specifically covering optional TE import

6. Coverage Assessment

Test Type	Exists	Covers Bug
Fix PR tests	NO	N/A
NMFW regression tests	NO	N/A

7. New Regression Test

NOT ADDED - Optional dependency import fix.

8. Conclusion

Verdict: PARTIAL COVERAGE - No dedicated test for optional TE import handling.