Repository navigation
Conversation
…norm core The next-token shift -- the base logits that predict the token at position i sit at i - 1 -- was spelled out at each of _teacher_logits' three callers: DFlash KD and DSpark as (safe_label_indices - 1).clamp(min=0), LiLiCorr as anchor + s for the token at anchor + 1 + s. A wrong shift trains without an error, just worse, so it now lives once: _teacher_logits takes the label positions and applies the shift itself, and the callers pass what they supervise. DFlash2's predecessor gather uses the same offset for a different quantity (the previous token, not the teacher) and keeps its own. _maybe_apply_base_final_norm took the producer's whole outputs dict, so _teacher_logits wrapped its one flag in a single-key dict to call it. The bool now has its own function, _apply_base_final_norm, which the dict reader delegates to; EAGLE keeps calling the dict reader, unchanged. Behavior is unchanged: against this branch's parent, offline DFlash KD, LiLiCorr (penalty and calibration) and DSpark give bitwise-equal losses and gradients in fp32 and bf16, with pre- and post-norm captures stored in float64. The teacher tests now pass labels and expect row label - 1, and fail without the shift; a new test pins that label 0 reads row 0 rather than wrapping to the last row. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…he DFlash family Two pieces of the family's forwards were written out once per model. - The step with nothing to supervise returns a zero loss that every trainable draft parameter feeds, so DDP's reduction still sees each of them. There were four copies in two forms: DFlash's indexed one element of each trainable parameter, while Domino's forward, DSpark's forward and DSpark's degenerate-loss branch summed every whole parameter, reading all draft weights to produce a zero. HFDFlashModel._draft_zero_loss now has the cheap form once. - The concatenation of the base layers the draft's fc reads, hidden_states[lid + 1] for each target layer, appeared in the DFlash, DSpark and Domino forwards. It is now HFDFlashModel._target_hidden, which also carries the "+ 1" (hidden_states[0] is the embedding output) that modify()'s check on negative ids refers to. Losses and gradients are unchanged; DSpark's DDP-coverage tests, which require a gradient on every draft parameter for a batch with no valid anchor and for one with zero total weight, pass on the shared helper. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
_build_generate_swa_mask hand-wrote the draft's visibility rule a second time -- context before the block, the window counted from each query's real position, block-causal under dflash_draft_attention="causal" -- and its docstring asked readers to keep it in step with training. It is the training rule, _draft_mask_mod, for the one block being drafted, anchored at ctx_len, so it now evaluates that. The grid evaluation and the additive conversion that the training mask's dense path already did become _dense_additive_mask, shared by both. A verbatim copy of the old builder stays in the tests as the reference: over contexts of 1, 5 and 16 tokens, no window, a window of one block and of two, and bidirectional and causal blocks, the masks are equal element for element (and None where the old one returned None). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Comment |
h-guo18
added this pull request to stack #2758
October 12, 2026 09:10
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## haoguo/dflash-training-perf #2757 +/- ##
===============================================================
- Coverage 78.20% 78.18% -0.03%
===============================================================
Files 647 647
Lines 71795 71714 -81
===============================================================
- Hits 56145 56067 -78
+ Misses 15650 15647 -3
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:
|
Which token each block position is trained on, which positions count, and how they are weighted were written out five times: DFlash _compute_loss, DFlash2 _compute_loss (for the selector), LiLiCorr _block_targets, DSpark _compute_dspark_loss and Domino _compute_domino_loss. The copies differed only in the label offset (DFlash, DFlash2 and LiLiCorr train position k on the token at anchor + k; DSpark and Domino on anchor + k + 1) and in whether slot 0 is skipped, so a fix to one never reached the others. - _block_targets returns the label token, its sequence position and the supervision mask, given the class constant _label_shift (0, or 1 for DSpark and Domino). The teacher reads its rows from those positions, including LiLiCorr's candidate logits. - _block_tokens is the teacher-forced token gather that DSpark's and Domino's heads and DFlash2's predecessor ids each wrote by hand. - _block_position_weights holds the D-PACE and decay weighting. DSpark and Domino pass "decay" explicitly, which is what they do today: they never read dflash_loss_objective. - _masked_top1 replaces five copies of the masked top-1 accuracy. The in-loss zero branches keep routing through the draft's output (flat_logits.sum() * 0.0) rather than _draft_zero_loss: with DFlashModule as its own FSDP2 unit (the MiniMax-M2.7 launcher layout), reading the resharded parameters directly skips the unit's backward collectives on that rank and the job aborts with a collective mismatch. Behavior is unchanged. Against the previous head, a train step of DFlash (KD and CE), DFlash2 (ce, tv and lambda, with the selector), LiLiCorr, DSpark (vanilla and rnn heads, confidence head on) and Domino, offline and online, with D-PACE and decay, and steps that take the zero-loss branch, gives bitwise-equal loss, gradients and metrics in fp32 and bf16. DFlash2's predecessor ids now come from the shared gather; they differ from the old formula only at positions whose label lies past the sequence end, which carry no selector weight. A new test pins each variant's labels and supervision on a hand-worked example; it fails if DSpark's shift is set to 0 or if the anchor slot stops being excluded. No existing test did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…Spark DFlash's and DSpark's pseudo_speculative_generate repeated the same first half: run the base model on the prefix, take its greedy next token, gather the target layers onto the draft's device, build the block (anchor then mask tokens), its position ids and mask, and run the draft. Only the decoding after that differs (DFlash takes the argmax of positions 1..bs-1; DSpark walks its Markov head over the block). - _generation_base returns the base model's next token and hidden states, and carries the offline check both copies made. - _draft_block drafts the block that token anchors. Its position ids now come from _build_position_ids at anchor = ctx_len -- the training layout -- the way the generation mask already comes from the training mask rule. The comment it replaces described a uniform-block layout training has not used since the anchors were introduced. - _target_hidden takes an optional device, so the last two hand-written hidden_states[lid + 1] picks go through it. The docstring's note that generation runs with attention_mask=None was stale since the generation mask was added and now says what is passed. Generated tokens are unchanged: DFlash (full attention, sliding window, causal blocks), DFlash2, LiLiCorr, DSpark and Domino online give bitwise-equal draft tokens in fp32 and bf16 against the previous commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
… a second default Thirteen variant fields were read as getattr(config, "dflash_...", <default>) although DFlashConfig declares every one of them with a default, so each default lived twice. One pair had already drifted: dflash_fp32_master_weights defaults to True in the config and to False in HFDFlashModel.modify's fallback, which a reader would take for the default. The fallbacks never run, because modify() always receives a DFlashConfig, so the fields are now read directly. convert_to_dflash_model chose the registry for dflash_architecture_config.projector_type with an if/elif chain and listed the accepted names again by hand in its error. A table maps each projector type to its registry, and the error lists the table's keys. No behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Contributor
|
…e skill tables The shared generation step raises one message for the whole family, "An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate". The known-failure rows in the DFlash and DSpark skill references still quoted the two per-variant messages it replaced, so a search for the actual error found no row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
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: Refactor, no behavior change
Stacked on #2279 (its branch is the base); kept in draft until #2279 lands. Rules the DFlash family wrote out more than once now live in one place each:
_block_targets,_block_tokensand_block_position_weights, driven by one class constant,_label_shift: DFlash, DFlash2 and LiLiCorr train positionkon the token atanchor + k, and DSpark and Domino onanchor + k + 1. DSpark and Domino still always use decay weighting, as today: they never readdflash_loss_objective. A follow-up makes them honor it.isit ati - 1. DFlash KD, DSpark and LiLiCorr each applied that shift before calling_teacher_logits. It now takes label positions and shifts them itself, and every caller passes the positions it supervises. A wrong shift trains silently, just worse._apply_base_final_norm(hidden, prenorm, norm)takes the flag as a bool._maybe_apply_base_final_normkeeps reading it from the producer's dict for EAGLE, and delegates._draft_zero_loss, in the cheap form (one element per parameter, not a fullsum()of every weight). The zero branches inside the loss keep routing through the draft's output (flat_logits.sum() * 0.0). WithDFlashModuleas its own FSDP2 unit (the MiniMax-M2.7 launcher layout), reading the resharded parameters directly there skips that unit's backward collectives on the rank and aborts the job._masked_top1.hidden_states[lid + 1]for each target layer, in three forwards and both generation paths, becomes_target_hidden.pseudo_speculative_generateshared their first half; it is now_generation_base+_draft_block. The block's position ids come from the training_build_position_idsatanchor = ctx_len._build_generate_swa_maskhand-wrote the training visibility rule a second time. It now evaluates_draft_mask_modfor its one block, anchored atctx_len, through the dense helper the training mask also uses.getattr(config, name, <default>)next toDFlashConfig's own default. One pair had drifted:dflash_fp32_master_weightsdefaults toTruein the config andFalsein the fallback. They are now read directly.projector_typedispatch goes through a table instead of an if/elif chain.Net change to source against #2279: +357 / −426 lines.
Usage
No API change.
Testing
tests/unit/torch/speculative(+tests/unit/torch/export/test_hf_spec_rope_export.py): 411 passed on transformers 5.5.4, 5.12.1 and 5.18.0 (5.5.4 and 5.18.0 fall within CI'stf_minandtf_latestpins).tests/gpu/torch/speculative/plugins: 14 passed, 1 skipped.fully_shardunit, one rank on the zero branch): both ranks finish, as before.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Follow-up to the design review of #2279; review after it.
🤖 Generated with Claude Code