Add an FSDP2-native Ulysses sequence parallelism backend - #4209
Open
qgallouedec wants to merge 14 commits into
Open
Add an FSDP2-native Ulysses sequence parallelism backend#4209qgallouedec wants to merge 14 commits into
qgallouedec wants to merge 14 commits into
Conversation
…ntion sequence mixers Findings from running recent architectures through `sp_backend="torch"` against an unsharded reference (logits, loss, gradients): - gpt-oss: attention sinks are one logit per head, so slice `s_aux` to the heads each rank holds after the all-to-all. Passes with flash attention. - Qwen2.5-VL: its attention layers only get the rotary embeddings, never `position_ids`. The wrapper now falls back to the `position_ids` seen at the model's forward (stashed by the pre-hook), taking the text axis of multimodal `[3, batch, seq]` positions. Text-only training passes. - LFM2 (short convolutions), Mamba hybrids, linear attention: refuse every layer type other than full and sliding-window attention with one message, since those layers mix tokens inside the local shard and would silently be cut at the shard boundaries. Verified passing: Llama 3.2, Qwen3, Phi-3, Qwen3-MoE, GLM-4-MoE, DeepSeek-V3 (MLA), Cohere2 and Gemma3 (sliding windows, flash attention), gpt-oss, Qwen2.5-VL, Qwen3-VL, Gemma3 (VLM) and Idefics3 on text. Qwen3-VL and Gemma3ForConditionalGeneration ignore `shift_labels`, tracked in huggingface/transformers#48491; the docs say to compute the loss by hand for those. The test script now handles vision-language checkpoints and syncs weights from the reference (tiny checkpoints can lack a weight and initialize it at random per load), and passes `use_cache=False` at forward time.
`sp_backend` no longer defaults to `deepspeed`. Left unset, it resolves to `deepspeed` under the DeepSpeed engine and `torch` otherwise, so `--parallelism_config_sp_size N` with FSDP2 just works. Explicit values are kept and still raise on a mismatch. The resolution runs in `AcceleratorState` once the distributed type is settled, before the device mesh is built, since the DeepSpeed backend builds none. The CLI flag defaults to `None` and only sets the env var when given. Also shortens the torch backend docs section and the module docstring.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
SunMarc
reviewed
Sep 3, 2026
SunMarc
left a comment
SunMarc
left a comment
Member
There was a problem hiding this comment.
Thanks a lot for this implementation ! Left a bunch of small comments but it should be rather easy to fix ! One question I had is does this integrate cleanly with GRPO or not. With my integration, it wasn't that practical due to the dataloader but maybe there were other blockers.
Both sides add a guard at the same point of `_validate_accelerator`: this branch refuses `sp_backend="torch"` under DeepSpeed, and main (#4180) refuses FSDP2 with no dimension to shard across. Both are kept as written.
#4180 refuses FSDP2 when neither `dp_shard` nor `cp` is enabled. `sp` is folded into `dp_shard_cp` the same way `cp` is, so `sp_size=2` alone gives FSDP a dimension to shard across, and the guard rejected it. Testing `dp_shard_cp_dim_names` directly is the condition the comment describes and the idiom the mesh builder already uses.
Transformers' flash attention path takes passed `cu_seq_lens_q/k` and `max_length_q/k` over `position_ids` and only rebuilds them when they are absent. A padding-free collator computes them for the shard it sees, so under Ulysses they described the local sequence while the kernel ran on the gathered one. On the packed FA2 test with such kwargs passed, 41.8% of the logits mismatched the reference (max abs diff 0.62 against a 0.02 tolerance) and nothing raised. The wrapper now drops the four kwargs after gathering `position_ids`, and the test passes them explicitly so the case stays covered.
`sp_backend` defaults to torch at init, so a handler exists as soon as `sp_size > 1`, and `_resolve_backends` switches it to DeepSpeed under that engine when the user left it unset. It runs first thing in `_validate_accelerator`, which now precedes the device mesh build: the DeepSpeed backend builds no mesh, so the backend has to be known before. The torch backend's FSDP2 requirement moves with it, as a single check on `accelerator.is_fsdp2`, and `state.py` no longer carries any of this.
`_prepare_sp_model` runs the four steps in order: refuse recurrent layers, refuse attention that cannot be computed over the gathered sequence, route attention through Ulysses, attach the position_ids hook. The validation that lived inside `_attach_sequence_parallel_hooks` is its own function, named after the recurrent-layer check beside it, and the hooks function only attaches the hook. Its body now says what the hook does: number the shard by where it sits in the sequence, and record the positions for attention calls that come without them.
Every attention call all-gathered its shard of `position_ids`: 36 layers with recompute is 72 collectives per step, latency-bound at 120-270 us each, so 9-20 ms per step on 2-8 GPUs. The model's forward pre-hook already saw `position_ids` once per forward to record them for attention calls that come without any, so it gathers them there and records the whole sequence; the wrapper reads that instead of gathering each layer's kwarg. Every layer's kwarg is its shard of what the hook saw, so nothing changes numerically: packed sdpa and flash-attn2 on 2 GPUs still match the single-GPU reference.
The four all-to-all call sites in the attention wrapper read `_SeqAllToAll.apply(group, x, 1, 2)` and meant opposite exchanges, since queries, keys and values arrive as `[batch, heads, local_seq, head_dim]` and the output as `[batch, seq, heads / sp_size, head_dim]`. `_gather_seq_scatter_heads` and `_gather_heads_scatter_seq` carry the direction in their name and state the layouts in their docstring. Both are thin wrappers over the one generic all-to-all, whose backward stays "swap scatter and gather".
The module's public entry point is the `sequence_parallel` context manager, the counterpart of torch's `context_parallel`, so the file is named after that rather than after the paper.
Member
Author
|
thks!
Not today, and I don't think it needs to be: TRL's GRPOConfig raises for But RL doesn't reach the lengths where SP matters. Without CP or SP, chunked loss plus activation offload already gets one GPU past 100k tokens, and RL tops out in the tens of thousands usually. 1M-token training is SFT / mid-training territory, and that's where the support is. If long-rollout RL ever gets there, the torch backend makes it tractable though!: it never touches the dataloader, it shards whatever the trainer wraps in the |
The test loaded its model with a hub flash-attention kernel, which needs a compatible `kernels` install that CI does not have, so it failed on the import before reaching the check it tests. The layer type is refused before any check on the attention implementation, so the default `sdpa` model covers it.
The prepare gate and the `maybe_context_parallel` gate checked `sp_enabled` and `sp_backend == "torch"` in opposite orders. Three comments described the per-layer gather the code no longer does; they now describe the single gather and the read.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds
sp_backend="torch": Ulysses sequence parallelism on FSDP2, no DeepSpeed. Redo of #4084. DeepSpeed backend unchanged.Why: torch CP only takes SDPA with a causal mask. Ulysses keeps the model's own attention call, so flash attention, packing and sliding windows work.
How:
spis a device mesh dim, folded into the FSDP sharding mesh likecp.ALL_ATTENTION_FUNCTIONSis wrapped with all-to-all (scatter heads, gather sequence) around the original attention. Only prepared modules are affected, a second model in the process keeps plain attention.maybe_context_parallel, so a CP loop works unchanged.sp_backendresolves todeepspeedunder DeepSpeed,torchotherwise.sp_sizenot dividing the heads.Tested against an unsharded forward/backward (logits, loss, grads), 4xH100:
Found on the way: huggingface/transformers#48491 (Qwen3-VL and Gemma3 VLM ignore
shift_labels).Differences with #4084, and why:
position_idsare gathered rather than dropped in favour ofcu_seqlens: sdpa cannot usecu_seqlens, so packing broke there. Packed sdpa and FA2 are now verified exact._prepare_deepspeedis untouched, so the DeepSpeed path is provably unchanged.Follow-ups: shared CP/SP batch adapter, kv-head replication.
Who can review?
@SunMarc