Repository navigation
Conversation
This was referenced Aug 2, 2026
Closed
yaoyu-33
added
area:training
Design proposal (RFC) for adding DSpark speculative-decoding draft training to Megatron-Bridge, tracking NVIDIA-NeMo#5230. It describes the DSpark draft (a semi-autoregressive parallel backbone plus Markov and confidence heads), its three-term CE / TV / confidence training objective and acceptance metric, and how it maps onto the existing ModelOpt distillation subsystem pattern (provider hook, custom forward step, loss, thin entry, HF weight mapping). It includes a phased implementation plan and references the NeMo AutoModel reference implementation. Signed-off-by: khazic <khazzz1c@gmail.com>
Requested in review: the Megatron-Bridge design doc does not need a breakdown of the AutoModel implementation; the pointer in the overview is enough. Signed-off-by: khazic <khazzz1c@gmail.com>
The design page carried a four-phase implementation plan and an open-questions list. Neither describes what this PR (or the repo) actually contains: they promise work that is not scheduled and has no owner, so they go stale the moment priorities move. The sequencing and the open decisions belong in the tracking issue (NVIDIA-NeMo#5230), which is where they can be updated. The remaining page is the design itself: the algorithm, where it fits in Megatron-Bridge, and the proposed file surface. The frozen-target supervision section keeps both capture options, since the trade-off between them is design content rather than a schedule. Signed-off-by: khazic <khazzz1c@gmail.com>
First increment of DSpark speculative-decoding draft training (NVIDIA-NeMo#5230; design in NVIDIA-NeMo#5233): the architecture-agnostic, pure-torch core. Adds src/megatron/bridge/training/post_training/dspark/: * heads.py: the sequential head (VanillaMarkov / GatedMarkovHead / RNNHead) that injects intra-block token dependency onto the parallel-backbone logits, plus the confidence head. * loss.py: the three-term CE + total-variation + confidence-BCE objective with per-position decay, and the analytical acceptance metrics (accept_rate = 1 - TV, tau = sum cumprod(accept) + 1). The data-parallel denominator reduction is an injectable dp_group, so the same math wires to a Megatron-Core forward step later. CPU unit tests cover the head math and the loss / acceptance math. The parallel-aware draft backbone and the forward-step / provider integration (which need Megatron-Core and multi-GPU validation) are follow-ups. Signed-off-by: khazic <khazzz1c@gmail.com>
The DSpark loss denominator reduction now runs only when a dp_group is passed; dp_group=None is local-only. Previously it fell back to the default process group whenever torch.distributed was initialized, whose members are not the DP replicas under Megatron-Core parallelism (TP/PP/CP), so the caller must pass the data-parallel group explicitly. Add a 2-GPU functional test (tests/functional_tests/test_groups/training/test_dspark_loss_dp.py) that spawns two NCCL ranks and asserts the DP-averaged gradient through a shared head matches the single-process full-batch gradient. This validates the dp_group reduction, which cannot be exercised on a single process, and needs a GPU node to run. Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
yaoyu-33
force-pushed
the
khazic/fix/dspark-tv-half-factor
branch
from
August 7, 2026 00:21
c505e09 to
3cf0cea
Compare
This was referenced Aug 11, 2026
Contributor
Author
|
Closing this: the change is in the wrong direction and @yaoyu-33 is right on the review in #5234. I checked it independently. Paper Eq. 10 minimizes the weighted raw L1 distance at coefficient 0.9. The factor So the implementation was already correct, and what was actually wrong was the docstring in Both descriptions are corrected in #5233 and #5234 instead, and #5234 adds a regression test that asserts the pair together (disjoint distributions give |
This was referenced Aug 13, 2026
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
3 participants
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
The trained TV term summed the raw probability L1 distance, which is twice the total-variation distance. With the default alphas the objective was effectively
0.1 * CE + 1.8 * TVinstead of the documented0.1 * CE + 0.9 * TV, and thel1_lossmetric reported twice the documented TV value.Fix
Move the L1-vs-TV convention into the helper itself:
_l1_distance_per_tokenbecomes_tv_distance_per_tokenand returns0.5 * ||p_draft - p_target||_1directly.accept_rate = 1 - tvandl1_num = (tv * weight).sum()then fall out with no inline0.5constants, so no future call site can accidentally grab the un-halved distance.Note: the
l1_lossmetric key is kept for continuity but now reports the TV value (half the raw L1), matching the module docstring and the design doc. Runs compared across this change will show the term drop by 2x; that is the fix, not a regression.Tests
New unit test
test_l1_loss_is_total_variation_not_full_l1builds two disjoint saturated distributions (raw L1 = 2.0, TV = 1.0) and asserts both the metric and the loss equal 1.0. The test fails on the parent commit and passes here. All 17 dspark unit tests pass; ruff check and format are clean.Merge note
This branch and the masked-input-robustness fix PR (same stack) both touch the
l1_distblock ofdspark_loss, so merging both into the DSpark core branch (#5234) will conflict there. Resolve by scrubbing first, then using the TV helper:tv_dist = _tv_distance_per_token(...),tv_dist = torch.where(eval_bool, tv_dist, 0.0), thenaccept_rate = (1.0 - tv_dist).clamp_(0.0, 1.0).