[shardformer]: fix llama causal mask with 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. Same class of bug as #6451 (Qwen2/Qwen3), in the LLaMA forward.

📝 What does this PR do?

With sequence parallelism enabled and flash attention disabled, LLaMA trains with the wrong attention:

1. Causal mask silently dropped (bidirectional attention). With SP the policy replaces the attention forward (get_llama_flash_attention_forward); its non-flash branch adds the mask only if it is not None. llama_model_forward builds that mask with transformers' _update_causal_mask, which for sdpa (the default attention implementation), in training and without padding returns None — it relies on SDPA's is_causal instead:

# transformers 4.51.3
LlamaModel(cfg(attn_implementation="sdpa")).train()._update_causal_mask(all_ones_mask, h, cache_position, None, False)   # -> None
LlamaModel(cfg(attn_implementation="eager")).train()._update_causal_mask(all_ones_mask, h, cache_position, None, False)  # -> (B, 1, S, S)

So every token attends to the whole sequence. Training runs without errors; the outputs and gradients are wrong. This is the common pretraining setup (packed samples, no padding, default sdpa), for every SP mode that goes through this path (checked all_to_all and split_gather).

2. Crash with pipeline + sequence parallelism. On stages after the first, hidden_states only holds the local part of the sequence while cache_position covers the full sequence, so _update_causal_mask fails:

RuntimeError: The size of tensor a (12) must match the size of tensor b (24) at non-singleton dimension 0   (llama.py:143)

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 nothing changes. All SP modes compute attention over the full sequence (q_len *= sp_size for the modes sharing the TP group, all-to-all for Ulysses), so the (B, 1, S, S) mask fits all of them.

The existing LLaMA SP configs all enable flash attention (explicitly or via enable_all_optimization), which is why the tests did not catch this. Added three 4-GPU configs without flash attention: Ulysses, Ulysses + PP, TP + SP split_gather.

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

LLaMA, flash attention off	main	this PR
sp 2, all_to_all	all ranks: outputs not close	passes
sp 2 (tp 2), split_gather	all ranks: outputs not close	passes
pp 2 + sp 2, all_to_all	stage 2 crashes (llama.py:143)	passes
sp 2, all_to_all, flash attention on (control)	passes	passes
all 13 configs of run_llama_test (incl. the 3 new ones), run one by one	-	10 pass; the 3 ring_attn configs were skipped because they need flash_attn, which I cannot install on my machine

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