Skip to content

fix: support QSA with sequence parallelism - #180

Merged
tastelikefeet merged 3 commits into
modelscope:mainfrom
taking-lying-flat:fix/qsa-sequence-parallel
Sep 5, 2026
Merged

fix: support QSA with sequence parallelism#180
tastelikefeet merged 3 commits into
modelscope:mainfrom
taking-lying-flat:fix/qsa-sequence-parallel

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Summary

  • project each local sequence-parallel shard with the replicated QSA indexer before communication
  • all-gather the compact indexer Q/K activations across the layer TP group
  • build the QSA mask from the restored full sequence so it matches TE's full-sequence Q/K/V
  • preserve the existing packed-sequence and context-parallel fallbacks

For Qwen3.8-Flash-Next, the gathered width is 640 ((4 + 1) * 128) instead of the 2560-wide hidden state, reducing the new communication volume by 75% compared with gathering hidden_states directly.

Fixes #178.

Validation

  • verified with a real 2-process Gloo collective that both TP ranks produce a mask identical to the non-SP full-sequence path, including a pooling block crossing the rank boundary
  • verified that globally short sequences still return None before projection or communication
  • verified the gather API and first-dimension concatenation semantics against Megatron-Core 0.16.0, 0.16.1, 0.17.0, 0.18.0, and 0.19.0
  • pre-commit run --all-files
  • python3 -m py_compile src/mcore_bridge/model/gpts/qwen4_exp.py src/mcore_bridge/model/modules/qsa_indexer.py

The validation harness was intentionally not committed; this PR contains production code only.

@tastelikefeet
tastelikefeet merged commit c419b4b into modelscope:main Sep 5, 2026
1 check passed
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.

QSAIndexer.select_mask incompatible with sequence_parallel: mask built on s/TP sequence crashes TE get_full_mask

3 participants