From 6ee408e61aa0fee110f6938e1c41c0356e259f14 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Tue, 22 Sep 2026 10:25:27 -0500 Subject: [PATCH 1/2] [shardformer] fix qwen2/qwen3 causal mask with pipeline + sequence parallelism With sequence parallelism the policy replaces the attention forward, whose non-flash path needs an explicit 4d causal mask and skips masking when it gets None. The Qwen2 / Qwen3 pipeline forward however built the mask for the transformers attention implementation; for sdpa and an all-ones padding mask that is None, so with pipeline + sequence parallelism and flash attention off, the attention silently became bidirectional. Always build the 4d causal mask when sequence parallelism is on without flash attention, which is exactly when the attention forward is replaced. Add a pp 2 + sp 2 (all_to_all) config without flash attention to the qwen2 and qwen3 tests. (The mask length on later pipeline stages, the other bug originally fixed here, has been fixed on main by 641cda4.) --- colossalai/shardformer/modeling/qwen2.py | 10 ++++++++++ colossalai/shardformer/modeling/qwen3.py | 10 ++++++++++ .../test_shardformer/test_model/test_shard_qwen2.py | 13 +++++++++++++ .../test_shardformer/test_model/test_shard_qwen3.py | 13 +++++++++++++ 4 files changed, 46 insertions(+) diff --git a/colossalai/shardformer/modeling/qwen2.py b/colossalai/shardformer/modeling/qwen2.py index 80dfb003dccf..c1a4ae161407 100644 --- a/colossalai/shardformer/modeling/qwen2.py +++ b/colossalai/shardformer/modeling/qwen2.py @@ -138,6 +138,16 @@ def qwen2_model_forward( q_padding_mask=attention_mask, is_causal=True, ) + elif shard_config.enable_sequence_parallelism: + # the attention forward is replaced for sequence parallelism (see the policy), and it takes + # a 4d causal mask, rather than the mask for the attention implementation of transformers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) else: if self.config._attn_implementation == "flash_attention_2": # 2d mask is passed through the layers diff --git a/colossalai/shardformer/modeling/qwen3.py b/colossalai/shardformer/modeling/qwen3.py index 9b6670ba051b..01982ef1e0d0 100644 --- a/colossalai/shardformer/modeling/qwen3.py +++ b/colossalai/shardformer/modeling/qwen3.py @@ -136,6 +136,16 @@ def qwen3_model_forward( q_padding_mask=attention_mask, is_causal=True, ) + elif shard_config.enable_sequence_parallelism: + # the attention forward is replaced for sequence parallelism (see the policy), and it takes + # a 4d causal mask, rather than the mask for the attention implementation of transformers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) else: if self.config._attn_implementation == "flash_attention_2": # 2d mask is passed through the layers diff --git a/tests/test_shardformer/test_model/test_shard_qwen2.py b/tests/test_shardformer/test_model/test_shard_qwen2.py index 865563adc625..193be8bce467 100644 --- a/tests/test_shardformer/test_model/test_shard_qwen2.py +++ b/tests/test_shardformer/test_model/test_shard_qwen2.py @@ -120,6 +120,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, "precision": "fp16", "initial_scale": 1, }, + { # Ulysess without flash attention, the replaced attention forward needs a 4d causal mask + "tp_size": 1, + "pp_size": 2, + "sp_size": 2, + "num_microbatches": 2, + "enable_sequence_parallelism": True, + "sequence_parallelism_mode": "all_to_all", + "enable_flash_attention": False, + "use_lazy_init": True, + "zero_stage": 1, + "precision": "fp16", + "initial_scale": 1, + }, { "tp_size": 2, "pp_size": 2, diff --git a/tests/test_shardformer/test_model/test_shard_qwen3.py b/tests/test_shardformer/test_model/test_shard_qwen3.py index 9670a5999e8c..7c76cd834f1b 100644 --- a/tests/test_shardformer/test_model/test_shard_qwen3.py +++ b/tests/test_shardformer/test_model/test_shard_qwen3.py @@ -116,6 +116,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, "precision": "fp16", "initial_scale": 1, }, + { # Ulysess without flash attention, the replaced attention forward needs a 4d causal mask + "tp_size": 1, + "pp_size": 2, + "sp_size": 2, + "num_microbatches": 2, + "enable_sequence_parallelism": True, + "sequence_parallelism_mode": "all_to_all", + "enable_flash_attention": False, + "use_lazy_init": True, + "zero_stage": 1, + "precision": "fp16", + "initial_scale": 1, + }, { "tp_size": 2, "pp_size": 2, From 15936542f84b69e7500e59517cb7d07a8816679f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:04:02 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_checkpoint_io/test_gemini_checkpoint_io.py | 2 +- tests/test_zero/test_gemini/test_optim.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index ba60de0bd1c1..3649b18ec8bb 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -23,7 +23,7 @@ ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False] try: - import apex.amp + pass APEX_AVAILABLE = True except ImportError: diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index b89c591882d6..a0970e43db9e 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -15,7 +15,7 @@ from tests.kit.model_zoo import model_zoo, run_fwd_bwd try: - import apex.amp + pass APEX_AMP_AVAILABLE = True except ImportError: