Repository navigation
Conversation
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>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to This change stops DSpark and Domino from training on padded tokens, and DFlash masking is reported unchanged. No actionable merge risk remains from the supplied evidence. 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 change: preventing padding from contributing to DSpark and Domino loss masks. |
| 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 pull request changes only the DFlash, DSpark, and Domino mask logic plus a unit test. The added code contains no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hard… |
✨ 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.
Contributor
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2761 +/- ##
==========================================
+ Coverage 69.47% 78.27% +8.79%
==========================================
Files 646 646
Lines 71672 71662 -10
==========================================
+ Hits 49796 56090 +6294
+ Misses 21876 15572 -6304
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:
|
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
1 participant
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