Repository navigation
Conversation
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Build the shared loss mask modelopt/torch/speculative/plugins/hf_dflash.py |
_build_loss_mask combines valid-label positions, attention_mask, and an optional loss_mask. |
Use and test shared masking modelopt/torch/speculative/plugins/hf_dflash.py, modelopt/torch/speculative/plugins/hf_domino.py, modelopt/torch/speculative/plugins/hf_dspark.py, tests/unit/torch/speculative/plugins/test_hf_dflash.py |
The three training paths call _build_loss_mask. A parameterized test checks that attention-masked padding positions with labels produce exactly zero loss. |
Priority:
Estimated code review effort: 2 (Simple) | ~10 minutes
Change: Bug fix
Merge Risk: 
ef127
-
The intended padding-loss issue is addressed without an established remaining merge-blocking risk.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the primary fix: excluding padding from DSpark and Domino loss masks during speculative training. Docstring Coverage Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 4 files. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The PR changes only three modelopt Python files and one test. The added code centralizes loss-mask construction and does not add torch.load(..., weights_only=False), numpy.load(..., allow_pickle… Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only three modelopt Python files and one test. The added code centralizes loss-mask construction and does not add torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec on external input, or # nosec. No pyproject.toml or requirements.txt changes add dependencies.
-
Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2761 +/- ##
==========================================
- Coverage 79.82% 71.25% -8.57%
==========================================
Files 661 661
Lines 74928 74918 -10
==========================================
- Hits 59810 53385 -6425
- Misses 15118 21533 +6415
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
HFDFlashModel.forward builds the training loss mask as (labels != -100) * attention_mask * loss_mask. The attention_mask term was added in #1975 because labels do not, in general, mark padding with -100: the online collators copy labels from the padded input_ids and write -100 only where answer_only_loss masks non-assistant tokens. HFDSparkModel.forward and HFDominoModel.forward kept an older copy of the rule that consults attention_mask only when labels are absent. With labels present they ignored it, so online training with answer_only_loss=false on a padded batch sampled anchors in the padded tail and trained on pad tokens. The shipped recipes set answer_only_loss=true, which gives padding -100 labels, and streaming and offline batches carry a loss_mask that is zero on padding, so no shipped configuration hit this. The rule lived in three copies and #1975 updated only one. Move it into one helper, HFDFlashModel._build_loss_mask, and call it from the DFlash, DSpark and Domino training forwards, so the next change to which positions are supervised reaches every variant. DFlash's mask is bit-identical to before (same ops in the same order). DFlash2 and LiLiCorr inherit the DFlash forward and need no change. Add a regression test that runs each variant's real training forward on a batch whose only labels sit on padded positions and expects nothing to be supervised. Without this fix it fails for DSpark and Domino. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
h-guo18
force-pushed
the
haoguo/dflash-variant-loss-mask
branch
from
October 24, 2026 06:57
2d38b37 to
ef1276c
Compare
This branch has not been deployed
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.
What does this PR do?
Type of change: Bug fix
DSpark and Domino online training could train on padding.
HFDFlashModel.forwardbuilds its loss mask as(labels != -100) * attention_mask * loss_mask. Theattention_maskterm was added in #1975 because labels do not, in general, mark padding with -100. The online collators (LanguageDataCollator,VisionLanguageDataCollator) copylabelsfrom the paddedinput_idsand write -100 only whereanswer_only_lossmasks non-assistant tokens.HFDSparkModel.forwardandHFDominoModel.forwardkept an older copy of the rule that consultsattention_maskonly whenlabelsis absent. With labels present they ignored it, so on a padded batch they sampled anchors in the padded tail and put pad tokens into the loss.The rule lived in three copies and #1975 updated only one. This PR moves DFlash's rule into a single helper,
HFDFlashModel._build_loss_mask, and calls it from the DFlash, DSpark and Domino training forwards. The next change to which positions are supervised then reaches every variant. DFlash's mask is bit-identical to before (same ops in the same order). DFlash2 and LiLiCorr inherit the DFlash forward and need no change.Affected: DSpark and Domino trained with
data.mode=onlineandtraining.answer_only_loss=falseon batches that contain padding, i.e. any sample shorter thantraining.training_seq_len.Not affected:
answer_only_loss: true, which gives padding -100 labels. This is why the bug went unnoticed.loss_maskthat is zero on padding.Usage
N/A. Bug fix with no API change.
Testing
test_hf_dflash.py::test_padding_is_never_supervised[dflash|dspark|domino]. It runs each variant's real training forward on a batch whose only non -100 labels sit on padded positions, and expects zero loss.dsparkanddominocases fail with loss 0.456 and 3.486.taskset -c 0-3).tests/unit/torch/speculativeon CPU with torch 2.10: 315 passed on transformers 5.5.4, 5.12.1 and 5.18.0. 5.5.4 and 5.18.0 fall within CI'stf_min(~=5.5.0) andtf_latest(~=5.18.0) pins. Main before this change: 312 passed on each.labels,attention_maskandloss_mask(absent, int, bool, float, bf16, fp64, fractional).loss_maskbatches. DSpark and Domino are bit-identical too, except on a padded batch whose padding keeps its labels, which is the fixed case.pre-commit run --files <changed files>: clean.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Follow-up to #1975, which fixed the rule in the DFlash forward only. DSpark and Domino still duplicate the rest of the DFlash training-forward prologue (base forward, anchor sampling, draft inputs). This PR consolidates only the loss-mask rule.
🤖 Generated with Claude Code
Summary by CodeRabbit