diff --git a/docs/index.md b/docs/index.md index 2b17a43a04..b50f8bf19b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -82,6 +82,7 @@ training/moe-optimization.md training/peft.md training/packed-sequences.md training/multi-token-prediction.md +training/dspark-speculative-decoding.md training/callbacks.md ``` diff --git a/docs/training/dspark-speculative-decoding.md b/docs/training/dspark-speculative-decoding.md new file mode 100644 index 0000000000..180f2e7201 --- /dev/null +++ b/docs/training/dspark-speculative-decoding.md @@ -0,0 +1,281 @@ +# DSpark Speculative Decoding (Draft Training) Design + +## Overview + +This page is a design proposal (RFC) for adding **DSpark** speculative-decoding +draft training to Megatron-Bridge. It tracks +[issue #5230](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/5230). + +DSpark ([arXiv:2607.05147](https://arxiv.org/abs/2607.05147)) is a +speculative-decoding method that trains a lightweight **draft** against a frozen +target. The draft proposes a block of tokens; the target verifies them. DSpark's +draft is *semi-autoregressive*: a parallel backbone predicts a whole block in one +pass, and a small sequential head restores the intra-block token dependency that a +purely parallel drafter loses (the "suffix decay" problem). A confidence head +predicts per-position acceptance so the serving engine can schedule its verify +budget. + +A working draft-training reference already exists in NeMo AutoModel at +`nemo_automodel/components/speculative/dspark/`. This document specifies how the +same training recipe maps onto Megatron-Bridge and Megatron-Core. It +intentionally does **not** ship the draft module itself: that is a Megatron-Core +neural module (see [Decision: Megatron-Core native](#decision-megatron-core-native)) +whose parallel/loss integration must be validated on multi-GPU runs, so it belongs +in later, separately reviewed PRs. + +## Background: the DSpark draft + +### Semi-autoregressive draft + +For a block of `block_size` positions (`gamma`, default 7 offline / 5 in +production), the draft factorizes the block distribution as +`P(x | x0) = prod_k p_k(x_k | x0, x_1..x_{k-1})`: + +- **Parallel backbone.** A small transformer predicts every position of the block + in one forward pass. Crucially, the backbone does not re-encode the prompt: it + attends to the frozen target's hidden states as the attention **context K/V**, + with the block's own "noise" slots as the queries. Each block query attends to + the context strictly before its anchor and bidirectionally within its own block. + This is the block-parallel (DFlash-style) half of the draft. +- **Sequential head.** A lightweight module adds intra-block dependency on top of + the parallel logits. DSpark provides a **Markov head** (a low-rank, + token-conditioned additive logit bias, rank 256) and an **RNN head** (a + GRU-like recurrence carrying prefix state across the block). At training time + the head is teacher-forced with the previous block tokens; only at inference is + it run serially. +- **Confidence head.** A small projection `c_k = sigmoid(w^T [h_k ; embed(x_{k-1})])` + predicts each position's acceptance probability. + +The draft **shares and freezes the target's token embedding and LM head**; there +is no separate (compressed) draft vocabulary or `d2t`/`t2d` remapping. The +target's selected decoder-layer hidden states are fused by a learned linear +projection (`fc`) into the draft hidden size and used as the attention context. + +### Training objective + +The draft is trained with three position-weighted terms (weights `exp(-k/gamma)`, +emphasizing earlier positions): + +- `L_ce` (weight 0.1): cross-entropy of the draft logits against the target's + teacher-forced next tokens. +- `L_l1` (weight 0.9): the **raw** probability L1 distance `||p_draft - p_target||_1` + to the target's next-token distribution. This is the dominant term and a direct + acceptance proxy. +- `L_conf` (weight 1.0): binary cross-entropy training the confidence head against + the analytical acceptance label `c_k* = 1 - 0.5 * ||p_draft - p_target||_1`. + +The factor `1/2` belongs to the acceptance label only. Paper +[Eq. 10](https://arxiv.org/html/2607.05147#S3.E10) minimizes the weighted **raw** +L1 distance at coefficient 0.9; the `1/2` that converts L1 into a total-variation +distance appears in [Eq. 8](https://arxiv.org/html/2607.05147#S3.E8), where the +analytical acceptance/confidence target is formed. Both the +[DeepSpec](https://github.com/deepseek-ai/DeepSpec/blob/main/deepspec/modeling/dspark/loss.py) +and [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/speculative/dspark/loss.py) +implementations make the same distinction. Training `0.5 * L1` instead would halve +the dominant term's gradient while leaving the published coefficient in place, so +the two must not be conflated. A regression test asserts the pair together +(disjoint distributions give `l1_loss == 2` and `accept_rate == 0`) so the +implementation and this page cannot drift apart. + +The target's teacher distribution is produced by passing the target's last hidden +state through the (shared, frozen) LM head and is always detached, so no gradient +flows into the target. + +### Acceptance and confidence + +Training measures acceptance analytically rather than by running rejection +sampling. Per position, `accept_rate = clamp(1 - 0.5 * ||p_draft - p_target||_1, 0, 1)` +(this is where the `1/2` of Eq. 8 lives). The expected accepted +prefix length of a block is `tau = sum_k cumprod(accept_rate)_k + 1` (a token +survives only if every earlier token in its block is accepted; the `+1` counts the +verified anchor token). `tau` is the training-time speedup proxy. The confidence +head regresses to `accept_rate`; the paper applies a post-hoc Sequential +Temperature Scaling calibration for serving. + +## Where DSpark fits in Megatron-Bridge + +### Current state + +Megatron-Bridge has no standalone speculative or draft-training subsystem today. +The only draft-adjacent training path is **Multi-Token Prediction (MTP)**, which +is a native Megatron-Core module that Bridge integrates thinly (provider flags, +batch plumbing in `training/gpt_step.py`, and per-model HF weight mappings such as +`mtp.*` in `models/qwen/qwen3_bridge.py`). EAGLE and Medusa appear only in the +ModelOpt **export** path, not in training. + +DSpark's draft (semi-autoregressive block plus Markov and confidence heads) is a +new neural module, so the first thing this RFC has to settle is which repository +owns it end to end. The two candidates lead to genuinely different contracts, and +picking one is a prerequisite for every section below. + +1. **Megatron-Core native** (as MTP does in + `megatron/core/transformer/multi_token_prediction.py`): Megatron-Core owns the + parallel module, the loss folding, and the sharded checkpoint; Bridge owns + configuration and the HF mapping. Requires a Megatron-LM change and a + Bridge-defined HF export schema. +2. **ModelOpt-injected** (as EAGLE/Medusa are, via `modelopt.torch.speculative`): + ModelOpt owns the transformation, module, loss, and export; Bridge only + configures it. Requires nothing new in Megatron-LM, but ModelOpt has to grow a + Megatron-side DSpark plugin first. + +### Decision: Megatron-Core native + +**This RFC selects option 1.** ModelOpt's DSpark support is currently +HF-path-only: `modelopt/torch/speculative/plugins/` ships `hf_dspark.py` and +`modeling_dspark.py`, and `modelopt/torch/export/plugins/hf_spec_export.py` ships +`DSparkExporter`, but the Megatron-side plugins are `megatron_eagle.py` and +`megatron_medusa.py` only. There is no `megatron_dspark.py`. Choosing option 2 +would therefore mean landing a Megatron-native DSpark in Model-Optimizer first and +then waiting on Bridge's ModelOpt pin (currently `nvidia-modelopt==0.46.0rc1`) to +pick it up, which puts a cross-repo dependency in front of every Bridge-side +milestone. Option 1 keeps the critical path inside Megatron-LM and Bridge, where +MTP already demonstrates the pattern end to end. + +Consequences of the decision, which the rest of this document then specifies: + +- The parallel draft module, its loss folding, and its sharded checkpoint are a + **Megatron-Core** contribution, not a Bridge one. +- Bridge owns the **HF export schema** for trained drafts, because ModelOpt's + `DSparkExporter` covers its own HF-path artifacts and not an MCore-native + checkpoint. +- The architecture-agnostic pieces that depend only on `torch` (the sequential and + confidence heads, and the CE / L1 / confidence objective) can be developed and + unit-tested in Bridge, and move to Megatron-Core when the parallel module lands + there. They are not a substitute for that module. + +## Proposed design + +The construction path is `ModelConfig` + `ModelBuilder`, **not** a +`ModelProviderMixin` subclass: `ModelProviderMixin`-based configuration is +deprecated in current Bridge main (`models/model_provider.py` raises a +`DeprecationWarning` directing callers to `ModelConfig` + `ModelBuilder`), so a new +provider subclass could not be the stable API this RFC promises. The DSpark config +subclasses `GPTModelConfig` and names a `DSparkModelBuilder`, mirroring how +`MambaModelConfig` names `MambaModelBuilder`. + +| Piece | Location (proposed) | Mirrors | +|---|---|---| +| Draft config + builder | `models/dspark/dspark_builder.py` (`DSparkModelConfig`, `DSparkModelBuilder`) | `models/mamba/mamba_builder.py` (`MambaModelConfig` naming `MambaModelBuilder`) | +| Parallel draft module | Megatron-Core `megatron/core/transformer/` | `multi_token_prediction.py` | +| Custom forward + loss dispatch | `training/gpt_step.py` (`forward_step_dspark`) | `forward_step_modelopt` | +| DSpark loss | `training/post_training/dspark/` | `post_training/distillation.py` (the three-term CE/L1/confidence loss and the `tau` acceptance metric) | +| Thin entry point | `training/dspark.py` | `training/distill.py` | +| HF to/from Megatron weight mapping for `dspark.*` params | `models//_bridge.py` mapping registry | the `mtp.*` mappings in `qwen3_bridge.py` | +| Recipe + example | `recipes//…` + `examples/dspark/` | the distillation recipe/example | +| Docs | this page | `multi-token-prediction.md` | + +### Loss return contract + +`dspark_loss` follows the Megatron-Core 3-tuple contract used by +`training/losses.py`: it returns the **unnormalized** micro-batch numerator, this +micro-batch's `num_tokens`, and a report mapping each metric to a reducible +`[numerator, denominator]` pair. Normalization happens once, against the globally +reduced token count. + +This requires `calculate_per_token_loss=True` on the model config, and DSpark +training must fail fast when it is unset. With the default `False`, Megatron-Core +divides each micro-batch by its own `num_tokens` before averaging over +micro-batches, which trains a mean of micro-batch means. That is not a rounding +detail: with two ranks running two micro-batches each at uneven supervised-token +counts, the per-micro-batch normalization produces a gradient **37% away** from +the full-batch reference, while the global-token normalization matches it to +`1.8e-07`. + +Reporting pairs exist for the same reason. A logged ratio must be formed as +`sum(numerator) / sum(denominator)` over the log window and the data-parallel +group; averaging per-micro-batch ratios reintroduces the same bias in the metrics. + +### Memory contract for the full-vocabulary terms + +Both the cross-entropy and the L1 term need the draft's normalized distribution +over the whole vocabulary. At reference scale (`num_blocks=512`, `block_size=7`, +`vocab≈150k`) one FP32 distribution tensor is 2.03 GiB, and computing the terms +separately keeps several live at once: the L1 needs the draft probabilities, the +target probabilities, and their difference, while the cross-entropy saves its own +`log_softmax` for the backward. Measured on an 80 GB card, an unchunked pass costs +**10.15 GiB** of forward-plus-backward overhead above its inputs. + +The implementation therefore derives both terms from a single FP32 `log_softmax`, +chunks the flattened rows, and recomputes each chunk in the backward. That brings +the overhead to **2.06 GiB** (4.9x) at identical values for both terms. Any owner +of this module inherits the same constraint: either chunk, or state and enforce a +much smaller anchor-count contract with measurements. Chunking only the L1 is not +enough, because the cross-entropy's saved `log_softmax` is a full-size tensor of +the same order. + +### Frozen weights: replicas, not aliases + +The draft does **not** physically alias the target's parameters. The AutoModel +reference creates draft-owned `embed_tokens` and `lm_head`, copies the target +weights into them, freezes the replicas, and controls their trainability +separately. This RFC adopts the same storage model, because it is what makes the +remaining contracts well-defined: + +- **Pipeline placement.** Replicas let the draft's embedding and LM head sit on + whatever pipeline stage hosts the draft, instead of forcing the draft to share a + stage with the target's first and last stages. +- **Checkpointing.** The frozen replicas are ordinary sharded-checkpoint entries. + They are saved (not reconstructed) so a resumed run is bit-identical without + needing the target checkpoint present. +- **Trainability.** "Frozen" is a `requires_grad=False` property of draft-owned + parameters, not an aliasing side effect, so unfreezing the LM head later is a + config change rather than a structural one. + +### Frozen-target supervision + +The draft trains against features captured from the frozen target: the selected +decoder-layer hidden states (fused by `fc` into the draft's context K/V) and the +final hidden state (which drives the teacher distribution through the shared LM +head). Two capture modes are supported, with different schemas: + +- **Colocated capture** during the same forward (analogous to how MTP consumes + extra token IDs), keeping the target and draft in one model. No cache format; + costs target memory. +- **Offline precomputation** of target features, streamed to a draft-only training + job (as AutoModel's `precompute_dspark.py` does). Sidesteps holding the target + in memory, and requires a versioned cache schema recording the target checkpoint + id, the selected layer indices, the hidden dtype, the anchor sampling seed, and + the block size, so a cache can never be silently paired with a different target. + +Both modes must respect tensor, pipeline, and context parallelism: + +- **Tensor parallel.** The teacher and draft logits are vocab-sharded, so the L1 + distance must be computed shard-locally and summed across the TP group rather + than gathered: `||p_d - p_t||_1` is the sum over TP shards of the local absolute + difference, with the softmax denominators reduced across the group first. The + exactness of this term is what the acceptance label depends on, so an + approximate (top-k, or gathered in lower precision) path is not acceptable here. +- **Pipeline parallel.** The selected decoder-layer hidden states and the final + hidden state must be transported to the draft's stage. When the draft shares the + last stage, this is a local read; otherwise the selected hiddens are additional + pipeline payload, and their dtype and layout are part of the module contract. +- **Context parallel.** Anchor blocks must not straddle CP shard boundaries. + Anchors are sampled per CP shard from positions whose whole block is local, and + the loss mask zeroes any block that would cross a boundary, so no cross-shard + attention or gather is needed for the draft. + +### Checkpoint and export ownership + +- **Training checkpoint.** MCore-native sharded checkpoint, owned by the + Megatron-Core module, covering the draft backbone, `fc`, `markov_head`, + `confidence_head`, and the frozen embedding and LM-head replicas. +- **HF export.** Owned by Bridge. ModelOpt's `DSparkExporter` targets ModelOpt's + own HF-path artifacts, so an MCore-native checkpoint needs its own HF schema + before `AutoBridge` can round-trip it. The bridge mapping registry gains + `dspark.*` entries (QKV and gated-MLP splits for the draft backbone, plain + linears for `fc` and the heads); the frozen replicas are exported as ordinary + weights so a served draft is self-contained. The `mtp.*` mappings in + `qwen3_bridge.py` are the closest existing template. + +These mappings are specified only once the storage model above is fixed; listing +them earlier would have been guesswork about whether the shared tensors are stored +or reconstructed. + +## References + +- DSpark paper: [arXiv:2607.05147](https://arxiv.org/abs/2607.05147), + *Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation*. +- Draft-training reference implementation: NeMo AutoModel, + [`nemo_automodel/components/speculative/dspark/`](https://github.com/NVIDIA-NeMo/Automodel/tree/main/nemo_automodel/components/speculative/dspark). +- Tracking issue: [#5230](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/5230). +- Related in-repo training: [Multi-Token Prediction](multi-token-prediction.md).