cannot import name 'check_set_window_size' from 'transformer_engine.pytorch.attention'

[njw1123]

transformer_engine=2.4.0

  if depth == 0 or not hasattr(obj, "__dict__") or not isinstance(obj, types.ModuleType):
Traceback (most recent call last):
  File "/jizhicfs/marvinhjia/njw1123/megatron_ori/pretrain_gpt.py", line 12, in <module>
    from megatron.training import get_args
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/training.py", line 27, in <module>
    from megatron.utils import lazy_gc, average_losses_across_data_parallel_group, get_all_moe_loss
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/utils.py", line 37, in <module>
    from megatron.model.module import param_is_not_shared
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/model/__init__.py", line 21, in <module>
    from .bert_model import BertModel
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/model/bert_model.py", line 23, in <module>
    from megatron.model.language_model import parallel_lm_logits
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/model/language_model.py", line 26, in <module>
    from megatron.model.transformer import ParallelTransformer
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/model/transformer.py", line 35, in <module>
    from megatron.model.mamba_utils import allocate_layers, mamba_select_layers_for_pipeline_parallel, _init_weights 
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/model/mamba_utils.py", line 17, in <module>
    from megatron.custom_layers.megatron_layers import TENorm
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/custom_layers/__init__.py", line 2, in <module>
    from .custom_transformer import CustomTransformerLayer, MOETransformerLayer
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/custom_layers/custom_transformer.py", line 18, in <module>
    from .custom_attention import CustomMultiheadAttention, DeepSeekV2Attention
  File "/usr/local/python3/lib/python3.10/site-packages/megatron_lm-1.1.5-py3.10.egg/megatron/custom_layers/custom_attention.py", line 12, in <module>
    from transformer_engine.pytorch.attention import check_set_window_size, _SplitAlongDim, InferenceParams, apply_rotary_pos_emb, DotProductAttention
ImportError: cannot import name 'check_set_window_size' from 'transformer_engine.pytorch.attention' (/usr/local/python3/lib/python3.10/site-packages/transformer_engine/pytorch/attention/__init__.py)

[KshitijLakhani]

@njw1123 This most likely seems to be due to an incorrect import.
The attention code went through refactoring: Part 1 and Part 2.

I'd suggest to try importing check_set_window_size the way it is done in TransformerEngine. Here's an example of the same:

TransformerEngine/tests/pytorch/fused_attn/test_fused_attn.py

Line 23 in 1ae1d22

check_set_window_size,

You might also want to correct any other imports in your application that may have been affected due to this refactoring, if any. The test directory would be the best place to search how the imports are being used within TransformerEngine and you could try replicating the same.

The two refactoring PRs do have a detailed description of the directory/module change in case you are interested.

Hope this helps !