Skip to content

fix(sp): exclude internal ring padding from MoE auxiliary loss - #10181

Open
Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/ring-moe-padding
Open

Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/ring-moe-padding

Conversation

@Excelius-Wang

Copy link
Copy Markdown
Contributor

Ring sequence parallelism pads each packed sequence independently. The MoE output hook currently trims gathered router logits to the original total length, retaining internal padding and dropping valid tokens from later sequences. This changes the load-balancing loss and its gradients when router auxiliary loss is enabled.

Filter the gathered router logits using the padded position-ID mask in the ring path, reusing the existing padding layout. Keep the Ulysses-only trimming behavior.

Validation:

  • CPU Gloo regression tests cover ring-only, ring + Ulysses, and Ulysses-only groups; compare router logits, Qwen3-MoE auxiliary loss, and local gradients against an unpadded reference across five packing layouts, including a single sequence; also cover disabled sequence parallelism. Before: two ring configurations fail; after: all four configurations pass.
  • Related tests: 46 passed, 4 CUDA-only tests skipped.
  • pre-commit run --all-files: passed.

No full GPU distributed-attention training run was performed.

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.

1 participant