fix(gkd): preserve student gradients for empty token partitions - #10085
Open
Excelius-Wang wants to merge 1 commit into
Open
fix(gkd): preserve student gradients for empty token partitions#10085Excelius-Wang wants to merge 1 commit into
Excelius-Wang wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When a context-parallel partition has no active response tokens,
jsd_lossreturns a detached zero.cp_reducecan 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.