fix mamba_hybrid model pretrain in recompute mode · Pull Request !3094 · Ascend/MindSpeed-LLM

`--recompute-method uniform \
--recompute-granularity full \
--recompute-num-layers 1 \
`
mamba混合模型，设置上述recompute参数之后，训练过程中会报错
```
Traceback (most recent call last):
  File "/home/ma-user/MindSpeed-LLM/pretrain_mamba.py", line 243, in <module>
    main()
  File "/home/ma-user/MindSpeed-LLM/pretrain_mamba.py", line 235, in main
    pretrain(train_valid_test_datasets_provider,
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/training/training.py", line 420, in pretrain
    iteration, num_floating_point_operations_so_far = train(*train_args)
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/training/training.py", line 569, in train
    train_step(forward_step_func,
  File "/home/ma-user/MindSpeed-LLM/megatron/training/training.py", line 590, in train_step
    losses_reduced = forward_backward_func(
  File "/home/ma-user/MindSpeed-LLM/megatron/core/pipeline_parallel/schedules.py", line 439, in forward_backward_no_pipelining
    output_tensor, num_tokens = forward_step(
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/pipeline_parallel/schedules.py", line 78, in wrapper
    output, num_tokens = fn(
  File "/home/ma-user/MindSpeed-LLM/megatron/core/pipeline_parallel/schedules.py", line 264, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
  File "/home/ma-user/MindSpeed-LLM/pretrain_mamba.py", line 170, in forward_step
    output_tensor = model(tokens, position_ids, attention_mask,
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ma-user/MindSpeed-LLM/megatron/core/distributed/distributed_data_parallel.py", line 204, in forward
    return self.module(*inputs, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ma-user/MindSpeed-LLM/megatron/legacy/model/module.py", line 190, in forward
    outputs = self.module(*inputs, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ma-user/MindSpeed-LLM/megatron/core/models/mamba/mamba_model.py", line 188, in forward
    hidden_states = self.decoder(
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/ssm/mamba_block.py", line 29, in mamba_block_forward
    hidden_states = self._mamba_block_method_checkpointed_forward_func(
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/ssm/mamba_block.py", line 119, in _mamba_block_method_checkpointed_forward_func
    hidden_states = checkpoint_handler(
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/ssm/mamba_block.py", line 105, in checkpoint_handler
    return tensor_parallel.checkpoint(
  File "/home/ma-user/MindSpeed-LLM/MindSpeed/mindspeed/core/tensor_parallel/random.py", line 291, in wrapper
    return checkpoint(function, distribute_saved_activations, *args)
  File "/home/ma-user/MindSpeed-LLM/megatron/core/tensor_parallel/random.py", line 301, in checkpoint
    return CheckpointFunction.apply(function, distribute_saved_activations, *args)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/tensor_parallel/layers.py", line 231, in wrapper
    return fn(ctx, run_function, distribute_saved_activations, *args)
  File "/home/ma-user/MindSpeed-LLM/megatron/core/tensor_parallel/random.py", line 240, in forward
    outputs = run_function(*args)
  File "/home/ma-user/MindSpeed-LLM/mindspeed_llm/core/ssm/mamba_block.py", line 78, in custom_forward
    hidden_states = layer(
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ma-user/MindSpeed-LLM/megatron/core/ssm/mamba_layer.py", line 67, in forward
    hidden_states = hidden_states.to(dtype=self.config.params_dtype)
AttributeError: 'tuple' object has no attribute 'to'
```
原因是开启`self.config.recompute_granularity == 'full'`之后，会走`_mamba_block_method_checkpointed_forward_func`函数。其中Attention层的forward返回值为tuple，需要在拿到`hidden_status`之后进行解tuple。

GVPAscend/MindSpeed-LLM

搜索帮助

GVP Ascend/MindSpeed-LLM