[bugfix] fix GDN sequence parallel and CP-aware cu_seqlens resolution - #162
Conversation
| tensor_parallel_output_grad=False, | ||
| group=tp_group, | ||
| ) | ||
| saved_linear_sp = self._set_linear_sequence_parallel(False) |
There was a problem hiding this comment.
Consider wrapping this in try/finally to make the temporary attribute mutation safe.
| 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 |
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
|
thanks! |
| tp_group = self.pg_collection.tp | ||
| saved_linear_sp = {} | ||
| if use_sp: | ||
| hidden_states = gather_from_sequence_parallel_region( |
There was a problem hiding this comment.
This communication does not consider the linear layer with the LoRA adapter. Training a Qwen 3.5 model with PEFT using this code snippet will result in a shape mismatch.
There was a problem hiding this comment.
We indeed did not consider LoRA compatibility previously. You are welcome to fix it if you are interested. If you would like me to resolve this issue, I can provide a LoRA‑compatible version later.
There was a problem hiding this comment.
Thank you, Sunyi. I'm already working on it and would be happy to submit a pull request.
4 participants
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.