Clean up remaining HAVE_TE gating in GPT experimental specs and inference controller

[returnL]

What does this PR do ?

Follow up on #3763 by cleaning up two remaining low-risk HAVE_TE gating inconsistencies.

• use shared megatron.core.extensions.transformer_engine.HAVE_TE in
  megatron/core/models/gpt/experimental_attention_variant_module_specs.py
• remove an unused local TE detection block in
  megatron/core/inference/text_generation_controllers/text_generation_controller.py

This PR intentionally does not touch files where local TE detection is still required or was previously reverted.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[returnL]

Hi @mcore-oncall , this is a follow-up cleanup PR to #3763 that removes remaining HAVE_TE gating inconsistencies.

The changes are minimal and low-risk:

• Use shared HAVE_TE in experimental_attention_variant_module_specs.py
• Remove unused local TE detection in text_generation_controller.py

Could you please help review or assign reviewers? Thanks!

[dimapihtar]

Hi @returnL

Thanks a lot for the contribution! Could you resolve merge conflicts, please?

[returnL]

Hi @dimapihtar,

Thanks for the reminder. I have rebased this PR on the latest main, resolved the conflicts in the TE gating/import changes, and ran tools/autoformat.sh.

The updated branch has been pushed. Please let me know if anything else is needed.

[greptile-apps]

Greptile Summary

This PR removes a duplicated local transformer_engine import-detection block from text_generation_controller.py (where HAVE_TE was never used) and migrates experimental_attention_variant_module_specs.py to the shared HAVE_TE flag from megatron.core.extensions.transformer_engine, replacing the ad-hoc try/except ImportError pattern.

• text_generation_controller.py: Removes the orphaned try: import transformer_engine … HAVE_TE = True/False block entirely and reorders the displaced imports back into their natural alphabetical position—no behavioural change.
• experimental_attention_variant_module_specs.py: Replaces the local TE detection with from megatron.core.extensions.transformer_engine import HAVE_TE, then conditionally imports TESpecProvider (or sets it to None) under if HAVE_TE / else—consistent with how other files in the codebase already gate TE-specific symbols.

Confidence Score: 5/5

Safe to merge — both changes are pure cleanup with no logic alterations.

The removed try/except block in text_generation_controller.py was verified to be entirely unused (HAVE_TE never referenced after definition). The migration to the shared HAVE_TE flag in experimental_attention_variant_module_specs.py is consistent with how every other file in the extensions layer handles TE availability, and the TESpecProvider = None fallback mirrors the pre-existing NameError-equivalent behaviour when TE is absent.

No files require special attention.

Important Files Changed

Filename	Overview
megatron/core/inference/text_generation_controllers/text_generation_controller.py	Removes an unused HAVE_TE detection block and moves displaced imports to their canonical top-of-file position; no logic change.
megatron/core/models/gpt/experimental_attention_variant_module_specs.py	Switches from a local try/except TE detection to the shared HAVE_TE flag; TESpecProvider falls back to None when TE is absent, matching the pattern used elsewhere in the repo.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Module Import] --> B{HAVE_TE?\nfrom megatron.core.extensions.transformer_engine}
    B -- Yes --> C[import TESpecProvider\nfrom transformer_engine_spec_provider]
    B -- No --> D[TESpecProvider = None]
    C --> E[_get_backend_spec_provider]
    D --> E
    E --> F{config.use_kitchen?}
    F -- Yes --> G[KitchenSpecProvider\nfallback=TESpecProvider]
    F -- No --> H[TESpecProvider]

Loading

_{Reviews (1): Last reviewed commit: "Merge branch 'main' into unify-te-gating..." | Re-trigger Greptile}