Fix import error of megatron.core for tests
State: closed | Pull request: yes
Labels: name: stale

name: NLP

# What does this PR do ?

Installing NeMo/TransformerEngine and checkout Megatron-LM code by using its `main` branch.
Then run pytest
```python
PYTHONPATH=/home/robin/Megatron-LM/ pytest --tb=native tests/
```
It directly reported error:
```
ImportError while loading conftest '/home/robin/try1/NeMo/tests/conftest.py'.                                                                                                                                                               
tests/conftest.py:28: in <module>                                                                                                                                                                                                           
    from tests.fixtures.tts import *                                                                                                                                                                                                        
tests/fixtures/tts.py:21: in <module>
    from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
nemo/collections/asr/__init__.py:15: in <module>
    from nemo.collections.asr import data, losses, models, modules
nemo/collections/asr/models/__init__.py:36: in <module>
    from nemo.collections.asr.models.transformer_bpe_models import EncDecTransfModelBPE
nemo/collections/asr/models/transformer_bpe_models.py:52: in <module>
    from nemo.collections.nlp.modules.common import TokenClassifier
nemo/collections/nlp/__init__.py:15: in <module>
    from nemo.collections.nlp import data, losses, models, modules
nemo/collections/nlp/models/__init__.py:28: in <module>
    from nemo.collections.nlp.models.language_modeling import MegatronGPTPromptLearningModel
nemo/collections/nlp/models/language_modeling/__init__.py:16: in <module>
    from nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model import (
nemo/collections/nlp/models/language_modeling/megatron_gpt_prompt_learning_model.py:31: in <module>
    from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py:78: in <module>
    from megatron.core.models.gpt import GPTModel as MCoreGPTModel
../../Megatron-LM/megatron/core/models/gpt/__init__.py:1: in <module>
    from .gpt_model import GPTModel
../../Megatron-LM/megatron/core/models/gpt/gpt_model.py:15: in <module>
    from megatron.core.transformer.transformer_block import TransformerBlock
../../Megatron-LM/megatron/core/transformer/transformer_block.py:10: in <module>
    from megatron.core.transformer.custom_layers.transformer_engine import TENorm
../../Megatron-LM/megatron/core/transformer/custom_layers/transformer_engine.py:67: in <module>
    class TELinear(te.pytorch.Linear):
E   AttributeError: module 'transformer_engine' has no attribute 'pytorch'
```
The reason for this error is that GPTModel hasn't been import correctly.

# Changelog 
- Fix import error of megatron.core for tests

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

@XuesongYang @MaximumEntropy, @ericharper, @ekmb, @yzhang123, @VahidooX, @vladgets, or @okuchaiev