[bugfix] fix GDN sequence parallel and CP-aware cu_seqlens resolution - #162
Merged
Conversation
addsubmuldiv
approved these changes
Aug 10, 2026
hjh0119
requested changes
Aug 10, 2026
| tensor_parallel_output_grad=False, | ||
| group=tp_group, | ||
| ) | ||
| saved_linear_sp = self._set_linear_sequence_parallel(False) |
Collaborator
There was a problem hiding this comment.
Consider wrapping this in try/finally to make the temporary attribute mutation safe.
Comment on lines
+150
to
+162
| if cu_seqlens.numel() > 0 and int(cu_seqlens[0].item()) != 0: | ||
| total_cu = int(cu_seqlens[-1].item()) | ||
| if total_cu == total_seq_len: | ||
| cu_seqlens = torch.cat([ | ||
| torch.zeros(1, dtype=cu_seqlens.dtype, device=cu_seqlens.device), | ||
| cu_seqlens, | ||
| ]) | ||
| elif total_cu - int(cu_seqlens[0].item()) == total_seq_len: | ||
| cu_seqlens = cu_seqlens - cu_seqlens[0] | ||
| seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] | ||
| if not bool(cu_seqlens[-1].eq(total_seq_len) & (seq_lengths % cp_size).eq(0).all()): | ||
| return None | ||
| return cu_seqlens |
Collaborator
There was a problem hiding this comment.
int() / bool() here force a device sync, once per GDN layer per forward. These
comparisons work directly on tensors, e.g. cu_seqlens[0] != 0
Collaborator
|
thanks! |
hjh0119
approved these changes
Aug 10, 2026
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.
Summary
Fix modelscope/ms-swift#9791
GatedDeltaNet.forward: when SP is enabled, gather the complete sequence first, temporarily disable thesequence_parallelflag of linear layers, and scatter the output after computation completes._resolve_cu_seqlensto validate and normalizecu_seqlensunder CP + packed sequence (THD) scenarios (supporting padded offsets, leading‑zero padding, and CP alignment checks for sequence lengths); fall back to the originalcu_seqlens_qwhen parsing fails._set_linear_sequence_parallel/_restore_linear_sequence_parallelto centrally manage the SP flags ofin_proj/in_proj_qkvz/in_proj_ba/out_proj, preventing behavioral conflicts between full‑sequence computation after gather and linear‑layer SP.seq_lenunder SP mode: multiplyseq_lenbycp_sizewhen SP is enabled; keepseq_len *= sp_size * cp_sizewhen SP is disabled.