开启ulysses CP并行训练报错

一、问题现象（附报错日志上下文）：
File "/home/atlas/y00880253/modellink-1.0.rc3/megatron/core/transformer/attention.py", line 315, in forward
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    core_attn_out = self.core_attention(
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/atlas/y00880253/MindSpeed-1.0.RC3_core_r0.6.0/mindspeed/core/context_parallel/ulysses_context_parallel.py", line 92, in forward
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx)
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/autograd/function.py", line 539, in apply
    result = forward_call(*args, **kwargs)
  File "/home/atlas/y00880253/MindSpeed-1.0.RC3_core_r0.6.0/mindspeed/core/context_parallel/ulysses_context_parallel.py", line 92, in forward
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/atlas/y00880253/MindSpeed-1.0.RC3_core_r0.6.0/mindspeed/core/context_parallel/ulysses_context_parallel.py", line 49, in forward
    key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx)
  File "/root/anaconda3/envs/mlrc3/lib/python3.9/site-packages/torch/autograd/function.py", line 539, in apply
        return single_all_to_all(input_, scatter_idx, gather_idx, group)result = forward_call(*args, **kwargs)

File "/home/atlas/y00880253/MindSpeed-1.0.RC3_core_r0.6.0/mindspeed/core/context_parallel/ulysses_context_parallel.py", line 23, in single_all_to_all
  File "/home/atlas/y00880253/MindSpeed-1.0.RC3_core_r0.6.0/mindspeed/core/context_parallel/ulysses_context_parallel.py", line 92, in forward
    input_t = input_.reshape(
RuntimeError: shape '[-1, 8, 0, 128]' is invalid for input of size 1048576

二、软件版本:
-- CANN 版本 (e.g., CANN 3.0.x，5.x.x): 8.0.RC3
--Tensorflow/Pytorch/MindSpore 版本: 2.1
--Python 版本 (e.g., Python 3.7.5): 3.9
--操作系统版本 (e.g., Ubuntu 18.04): eulerosv2r12

三、测试步骤：
使用默认mcore训练脚本
bash examples/mcore/chatglm3/pretrain_chatglm3_6B_32K.sh

Ascend/MindSpeed

内容风险标识

评论 (1)

Ascend/MindSpeed .gitee-modal { width: 500px !important; }

内容风险标识