Skip to content

fix(gkd): preserve student gradients for empty token partitions - #10085

Open
Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/gkd-empty-token-backward
Open

fix(gkd): preserve student gradients for empty token partitions#10085
Excelius-Wang wants to merge 1 commit into
modelscope:mainfrom
Excelius-Wang:fix/gkd-empty-token-backward

Conversation

@Excelius-Wang

Copy link
Copy Markdown
Contributor

When a context-parallel partition has no active response tokens, jsd_loss returns a detached zero. cp_reduce can give that rank a nonzero reduced loss, but backward still fails because the local student graph is disconnected. The same empty path is reached when all teacher top-k rows are uncovered.

Return the sum of the empty student logits so the loss remains zero and backward produces zero student gradients. Nonempty-token computation is unchanged.

Validation: CPU regression tests cover all three beta branches, full-vocabulary and top-k masking, uncovered top-k rows, and backward through a tiny linear model. The expanded suite covers FP16/BF16/FP32/FP64, vocabulary-size alignment, no-grad evaluation, and two AdamW updates with an empty accumulation microbatch. Restoring baseline jsd_loss gives 23 errors; all five test methods pass with the fix. Nonempty loss and gradient comparisons also pass. A two-process Gloo probe starts with a globally valid prompt/response label sequence and executes the production label shift, CP slicing and Megatron trainer loss method. The empty rank fails on the baseline and both ranks pass after the fix. The probe uses a tiny linear model and CPU-adapted index allocation, not a full Megatron runtime. Applicable pre-commit hooks pass. No GPU/Megatron end-to-end run.

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