[shardformer]: fix attention mask of qwen2/qwen3 with pipeline + sequence parallelism

[LiRunGuo]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install (ran the pinned hooks manually instead: black 24.10.0 --line-length=120, isort 5.13.2 --profile black, autoflake 2.3.1, all clean)

🚨 Issue number

No existing issue; described below.

📝 What does this PR do?

Fixes two bugs in the pipeline forward of Qwen2 and Qwen3 (Qwen2PipelineForwards.qwen2_model_forward, Qwen3PipelineForwards.qwen3_model_forward, identical code) when pipeline parallelism is combined with sequence parallelism:

1. Wrong key length of the attention mask on later stages. For stages after the first, the sequence is split along the SP group, so the forward does seq_length *= sp_size — but seq_length_with_past was computed before that and is not updated. The mask is built as (B, 1, S, S / sp_size):

RuntimeError: The expanded size of the tensor (24) must match the existing size (12) at non-singleton dimension 3.
Target sizes: [2, 8, 24, 24].

This happens with enable_flash_attention=True whenever ColoAttention falls back to the SDPA kernel (no flash_attn installed); the Dao kernel path ignores the mask shape, which is probably why CI did not catch it. The existing 4-GPU config "Ulysess + Flash attention" in test_shard_qwen3.py fails this way on main without flash_attn (I ran it as is); Qwen2 has the same code and fails the same way with the equivalent config (table below).

2. Causal mask dropped without flash attention (silently wrong results). With SP the policy replaces the attention forward (get_qwen{2,3}_flash_attention_forward), whose non-flash branch needs an explicit 4d causal mask and skips masking when it gets None. The pipeline forward however builds the mask for the transformers attention implementation; for sdpa (the default) and an all-ones padding mask, _prepare_4d_causal_attention_mask_for_sdpa returns None, so attention becomes bidirectional. Training runs, but the outputs are wrong.

Fix.

1. Recompute seq_length_with_past after seq_length *= sp_size (as the llama forward does).
2. When SP is enabled (and flash attention is not), always build the 4d causal mask, matching the condition under which the policy replaces the attention forward. Without SP the transformers attention is used, and its mask handling is unchanged.

Test. Added a "Ulysses without flash attention" config (tp 1, pp 2, sp 2, all_to_all, enable_flash_attention=False) to run_qwen2_test and run_qwen3_test.

Results (4x A800, torch 2.5.1+cu124, transformers 4.51.3, no flash_attn; the repo's own check_forward_backward, which compares loss, output and grads with the unsharded model)

pp 2 + sp 2 (all_to_all)	main	this PR
Qwen3, flash attention on	stage 2 fails: mask [.., 24, 12] vs [.., 24, 24]	passes
Qwen3, flash attention off	all ranks: outputs not close (bidirectional attention)	passes
Qwen2, flash attention on	stage 2 fails: mask size mismatch	passes
Qwen2, flash attention off	all ranks: outputs not close	passes
full test_shard_qwen3.py::test_qwen3 (all 4-GPU configs, incl. the new one)	fails on the existing "Ulysess + Flash attention" config	passes
full test_shard_qwen2.py::test_qwen2 (all 4-GPU configs, incl. the new one)	not run	passes

Not in this PR

LLaMA with the same config (pp 2 + sp 2, all_to_all, flash attention off) also fails, but differently (The size of tensor a (12) must match the size of tensor b (24) at llama.py:143), so it is a separate issue; I have not changed LLaMA here.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.