Update the TE support to ModelOPT-PEFT

[jingyu-ml]

What does this PR do?

Type of change: new feature

Overview:

Added optional Transformer Engine (TE) support in modelopt/torch/peft/lora/plugins/megatron.py by importing TEColumnParallelLinear and TERowParallelLinear behind a try/except guard (HAS_TE).

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: Yes
• Did you add or update any necessary documentation?: Yes
• Did you update Changelog?:No

Additional Information

Summary by CodeRabbit

• New Features

  • Added Transformer Engine support for LoRA adapters
  • Extended LoRA to support NVFP4 quantization

• Tests

  • Added comprehensive tests validating Transformer Engine with LoRA and quantization
  • Added test coverage for combined quantization and LoRA workflows

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

• 🔍 Trigger a full review

📝 Walkthrough

Walkthrough

This PR adds conditional Transformer Engine (TE) support to the Megatron LoRA plugin by introducing TE-based LoRA wrapper classes and corresponding quantized variants with proper registry registration. Comprehensive tests validate LoRA and quantization behavior on TE-integrated Megatron models across multiple GPU workers.

Changes

Cohort / File(s)	Summary
TE LoRA Support Implementation modelopt/torch/peft/lora/plugins/megatron.py	Adds HAS_TE flag for conditional TE availability. Introduces four new classes: _LoRATEMCoreColumnParallelLinear, _LoRATEMCoreRowParallelLinear, _QuantLoRATEMCoreColumnParallelLinear, and _QuantLoRATEMCoreRowParallelLinear. Registers TE-based linear module variants with LoRAModuleRegistry and QuantModuleRegistry when TE imports succeed.
TE Integration Tests tests/gpu/torch/peft/test_megatron_peft_te.py	New test module validating TE-backed Megatron GPT models with LoRA adapters. Includes model provider function, test functions for LoRA forward passes, quantization-then-LoRA workflows, and LoRA-then-quantization workflows. Uses multiprocessing to validate behavior across GPU workers with NCCL backend.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 21.43% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: adding Transformer Engine (TE) support to ModelOPT-PEFT, which is reflected in the actual code changes that introduce TE-based LoRA wrapper classes and registration.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch jingyux/modelopt-peft-te

Important

Action Needed: IP Allowlist Update

If your organization protects your Git platform with IP whitelisting, please add the new CodeRabbit IP address to your allowlist:

• ✨ 136.113.208.247/32 (new)
• 34.170.211.100/32
• 35.222.179.152/32

Reviews will stop working after February 8, 2026 if the new IP is not added to your allowlist.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.73%. Comparing base (2a46753) to head (046285f).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #835      +/-   ##
==========================================
+ Coverage   73.44%   73.73%   +0.28%     
==========================================
  Files         194      196       +2     
  Lines       20034    20412     +378     
==========================================
+ Hits        14714    15050     +336     
- Misses       5320     5362      +42     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Edwardf0t1]

LGTM.

@jingyu-ml Could you add more context about why we need this change and what use cases are supported?

[jingyu-ml]

LGTM.

@jingyu-ml Could you add more context about why we need this change and what use cases are supported?

This is because not every model uses Megatron’s non-TE parallel linear layer. We need to make PEFT easier to use when a model relies on the TE parallel linear layer.