Repository navigation
fix(training): exclude dropped anchor blocks from the DSpark loss and metrics - #5242
Conversation
|
/ok to test c4dae91 |
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>
2622b35 to
eebe7bb
Compare
…racts Select Megatron-Core native ownership and give the reasons. ModelOpt's DSpark support is HF-path only (hf_dspark and modeling_dspark plugins plus DSparkExporter, with megatron_eagle and megatron_medusa but no megatron_dspark), so routing through ModelOpt would put a Model-Optimizer contribution and a pin bump in front of every Bridge milestone. Replace the ModelProviderMixin recommendation with ModelConfig + ModelBuilder. Provider-based configuration is deprecated in current main, so it could not be the stable API this page promises. Correct the objective: the training term is the raw weighted L1 of paper Eq. 10, and the factor 1/2 belongs only to the Eq. 8 acceptance label. Training half the L1 at the published coefficient would halve the dominant gradient. Define the contracts the previous revision left open: replica-not-alias storage for the frozen embedding and LM head, shard-local exact L1 under tensor parallelism, selected-hidden transport under pipeline parallelism, block-aligned anchor sampling under context parallelism, the versioned offline feature-cache schema, the Megatron-Core loss return contract with its calculate_per_token_loss requirement, the measured memory contract for the probability L1, and Bridge as the owner of the HF export schema. 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>
Return the unnormalized micro-batch numerator plus num_tokens and reducible [numerator, denominator] reporting pairs, instead of dividing by a locally reduced denominator. Megatron-Core calls the loss once per micro-batch, so a locally normalized value accumulates a mean of micro-batch means and biases the objective whenever micro-batches carry different supervised-token counts. The unbiased global mean needs calculate_per_token_loss=True, which assert_per_token_loss_config now enforces at the integration point. Chunk the FP32 probability L1 so the softmax temporaries are bounded by chunk_tokens * vocab rather than the full [batch, num_blocks, block_size, vocab] tensor, recomputing each chunk in the backward. Three such tensors at reference scale are several GiB of forward temporaries before autograd storage. Correct the objective's documentation: the training term is the raw weighted L1 of paper Eq. 10, and the factor 1/2 belongs only to the Eq. 8 acceptance label. The implementation already did this; the docstring and the design page did not. A disjoint-distribution regression now asserts l1_loss == 2 with accept_rate == 0 so the two cannot drift apart again. Rewrite the 2-GPU functional test around the new contract: two ranks, two micro-batches each, uneven supervised-token counts, asserting both that the global-token gradient matches the single-process reference and that the per-micro-batch normalization does not. Add a dedicated L0 launch script so CI actually selects the test, using a single pytest coordinator because the test self-spawns its ranks. Signed-off-by: khazic <khazzz1c@gmail.com>
|
Status note: this PR needs a rebase and is not reviewable as it stands. It was stacked on #5241, which I have now closed. Paper Eq. 10 trains the weighted raw L1 and the factor The change this PR actually contributes is still valid and I am keeping it. #5234 has meanwhile been reworked fairly heavily (Megatron-Core 3-tuple return contract, no |
… metrics block_keep_mask marked anchors as dropped but nothing acted on it: eval_mask alone drove the loss weight, so a dropped anchor still contributed gradient, and it still entered num_tokens and the acceptance diagnostics. Combine the two into one supervision mask, eval_mask AND the kept-anchor flag, and drive the loss weight, the token count, and the acceptance metrics from it. Dropping an anchor is now exactly equivalent to zeroing its eval_mask, which a test asserts directly, and dropped anchors leave the metric denominators rather than sitting in them at weight zero. Signed-off-by: khazic <khazzz1c@gmail.com>
eebe7bb to
f7fd089
Compare
|
Rebased onto the reworked #5234 in f7fd089. The closed #5241 is out of this branch's history, so the diff no longer carries the Reapplied rather than replayed, because #5234 changed underneath it: Tests: the two original cases are kept, with the equivalence check extended to |
…string The docs are built with MyST, where a line of dashes under text is a setext H2. The generated apidocs page for this module therefore started at H2 and tripped "Document headings start at H2, not H1 [myst.header]", which fails the docs build because it treats warnings as errors. Use bold text instead; this was the only RST underline under src/megatron/bridge. Signed-off-by: khazic <khazzz1c@gmail.com>
…azic/fix/dspark-block-keep-mask
|
Merged the current head of #5234 in, which is what this branch stacks on. That also brings over the docs fix for the module docstring in No change to the fix itself. |
# Conflicts: # src/megatron/bridge/training/post_training/dspark/loss.py # tests/unit_tests/training/post_training/test_dspark_loss.py
|
/ok to test 39043f8 |
|
Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing: Add |
|
This PR addresses the reported mask mismatch in NVBug 6651438: acceptance rate and tau must describe the same kept-anchor population. The combined supervision mask and dropped-anchor regression tests cover the reported mechanism. Tracking this as pending QA verification; this comment does not claim a completed release validation. |
cuichenx
removed
the
needs-review
|
/ok to test f9b99d3 |
cuichenx
removed
the
needs-more-tests
|
The link-check rerun now fails on retired documentation links (404s), not the earlier GitHub 503s. Those references are corrected on main; could you bring in current main and rerun validation? |
This branch was successfully deployed
6 participants
What
block_keep_maskwas consumed only by thetaudiagnostic. The loss weight was built fromeval_maskalone, so a dropped anchor block (keep = 0) whose eval prefix was still set received full CE/TV/confidence gradient, and its positions were averaged intoaccept_rate, biasing both the objective and the diagnostics.Fix
Define the supervised-position mask once in
dspark_loss(supervised = eval_mask * block_keep_mask.unsqueeze(-1)), derive the decayed loss weight from it, and pass it into_acceptance_metrics(now keyword-only per the repo style rule on same-type parameters). One mask drives the loss weight,accept_rate, and thetaublock validity, so the three sites cannot drift.Tests
Two new unit tests: dropped blocks produce exactly zero gradient and the loss equals the eval-mask-zeroed equivalent;
accept_rate/tauignore a dropped block that deliberately mismatches the teacher. Both fail on the parent commit and pass here. All 18 dspark unit tests pass; ruff is clean.Merge note
the masked-input-robustness fix PR (same stack) sanitizes padded inputs keyed on
eval_maskalone. Whichever of the two branches merges into the DSpark core branch (#5234) second should rebase that sanitization onto the combinedsupervisedmask (one-line change: buildsupervisedfirst and usesupervised.bool()as the scrub mask), so that positions of dropped anchors are scrubbed as well as zero-weighted.