Repository navigation
docs(training): add DSpark speculative decoding design - #5233
kamran-nvidia merged 4 commits into
Conversation
yaoyu-33
added
area:training
| head regresses to `accept_rate`; the paper applies a post-hoc Sequential | ||
| Temperature Scaling calibration for serving. | ||
|
|
||
| ### Reference implementation (NeMo AutoModel) |
There was a problem hiding this comment.
Removed that section. The pointer to the AutoModel reference in the overview stays.
aroshanghias-nvd
added
the
waiting-on-maintainers
yaoyu-33
removed
the
needs-review
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>
097efcf to
a4047be
Compare
yaoyu-33
left a comment
yaoyu-33
left a comment
There was a problem hiding this comment.
Architecture disposition: changes requested. This is a useful docs-only RFC, but it does not yet define an implementable stable contract. Before merge, it should correct the objective scaling, choose one end-to-end owner (or specify two complete alternative contracts), and define the frozen-weight, distributed-loss, checkpoint, and export boundaries. No DSpark core implementation is requested in this PR.
|
|
||
| - `L_ce` (weight 0.1): cross-entropy of the draft logits against the target's | ||
| teacher-forced next tokens. | ||
| - `L_tv` (weight 0.9): total-variation distance `0.5 * ||p_draft - p_target||_1` |
There was a problem hiding this comment.
[High] This halves the dominant loss relative to the cited DSpark objective. Paper Eq. 10 minimizes the weighted raw L1 distance ||p_draft - p_target||_1 with coefficient 0.9; the factor 1/2 appears in Eq. 8 only when forming the analytical acceptance/confidence target. The primary DeepSpec implementation and current NeMo AutoModel implementation make the same distinction. With the documented coefficient, an implementation following this line gets half the intended training gradient (the proposed follow-up #5241 makes exactly that change). Please define the training term as weighted raw L1 and retain 0.5 * L1 only for the acceptance/confidence target, or explicitly document and validate a different coefficient contract.
| Plus an `examples/distillation/` recipe and a `docs/training/` page. This | ||
| provider-hook plus custom-forward-step plus custom-loss plus thin-entry pattern is | ||
| the established, low-friction way to add a training subsystem here, and it is the | ||
| recommended shape for DSpark. |
There was a problem hiding this comment.
[High] The recommended extension surface is already deprecated and conflates two different owners. Current Bridge main warns that ModelProviderMixin-based construction must migrate to ModelConfig + ModelBuilder, so a new dynamic provider subclass cannot be the promised stable API. The two choices above also lead to different contracts: an MCore-native design makes MCore own the parallel module/loss/sharded checkpoint, while a ModelOpt-injected design makes ModelOpt own the transformation/module/loss/export and Bridge only configures it. Current ModelOpt main already has an HF DSpark plugin and DSparkExporter, while Bridge currently pins ModelOpt 0.46.0rc1, which predates that support. Please choose one end-to-end owner, or give each alternative its own complete module/loss/construction/checkpoint/export contract, and make the Bridge hook builder-backed rather than prescribing the legacy provider path.
|
|
||
| Because the draft reuses the target's embedding and LM head and adds `fc`, | ||
| `markov_head`, `confidence_head`, and the draft decoder layers, the bridge mapping | ||
| registry must gain `dspark.*` entries (QKV and gated-MLP splits for the draft |
There was a problem hiding this comment.
[High] dspark.* Bridge mappings are not implied until the frozen-weight and export schema are defined. The references do not physically alias target parameters: AutoModel creates draft-owned embed_tokens/lm_head, copies the target weights, and freezes the replicas, then controls their trainability separately. That distinction determines PP placement, sharded save/resume, and whether serving export stores or reconstructs them. It also differs by owner: ModelOpt already owns its DSpark deployment schema/exporter, whereas an MCore-native checkpoint would need a separately specified HF schema before AutoBridge can map it. Please define alias-versus-replica storage, TP/vocab-sharded exact-L1 computation, selected/final-hidden PP transport, CP masking, online checkpoint versus offline-cache schema, and the repository that owns export; then list only the Bridge mappings that remain necessary.
|
Thanks @yaoyu-33, all three are addressed in 02156f9. I verified each against the sources you cited before changing anything. Objective scaling. You are right and the page was wrong. Eq. 10 minimizes the weighted raw L1 at coefficient 0.9, and the Owner. Settled on Megatron-Core native, with the reasoning written into a new "Decision" section rather than left implicit. The deciding fact is that ModelOpt's DSpark is HF-path only today: Deprecated extension surface. Dropped. The construction path is now Contracts that were missing. The page now specifies replica-not-alias storage for the frozen embedding and LM head and why that choice drives PP placement, sharded save/resume, and export; shard-local exact L1 under TP (with gathered or top-k paths ruled out, since the acceptance label depends on its exactness); selected-hidden transport under PP; block-aligned anchor sampling under CP; the versioned offline feature-cache schema; and the checkpoint/export ownership split. It also documents the Megatron-Core loss return contract and its |
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>
…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>
02156f9 to
ac09592
Compare
kamran-nvidia
added
needs-review
yaoyu-33
added
ready-to-merge
|
/ok to test ac09592 |
5 participants
What
Adds a design proposal (RFC) page,
docs/training/dspark-speculative-decoding.md,for DSpark speculative-decoding draft training in Megatron-Bridge. Tracks #5230.
Why a design PR first
Megatron-Bridge has no draft-training subsystem today (MTP is the only
draft-adjacent path, and it is a Megatron-Core native module). DSpark's draft (a
semi-autoregressive parallel backbone plus Markov and confidence heads) is a new
neural module that must live either in Megatron-Core (like MTP) or as a
ModelOpt-injected module, and its parallel/loss integration can only be validated
on multi-GPU. So the draft module and its trainer belong in later, separately
reviewed PRs; this PR establishes the algorithm spec, the Bridge integration
shape, and the plan.
Contents
the frozen target's hidden states as context K/V) plus a Markov/RNN sequential
head and a confidence head; the draft shares and freezes the target embedding
and LM head.
confidence BCE (1.0); acceptance =
1 - TV,tau = sum cumprod(accept) + 1.hook,
forward_step_dspark, apost_training/dspark.pyloss, a thintraining/dspark.pyentry, anddspark.*HF weight mappings).Reference implementation: NeMo AutoModel
nemo_automodel/components/speculative/dspark/. Paper:arXiv:2607.05147.
Docs-only; the page is registered in the "Training and Customization" toctree.