Skip to content

fix: align channel loss with sequence parallel labels - #9977

Open
taking-lying-flat wants to merge 2 commits into
modelscope:mainfrom
taking-lying-flat:fix/issue-9973-channel-loss-sp
Open

fix: align channel loss with sequence parallel labels#9977
taking-lying-flat wants to merge 2 commits into
modelscope:mainfrom
taking-lying-flat:fix/issue-9973-channel-loss-sp

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

Fixes #9973.

When sequence parallelism is enabled, per_token_loss_func_sp gathers the per-token loss into the full-sequence layout, while channel-loss metrics still use labels from the local sequence shard. This causes an IndexError and applies the causal label shift twice.

This PR:

  • preserves the gathered labels returned by GatherLoss;
  • removes SP padding from the gathered loss and labels using the same mask;
  • reuses the already-shifted SP labels for channel masking;
  • flattens the aligned full-sequence loss and mask before channel slicing;
  • keeps packing/padding-free channel boundaries aligned with the full-sequence cu_seqlens;
  • flattens the gathered loss before applying per-token loss_scale, covering batch sizes greater than one when no SP padding is needed.

@taking-lying-flat
taking-lying-flat force-pushed the fix/issue-9973-channel-loss-sp branch from cc5b50d to 3e06621 Compare August 24, 2026 12:04
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.

enable_channel_loss crashes with IndexError when combined with sequence parallelism (sequence_parallel_size > 1)

1 participant