Checklist / 检查清单
Bug Description / Bug 描述
使用现有的 swift 训练配置：
"sequence_parallel": true,
搭配 lora 训练会报错
环境：mcore-bridge 1.6.2，Qwen3.5-4B（GDN: 16 key heads / 32 value heads / head_dim 128），TP=4 + sequence_parallel=true，LoRA target_modules=all-linear，padding_free=false，CP=1。
报错原因：
1.6.2 当中，
GatedDeltaNet 会自己去执行对 hidden state 的 all gather:
此处会把 tp + sp 切分的 Sequence / N 的数据做一个 all gather， 每一张卡都有所有的token
hidden_states = gather_from_sequence_parallel_region(
hidden_states,
tensor_parallel_output_grad=False,
group=tp_group,
)
然后，他会关闭 in proj 的 all gather ( 避免 in proj 再次做 all gather )
saved_linear_sp = self._set_linear_sequence_parallel(False)
但是，_set_linear_sequence_parallel 的方式是 ：
修改 ： module.sequence_parallel = False
而当 lora 训练的时候， 这个 module （ self.in_proj）实际上是一个 lora layer , module.base 才是实际 forward 的类 ：
LoraParallelLinear 自己只是个壳，真正做 SP gather 的是它里面的 base_layer（TEColumnParallelLinear）。而 #162 的 _set_linear_sequence_parallel 用 setattr(module, 'sequence_parallel', False) 改的是包装器上那份拷贝——base_layer 自己的 sequence_parallel 依然是 True，前向时照旧对输入做 all-gather。
因此这会导致对 in proj 设置的 sequence_parallel 对 base_layer 其实没有生效， forward 的过程中当中还会**次对数据做一次 all gather **
[64,1,2560] ──#162手动gather──> [256,1,2560] ──in_proj──> 应得 [256,1,3088]
实得 [1024,1,3088] ★序列×4，宽度正确★
所以，需要对 lora 的情况做特殊的处理
How to Reproduce / 如何复现
如上所示
Additional Information / 补充信息
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 445, in forward
hidden_states, context = self._layer_forward(
^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
return layer(hidden_states=hidden_states, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
return super().call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_layer.py", line 342, in forward
hidden_states, context = self._forward_attention(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/megatron/core/transformer/transformer_layer.py", line 652, in _forward_attention
attention_output_with_bias = self.self_attention(
^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/gated_delta_net.py", line 355, in forward
query, key, value = torch.split(
^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/functional.py", line 173, in split
return tensor.split(split_size_or_sections, dim)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/_tensor.py", line 1066, in split
return torch._VF.split_with_sizes(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 2048 (input tensor's size at dimension -1), but got split_sizes=[128, 128, 256]
python-BaseException
python-BaseException
Checklist / 检查清单
Bug Description / Bug 描述
使用现有的 swift 训练配置：
"sequence_parallel": true,
搭配 lora 训练会报错
环境：mcore-bridge 1.6.2，Qwen3.5-4B（GDN: 16 key heads / 32 value heads / head_dim 128），TP=4 + sequence_parallel=true，LoRA target_modules=all-linear，padding_free=false，CP=1。
报错原因：
1.6.2 当中，
GatedDeltaNet 会自己去执行对 hidden state 的 all gather:
此处会把 tp + sp 切分的 Sequence / N 的数据做一个 all gather， 每一张卡都有所有的token
hidden_states = gather_from_sequence_parallel_region(
hidden_states,
tensor_parallel_output_grad=False,
group=tp_group,
)
然后，他会关闭 in proj 的 all gather ( 避免 in proj 再次做 all gather )
saved_linear_sp = self._set_linear_sequence_parallel(False)
但是，_set_linear_sequence_parallel 的方式是 ：
修改 ： module.sequence_parallel = False
而当 lora 训练的时候， 这个 module （ self.in_proj）实际上是一个 lora layer , module.base 才是实际 forward 的类 ：
LoraParallelLinear 自己只是个壳，真正做 SP gather 的是它里面的 base_layer（TEColumnParallelLinear）。而 #162 的 _set_linear_sequence_parallel 用 setattr(module, 'sequence_parallel', False) 改的是包装器上那份拷贝——base_layer 自己的 sequence_parallel 依然是 True，前向时照旧对输入做 all-gather。因此这会导致对 in proj 设置的 sequence_parallel 对 base_layer 其实没有生效， forward 的过程中当中还会**次对数据做一次 all gather **
[64,1,2560] ──#162手动gather──> [256,1,2560] ──in_proj──> 应得 [256,1,3088]
实得 [1024,1,3088] ★序列×4，宽度正确★
所以，需要对 lora 的情况做特殊的处理
How to Reproduce / 如何复现
如上所示
Additional Information / 补充信息
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 445, in forward
hidden_states, context = self._layer_forward(
^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
return layer(hidden_states=hidden_states, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
return super().call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_layer.py", line 342, in forward
hidden_states, context = self._forward_attention(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/megatron/core/transformer/transformer_layer.py", line 652, in _forward_attention
attention_output_with_bias = self.self_attention(
^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/mcore_bridge/model/modules/gated_delta_net.py", line 355, in forward
query, key, value = torch.split(
^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/functional.py", line 173, in split
return tensor.split(split_size_or_sections, dim)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/envs/megatron2/lib/python3.12/site-packages/torch/_tensor.py", line 1066, in split
return torch._VF.split_with_sizes(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 2048 (input tensor's size at dimension -1), but got split_sizes=[128, 128, 256]
python-BaseException
python-BaseException