Repository navigation
Conversation
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>
07f251e to
f334076
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>
…DSpark loss Padded teacher-forcing layouts fill unsupervised positions with sentinel target ids (-1 / -100) and non-finite logits. A sentinel id indexes out of bounds in the gather, and zero-weighting a non-finite row does not save the step: 0 * NaN is NaN in the forward, and the backward of a non-finite row poisons the gradient regardless of its weight. Replace the inputs at unsupervised positions instead of masking the outputs. torch.where on draft_logits selects exact zeros into those slots in the backward rather than multiplying by zero, so no NaN reaches draft_logits.grad; target_ids are filled with a valid row; the reduced L1 and the confidence logits are scrubbed for non-finite teacher rows without copying the full-vocab teacher. Keying off the combined supervision mask means dropped anchors get the same protection as eval-masked positions. A poisoned batch now produces bit-identical losses, metrics, and gradients to the equivalent clean one, which the tests assert term by term. Note in the head's docstring that the embedding lookup stays strict: it is the producer's job to replace sentinels, and clamping there would silently train on row 0 rather than fail. Signed-off-by: khazic <khazzz1c@gmail.com>
f334076 to
ba3eede
Compare
|
Rebased onto #5242 (itself now on the reworked #5234) in ba3eede. The closed #5241 is out of this branch's history, so the diff no longer carries the Reapplied rather than replayed. The sanitizing is the same idea and the reasoning is now written into the code, because it is the non-obvious part: masking the outputs is not enough, since Two things changed with the new base. The scrub keys off #5242's combined supervision mask, so a poisoned position inside a dropped anchor is covered too, which a new test asserts. And the reduced-L1 scrub now sits ahead of the fused chunk output rather than the old Tests: the original poisoned-versus-clean case is kept and strengthened to compare every reporting pair's numerator and denominator rather than a fixed metric list. 28 unit tests pass, and the 2-GPU functional test passes under |
|
/ok to test ba3eede |
…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
…rk-masked-input-robustness
|
Merged the updated #5242 in (which itself now carries the current #5234 head), so this branch is back on top of the stack it belongs to. That clears the docs failure: The other red checks on the previous run were not code failures: No change to the sanitization fix itself. |
# Conflicts: # src/megatron/bridge/training/post_training/dspark/heads.py # src/megatron/bridge/training/post_training/dspark/loss.py # tests/unit_tests/training/post_training/test_dspark_loss.py
This branch was previously deployed
4 participants
What
dspark_losscomputed CE, softmax, and BCE over every position and only zeroed masked positions by multiplying with the zero weight afterwards. Two failure modes survive that scheme once real padded batches arrive:target_idspadded with anything other than-100(e.g.-1orvocab_size) at masked positions crashesF.cross_entropywith a device-side assert.0 * NaN == NaN, and a non-finite forward row also produces NaN gradients through cross_entropy's backward even when the upstream gradient is zero, so output-side masking alone cannot fix it.Fix
Sanitize the inputs at masked positions before any op:
draft_logitsrows are zeroed viatorch.where,target_idssentinels are replaced with 0 viamasked_fill, the per-position TV distance is scrubbed on the reduced[B, N, L]tensor (cheaper than copying the full-vocab teacher, which is detached and only consumed through this reduction), andconfidence_predis zeroed at masked positions. The padded-input contract is documented onDSparkForwardOutput.eval_mask.The draft-side
torch.wherecosts one extra logits-sized allocation per step. This is a deliberate robustness trade: it is the only placement that guarantees finite gradients regardless of how the future producer pads, and the loss already materializes two fp32 softmax copies of the same size. If a future producer is guaranteed to pad the draft with finite values, it can be dropped.heads.get_prev_embeddingsintentionally stays a strict lookup. Clamping sentinels there would silently map genuinely corrupt negative ids at supervised positions to row 0 (and corrupt the RNN head's recurrent state for later positions in the block), so the docstring now states the producer contract instead: replace sentinel pads before calling, exactly asdspark_losssanitizestarget_ids.Tests
New unit test poisons masked positions with
-1target ids,-infdraft rows, and NaN teacher/confidence values, then asserts the loss, all metrics, and the gradients are finite and exactly equal to the clean-input reference, with zero gradient at masked positions. It fails (crashes) on the parent commit and passes here. All 17 dspark unit tests pass; ruff is clean.Merge note
The sanitize mask is
eval_mask, matching this branch's loss weighting. the block-keep-mask fix PR (same stack) redefines the weight aseval_mask AND block_keep_mask; whichever branch merges into the DSpark core branch (#5234) second should key the sanitization on that combinedsupervisedmask so dropped-anchor padding is scrubbed too. This branch also conflicts textually with the TV-half-factor fix PR (same stack) in thel1_distblock; resolve by scrubbing first, then computing TV via the helper from that branch.