Skip to content

QSAIndexer.select_mask incompatible with sequence_parallel: mask built on s/TP sequence crashes TE get_full_mask #178

Closed
#180
Closed
QSAIndexer.select_mask incompatible with sequence_parallel: mask built on s/TP sequence crashes TE get_full_mask#178
#180
Labels
bugSomething isn't working

Description

@haha916
haha916
opened on Aug 31, 2026

Checklist / 检查清单

  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

Bug Description

Training Qwen/Qwen3.8-Flash-Next (qwen4_exp) with --sequence_parallel true (TP=2) crashes on the first forward pass with long sequences (> indexer_budget) when the QSA sparse-attention path is active:
RuntimeError: The size of tensor a (8886) must match the size of tensor b (4443) at non-singleton dimension 3

[INFO:swift] The training of Epoch 0 starts...
/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py:2099: UserWarning: window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=arbitrary
warnings.warn(
/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py:2099: UserWarning: window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=arbitrary
warnings.warn(
/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py:2099: UserWarning: window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=arbitrary
warnings.warn(
/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py:2099: UserWarning: window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=arbitrary
warnings.warn(
[rank0]: Traceback (most recent call last):
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/cli/_megatron/sft.py", line 7, in
[rank0]: megatron_sft_main()
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 93, in megatron_sft_main
[rank0]: return MegatronSft(args).main()
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank0]: result = self.run()
[rank0]: ^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 68, in run
[rank0]: trainer.train(train_dataset, val_dataset)
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 790, in train
[rank0]: self.run_train_step(train_data_iterator, val_data_iterator)
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 732, in run_train_step
[rank0]: metrics, grad_norm, update_successful = self.train_step(train_data_iterator)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 998, in train_step
[rank0]: metrics = forward_backward_func(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 2263, in forward_backward_pipelining_without_interleaving
[rank0]: output_tensor, num_tokens = forward_step(
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 437, in forward_step
[rank0]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/trainer.py", line 122, in forward_step
[rank0]: output_tensor = model(**data)
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank0]: return self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 489, in forward
[rank0]: outputs = self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/mm_gpt_model.py", line 104, in forward
[rank0]: return self.language_model(
[rank0]: ^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpt_model.py", line 316, in forward
[rank0]: decoder_output = self.decoder(
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py", line 481, in call
[rank0]: return super().call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank0]: return super().call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 410, in forward
[rank0]: checkpointed_result = self._checkpointed_forward(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 206, in _checkpointed_forward
[rank0]: hidden_states, context = checkpoint_handler(custom(layer_idx, chunk_end))
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 189, in checkpoint_handler
[rank0]: return tensor_parallel.checkpoint(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 649, in checkpoint
[rank0]: return CheckpointFunction.apply(function, distribute_saved_activations, *args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
[rank0]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 581, in forward
[rank0]: outputs = run_function(*args)
[rank0]: ^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 165, in wrapped_forward
[rank0]: return forward_func(
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 135, in custom_forward
[rank0]: hidden_states, context = self._layer_forward(
[rank0]: ^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
[rank0]: return layer(hidden_states=hidden_states, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank0]: return super().call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen4_exp.py", line 90, in forward
[rank0]: hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen3_next.py", line 332, in forward
[rank0]: core_attn_out = self.core_attention(
[rank0]: ^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py", line 1858, in forward
[rank0]: core_attn_out = super().forward(query, key, value, attention_mask, **_fa_kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/jit.py", line 67, in wrapper
[rank0]: return disabled_f(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 227, in nonrecursive_disable_wrapper
[rank0]: return fn(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py", line 1587, in forward
[rank0]: return self.unfused_attention(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/backends.py", line 393, in forward
[rank0]: dpa_utils.get_full_mask(
[rank0]: File "/usr/local/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py", line 1398, in get_full_mask
[rank0]: attention_mask = torch.logical_or(swa_mask, attention_mask)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: The size of tensor a (8886) must match the size of tensor b (4443) at non-singleton dimension 3
[rank1]: Traceback (most recent call last):
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/cli/_megatron/sft.py", line 7, in
[rank1]: megatron_sft_main()
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 93, in megatron_sft_main
[rank1]: return MegatronSft(args).main()
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank1]: result = self.run()
[rank1]: ^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 68, in run
[rank1]: trainer.train(train_dataset, val_dataset)
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 790, in train
[rank1]: self.run_train_step(train_data_iterator, val_data_iterator)
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 732, in run_train_step
[rank1]: metrics, grad_norm, update_successful = self.train_step(train_data_iterator)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 998, in train_step
[rank1]: metrics = forward_backward_func(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 2263, in forward_backward_pipelining_without_interleaving
[rank1]: output_tensor, num_tokens = forward_step(
[rank1]: ^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 437, in forward_step
[rank1]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/trainer.py", line 122, in forward_step
[rank1]: output_tensor = model(**data)
[rank1]: ^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank1]: return self.module(*inputs, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 489, in forward
[rank1]: outputs = self.module(*inputs, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/mm_gpt_model.py", line 104, in forward
[rank1]: return self.language_model(
[rank1]: ^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpt_model.py", line 316, in forward
[rank1]: decoder_output = self.decoder(
[rank1]: ^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py", line 481, in call
[rank1]: return super().call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank1]: return super().call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 410, in forward
[rank1]: checkpointed_result = self._checkpointed_forward(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 206, in _checkpointed_forward
[rank1]: hidden_states, context = checkpoint_handler(custom(layer_idx, chunk_end))
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 189, in checkpoint_handler
[rank1]: return tensor_parallel.checkpoint(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 649, in checkpoint
[rank1]: return CheckpointFunction.apply(function, distribute_saved_activations, *args)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
[rank1]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 581, in forward
[rank1]: outputs = run_function(*args)
[rank1]: ^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 165, in wrapped_forward
[rank1]: return forward_func(
[rank1]: ^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 135, in custom_forward
[rank1]: hidden_states, context = self._layer_forward(
[rank1]: ^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
[rank1]: return layer(hidden_states=hidden_states, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank1]: return super().call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen4_exp.py", line 90, in forward
[rank1]: hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen3_next.py", line 332, in forward
[rank1]: core_attn_out = self.core_attention(
[rank1]: ^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py", line 1858, in forward
[rank1]: core_attn_out = super().forward(query, key, value, attention_mask, **_fa_kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/jit.py", line 67, in wrapper
[rank1]: return disabled_f(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 227, in nonrecursive_disable_wrapper
[rank1]: return fn(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py", line 1587, in forward
[rank1]: return self.unfused_attention(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/backends.py", line 393, in forward
[rank1]: dpa_utils.get_full_mask(
[rank1]: File "/usr/local/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
[rank1]: return func(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py", line 1398, in get_full_mask
[rank1]: attention_mask = torch.logical_or(swa_mask, attention_mask)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: The size of tensor a (8886) must match the size of tensor b (4443) at non-singleton dimension 3
[rank3]: Traceback (most recent call last):
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/cli/_megatron/sft.py", line 7, in
[rank3]: megatron_sft_main()
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 93, in megatron_sft_main
[rank3]: return MegatronSft(args).main()
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank3]: result = self.run()
[rank3]: ^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 68, in run
[rank3]: trainer.train(train_dataset, val_dataset)
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 790, in train
[rank3]: self.run_train_step(train_data_iterator, val_data_iterator)
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 732, in run_train_step
[rank3]: metrics, grad_norm, update_successful = self.train_step(train_data_iterator)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 998, in train_step
[rank3]: metrics = forward_backward_func(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 2263, in forward_backward_pipelining_without_interleaving
[rank3]: output_tensor, num_tokens = forward_step(
[rank3]: ^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 437, in forward_step
[rank3]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/trainer.py", line 122, in forward_step
[rank3]: output_tensor = model(**data)
[rank3]: ^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank3]: return self.module(*inputs, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 489, in forward
[rank3]: outputs = self.module(*inputs, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/mm_gpt_model.py", line 104, in forward
[rank3]: return self.language_model(
[rank3]: ^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpt_model.py", line 316, in forward
[rank3]: decoder_output = self.decoder(
[rank3]: ^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py", line 481, in call
[rank3]: return super().call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank3]: return super().call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 410, in forward
[rank3]: checkpointed_result = self._checkpointed_forward(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 206, in _checkpointed_forward
[rank3]: hidden_states, context = checkpoint_handler(custom(layer_idx, chunk_end))
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 189, in checkpoint_handler
[rank3]: return tensor_parallel.checkpoint(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 649, in checkpoint
[rank3]: return CheckpointFunction.apply(function, distribute_saved_activations, *args)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
[rank3]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 581, in forward
[rank3]: outputs = run_function(*args)
[rank3]: ^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 165, in wrapped_forward
[rank3]: return forward_func(
[rank3]: ^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 135, in custom_forward
[rank3]: hidden_states, context = self._layer_forward(
[rank3]: ^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
[rank3]: return layer(hidden_states=hidden_states, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank3]: return super().call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen4_exp.py", line 90, in forward
[rank3]: hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen3_next.py", line 332, in forward
[rank3]: core_attn_out = self.core_attention(
[rank3]: ^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py", line 1858, in forward
[rank3]: core_attn_out = super().forward(query, key, value, attention_mask, **_fa_kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/jit.py", line 67, in wrapper
[rank3]: return disabled_f(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 227, in nonrecursive_disable_wrapper
[rank3]: return fn(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py", line 1587, in forward
[rank3]: return self.unfused_attention(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/backends.py", line 393, in forward
[rank3]: dpa_utils.get_full_mask(
[rank3]: File "/usr/local/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
[rank3]: return func(*args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py", line 1398, in get_full_mask
[rank3]: attention_mask = torch.logical_or(swa_mask, attention_mask)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: RuntimeError: The size of tensor a (11888) must match the size of tensor b (5944) at non-singleton dimension 3
[rank2]: Traceback (most recent call last):
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/cli/_megatron/sft.py", line 7, in
[rank2]: megatron_sft_main()
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 93, in megatron_sft_main
[rank2]: return MegatronSft(args).main()
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank2]: result = self.run()
[rank2]: ^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/pipelines/train/sft.py", line 68, in run
[rank2]: trainer.train(train_dataset, val_dataset)
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 790, in train
[rank2]: self.run_train_step(train_data_iterator, val_data_iterator)
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 732, in run_train_step
[rank2]: metrics, grad_norm, update_successful = self.train_step(train_data_iterator)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/base.py", line 998, in train_step
[rank2]: metrics = forward_backward_func(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 2263, in forward_backward_pipelining_without_interleaving
[rank2]: output_tensor, num_tokens = forward_step(
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/pipeline_parallel/schedules.py", line 437, in forward_step
[rank2]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/swift/megatron/trainers/trainer.py", line 122, in forward_step
[rank2]: output_tensor = model(**data)
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank2]: return self.module(*inputs, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 489, in forward
[rank2]: outputs = self.module(*inputs, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/mm_gpt_model.py", line 104, in forward
[rank2]: return self.language_model(
[rank2]: ^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpt_model.py", line 316, in forward
[rank2]: decoder_output = self.decoder(
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py", line 481, in call
[rank2]: return super().call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank2]: return super().call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 410, in forward
[rank2]: checkpointed_result = self._checkpointed_forward(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 206, in _checkpointed_forward
[rank2]: hidden_states, context = checkpoint_handler(custom(layer_idx, chunk_end))
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 189, in checkpoint_handler
[rank2]: return tensor_parallel.checkpoint(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 649, in checkpoint
[rank2]: return CheckpointFunction.apply(function, distribute_saved_activations, *args)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
[rank2]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/tensor_parallel/random.py", line 581, in forward
[rank2]: outputs = run_function(*args)
[rank2]: ^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 165, in wrapped_forward
[rank2]: return forward_func(
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 135, in custom_forward
[rank2]: hidden_states, context = self._layer_forward(
[rank2]: ^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/modules/transformer_block.py", line 251, in _layer_forward
[rank2]: return layer(hidden_states=hidden_states, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 352, in call
[rank2]: return super().call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen4_exp.py", line 90, in forward
[rank2]: hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/mcore_bridge/model/gpts/qwen3_next.py", line 332, in forward
[rank2]: core_attn_out = self.core_attention(
[rank2]: ^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py", line 1858, in forward
[rank2]: core_attn_out = super().forward(query, key, value, attention_mask, **_fa_kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/jit.py", line 67, in wrapper
[rank2]: return disabled_f(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 227, in nonrecursive_disable_wrapper
[rank2]: return fn(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py", line 1587, in forward
[rank2]: return self.unfused_attention(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/backends.py", line 393, in forward
[rank2]: dpa_utils.get_full_mask(
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
[rank2]: return func(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/transformer_engine/pytorch/attention/dot_product_attention/utils.py", line 1398, in get_full_mask
[rank2]: attention_mask = torch.logical_or(swa_mask, attention_mask)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: RuntimeError: The size of tensor a (11888) must match the size of tensor b (5944) at non-singleton dimension 3
[rank0]:[W831 19:34:05.375707529 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
W0831 19:34:07.007000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 878 closing signal SIGTERM
W0831 19:34:07.007000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 880 closing signal SIGTERM
W0831 19:34:07.007000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 881 closing signal SIGTERM
W0831 19:34:07.007000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 882 closing signal SIGTERM
W0831 19:34:07.007000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 883 closing signal SIGTERM
W0831 19:34:07.008000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 884 closing signal SIGTERM
W0831 19:34:07.009000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 885 closing signal SIGTERM
E0831 19:34:09.040000 812 site-packages/torch/distributed/elastic/multiprocessing/api.py:986] failed (exitcode: 1) local_rank: 1 (pid: 879) of binary: /usr/local/bin/python
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/usr/local/lib/python3.12/site-packages/torch/distributed/run.py", line 994, in
main()
File "/usr/local/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 362, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/torch/distributed/run.py", line 990, in main
run(args)
File "/usr/local/lib/python3.12/site-packages/torch/distributed/run.py", line 981, in run
elastic_launch(
File "/usr/local/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 170, in call
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 317, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

/usr/local/lib/python3.12/site-packages/swift/cli/_megatron/sft.py FAILED

Failures:
[1]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 0 (local_rank: 0)
exitcode : 1 (pid: 878)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[2]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 2 (local_rank: 2)
exitcode : 1 (pid: 880)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[3]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 3 (local_rank: 3)
exitcode : 1 (pid: 881)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[4]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 4 (local_rank: 4)
exitcode : -15 (pid: 882)
error_file: <N/A>
traceback : Signal 15 (SIGTERM) received by PID 882
[5]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 5 (local_rank: 5)
exitcode : -15 (pid: 883)
error_file: <N/A>
traceback : Signal 15 (SIGTERM) received by PID 883
[6]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 6 (local_rank: 6)
exitcode : -15 (pid: 884)
error_file: <N/A>
traceback : Signal 15 (SIGTERM) received by PID 884
[7]:
time : 2026-08-31_19:34:09
host : nb-had3tv4g5r-0
rank : 7 (local_rank: 7)
exitcode : -15 (pid: 885)
error_file: <N/A>
traceback : Signal 15 (SIGTERM) received by PID 885

Root Cause (first observed failure):
[0]:
time : 2026-08-31_19:34:07
host : nb-had3tv4g5r-0
rank : 1 (local_rank: 1)
exitcode : 1 (pid: 879)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html

How to Reproduce / 如何复现

megatron sft
--model Qwen/Qwen3.8-Flash-Next
--dataset <long function-calling jsonl, samples ~9k-12k tokens>
--tensor_model_parallel_size 2
--expert_model_parallel_size 4
--pipeline_model_parallel_size 2
--sequence_parallel true
--max_length 16384
--padding_free false
...

Crash trace: Qwen4ExpLayer.forward -> _qsa_select_mask (mask returned) -> Qwen3NextSelfAttention.forward -> TE DotProductAttention -> transformer_engine/pytorch/attention/dot_product_attention/utils.py get_full_mask -> torch.logical_or(swa_mask, attention_mask).

Root Cause

  1. QSA activates only when s // compress_ratio > block_topk, i.e. sequence length exceeds indexer_budget — long samples trigger it.
  2. With sequence_parallel=true, the layer input hidden_states is the SP-split tensor [s/TP, b, h] (=4443 for s=8886, TP=2). QSAIndexer.select_mask builds the selection mask on this split sequence -> [b, 1, s/TP, s/TP].
  3. Qwen3NextSelfAttention.forward then gathers hidden_states back to the full sequence (gather_from_sequence_parallel_region, line ~465), so Q/K/V are full-length (8886), but the QSA mask is never gathered.
  4. TE with window_size set + attn_mask_type=arbitrary builds swa_mask [1,1,8886,8886] and does logical_or with the [1,1,4443,4443] mask -> shape mismatch.

_qsa_select_mask guards against packing/padding-free and context_parallel_size > 1, but not against sequence_parallel.

Why the official example does not reproduce

examples/models/qwen4_exp/megatron_sft.sh uses --max_length 2048 + swift/self-cognition#500 (short samples). With s <= indexer_budget, select_mask hits the no-op guard (max_blocks <= block_topk) and returns None, so the sparse path never produces a mask. Any dataset with samples longer than the indexer budget reproduces the crash.

Suggested Fix

  • In QSAIndexer.select_mask (or Qwen4ExpLayer._qsa_select_mask): if config.sequence_parallel and tensor_model_parallel_size > 1, either gather_from_sequence_parallel_region(hidden_states) before building the mask, or skip QSA (return None, fall back to full attention) like the packing/CP guards.
  • Alternatively all-gather the produced mask back to the full sequence before passing it to attention.

Workaround

--sequence_parallel false avoids the crash (verified). Note the memory impact is small when --recompute_granularity full is used.

Environment

  • mcore-bridge: git main (1.7.0.dev0)
  • ms-swift: git main
  • megatron-core: 0.18.0
  • transformers: git main (>=5.16.0)
  • torch 2.11.0 / CUDA 13.0.3 / 8×141GB

Additional Information / 补充信息

No response

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions