[QUESTION]I cannot figure out TE(transformer_engine)

[Guodanding]

I have not installed nvidia transformer_engine. I try running some example scripts (e.g train_mixtral_8x7b_distributed.sh) but fail with the TE error:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/pretrain_gpt.py", line 246, in <module>
[rank0]:     pretrain(
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/megatron/training/training.py", line 270, in pretrain
[rank0]:     model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/megatron/training/training.py", line 581, in setup_model_and_optimizer
[rank0]:     model = get_model(model_provider_func, model_type)
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/megatron/training/training.py", line 455, in get_model
[rank0]:     model = model_provider_func(
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/pretrain_gpt.py", line 74, in model_provider
[rank0]:     transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(args.num_experts, args.moe_grouped_gemm, args.qk_layernorm)
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/megatron/core/models/gpt/gpt_layer_specs.py", line 63, in get_gpt_layer_with_transformer_engine_spec
[rank0]:     mlp = _get_mlp_module_spec(
[rank0]:   File "/share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/megatron/core/models/gpt/gpt_layer_specs.py", line 154, in _get_mlp_module_spec
[rank0]:     linear_fc1 = TEColumnParallelGroupedLinear
[rank0]: NameError: name 'TEColumnParallelGroupedLinear' is not defined

I have read Megatron-LM/megatron/core/models/gpt/gpt_layer_specs.py and find that HAVE_TE has been set to False:

try:
    from megatron.core.transformer.custom_layers.transformer_engine import (
        TEColumnParallelGroupedLinear,
        TEDotProductAttention,
        TELayerNormColumnParallelLinear,
        TENorm,
        TERowParallelGroupedLinear,
        TERowParallelLinear,
    )

    HAVE_TE = True
except ImportError:

    HAVE_TE = False

But it is useless because the code is judged by:
use_te = args.transformer_impl == "transformer_engine"
or
if args.transformer_impl == "transformer_engine" then
I don't know how to set args.transformer_impl and what choices can be set.
Or Megatron-LM must be run with nvidia transformer_engine?

[clarence-lee-sheng]

You can set a --transformer-impl local flag

There is a --transformer-impl argument in megatron/training/arguments.py L707
(

Megatron-LM/megatron/training/arguments.py

Line 707 in b76a7d3

group.add_argument('--transformer-impl', default='transformer_engine',

)

[Guodanding]

You can set a --transformer-impl local flag

There is a --transformer-impl argument in megatron/training/arguments.py L707 (

Megatron-LM/megatron/training/arguments.py

Line 707 in b76a7d3

group.add_argument('--transformer-impl', default='transformer_engine',
)

Thanks!
I know i can set default to 'local' to achieve my goal. But is there 'standard' way to set the args? I have try this in train_mixtral_8x7b_distributed.sh:

ENGINE_ARGS=(
    --transformer_impl local \
)

torchrun ${DISTRIBUTED_ARGS[@]} /share/home/wenjunyi/guokb/workspace/imple-moe/Megatron-LM/pretrain_gpt.py \
    ${MODEL_ARGS[@]} \
    ${MOE_ARGS[@]} \
    ${DATA_ARGS[@]} \
    ${TRAINING_ARGS[@]} \
    ${MODEL_PARALLEL_ARGS[@]} \
    ${LOGGING_ARGS[@]} \
    ${ENGINE_ARGS[@]}

but fail:
pretrain_gpt.py: error: unrecognized arguments: --transformer_impl local

[clarence-lee-sheng]

For your arguments.py, is there the --transformer_impl argument?

[Guodanding]

For your arguments.py, is there the --transformer_impl argument?

yes.

[1195343015]

I have installed transformer_engine, but the same error.

[Guodanding]

Is Apex needed?

[1195343015]

I think it's not necessary

[Guodanding]

ok.
In fact, after i set --transformer_impl argument to local, i meet another error:
[rank1]: AssertionError: (RMSNorm) is not supported in by torch Layernorm when instantiating WrappedTorchLayerNorm
So i am wondeing if i miss some needed packages.

[1195343015]

You'd better use PyTorch NGC container as your environment (which contains te and apex and so on), but I confirm that there's some bugs with megatron here.

[Guodanding]

ok. Thanks.

[Guodanding]

I confirm that there's some bugs with megatron here.

agree.

[Guodanding]

After i install apex, the error changes to:
[rank0]: AssertionError: (RMSNorm) is not supported in FusedLayerNorm when instantiating FusedLayerNorm