Skip to content

fix sequence-parallel split for uneven lengths - #6440

Open
zjn20030811 wants to merge 3 commits into
hpcaitech:mainfrom
zjn20030811:hotfix/uneven-sequence-split
Open

fix sequence-parallel split for uneven lengths#6440
zjn20030811 wants to merge 3 commits into
hpcaitech:mainfrom
zjn20030811:hotfix/uneven-sequence-split

Conversation

@zjn20030811

@zjn20030811 zjn20030811 commented Sep 4, 2026

Copy link
Copy Markdown

Summary

  • Allow sequence-parallel split-forward communication when the split dimension is not divisible by the process-group size by padding the collective buffer with zeros.
  • Keep the Llama hidden states, attention mask, position IDs, and cache positions on the same padded extent, then trim gathered model outputs back to their logical sequence length.
  • Align distributed causal-language-model labels with padded local logits and ignore synthetic tail positions in both loss and gradient computation.
  • Preserve original extents in the split/gather autograd operations so padded slots cannot leak gradients into valid tokens.
  • Add a two-rank CPU gloo regression covering collective values, forward and backward shapes, distributed loss equivalence, and a tiny Llama all-to-all forward/backward pass.

Fixes #6381

Validation

  • python -m pytest tests/test_shardformer/test_layer/test_uneven_sequence_split.py -q -s (1 passed)
  • python -m py_compile colossalai/shardformer/layer/_operation.py colossalai/shardformer/layer/loss.py colossalai/shardformer/modeling/llama.py tests/test_shardformer/test_layer/test_uneven_sequence_split.py
  • ruff check tests/test_shardformer/test_layer/test_uneven_sequence_split.py

@zjn20030811
zjn20030811 requested a review from a team as a code owner September 4, 2026 04:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG]: colossal cannot split tensor evenly when using Sequential Parallelism in hybirdplugin

1 participant