[shardformer]: fix qwen2/qwen3 causal mask with pipeline + sequence parallelism and no flash attention

[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.

Update (rebased onto 71c8abe): this PR originally fixed two bugs. The first one (wrong key length of the attention mask on later pipeline stages) was fixed on main by 641cda4 (2026-10-09) in the same way, so I dropped it. The second one, below, is still present on main (re-verified on 71c8abe), and is what this PR fixes now.

📝 What does this PR do?

Bug: with pipeline + sequence parallelism and flash attention off, Qwen2 / Qwen3 silently use bidirectional attention.

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 (Qwen{2,3}PipelineForwards.qwen{2,3}_model_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 every token attends to the whole sequence. Training runs without errors; outputs and gradients are wrong.

(641cda4 fixed the same class of bug for LLaMA by always building the 4d mask in the non-flash path; the Qwen2/Qwen3 pipeline forwards were not changed.)

Fix. When SP is enabled and flash attention is not — exactly when the policy replaces the attention forward — build the 4d causal mask over the full sequence with _prepare_4d_causal_attention_mask. 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. The existing SP configs all enable flash attention, which is why this was not caught.

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

pp 2 + sp 2 (all_to_all), flash attention off	main (71c8abe)	this PR
Qwen3	all 4 ranks: outputs not close	passes
Qwen2	all 4 ranks: outputs not close	passes
full test_shard_qwen3.py::test_qwen3 (incl. the new config)	–	passes
full test_shard_qwen2.py::test_qwen2 (incl. the new config)	–	passes

💥 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.

[LiRunGuo]

cc @richardoo-707: 641cda4 fixed the mask-length part of this PR in the same way (thanks!), so I rebased and dropped it. What remains is the Qwen2/Qwen3 counterpart of the LLaMA non-flash mask change in 641cda4: with pipeline + sequence parallelism and flash attention off, the Qwen pipeline forward still passes None to the replaced attention, which then runs without a causal mask. I re-verified on 71c8abe: all 4 ranks fail check_forward_backward on main and pass with this PR.