ImportError: cannot import name 'get_args' from 'megatron'

[peter-fei]

No functions named get_args() in megatron
The funxitons in init doesn't include get_args

r"""
A set of modules to plugin into Megatron-LM with FastMoE
"""
from .utils import add_fmoe_args

from .layers import MegatronMLP
from .layers import fmoefy

from .checkpoint import save_checkpoint
from .checkpoint import load_checkpoint

from .distributed import DistributedDataParallel

from .balance import reset_gate_hook
from .balance import get_balance_profile
from .balance import generate_megatron_gate_hook
from .balance import add_balance_log

from .patch import patch_forward_step
from .patch import patch_model_provider

[laekov]

What is the version of your Megatron-LM? They have being significantly changing the API in their recent versions.

[peter-fei]

My version 1.1.0

[laekov]

1.1.0 may be too old. We have verivied support for 2.2, 2.5 and 3.0.2 . If you have to use 1.1.0, you need to either modify megatron or fastmoe.

[peter-fei]

Thanks. By the way, How can I get a higher version?

[laekov]

Check out the tags here