出现 forward() missing 1 required positional argument: 'memory_efficient'

[TongLi3701]

您好，运行训练的script 出来以下问题，请问是哪里有问题呢？

Traceback (most recent call last):
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/pretrain_llama.py", line 151, in <module>
    pretrain(train_valid_test_datasets_provider, model_provider,
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/training.py", line 153, in pretrain
    iteration = train(forward_step_func,
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/training.py", line 711, in train
    train_step(forward_step_func,
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/training.py", line 426, in train_step
    losses_reduced = forward_backward_func(
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/core/pipeline_parallel/schedules.py", line 363, in forward_backward_no_pipelining
    output_tensor = forward_step(forward_step_func, data_iterator,
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/core/pipeline_parallel/schedules.py", line 218, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/pretrain_llama.py", line 93, in forward_step
    output_tensor = model(tokens, position_ids, attention_mask,
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/distributed.py", line 58, in forward
    return self.module(*inputs, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/module.py", line 183, in forward
    outputs = self.module(*inputs, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/llama_model.py", line 114, in forward
    lm_output = self.language_model(
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/language_model.py", line 520, in forward
    encoder_output = self.encoder(
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/transformer.py", line 1303, in forward
    hidden_states = layer(
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/transformer.py", line 789, in forward
    layernorm_output = self.input_layernorm(hidden_states)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/litong/workspace/MyDev/Megatron-LLaMA/megatron/model/fused_layer_norm.py", line 90, in forward
    return FusedRMSNormAffineFunction.apply(input, weight, self.normalized_shape, self.eps)
  File "/home/litong/.conda/envs/megatron-llama/lib/python3.9/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
TypeError: forward() missing 1 required positional argument: 'memory_efficient'