diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f60fad5cc15..88d7a89466d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -119,6 +119,7 @@ repos: modelopt/torch/quantization/plugins/attention.py| modelopt/torch/sparsity/attention_sparsity/methods/vsa_utils.py| modelopt/torch/speculative/eagle/utils.py| + modelopt/torch/speculative/plugins/dflash_losses.py| modelopt/torch/speculative/plugins/hf_domino.py| modelopt/torch/speculative/plugins/modeling_domino.py| modelopt/torch/speculative/plugins/hf_dflash.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6397b993a0..294bcc88e82 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -36,6 +36,7 @@ Changelog - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. +- The DFlash2 and LiLiCorr training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics`` and ``lilicorr_metrics``. **Deprecations** @@ -46,6 +47,7 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization. - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format. +- Fix Domino AR validation (``pseudo_speculative_generate``), which scored the backbone alone at the wrong block positions, and a DSpark AR-validation crash when ``device_map="auto"`` places a tied ``lm_head`` on another GPU than the draft. DSpark and Domino now reject ``dflash_architecture_config.shift_label: false`` at conversion (DSpark trained with the shift regardless) and export the ``shift_label`` they were trained with; checkpoints saved with ``shift_label: false`` still load. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 92ce9121b53..fdd76b37d8b 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -521,19 +521,19 @@ def _export_config(self): """Extend the DFlash config with the Domino head fields.""" config = super()._export_config() draft_config = self.model.dflash_config + # Read off the built head rather than the config dict, so the export describes the + # head these weights belong to. + head = self.model.dflash_module - # Present because HFDominoModel.modify validates them at convert time. - emb_dim = draft_config.emb_dim - gru_hidden_dim = draft_config.gru_hidden_dim # Mirror the reference checkpoint: emb_dim also appears at the top level. - config["emb_dim"] = emb_dim + config["emb_dim"] = head.emb_dim config["dflash_config"].update( { "projector_type": getattr(draft_config, "projector_type", "domino"), - "shift_label": getattr(draft_config, "shift_label", True), - "pure_draft_prefix_len": getattr(draft_config, "pure_draft_prefix_len", 1), - "gru_hidden_dim": gru_hidden_dim, - "emb_dim": emb_dim, + "shift_label": bool(self.model._label_shift), + "pure_draft_prefix_len": head.pure_draft_prefix_len, + "gru_hidden_dim": head.gru_hidden_dim, + "emb_dim": head.emb_dim, } ) return config @@ -615,14 +615,17 @@ def _export_config(self): """Extend the DFlash config with the DSpark head fields.""" config = super()._export_config() draft_config = self.model.dflash_config + # Read off the built head, not the config dict: DSparkModule casts markov_rank to int + # and lower-cases markov_head_type, and the export must describe what was trained. + head = self.model.dflash_module config["dflash_config"].update( { "projector_type": getattr(draft_config, "projector_type", "dspark"), - "shift_label": getattr(draft_config, "shift_label", True), - "markov_rank": draft_config.markov_rank, - "markov_head_type": self.model.dflash_module.markov_head_type, - "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)), + "shift_label": bool(self.model._label_shift), + "markov_rank": head.markov_rank, + "markov_head_type": head.markov_head_type, + "use_confidence_head": head.use_confidence_head, } ) if config["dflash_config"]["markov_head_type"] == "vanilla": @@ -634,7 +637,7 @@ def _export_config(self): "mask_token_id": config["dflash_config"]["mask_token_id"], "target_layer_ids": config["dflash_config"]["target_layer_ids"], "markov_rank": config["dflash_config"]["markov_rank"], - "sample_from_anchor": True, + "sample_from_anchor": config["dflash_config"]["shift_label"], "enable_confidence_head": config["dflash_config"]["use_confidence_head"], "confidence_head_with_markov": config["dflash_config"]["use_confidence_head"], } diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index f49c38254d6..1a53b21e4d4 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -340,7 +340,7 @@ class DFlashConfig(ModeloptBaseConfig): "absolute — there is no outer multiplier scaling the three terms as a group — so " "each is the coefficient with which its term enters the total, and " "`loss == origin_loss + lilicorr_loss` holds exactly. Both halves and all three " - "weights are reported in the `lilicorr_metrics` dict the model attaches to its " + "weights are reported in the `draft_metrics` dict the model attaches to its " "forward output, so a consumer of those metrics can check the identity per step " "and read which composition produced a checkpoint. The three weights are validated " "all-or-nothing (a negative value means unset): a config that sets some but not " diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 859d960d0b2..c250ca48298 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/modelopt/torch/speculative/dflash/conversion.py @@ -95,4 +95,9 @@ def restore_dflash_model( # Warm start is a one-time training-time init; the restored weights would overwrite it # anyway, and replaying it on the meta device would raise. config = config.model_copy(update={"dflash_init_checkpoint": None}) + # shift_label: false is refused at convert only: DSpark drafters saved with it before the + # refusal trained next-token alignment regardless, and must still load. + config.dflash_architecture_config = { + k: v for k, v in config.dflash_architecture_config.items() if k != "shift_label" + } return convert_to_dflash_model(model, config)[0] diff --git a/modelopt/torch/speculative/plugins/dflash_losses.py b/modelopt/torch/speculative/plugins/dflash_losses.py new file mode 100644 index 00000000000..b2e8279c3c3 --- /dev/null +++ b/modelopt/torch/speculative/plugins/dflash_losses.py @@ -0,0 +1,279 @@ +# Adapted from https://github.com/sgl-project/SpecForge/blob/8ea5ca6/specforge/core/dflash.py +# and https://github.com/deepseek-ai/DeepSpec/blob/add63ba/deepspec/modeling/dspark/loss.py +# Copyright (c) 2025 sgl-project +# Copyright (c) 2026 The DeepSpec Authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Loss components of the DFlash speculative-decoding family, independent of any variant. + +A training step's supervision: which token each block position is trained on and where +(:func:`_block_targets`), how the positions are weighted (:func:`_block_position_weights`), +and :class:`LossContext`, which carries them with the draft's logits to the loss. +""" + +import functools +from typing import NamedTuple + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint + +__all__ = ["DraftLogits", "LossContext", "LossOutput"] + + +def _block_positions(anchor_positions, block_size, offset, seq_len): + """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. + + Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end. + """ + k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1) + unclamped = anchor_positions.unsqueeze(-1) + k + offset + return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len + + +def _block_gather(x, positions): + """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs].""" + return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions) + + +def _block_tokens(input_ids, anchor_positions, block_size, offset): + """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs].""" + positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1]) + return _block_gather(input_ids, positions) + + +class _BlockTargets(NamedTuple): + """What each block position is trained on: the token at ``anchor + k + label_shift``.""" + + token_ids: torch.Tensor # [B, N, bs] + positions: torch.Tensor # [B, N, bs] their sequence indices, clamped onto the sequence + mask: torch.Tensor # [B, N, bs] 1.0 where the position is supervised + + +def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift): + """Labels of every block position and where they are supervised. + + A position is supervised when its block is kept and its label lies inside the sequence + and the loss mask, and is not the anchor itself, which the draft is given as input. + """ + positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1]) + k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1) + mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() + mask = mask * inside.float() + mask = mask * (k + shift > 0).float() + mask = mask * _block_gather(loss_mask, positions) + return _BlockTargets(_block_gather(input_ids, positions), positions, mask) + + +def _dpace_position_weights( + confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None +) -> torch.Tensor: + """Compute detached D-PACE per-position weights from draft confidences. + + Derived from D-PACE (arXiv:2605.18810). The paper factorizes the per-position + weight (Fig. 2 / Eq. 8) into a *cumulative confidence* times a *continuation + value*, which is equivalently the suffix sum of the cumulative confidences:: + + C_j = prod_{i<=j} q~_i # cumulative confidence (Eq. 8) + w_j = sum_{m>=j} C_m # = C_j * continuation value f~_j + + Each confidence is asymmetrically smoothed toward 1 (Eq. 7):: + + q~_i = (1 - alpha) * q_i + alpha, alpha in (0, 1], + + so the floor ``q~_i >= alpha`` keeps every cumulative product (hence every + weight) strictly positive. We evaluate the suffix sum from its definition as + ``total - exclusive_prefix_sum`` of ``C`` rather than reversing the tensor. + Positions with ``valid_mask == 0`` are multiplicative no-ops in ``C`` and + contribute nothing to the sum, matching the per-token loss mask. Weights are + detached (Eq. 9): they reweight the cross-entropy without adding gradient. + + Args: + confidences: ``[..., L]`` draft confidence ``q_i = exp(-CE_i)`` per position. + alpha: smoothing factor in (0, 1]; raises if outside that range. + valid_mask: optional ``[..., L]`` 0/1 mask; ``None`` treats all positions valid. + + Returns: + Detached weights with the same shape and dtype as ``confidences``. + """ + if not 0.0 < alpha <= 1.0: + raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {alpha}") + + with torch.no_grad(): + smoothed = alpha + (1.0 - alpha) * confidences.float() # Eq. 7 + if valid_mask is not None: + keep = valid_mask.to(torch.bool) + smoothed = torch.where(keep, smoothed, torch.ones_like(smoothed)) + cum_conf = torch.cumprod(smoothed, dim=-1) # Eq. 8 cumulative confidence C_j + if valid_mask is not None: + cum_conf = cum_conf * keep.to(cum_conf.dtype) + # Suffix sum w_j = sum_{m>=j} C_m, written as total minus the exclusive + # prefix sum so no axis reversal is needed (Eq. 8). + inclusive = torch.cumsum(cum_conf, dim=-1) + weights = inclusive[..., -1:] - inclusive + cum_conf + return weights.to(dtype=confidences.dtype) + + +def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None): + """``mask`` weighted per block position, by D-PACE or by exponential decay. + + ``first`` is the first predicted block position. ``target_ce`` returns the draft's + cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. + """ + block_size = mask.shape[-1] + if scheme == "dpace" and block_size > 1: + with torch.no_grad(): + # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9). + confidences = torch.exp(-target_ce()[..., first:].float()) + dpace = torch.ones_like(mask) + dpace[..., first:] = _dpace_position_weights( + confidences, dpace_alpha, valid_mask=mask[..., first:] + ) + return mask * dpace + if decay_factor > 0: + k = torch.arange(block_size, device=mask.device).view(1, 1, -1) + return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor) + return mask + + +def _masked_top1(logits, target_ids, mask): + """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" + correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) + return correct.sum().float() / (mask.sum() + 1e-6) + + +def _tvd_chunk(a, b): + """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``.""" + return (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) + + +def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None): + """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean. + + Materializing both [N, vocab] float32 softmax tensors at once OOMs at large + N (num_anchors*block_size) and a 150k vocab. We compute it in row chunks and + gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward + rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math + is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. + + Chunks with ``Tensor.split``, not slicing (each slice's backward zero-fills a full + [N, vocab] tensor) nor ``torch.chunk`` (different chunk shapes, so recompiles). + """ + _chunk = chunk_fn or _tvd_chunk + outs = [] + for a, b in zip( + final_logits.split(chunk_size, dim=0), + teacher_logits.split(chunk_size, dim=0), + strict=True, + ): + if torch.is_grad_enabled() and a.requires_grad: + outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False)) + else: + outs.append(_chunk(a, b)) + return torch.cat(outs, dim=0) + + +class DraftLogits(NamedTuple): + """The draft's logits per block position: [B, N, bs, V].""" + + final: torch.Tensor # the distribution the draft serves + backbone: torch.Tensor # before any head; ``final`` itself for a head-less variant + + +class LossOutput(NamedTuple): + """A training step's loss, the accuracy ``train_acc`` reports, and its metrics.""" + + loss: torch.Tensor + accuracy: float + metrics: dict # name -> detached 0-dim tensor + + +class LossContext: + """One training step's supervision; costly per-token quantities are computed lazily, once. + + ``weights`` are the per-position weights after decay or D-PACE, which reads the draft's + confidence off ``ce("final")``. ``ce`` keeps its gradient only when ``ce_grad`` says the + objective backpropagates through it; otherwise it runs under ``no_grad``. + """ + + def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad): + """Weigh the step's block positions. + + Args: + batch: The forward's ``DraftBatch``. + logits: The draft's ``DraftLogits``. + targets: The labels and their mask, from ``_block_targets``. + weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha. + teacher: Returns the base model's logits for every label position. + ce_grad: Whether the objective backpropagates through ``ce``. + """ + self.batch = batch + self.logits = logits + self.targets = targets + self._teacher_fn = teacher + self._ce_grad = ce_grad and torch.is_grad_enabled() + self._ce = {} + self._teacher = None + self.weights = _block_position_weights( + targets.mask, + **weighting, + target_ce=lambda: self.ce("final").detach().view(targets.mask.shape), + ) + self.weight_sum = self.weights.sum() + 1e-6 + + @functools.cached_property + def trainable(self) -> bool: + """Whether the weights leave anything to train on (reads the device once).""" + return bool(self.weight_sum > 1.0) + + def flat_logits(self, source): + """``logits.`` as [B * N * bs, V].""" + logits = getattr(self.logits, source) + return logits.reshape(-1, logits.shape[-1]) + + def ce(self, source): + """Per-token cross-entropy of ``logits.`` on the labels: [B * N * bs].""" + if source not in self._ce: + with torch.set_grad_enabled(self._ce_grad): + self._ce[source] = F.cross_entropy( + self.flat_logits(source), self.targets.token_ids.view(-1), reduction="none" + ) + return self._ce[source] + + def teacher(self): + """The base model's logits for every label position: [B, N, bs, V].""" + if self._teacher is None: + self._teacher = self._teacher_fn() + return self._teacher diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index d4d76d7e550..93965b2b6e4 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -71,12 +71,12 @@ off meta; ``_maybe_init_rotary_emb`` stays idempotent for the lazy pattern. """ +import functools import logging from pathlib import Path from typing import NamedTuple import torch -import torch.nn.functional as F from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -85,6 +85,7 @@ from ..dflash.conversion import DFlashDMRegistry from ..dflash.dflash_model import DFlashModel from .dflash_flex_attention import build_draft_block_mask +from .dflash_losses import DraftLogits, LossContext, LossOutput, _block_targets, _masked_top1 from .modeling_dflash import ( # noqa: F401 DFlashAttention, DFlashBaseModelOutput, @@ -133,78 +134,6 @@ def _gather_rows(x, positions): return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1)) -def _block_positions(anchor_positions, block_size, offset, seq_len): - """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. - - Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end. - """ - k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1) - unclamped = anchor_positions.unsqueeze(-1) + k + offset - return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len - - -def _block_gather(x, positions): - """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs].""" - return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions) - - -def _block_tokens(input_ids, anchor_positions, block_size, offset): - """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs].""" - positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1]) - return _block_gather(input_ids, positions) - - -class _BlockTargets(NamedTuple): - """What each block position is trained on: the token at ``anchor + k + label_shift``.""" - - token_ids: torch.Tensor # [B, N, bs] - positions: torch.Tensor # [B, N, bs] their sequence indices, clamped onto the sequence - mask: torch.Tensor # [B, N, bs] 1.0 where the position is supervised - - -def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift): - """Labels of every block position and where they are supervised. - - A position is supervised when its block is kept and its label lies inside the sequence - and the loss mask, and is not the anchor itself, which the draft is given as input. - """ - positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1]) - k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1) - mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() - mask = mask * inside.float() - mask = mask * (k + shift > 0).float() - mask = mask * _block_gather(loss_mask, positions) - return _BlockTargets(_block_gather(input_ids, positions), positions, mask) - - -def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None): - """``mask`` weighted per block position, by D-PACE or by exponential decay. - - ``first`` is the first predicted block position. ``target_ce`` returns the draft's - cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. - """ - block_size = mask.shape[-1] - if scheme == "dpace" and block_size > 1: - with torch.no_grad(): - # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9). - confidences = torch.exp(-target_ce()[..., first:].float()) - dpace = torch.ones_like(mask) - dpace[..., first:] = _dpace_position_weights( - confidences, dpace_alpha, valid_mask=mask[..., first:] - ) - return mask * dpace - if decay_factor > 0: - k = torch.arange(block_size, device=mask.device).view(1, 1, -1) - return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor) - return mask - - -def _masked_top1(logits, target_ids, mask): - """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" - correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) - return correct.sum().float() / (mask.sum() + 1e-6) - - def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal): """The draft's attention visibility rule, as a FlexAttention ``mask_mod``. @@ -247,68 +176,15 @@ def _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device): return attn_mask -def _dpace_position_weights( - confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None -) -> torch.Tensor: - """Compute detached D-PACE per-position weights from draft confidences. +class DraftBatch(NamedTuple): + """What a training step hands its loss: the batch, its blocks and the draft's output.""" - Derived from D-PACE (arXiv:2605.18810). The paper factorizes the per-position - weight (Fig. 2 / Eq. 8) into a *cumulative confidence* times a *continuation - value*, which is equivalently the suffix sum of the cumulative confidences:: - - C_j = prod_{i<=j} q~_i # cumulative confidence (Eq. 8) - w_j = sum_{m>=j} C_m # = C_j * continuation value f~_j - - Each confidence is asymmetrically smoothed toward 1 (Eq. 7):: - - q~_i = (1 - alpha) * q_i + alpha, alpha in (0, 1], - - so the floor ``q~_i >= alpha`` keeps every cumulative product (hence every - weight) strictly positive. We evaluate the suffix sum from its definition as - ``total - exclusive_prefix_sum`` of ``C`` rather than reversing the tensor. - Positions with ``valid_mask == 0`` are multiplicative no-ops in ``C`` and - contribute nothing to the sum, matching the per-token loss mask. Weights are - detached (Eq. 9): they reweight the cross-entropy without adding gradient. - - Args: - confidences: ``[..., L]`` draft confidence ``q_i = exp(-CE_i)`` per position. - alpha: smoothing factor in (0, 1]; raises if outside that range. - valid_mask: optional ``[..., L]`` 0/1 mask; ``None`` treats all positions valid. - - Returns: - Detached weights with the same shape and dtype as ``confidences``. - """ - if not 0.0 < alpha <= 1.0: - raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {alpha}") - - with torch.no_grad(): - smoothed = alpha + (1.0 - alpha) * confidences.float() # Eq. 7 - if valid_mask is not None: - keep = valid_mask.to(torch.bool) - smoothed = torch.where(keep, smoothed, torch.ones_like(smoothed)) - cum_conf = torch.cumprod(smoothed, dim=-1) # Eq. 8 cumulative confidence C_j - if valid_mask is not None: - cum_conf = cum_conf * keep.to(cum_conf.dtype) - # Suffix sum w_j = sum_{m>=j} C_m, written as total minus the exclusive - # prefix sum so no axis reversal is needed (Eq. 8). - inclusive = torch.cumsum(cum_conf, dim=-1) - weights = inclusive[..., -1:] - inclusive + cum_conf - return weights.to(dtype=confidences.dtype) - - -class _DFlashLossTerms(NamedTuple): - """The unreduced pieces behind the block loss, for variants that re-derive it. - - ``ce_per_token`` is ``None`` on the KD path, which never forms a per-position - cross-entropy. ``weights`` carries the position weighting (decay or D-PACE) and - normalizes by ``weight_sum``; ``supervised_mask`` is the unweighted mask the - reported accuracy uses. - """ - - ce_per_token: torch.Tensor | None - weights: torch.Tensor - weight_sum: torch.Tensor - supervised_mask: torch.Tensor + input_ids: torch.Tensor # [B, seq] + loss_mask: torch.Tensor # [B, seq] 1.0 where a token may be a label + anchors: torch.Tensor # [B, N] each block's anchor position + keep: torch.Tensor # [B, N] whether each block is trained + base_outputs: DFlashBaseModelOutput + hidden: torch.Tensor # [B, N * bs, H] the draft's hidden states @DFlashDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) @@ -382,6 +258,16 @@ def modify(self, config): """Initialize DFlash draft module.""" super().modify(config) + # shift_label names a next-token variant's alignment in its export; the alignment is + # fixed by _label_shift, so a config turning it off is refused rather than ignored. + arch_config = config.dflash_architecture_config + if self._label_shift and not arch_config.get("shift_label", True): + raise ValueError( + f"dflash_architecture_config.shift_label={arch_config['shift_label']!r} is not " + f"supported: projector_type {arch_config.get('projector_type')!r} trains only " + "next-token alignment (block position k predicts the token at anchor + k + 1)." + ) + self.dflash_fp32_master_weights = config.dflash_fp32_master_weights base_config = self._base_llm_config @@ -667,6 +553,24 @@ def get_exporter(self): return DFlashExporter(self) + @staticmethod + def _build_loss_mask(input_ids, attention_mask, labels, loss_mask): + """Return the float ``[B, seq_len]`` mask of positions anchor sampling and the loss may use. + + Every DFlash-family training forward calls this so the variants mask the same positions. + """ + mask = torch.ones(input_ids.shape, device=input_ids.device) + if labels is not None: + mask = mask * (labels != LabelSmoother.ignore_index).float() + # Labels carry optional answer-only masking but do not in general mark padding with -100 + # (collators copy them from padded input_ids), so always intersect with attention_mask. + if attention_mask is not None: + mask = mask * attention_mask.float() + # In offline training, assistant mask is dumped and passed as kwarg. + if loss_mask is not None: + mask = mask * loss_mask + return mask + def _sample_anchor_positions(self, seq_len, loss_mask, device): """Randomly sample anchor positions per sample. @@ -917,104 +821,112 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None): bias = getattr(lm_head, "bias", None) return logits if bias is None else logits + bias[token_ids] - def _compute_loss( + def _online_base_outputs( self, - logits, input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - return_terms=False, + attention_mask, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + output_attentions=None, + cache_position=None, + **kwargs, ): - """Compute weighted cross-entropy (or KD) loss and accuracy. - - Args: - logits: Draft model output [B, N*block_size, vocab]. - input_ids: Original input token IDs [B, seq_len]. - anchor_positions: Anchor positions per block [B, N]. - block_keep_mask: Valid block mask [B, N]. - loss_mask: Token-level loss mask [B, seq_len]. - draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``. - Unused here; passed for variants whose head consumes them. - base_outputs: The step's ``DFlashBaseModelOutput``: the KD teacher under - ``dflash_self_logit_distillation``, and read by variants. - return_terms: Also return the unreduced pieces behind the loss, so a variant - can recompose the block objective from a different divergence without - rebuilding the target alignment and position weighting. - TODO: promote this into a shared divergence seam when the DFlash-family - loss code is refactored; DFlash2 is the only consumer today. + """Run the frozen base model: the target hidden states, and its logits for KD.""" + # Multimodal models need the top-level conditional-generation forward so their + # image/video features are inserted before the language model runs. Keep the + # long-standing narrow call for text-only models. + base_forward_kwargs = _multimodal_forward_kwargs(kwargs) + with torch.no_grad(): + if base_forward_kwargs: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=False, + output_attentions=output_attentions, + output_hidden_states=True, + cache_position=cache_position, + return_dict=True, + **base_forward_kwargs, + ) + else: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) - Returns: - (loss, accuracy) tuple. - """ - block_size = self.dflash_block_size - targets = _block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift + if not getattr(raw_outputs, "hidden_states", None): + raise RuntimeError( + "The base model did not return hidden states required for DFlash training. " + "Ensure its top-level multimodal forward supports output_hidden_states=True." + ) + return DFlashBaseModelOutput( + target_hidden=self._target_hidden(raw_outputs.hidden_states), logits=raw_outputs.logits ) - binary_eval_mask = targets.mask.view(-1) - flat_logits = logits.view(-1, logits.size(-1)) - flat_targets = targets.token_ids.view(-1) + def _draft_logits(self, batch): + """The draft's logits per block position; a variant with a head applies it here.""" + logits = self._base_model_lm_head(batch.hidden) + logits = logits.view(*batch.anchors.shape, self.dflash_block_size, -1) + return DraftLogits(final=logits, backbone=logits) - kd = self.dflash_self_logit_distillation - if kd and base_outputs is None: - raise ValueError( - "dflash_self_logit_distillation distills from base_outputs, but none was " - "passed; an override of _compute_loss must forward it." - ) - # Non-KD loss is per-token cross-entropy; compute it once (grad enabled) so the - # D-PACE confidences below can reuse it instead of a second CE pass. The KD path - # optimizes KL, so its confidences need a dedicated no_grad CE pass. - loss_per_token = None - if not kd: - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - - weight_mask = _block_position_weights( - targets.mask, - self.dflash_loss_objective, - 1 - self._label_shift, - self.dflash_loss_decay_factor, - self.dflash_dpace_alpha, - target_ce=lambda: ( - loss_per_token.detach() - if loss_per_token is not None - else F.cross_entropy(flat_logits, flat_targets, reduction="none") - ).view(targets.mask.shape), + @property + def _position_weighting(self): + """How block positions are weighted in the loss: ``"dpace"`` or ``"decay"``.""" + return self.dflash_loss_objective + + def _loss_context(self, batch): + """The step's labels, the draft's logits and the position weights.""" + targets = _block_targets( + batch.input_ids, + batch.anchors, + batch.keep, + batch.loss_mask, + self.dflash_block_size, + self._label_shift, + ) + return LossContext( + batch, + self._draft_logits(batch), + targets, + weighting={ + "scheme": self._position_weighting, + "first": 1 - self._label_shift, + "decay_factor": self.dflash_loss_decay_factor, + "dpace_alpha": self.dflash_dpace_alpha, + }, + teacher=functools.partial(self._teacher_logits, batch.base_outputs, targets.positions), + # The KD path optimizes the soft target; its D-PACE confidences need no gradient. + ce_grad=not self.dflash_self_logit_distillation, ) - flat_weights = weight_mask.view(-1) - valid_count = flat_weights.sum() + 1e-6 - - if valid_count > 1.0: - if kd: - flat_teacher = self._teacher_logits(base_outputs, targets.positions) - flat_teacher = flat_teacher.reshape(flat_logits.shape) - target_soft = torch.softmax(flat_teacher, dim=-1) - draft_logsoft = torch.log_softmax(flat_logits, dim=-1) - kd_loss = -(target_soft * draft_logsoft).sum(dim=-1) - loss = (kd_loss * flat_weights).sum() / valid_count - else: - loss = (loss_per_token * flat_weights).sum() / valid_count + def _compute_loss(self, ctx): + """The weighted cross-entropy (or KD) loss on the served logits, and its accuracy. - with torch.no_grad(): - accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item() - else: + A step whose position weights leave nothing to train returns a zero loss and + accuracy 0.0. + """ + final = ctx.logits.final + if not ctx.trainable: # Through the draft's output, so an FSDP unit around the draft still runs its backward. - loss = flat_logits.sum() * 0.0 - accuracy = 0.0 - - if return_terms: - terms = _DFlashLossTerms( - ce_per_token=loss_per_token, - weights=flat_weights, - weight_sum=valid_count, - supervised_mask=binary_eval_mask, - ) - return loss, accuracy, terms - return loss, accuracy + return LossOutput(final.sum() * 0.0, 0.0, {}) + + if self.dflash_self_logit_distillation: + flat_logits = ctx.flat_logits("final") + target_soft = torch.softmax(ctx.teacher().reshape(flat_logits.shape), dim=-1) + draft_logsoft = torch.log_softmax(flat_logits, dim=-1) + per_token = -(target_soft * draft_logsoft).sum(dim=-1) + else: + per_token = ctx.ce("final") + loss = (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum + + with torch.no_grad(): + accuracy = _masked_top1(final, ctx.targets.token_ids, ctx.targets.mask) + return LossOutput(loss, accuracy.item(), {}) def forward( self, @@ -1030,13 +942,16 @@ def forward( cache_position=None, **kwargs, ): - """Training forward with random anchor sampling. + """Training forward with random anchor sampling, shared by every DFlash-family variant. - Random anchor sampling instead of uniform block division - Bidirectional intra-block attention (no causal constraint) - Context sees strictly before anchor position - - Label alignment: position k predicts token at anchor+k - - Optional loss decay weighting + - Label alignment: position k predicts the token at anchor + k + ``_label_shift`` + - Position weighting by D-PACE or exponential decay + + A variant changes the step only through its hooks: ``_online_base_outputs``, + ``_draft_logits`` (its head), ``_position_weighting`` and ``_compute_loss``. """ if not self.training: if self.dflash_offline: @@ -1060,7 +975,7 @@ def forward( **kwargs, ) - bsz, seq_len = input_ids.shape + seq_len = input_ids.shape[1] block_size = self.dflash_block_size device = input_ids.device @@ -1081,58 +996,23 @@ def forward( and base_outputs.base_hidden is None ): raise KeyError("base_model_hidden_states") - target_hidden = base_outputs.target_hidden else: - # Multimodal models need the top-level conditional-generation forward so their - # image/video features are inserted before the language model runs. Keep the - # long-standing narrow call for text-only models. - base_forward_kwargs = _multimodal_forward_kwargs(kwargs) - use_top_level_forward = bool(base_forward_kwargs) - with torch.no_grad(): - if use_top_level_forward: - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=False, - output_attentions=output_attentions, - output_hidden_states=True, - cache_position=cache_position, - return_dict=True, - **base_forward_kwargs, - ) - else: - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - - if not getattr(raw_outputs, "hidden_states", None): - raise RuntimeError( - "The base model did not return hidden states required for DFlash training. " - "Ensure its top-level multimodal forward supports output_hidden_states=True." - ) - target_hidden = self._target_hidden(raw_outputs.hidden_states) - base_outputs = DFlashBaseModelOutput( - target_hidden=target_hidden, logits=raw_outputs.logits + base_outputs = self._online_base_outputs( + input_ids, + attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + cache_position=cache_position, + **kwargs, ) + target_hidden = base_outputs.target_hidden - # 2. Build loss mask. Labels carry optional answer-only masking, but do - # not in general mark padded tokens with -100 (the VLM collator creates - # them from padded input_ids). Always intersect with attention_mask so - # anchor sampling and loss never include the padded tail. - loss_mask = torch.ones(bsz, seq_len, device=device) - if labels is not None: - loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float() - if attention_mask is not None: - loss_mask = loss_mask * attention_mask.float() - - # In offline training, assistant mask is dumped and passed as kwarg. - if kwargs.get("loss_mask") is not None: - loss_mask = loss_mask * kwargs["loss_mask"] + # 2. Build loss mask. + loss_mask = self._build_loss_mask( + input_ids, attention_mask, labels, kwargs.get("loss_mask") + ) # 3. Random anchor sampling anchor_positions, block_keep_mask = self._sample_anchor_positions( @@ -1170,21 +1050,15 @@ def forward( ) # 6. Compute loss and accuracy - logits = self._base_model_lm_head(hidden) - loss, accuracy = self._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - draft_hidden=hidden, - base_outputs=base_outputs, + batch = DraftBatch( + input_ids, loss_mask, anchor_positions, block_keep_mask, base_outputs, hidden ) - + out = self._compute_loss(self._loss_context(batch)) return ModelOutput( - loss=loss, + loss=out.loss, logits=base_outputs.logits, - train_acc=[[accuracy]], + train_acc=[[out.accuracy]], + draft_metrics=out.metrics, ) @torch.no_grad() @@ -1219,36 +1093,42 @@ def pseudo_speculative_generate(self, input_ids, steps=1): The draft model's KV injection concatenates projected target_hidden as K/V with the block's own K/V, enabling the draft to "see" the target's context. - 6. **Decode**: Apply ``lm_head`` to draft hidden states at positions 1..block_size-1 - (skip position 0 which is the known anchor). Argmax gives draft tokens. + 6. **Decode**: Apply ``lm_head`` to draft hidden states from position + ``1 - _label_shift`` on, the first one that predicts a token after the anchor + (position 1 for DFlash, whose position 0 is the known anchor). Argmax gives draft + tokens. 7. **Return**: ``(base_token, draft_tokens[:steps])`` — base token is always - returned; draft tokens are truncated to ``steps`` (default: block_size-1). + returned; draft tokens are truncated to ``steps``. Note: This method re-runs the full target model from scratch on each call (no KV cache). For AR validation, it is called repeatedly with growing ``input_ids`` by ``AcceptanceRateValidation.validate()``. The ``steps`` - parameter should be set to ``block_size - 1`` for full block evaluation. + parameter should be set to ``block_size - 1 + _label_shift`` for full block + evaluation. Args: input_ids: Input token IDs [B, seq_len]. - steps: Number of draft tokens to return (capped at block_size-1). + steps: Number of draft tokens to return (capped at block_size - 1 + _label_shift). Returns: base_token: Next token from base model [B, 1]. - draft_tokens: Draft tokens [B, min(steps, block_size-1)] or None if steps < 1. + draft_tokens: Draft tokens [B, min(steps, block_size - 1 + _label_shift)] or None + if steps < 1. """ base_token, hidden_states = self._generation_base(input_ids) if steps < 1: return base_token, None draft_hidden = self._draft_block(base_token, hidden_states) - # Logits on positions 1..block_size-1 (skip anchor at position 0) - draft_logits = self._base_model_lm_head(draft_hidden[:, 1:, :]) - draft_tokens = draft_logits.argmax(dim=-1) # [B, block_size-1] + # Position k predicts the token at anchor + k + _label_shift, so this one drafts the + # first token after the anchor. + first = 1 - self._label_shift + draft_logits = self._base_model_lm_head(draft_hidden[:, first:, :]) + draft_tokens = draft_logits.argmax(dim=-1) # [B, block_size - first] # Return up to `steps` tokens. base_token already follows input_ids; draft_tokens is # produced on the draft's device, and callers cat both onto the running sequence. - num_tokens = min(steps, self.dflash_block_size - 1) + num_tokens = min(steps, self.dflash_block_size - first) return base_token, draft_tokens[:, :num_tokens].to(input_ids.device) diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 6e4171c9210..5d8760133e5 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -39,7 +39,8 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens +from .dflash_losses import LossOutput, _block_tokens +from .hf_dflash import HFDFlashModel from .modeling_dflash2 import DFlash2Module __all__ = ["HFDFlash2Model"] @@ -83,22 +84,6 @@ def modify(self, config): "optimizes a soft target instead. Set dflash_self_logit_distillation=false, " "or dflash_lk_loss_type='ce' to keep distillation." ) - self._selector_metrics = None - - def forward(self, *args, **kwargs): - """Run the DFlash training forward and attach the candidate-selector metrics. - - The variant is the objective, so the pipeline above it is inherited verbatim; - this only carries out what ``_compute_loss`` produced. Without it - ``selector_coverage`` -- the only signal distinguishing a selector that is - choosing from one being handed the gold token -- never reaches the logs. - """ - self._selector_metrics = None - outputs = super().forward(*args, **kwargs) - if self._selector_metrics is not None: - outputs["selector_metrics"] = self._selector_metrics - self._selector_metrics = None - return outputs def get_exporter(self): """Get the exporter for the DFlash2 draft model.""" @@ -163,7 +148,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas # every training step. The trainer converts them at the logging boundary. return loss, accuracy.detach(), coverage.detach() - def _lk_loss(self, terms): + def _lk_loss(self, ctx): """Re-weight the block objective between cross-entropy and acceptance. Both terms are read off the same per-position cross-entropy the backbone loss @@ -181,66 +166,34 @@ def _lk_loss(self, terms): The blend weight is detached, so it reshapes the objective without adding a gradient path of its own. """ - ce, weights, weight_sum = terms.ce_per_token, terms.weights, terms.weight_sum - assert ce is not None, ( - "the KD path produced no per-position cross-entropy, so q(gold) is " - "unavailable and the blend would silently fall back to it; modify() is " - "supposed to have rejected this combination at convert time" - ) + ce, weights, weight_sum = ctx.ce("final"), ctx.weights.view(-1), ctx.weight_sum gold_probability = torch.exp(-ce) tv_loss = ((1.0 - gold_probability) * weights).sum() / weight_sum if self.dflash_lk_loss_type == "tv": return tv_loss ce_loss = (ce * weights).sum() / weight_sum - mask = terms.supervised_mask + mask = ctx.targets.mask.view(-1) acceptance = (gold_probability.detach() * mask).sum() / (mask.sum() + 1e-6) ce_share = self.dflash_lk_ce_scale * torch.exp(-self.dflash_lk_ce_decay * acceptance) return ce_share * ce_loss + (1.0 - ce_share) * tv_loss - def _compute_loss( - self, - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - ): + def _compute_loss(self, ctx): """Backbone DFlash loss plus the candidate-selector cross-entropy. - Reuses ``HFDFlashModel._compute_loss`` for the backbone term, then rebuilds - the same target/weight alignment for the selector term. Reported accuracy - stays the backbone's top-1, so DFlash and DFlash2 runs remain comparable; - the selector's own accuracy is logged separately. + Reported accuracy stays the backbone's top-1, so DFlash and DFlash2 runs remain + comparable; the selector's own accuracy and coverage are added to the metrics. """ - loss, accuracy, terms = super()._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - draft_hidden=draft_hidden, - base_outputs=base_outputs, - return_terms=True, - ) + out = super()._compute_loss(ctx) + loss = out.loss if self.dflash_lk_loss_type != "ce": - loss = self._lk_loss(terms) - if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: - return loss, accuracy + loss = self._lk_loss(ctx) + if self.dflash_selector_loss_alpha <= 0: + return out._replace(loss=loss) - bsz = input_ids.shape[0] + batch, targets = ctx.batch, ctx.targets + bsz, n_blocks = batch.anchors.shape block_size = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - - # The backbone's supervision mask, without its position weighting (decay/D-PACE): - # that shapes *where* the backbone spends capacity, while the selector should learn - # every position's transition equally. - targets = _block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift - ) # Teacher-forced predecessor of block position k is the real token at anchor+k-1. # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the @@ -251,18 +204,24 @@ def _compute_loss( # from the anchor token, so ITS position 0 corresponds to block offset 1. A caller # that hands greedy_path the full 0..block_size-1 candidate set is off by one. predecessor_ids = _block_tokens( - input_ids, anchor_positions, block_size, self._label_shift - 1 + batch.input_ids, batch.anchors, block_size, self._label_shift - 1 ) + # The backbone's supervision mask, without its position weighting (decay/D-PACE): + # that shapes *where* the backbone spends capacity, while the selector should learn + # every position's transition equally. selector_loss, selector_accuracy, selector_coverage = self._selector_loss( - logits.reshape(bsz, n_blocks, block_size, -1), + ctx.logits.backbone, targets.token_ids, - draft_hidden.reshape(bsz, n_blocks, block_size, -1), + batch.hidden.reshape(bsz, n_blocks, block_size, -1), predecessor_ids, targets.mask, ) - self._selector_metrics = { + metrics = { + **out.metrics, "selector_accuracy": selector_accuracy, "selector_coverage": selector_coverage, } - return loss + self.dflash_selector_loss_alpha * selector_loss, accuracy + return LossOutput( + loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics + ) diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index fc33a5155b6..2e924b121af 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -68,13 +68,8 @@ from transformers.utils import ModelOutput from ..dflash.conversion import DominoDMRegistry -from .hf_dflash import ( - HFDFlashModel, - _block_position_weights, - _block_targets, - _block_tokens, - _masked_top1, -) +from .dflash_losses import _block_position_weights, _block_targets, _block_tokens, _masked_top1 +from .hf_dflash import HFDFlashModel from .modeling_dflash import DFlashBaseModelOutput from .modeling_domino import DominoModule @@ -132,10 +127,6 @@ def modify(self, config): self.dflash_lambda_base_start = config.dflash_lambda_base_start self.dflash_lambda_base_decay_ratio = config.dflash_lambda_base_decay_ratio self._lambda_base = self.dflash_lambda_base_start - if not getattr(self.dflash_module, "shift_label", True): - raise NotImplementedError( - "Domino currently supports shift_label=True (next-token alignment) only." - ) def get_exporter(self): """Get the exporter for the Domino draft model.""" @@ -266,14 +257,6 @@ def forward( two-term loss. Eval/offline-eval is delegated to the DFlash parent. """ if not self.training: - # Eval delegates to the DFlash backbone; the correction head is not - # applied yet, so warn once that acceptance rates are backbone-only. - if not getattr(self, "_warned_eval_head_bypass", False): - logger.warning( - "Domino eval uses the DFlash backbone only (correction head not " - "applied yet); reported acceptance rates are backbone-only." - ) - self._warned_eval_head_bypass = True return super().forward( input_ids=input_ids, attention_mask=attention_mask, @@ -365,6 +348,38 @@ def forward( return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], domino_metrics=metrics) + @torch.no_grad() + def pseudo_speculative_generate(self, input_ids, steps=1): + """Draft the block left to right, adding the correction head as training does. + + Block position ``k`` predicts the token after the anchor and the ``k`` tokens drafted + before it. From ``pure_draft_prefix_len`` on, its backbone logits get the correction + from the GRU state over exactly those tokens. + """ + base_token, hidden_states = self._generation_base(input_ids) + if steps < 1: + return base_token, None + draft_hidden = self._draft_block(base_token, hidden_states) + device = draft_hidden.device + # Under device_map="auto" the embedding table and a tied lm_head sit on the first GPU, + # not the draft's, and their outputs stay there. + backbone_logits = self._base_model_lm_head(draft_hidden).to(device) # [B, block_size, V] + + m = self.dflash_module + prev_token = base_token.squeeze(-1).to(device) + gru_state = None + draft_tokens = [] + for k in range(min(steps, self.dflash_block_size)): + prev_emb = self._base_model_embeddings(prev_token).to(device) + gru_out, gru_state = m.prefix_gru(prev_emb.unsqueeze(1), gru_state) + logits = backbone_logits[:, k, :] + if k >= m.pure_draft_prefix_len: + correction = m.embed_proj(torch.cat([draft_hidden[:, k, :], gru_out[:, 0, :]], -1)) + logits = logits + correction + prev_token = logits.argmax(dim=-1) + draft_tokens.append(prev_token) + return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device) + class DominoLambdaCallback(TrainerCallback): """Update the model's ``lambda_base`` from the HF Trainer global step. diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index b29d11584ce..0891f9a924b 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -63,19 +63,20 @@ import torch import torch.nn.functional as F -import torch.utils.checkpoint from transformers import PreTrainedModel from transformers.trainer_pt_utils import LabelSmoother from transformers.utils import ModelOutput from ..dflash.conversion import DSparkDMRegistry -from .hf_dflash import ( - HFDFlashModel, +from .dflash_losses import ( _block_position_weights, _block_targets, _block_tokens, _masked_top1, + _tvd_chunk, + _tvd_per_token, ) +from .hf_dflash import HFDFlashModel from .modeling_dflash import DFlashBaseModelOutput from .modeling_dspark import DSparkModule @@ -84,37 +85,6 @@ __all__ = ["HFDSparkModel"] -def _tvd_chunk(a, b): - """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``.""" - return (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) - - -def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None): - """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean. - - Materializing both [N, vocab] float32 softmax tensors at once OOMs at large - N (num_anchors*block_size) and a 150k vocab. We compute it in row chunks and - gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward - rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math - is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. - - Chunks with ``Tensor.split``, not slicing (each slice's backward zero-fills a full - [N, vocab] tensor) nor ``torch.chunk`` (different chunk shapes, so recompiles). - """ - _chunk = chunk_fn or _tvd_chunk - outs = [] - for a, b in zip( - final_logits.split(chunk_size, dim=0), - teacher_logits.split(chunk_size, dim=0), - strict=True, - ): - if torch.is_grad_enabled() and a.requires_grad: - outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False)) - else: - outs.append(_chunk(a, b)) - return torch.cat(outs, dim=0) - - @DSparkDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) class HFDSparkModel(HFDFlashModel): """DFlash model with the DSpark sequential (Markov) + confidence head. @@ -419,7 +389,8 @@ def pseudo_speculative_generate(self, input_ids, steps=1): if steps < 1: return base_token, None draft_hidden = self._draft_block(base_token, hidden_states) - backbone_logits = self._base_model_lm_head(draft_hidden) # [B, block_size, V] + # Under device_map="auto" a tied lm_head sits with the embedding table, not the draft. + backbone_logits = self._base_model_lm_head(draft_hidden).to(draft_hidden.device) # Autoregressive Markov sampling over the block. m = self.dflash_module diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index 908881f8c20..97c8f6aea25 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -77,7 +77,8 @@ from transformers import PreTrainedModel from ..dflash.conversion import LiLiCorrDMRegistry -from .hf_dflash import HFDFlashModel, _block_positions, _block_targets +from .dflash_losses import LossOutput, _block_positions +from .hf_dflash import HFDFlashModel from .modeling_lilicorr import LiLiCorrModule logger = logging.getLogger(__name__) @@ -173,7 +174,6 @@ def modify(self, config): self.dflash_lilicorr_w_cal, self.dflash_lilicorr_margin, ) - self._lilicorr_metrics = None @property def _needs_base_logits(self) -> bool: @@ -190,77 +190,33 @@ def get_exporter(self): return LiLiCorrExporter(self) - def _compute_loss( - self, - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - ): + def _compute_loss(self, ctx): """Add the LiLiCorr lattice terms to the DFlash block loss. The returned accuracy is the reranker's greedy path-prefix fraction — the quantity that tracks acceptance length — rather than the backbone's per-token accuracy, which is reported as ``origin_accuracy`` in the metrics instead. - - The two things this needs beyond the shared signature -- the target-layer hidden - states it anchors on, and the target distribution its penalty and calibration - terms read -- both come from ``base_outputs``, the container - ``HFDFlashModel.forward`` already builds for them. """ - target_hidden = base_outputs.target_hidden if base_outputs is not None else None - if draft_hidden is None or target_hidden is None: - raise ValueError( - "LiLiCorr requires draft_hidden and target_hidden in _compute_loss: the " - "reranker reads the draft's per-slot hidden states and anchors on a " - "committed target row." - ) - - origin_loss, origin_accuracy = super()._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - base_outputs=base_outputs, - ) + out = super()._compute_loss(ctx) + batch, targets = ctx.batch, ctx.targets # The backbone's own supervision, before its position weighting, so both losses # train the same positions. - targets = _block_targets( - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - self.dflash_block_size, - self._label_shift, - ) lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss( - logits=logits, + logits=ctx.logits.final, target_ids=targets.token_ids, # Block position 0 is the anchor -- the one row whose input embedding is a real # verified token -- so the lattice covers positions 1..bs-1. slot_mask=targets.mask[:, :, 1:] > 0.5, - anchor_positions=anchor_positions, - draft_hidden=draft_hidden, - target_hidden=target_hidden, - base_outputs=base_outputs, + anchor_positions=batch.anchors, + draft_hidden=batch.hidden, + target_hidden=batch.base_outputs.target_hidden, + base_outputs=batch.base_outputs, ) # The identity `loss == origin_loss + lilicorr_loss` is the cheap parity check # on this objective, so both halves are reported next to the total. - # Batched into one device-to-host sync rather than two per-step float() calls. - origin_scalars = torch.stack( - [ - origin_loss.detach().reshape(()).float(), - torch.as_tensor(origin_accuracy, dtype=torch.float32, device=origin_loss.device), - ] - ).tolist() - metrics["origin_loss"], metrics["origin_accuracy"] = origin_scalars - self._lilicorr_metrics = metrics - return origin_loss + lilicorr_loss, accuracy + metrics["origin_loss"] = out.loss.detach().float() + metrics["origin_accuracy"] = torch.full((), out.accuracy, device=out.loss.device) + return LossOutput(out.loss + lilicorr_loss, accuracy, metrics) def _compute_lilicorr_loss( self, @@ -277,7 +233,7 @@ def _compute_lilicorr_loss( Returns ``(loss, accuracy, metrics)``. ``loss`` already carries the absolute term weights, ``accuracy`` is the greedy path-prefix fraction, and ``metrics`` - holds Python floats read back in a single device sync. + holds detached 0-dim tensors. """ module = self.dflash_module head = module.lilicorr @@ -448,17 +404,18 @@ def _compute_lilicorr_loss( raw_metrics["lilicorr_penalty"] = penalty_loss raw_metrics["lilicorr_calibration"] = cal_loss - # One device sync for every scalar, rather than one per metric. - names = list(raw_metrics) - stacked = torch.stack([raw_metrics[name].detach().reshape(()).float() for name in names]) - metrics = dict(zip(names, stacked.tolist(), strict=True)) + metrics = {name: value.detach().reshape(()).float() for name, value in raw_metrics.items()} # Weights recorded alongside the terms so a run states its own composition - # without the reader expanding the config. - metrics["lilicorr_w_ce"] = self.dflash_lilicorr_w_ce - metrics["lilicorr_w_margin"] = self.dflash_lilicorr_w_margin - metrics["lilicorr_w_pen"] = self.dflash_lilicorr_w_pen - metrics["lilicorr_w_cal"] = self.dflash_lilicorr_w_cal - accuracy = metrics["lilicorr_selected_prefix"] / float(num_slots) + # without the reader expanding the config. float64 holds the configured values exactly. + weights = { + "lilicorr_w_ce": self.dflash_lilicorr_w_ce, + "lilicorr_w_margin": self.dflash_lilicorr_w_margin, + "lilicorr_w_pen": self.dflash_lilicorr_w_pen, + "lilicorr_w_cal": self.dflash_lilicorr_w_cal, + } + for name, weight in weights.items(): + metrics[name] = torch.full((), weight, dtype=torch.float64, device=loss.device) + accuracy = metrics["lilicorr_selected_prefix"].item() / float(num_slots) return loss, accuracy, metrics def _distractor_penalty( @@ -676,18 +633,3 @@ def block_mean(values): / supervised_count, } return metrics - - def forward(self, *args, **kwargs): - """Run the DFlash training forward and attach the LiLiCorr metrics. - - The whole variant is the objective, so the pipeline above it is inherited - verbatim rather than reimplemented; this only carries the metrics that - ``_compute_loss`` produced out to the caller, the way DSpark and Domino do - from their own forwards. - """ - self._lilicorr_metrics = None - outputs = super().forward(*args, **kwargs) - if self._lilicorr_metrics is not None: - outputs["lilicorr_metrics"] = self._lilicorr_metrics - self._lilicorr_metrics = None - return outputs diff --git a/modelopt/torch/speculative/plugins/modeling_domino.py b/modelopt/torch/speculative/plugins/modeling_domino.py index e6e9fc47c94..566967b3e7c 100644 --- a/modelopt/torch/speculative/plugins/modeling_domino.py +++ b/modelopt/torch/speculative/plugins/modeling_domino.py @@ -82,7 +82,6 @@ def __init__(self, config): f"pure_draft_prefix_len must be in [0, {self.block_size - 1}] " f"(block_size={self.block_size}), got {self.pure_draft_prefix_len}." ) - self.shift_label = getattr(config, "shift_label", True) # Causal state over the block's token embeddings. bias=False matches the # reference checkpoint (only weight_ih_l0 / weight_hh_l0 are stored). diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py index e6464380d17..11f1ff66788 100644 --- a/modelopt/torch/speculative/plugins/modeling_dspark.py +++ b/modelopt/torch/speculative/plugins/modeling_dspark.py @@ -89,10 +89,6 @@ def __init__(self, config): self.projector_type = getattr(config, "projector_type", "dspark") self.markov_rank = int(config.markov_rank) self.markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower() - # DSpark treats the anchor as the first prediction position (next-token - # alignment), matching Domino's shift_label path; kept as an attribute so - # the wrapper and exporter can read it uniformly. - self.shift_label = True if self.markov_rank <= 0: raise ValueError(f"DSpark requires markov_rank > 0, got {self.markov_rank}.") if self.markov_head_type not in ("vanilla", "gated", "rnn"): diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e125381ce6e..e7237b11b8f 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/modelopt_recipes/general/speculative_decoding/domino.yaml @@ -37,8 +37,6 @@ training: cp_size: 1 dp_shard_size: 1 disable_tqdm: true - # Keep off: eval runs the DFlash backbone only (correction head not applied - # yet), so AR would reflect the backbone alone, not the trained model. estimate_ar: false ar_validate_steps: 0 answer_only_loss: true diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index a071b071372..db69f3c4b9e 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml @@ -40,9 +40,6 @@ training: cp_size: 1 dp_shard_size: 1 disable_tqdm: true - # Keep off: eval runs the DFlash backbone only (Markov head not applied yet), - # so AR would reflect the backbone alone, not the trained model. Compare via - # export + the offline acceptance-length harness instead. estimate_ar: false ar_validate_steps: 0 answer_only_loss: true diff --git a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml index eb78e80c466..d96780e05d2 100644 --- a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml +++ b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml @@ -54,7 +54,6 @@ training: cp_size: 1 dp_shard_size: 1 disable_tqdm: true - # Eval runs the DFlash backbone only (no Markov head), so AR would understate the model. estimate_ar: false ar_validate_steps: 0 answer_only_loss: true diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 3f8434d1000..8494f6f4b76 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -28,16 +28,12 @@ The committed example is **online**, 2 tasks: `--full-conversations` matters: it keeps real assistant completions so `answer_only_loss` has assistant spans to mask. -**The inference side is intentionally not wired up yet.** The Domino correction head -is not applied in `pseudo_speculative_generate` or in the serving stack, so the -example ships no vLLM smoke test and no MT-Bench AR eval. Do not treat their absence -as a broken config. - -When that path lands, copy the smoke-test and AR-eval steps from -`hf_online_dflash.yaml` (its `task_1` and `task_2`) and append them to the Domino -config as **`task_2` and `task_3`** — Domino's `task_0`/`task_1` are already the -dataset build and training, so keep the source file's numbering and the destination's -distinct. +The example ships no vLLM smoke test or MT-Bench AR eval; do not treat their absence +as a broken config. Serving is not wired up yet, but HF AR validation +(`pseudo_speculative_generate`) applies the correction head: to measure AR, append +`hf_online_dflash.yaml`'s `task_2` as Domino's **`task_2`** with `--steps` equal to the +block size (Domino drafts the whole block: 16, not DFlash's 15). Add the smoke test +(`task_1` there) as `task_3` once serving lands. ## Recipe and training knobs @@ -87,11 +83,6 @@ export is the whole in-pipeline signal. ## Quality gate -**Do not trust in-training AR for Domino.** The recipe pins `estimate_ar: false` and -`ar_validate_steps: 0` deliberately: eval delegates to the DFlash backbone with the -correction head not applied, so reported acceptance rates are backbone-only. The code -logs this once as a warning — treat that warning as expected, not as a defect. - The training-regression gate from `dflash.md` applies and the Qwen3-8B example sets it: `MAX_FINAL_LOSS=5.0`, `MIN_FINAL_ACC=0.15`, checked by `check_regression.py` against `trainer_state.json`. Since no inference metric is produced, this is currently @@ -107,9 +98,8 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | Error pattern | Root cause | Fix | | --- | --- | --- | | `Domino (projector_type='domino') requires ['emb_dim', 'gru_hidden_dim'] in dflash_architecture_config` | GRU head dims missing | Set both in `dflash_architecture_config` | -| `Domino currently supports shift_label=True (next-token alignment) only` | `shift_label=false` | Leave it at `true` | +| `shift_label=False is not supported: projector_type 'domino' trains only next-token alignment` | `shift_label: false` | Leave it at `true` | | `pure_draft_prefix_len must be in [0, N] (block_size=B), got X` | Prefix length >= block size | Lower it below `block_size` | | `DominoLambdaCallback: state.max_steps unset (<=0); lambda_base curriculum disabled` (warning) | `training.max_steps` not set | Set `training.max_steps`, else the curriculum never runs | -| `Domino eval uses the DFlash backbone only ...` (warning) | Correction head not applied at eval | Expected — do not chase it; evaluate after export once the inference path lands | | DDP error about unused parameters | `ddp_find_unused_parameters` turned off | Keep it `true` | | Draft trains but quality is poor | Draft dims left at `Qwen3Config` defaults instead of matching the base | Set the GQA/FFN dims explicitly | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index f857d193d3f..bb2c6992ac5 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -117,15 +117,10 @@ the only in-pipeline evidence is training progress plus the export landing in ## Quality gate -**Do not trust in-training AR for DSpark.** The recipe pins `estimate_ar: false` and -`ar_validate_steps: 0` deliberately: eval runs the DFlash backbone only, with the -Markov head not applied, so any reported AR reflects the backbone alone rather than -the trained model. - -`pseudo_speculative_generate` *is* overridden for DSpark (unlike Domino), so a -non-offline model can generate correctly — but the offline/streaming path deletes base -layers and refuses. Evaluate by exporting and running the offline acceptance-length -harness separately. +`pseudo_speculative_generate` applies the Markov head, so AR validation measures the +trained drafter on an online model — but the offline/streaming path deletes base +layers and refuses. Evaluate those by exporting and running the offline +acceptance-length harness separately. Otherwise the training-regression gate from `dflash.md` (`MAX_FINAL_LOSS`, `MIN_FINAL_ACC` via `check_regression.py`) applies; none of the streaming examples set @@ -151,6 +146,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `DSpark (projector_type='dspark') requires 'markov_rank' (> 0) in dflash_architecture_config` | Markov head dimension missing | Set `dflash_architecture_config.markov_rank` | | `DSpark requires markov_rank > 0, got N` | Non-positive value | Set a positive rank | | `Unsupported markov_head_type: '...'. Expected 'vanilla', 'gated' or 'rnn'` | Typo or unsupported head | Use one of the three | +| `shift_label=False is not supported: projector_type 'dspark' trains only next-token alignment` | `dflash_architecture_config.shift_label: false` | Remove it; DSpark only trains next-token alignment | | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 | | `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export | | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) | diff --git a/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md b/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md index 8bde138c1b0..be7bba191f7 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md +++ b/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md @@ -63,12 +63,13 @@ produces an in-pipeline metric**, so what you check depends on the sheet: metric, and no launcher example sets one. A green benchmark task therefore says the run finished, not that the metric passed. Reporting PASS without having read the number off the log is a false PASS on the run's headline metric. -- **Sheet defines no inference metric** (currently Domino and DSpark — their eval path - runs the DFlash backbone with the new head bypassed, and Domino ships no benchmark - task at all) — do **not** go looking for a benchmark log. Report the training - regression gate instead, and state plainly that acceptance quality requires a - separate evaluation of the exported checkpoint. Never report a backbone-only - acceptance rate as the model's result. +- **Sheet defines no inference metric** (currently the Domino, LiLiCorr and DSpark + training examples and the online DFlash2 one, which ship no benchmark task) — do + **not** go looking for a benchmark log. Report the training regression gate instead, + and state plainly that acceptance quality requires a separate evaluation of the + exported checkpoint. Never report an AR measured by HF AR validation + (`pseudo_speculative_generate`) as the model's result for a variant whose head it + skips (LiLiCorr's reranker, DFlash2's selector). Where the gate is the training regression check (`check_regression.py` against `trainer_state.json`), confirm the `=== Regression Check ===` block is actually diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 5f2e8adc045..a0fc7301f39 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -22,6 +22,7 @@ import torch from _test_utils.torch.speculative.dflash import get_dflash_config from _test_utils.torch.transformers_models import get_tiny_llama +from accelerate import dispatch_model import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.plugins.dflash_flex_attention import is_block_mask @@ -296,6 +297,48 @@ def test_generate_gathers_sharded_hidden_states(self): assert draft_tokens.shape == (1, 2) +class TestLmHeadOffTheDraftDevice: + """``device_map="auto"`` keeps a tied ``lm_head`` with ``embed_tokens`` on the first GPU. + + The draft can land on another GPU, so a decode that adds the draft's own outputs to the + backbone logits must bring those logits to the draft first. + """ + + @pytest.mark.parametrize( + "head_fields", + [ + {}, + {"projector_type": "dspark", "markov_rank": 4}, + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + ], + ids=["dflash", "dspark", "domino"], + ) + def test_generate(self, head_fields): + config = get_dflash_config() + config["dflash_architecture_config"].update(head_fields) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + device_map = { + "model.embed_tokens": "cpu", + "lm_head": "cpu", + "model.layers": 0, + "model.norm": 0, + "model.rotary_emb": 0, + "dflash_module": 0, + } + # main_device="cpu" runs the "cpu" entries there rather than offloading them to the + # GPU, so the CPU stands in for the first GPU. + dispatch_model(model.eval(), device_map, main_device="cpu") + + input_ids = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + base_token, draft_tokens = model.pseudo_speculative_generate(input_ids, steps=2) + + assert base_token.device == input_ids.device + assert draft_tokens.device == input_ids.device + assert draft_tokens.shape == (1, 2) + + class TestDFlashFlexAttentionGPU: """FlexAttention path must match the dense-mask SDPA path it replaces. diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index d100b4a8f79..57487aaa522 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -22,7 +22,7 @@ import logging import os from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -39,13 +39,17 @@ import modelopt.torch.speculative as mtsp import modelopt.torch.speculative.plugins.dflash_flex_attention as dflash_flex_attention import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash +from modelopt.torch.speculative.plugins.dflash_losses import ( + DraftLogits, + _block_targets, + _dpace_position_weights, +) from modelopt.torch.speculative.plugins.hf_dflash import ( DFlashAttention, DFlashBaseModelOutput, DFlashModule, + DraftBatch, HFDFlashModel, - _block_targets, - _dpace_position_weights, build_target_layer_ids, ) from modelopt.torch.speculative.plugins.hf_dflash2 import HFDFlash2Model @@ -64,6 +68,15 @@ SEQ_LEN = 16 # must be multiple of BLOCK_SIZE +def _compute_loss(model, logits, input_ids, anchors, keep, loss_mask, base_outputs=None): + """``model._compute_loss`` on given draft logits [B, N * bs, V]: (loss, accuracy).""" + batch = DraftBatch(input_ids, loss_mask, anchors, keep, base_outputs, hidden=None) + block_logits = logits.view(*anchors.shape, BLOCK_SIZE, -1) + with patch.object(model, "_draft_logits", return_value=DraftLogits(block_logits, block_logits)): + out = model._compute_loss(model._loss_context(batch)) + return out.loss, out.accuracy + + class TestDFlashConvert: """Test DFlash model conversion.""" @@ -274,7 +287,7 @@ def _converted_model(self, objective, **overrides): def test_compute_loss_dpace_branch(self): """Default dpace objective produces a finite loss and valid accuracy.""" model = self._converted_model("dpace", dflash_self_logit_distillation=False) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = _compute_loss(model, *self._make_inputs()) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 @@ -283,7 +296,7 @@ def test_compute_loss_decay_branch(self): model = self._converted_model( "decay", dflash_loss_decay_factor=4.0, dflash_self_logit_distillation=False ) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = _compute_loss(model, *self._make_inputs()) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 @@ -293,15 +306,15 @@ def test_compute_loss_dpace_kd_branch(self): model = self._converted_model("dpace") inputs = self._make_inputs(vocab=vocab) teacher = DFlashBaseModelOutput(None, logits=torch.randn(1, SEQ_LEN, vocab)) - loss, acc = model._compute_loss(*inputs, base_outputs=teacher) + loss, acc = _compute_loss(model, *inputs, base_outputs=teacher) assert torch.isfinite(loss).item() assert 0.0 <= acc <= 1.0 - def test_kd_without_base_outputs_raises(self): + def test_kd_without_a_teacher_raises(self): """KD is decided by the config, so a missing teacher cannot fall back to CE.""" model = self._converted_model("dpace") - with pytest.raises(ValueError, match="base_outputs"): - model._compute_loss(*self._make_inputs()) + with pytest.raises(ValueError, match="base model's distribution"): + _compute_loss(model, *self._make_inputs(), base_outputs=DFlashBaseModelOutput(None)) class TestDFlashSaveRestore: @@ -1481,8 +1494,9 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self): anchors = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks] keep = torch.ones(bsz, n_blocks) - base_loss, base_acc = model._compute_loss(logits, input_ids, anchors, keep, loss_mask) - padded_loss, padded_acc = model._compute_loss( + base_loss, base_acc = _compute_loss(model, logits, input_ids, anchors, keep, loss_mask) + padded_loss, padded_acc = _compute_loss( + model, torch.cat([logits, torch.randn(bsz, pad * BLOCK_SIZE, vocab)], dim=1), input_ids, torch.cat([anchors, torch.zeros(bsz, pad, dtype=anchors.dtype)], dim=1), @@ -1529,6 +1543,36 @@ def test_block_targets_follow_the_label_shift(model_cls, expected): assert torch.equal(targets.mask, torch.tensor([mask], dtype=torch.float32)) +@pytest.mark.parametrize( + "head_fields", + [ + {"projector_type": "dspark", "markov_rank": 4}, + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + ], + ids=["dspark", "domino"], +) +def test_next_token_variants_refuse_shift_label_false_at_convert(head_fields): + """Left unset, as the recipes do, shift_label converts; set to false, it is refused. + + Restore does not refuse it: DSpark used to accept the field, ignore it in training and + save it into the checkpoint's modelopt_state. + """ + config = get_dflash_config() + config["dflash_architecture_config"].update(head_fields) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + + state = mto.modelopt_state(model) + saved_config = state["modelopt_state_dict"][0][1]["config"] + saved_config["dflash_architecture_config"]["shift_label"] = False + restored = mto.restore_from_modelopt_state(get_tiny_llama(num_hidden_layers=4), state) + assert restored._label_shift == 1 + + config["dflash_architecture_config"]["shift_label"] = False + with pytest.raises(ValueError, match="only next-token alignment"): + mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) + + class TestTeacherLogits: """The base distribution is projected only at the rows, and vocab entries, asked for. @@ -1622,9 +1666,9 @@ def test_kd_loss_matches_full_logits(self): hidden = self._hidden(model, bsz=1) inputs = TestDPaceLossIntegration._make_inputs(vocab=model.config.vocab_size) full = DFlashBaseModelOutput(None, logits=model._base_model_lm_head(hidden)) - loss_full, _ = model._compute_loss(*inputs, base_outputs=full) - loss, _ = model._compute_loss( - *inputs, base_outputs=DFlashBaseModelOutput(None, base_hidden=hidden) + loss_full, _ = _compute_loss(model, *inputs, base_outputs=full) + loss, _ = _compute_loss( + model, *inputs, base_outputs=DFlashBaseModelOutput(None, base_hidden=hidden) ) torch.testing.assert_close(loss, loss_full) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index c43677818c0..21ca9ca59e5 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -26,6 +26,7 @@ import json from copy import deepcopy +from unittest.mock import patch import pytest import torch @@ -312,7 +313,7 @@ def test_selector_metrics_reported(self): out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) # Carried out on the forward output, not left on a private attribute, so the # trainer can log them; and kept as tensors to avoid a per-step CPU-GPU sync. - metrics = out["selector_metrics"] + metrics = out["draft_metrics"] for key in ("selector_accuracy", "selector_coverage"): assert torch.is_tensor(metrics[key]) assert 0.0 <= metrics[key].item() <= 1.0 @@ -373,16 +374,16 @@ def test_overfits_a_single_batch(self): optimizer.step() assert out.train_acc[0][0] > 0.9 - assert out["selector_metrics"]["selector_accuracy"].item() > 0.9 - assert out["selector_metrics"]["selector_coverage"].item() > 0.9 + assert out["draft_metrics"]["selector_accuracy"].item() > 0.9 + assert out["draft_metrics"]["selector_coverage"].item() > 0.9 # Candidate selection is a topk over draft logits; check it in bf16, the training dtype. model.to(torch.bfloat16) torch.manual_seed(0) with torch.no_grad(): out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) - assert out["selector_metrics"]["selector_accuracy"].item() > 0.9 - assert out["selector_metrics"]["selector_coverage"].item() > 0.9 + assert out["draft_metrics"]["selector_accuracy"].item() > 0.9 + assert out["draft_metrics"]["selector_coverage"].item() > 0.9 class TestCandidateSelectorAlignment: @@ -540,6 +541,31 @@ def test_acceptance_term_is_a_probability(self): loss = self._loss(dflash_lk_loss_type="tv", dflash_selector_loss_alpha=0.0) assert 0.0 <= loss <= 1.0 + @pytest.mark.parametrize( + ("loss_type", "distill", "ce_grad"), + [("ce", False, True), ("tv", False, True), ("lambda", False, True), ("ce", True, False)], + ids=["ce", "tv", "lambda", "kd"], + ) + def test_cross_entropy_gradient_follows_the_objective(self, loss_type, distill, ce_grad): + """The per-token CE keeps its graph exactly when the objective backpropagates through it. + + D-PACE reads it first, under no_grad. A tv or lambda objective computed on that would + train the backbone through the selector alone; KD would hold a vocab-wide graph unused. + """ + config = _get_dflash2_config() + config.update( + dflash_lk_loss_type=loss_type, + dflash_self_logit_distillation=distill, + dflash_loss_objective="dpace", + ) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + model.float().train() + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + with patch.object(model, "_compute_loss", wraps=model._compute_loss) as compute_loss: + model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + assert compute_loss.call_args.args[0].ce("final").requires_grad is ce_grad + @pytest.mark.parametrize("loss_type", ["tv", "lambda"]) def test_distillation_conflict_is_rejected(self, loss_type): """Both terms read q(gold), which the KD path never forms.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_domino.py b/tests/unit/torch/speculative/plugins/test_hf_domino.py index 0030abe6f33..5f0fdf996d3 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_domino.py +++ b/tests/unit/torch/speculative/plugins/test_hf_domino.py @@ -24,6 +24,7 @@ import json from copy import deepcopy +import pytest import torch from _test_utils.torch.transformers_models import get_tiny_llama from safetensors.torch import load_file @@ -188,6 +189,41 @@ def spy(*args, **kwargs): assert windows == [6] +class TestDominoGenerate: + """AR validation drafts with the model that is trained and exported.""" + + @pytest.mark.parametrize("pure_draft_prefix_len", [0, 1]) + def test_drafts_are_the_training_head_argmax(self, pure_draft_prefix_len): + """Greedy drafts equal the training head's argmax when it is teacher-forced on them.""" + config = _get_domino_config() + config["dflash_architecture_config"]["pure_draft_prefix_len"] = pure_draft_prefix_len + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + model.float().eval() + head = model.dflash_module + torch.manual_seed(0) + with torch.no_grad(): + # Large enough to overturn the backbone's argmax, so a draft that skips the head + # or feeds it the wrong prefix disagrees. + for p in [*head.prefix_gru.parameters(), *head.embed_proj.parameters()]: + p.normal_() + input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN)) + + base_token, drafts = model.pseudo_speculative_generate(input_ids, steps=BLOCK_SIZE) + + with torch.no_grad(): + _, hidden_states = model._generation_base(input_ids) + hidden = model._draft_block(base_token, hidden_states) + final_logits = model._apply_domino_head( + hidden, + model._base_model_lm_head(hidden), + torch.cat([base_token, drafts], dim=1), # the anchor, then the drafts + torch.zeros(2, 1, dtype=torch.long), + 1, + ) + assert torch.equal(final_logits.argmax(dim=-1), drafts) + + class TestLambdaSchedule: """Test the lambda_base curriculum schedule.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 2dabbe37d57..66fc64d37db 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -36,8 +36,9 @@ import modelopt.torch.quantization as mtq import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG +from modelopt.torch.speculative.plugins.dflash_losses import _tvd_chunk, _tvd_per_token from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel -from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel, _tvd_chunk, _tvd_per_token +from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel from modelopt.torch.speculative.plugins.modeling_dflash import ( DFlashBaseModelOutput, DFlashModule, @@ -60,6 +61,7 @@ def _get_dspark_config( confidence_alpha=0.0, block_size=BLOCK_SIZE, num_layers=NUM_DRAFT_LAYERS, + markov_rank=MARKOV_RANK, ): """Create a DSpark config for testing (dflash mode + projector_type=dspark).""" config = deepcopy(DFLASH_DEFAULT_CFG["config"]) @@ -71,7 +73,7 @@ def _get_dspark_config( config["dflash_architecture_config"] = { "num_hidden_layers": num_layers, "projector_type": "dspark", - "markov_rank": MARKOV_RANK, + "markov_rank": markov_rank, "markov_head_type": head_type, "use_confidence_head": use_confidence_head, "pure_draft_prefix_len": 1, @@ -270,7 +272,9 @@ def spy(*args, **kwargs): class TestDSparkExporter: """Test vanilla serving exports and the native gated/recurrent formats.""" - def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): + def _export( + self, tmp_path, head_type="vanilla", use_confidence_head=False, markov_rank=MARKOV_RANK + ): model = get_tiny_llama(num_hidden_layers=4) mtsp.convert( model, @@ -278,7 +282,9 @@ def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): ( "dflash", _get_dspark_config( - head_type=head_type, use_confidence_head=use_confidence_head + head_type=head_type, + use_confidence_head=use_confidence_head, + markov_rank=markov_rank, ), ) ], @@ -312,8 +318,12 @@ def test_export_weight_keys_match_reference(self, tmp_path, head_type, use_confi @pytest.mark.parametrize("use_confidence_head", [False, True]) def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confidence_head): """Vanilla exports select the runtime and preserve the trained head settings.""" + # DSparkModule casts markov_rank to int; vLLM sizes its Markov table from the export. export_dir = self._export( - tmp_path, head_type=head_type, use_confidence_head=use_confidence_head + tmp_path, + head_type=head_type, + use_confidence_head=use_confidence_head, + markov_rank=float(MARKOV_RANK), ) with open(export_dir / "config.json") as f: cfg = json.load(f) @@ -323,6 +333,7 @@ def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confiden dc = cfg["dflash_config"] assert dc["projector_type"] == "dspark" assert dc["markov_rank"] == MARKOV_RANK + assert type(dc["markov_rank"]) is int assert dc["markov_head_type"] == head_type.lower() assert dc["use_confidence_head"] is use_confidence_head assert dc["shift_label"] is True @@ -331,6 +342,7 @@ def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confiden if vanilla: assert cfg["model_type"] == "qwen3" assert cfg["markov_rank"] == MARKOV_RANK + assert type(cfg["markov_rank"]) is int assert cfg["sample_from_anchor"] is True assert cfg["enable_confidence_head"] is use_confidence_head assert cfg["confidence_head_with_markov"] is use_confidence_head diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py index 0d132f9ce49..301ff60946e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py +++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py @@ -269,7 +269,7 @@ def test_forward_loss_and_metrics(self, variant): assert out.loss.dim() == 0 assert torch.isfinite(out.loss) - metrics = out.lilicorr_metrics + metrics = out.draft_metrics for key in ( "origin_loss", "lilicorr_loss", @@ -291,9 +291,9 @@ def test_loss_is_exactly_origin_plus_lilicorr(self, variant): model = _converted(variant=variant) model.train() out = model(**_make_batch(model.dflash_config.vocab_size)) - metrics = out.lilicorr_metrics + metrics = out.draft_metrics assert out.loss.item() == pytest.approx( - metrics["origin_loss"] + metrics["lilicorr_loss"], abs=1e-5 + (metrics["origin_loss"] + metrics["lilicorr_loss"]).item(), abs=1e-5 ) def test_hinge_only_active_in_the_margin_variant(self): @@ -301,20 +301,20 @@ def test_hinge_only_active_in_the_margin_variant(self): base = _converted(variant="base") base.train() assert ( - base(**_make_batch(base.dflash_config.vocab_size)).lilicorr_metrics["lilicorr_margin"] + base(**_make_batch(base.dflash_config.vocab_size)).draft_metrics["lilicorr_margin"] == 0.0 ) margin = _converted(variant="margin") margin.train() - metrics = margin(**_make_batch(margin.dflash_config.vocab_size)).lilicorr_metrics + metrics = margin(**_make_batch(margin.dflash_config.vocab_size)).draft_metrics # The hinge is relu(margin - gap) on an untrained head, so it is active. assert metrics["lilicorr_margin"] > 0.0 def test_penalty_term_is_active(self): model = _converted() model.train() - metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics + metrics = model(**_make_batch(model.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_penalty"] > 0.0 @pytest.mark.parametrize("w_cal", [0.0, 0.5]) @@ -324,14 +324,16 @@ def test_calibration_enters_the_loss_with_its_weight(self, w_cal): # term has supervised slots to be positive on. model = _converted(topk=VOCAB_SIZE, dflash_lilicorr_w_cal=w_cal) model.train() - metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics + metrics = model(**_make_batch(model.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_w_cal"] == w_cal - assert (metrics["lilicorr_calibration"] > 0.0) is (w_cal > 0.0) + assert bool(metrics["lilicorr_calibration"] > 0.0) is (w_cal > 0.0) weights = VARIANTS["base"] - assert metrics["lilicorr_loss"] == pytest.approx( - weights["w_ce"] * metrics["lilicorr_ce"] - + weights["w_pen"] * metrics["lilicorr_penalty"] - + w_cal * metrics["lilicorr_calibration"], + assert metrics["lilicorr_loss"].item() == pytest.approx( + ( + weights["w_ce"] * metrics["lilicorr_ce"] + + weights["w_pen"] * metrics["lilicorr_penalty"] + + w_cal * metrics["lilicorr_calibration"] + ).item(), abs=1e-5, ) @@ -423,9 +425,9 @@ def test_offline_forward_matches_online(self, prenorm): assert offline_out.loss.item() == pytest.approx(online_out.loss.item(), abs=1e-5) for term in ("lilicorr_penalty", "lilicorr_calibration"): - assert online_out.lilicorr_metrics[term] > 0.0 - assert offline_out.lilicorr_metrics[term] == pytest.approx( - online_out.lilicorr_metrics[term], abs=1e-5 + assert online_out.draft_metrics[term] > 0.0 + assert offline_out.draft_metrics[term].item() == pytest.approx( + online_out.draft_metrics[term].item(), abs=1e-5 ) @pytest.mark.parametrize( @@ -449,8 +451,8 @@ def test_degenerate_lattice_reports_finite_metrics(self, kwargs, degenerate): out = model(**_make_batch(model.dflash_config.vocab_size)) assert torch.isfinite(out.loss), degenerate - for name, value in out.lilicorr_metrics.items(): - assert math.isfinite(value), f"{name} is {value} ({degenerate})" + for name, value in out.draft_metrics.items(): + assert torch.isfinite(value), f"{name} is {value} ({degenerate})" def test_gradient_reaches_head_and_drafter_body(self): """Co-training: the lattice terms must reach the backbone, not only the head. @@ -492,14 +494,14 @@ def test_supervision_follows_lattice_coverage(self): """ wide = _converted(topk=VOCAB_SIZE) wide.train() - metrics = wide(**_make_batch(wide.dflash_config.vocab_size)).lilicorr_metrics + metrics = wide(**_make_batch(wide.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_slot_gt_in_lattice"] == 1.0 assert metrics["lilicorr_oracle_prefix"] == float(BLOCK_SIZE - 1) assert metrics["lilicorr_zero_prefix_rate"] == 0.0 narrow = _converted(topk=CANDIDATE_TOPK) narrow.train() - metrics = narrow(**_make_batch(narrow.dflash_config.vocab_size)).lilicorr_metrics + metrics = narrow(**_make_batch(narrow.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_slot_gt_in_lattice"] < 1.0 assert metrics["lilicorr_zero_prefix_rate"] > 0.0 @@ -531,7 +533,7 @@ def test_eval_forward_delegates_to_the_base_model(self): model.eval() batch = _make_batch(model.dflash_config.vocab_size) out = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]) - assert not hasattr(out, "lilicorr_metrics") or out.get("lilicorr_metrics") is None + assert not hasattr(out, "draft_metrics") or out.get("draft_metrics") is None def test_pseudo_speculative_generate_still_runs(self): """AR validation falls back to the backbone; it must not raise.""" @@ -658,7 +660,7 @@ def test_single_batch_overfit(self): out = model(**batch) out.loss.backward() optimizer.step() - losses.append(out.lilicorr_metrics["lilicorr_loss"]) + losses.append(out.draft_metrics["lilicorr_loss"].item()) assert all(math.isfinite(loss) for loss in losses), losses # A trend, not two endpoints. The DFlash forward draws its own masks, so successive diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml index f3095079a15..7adf3400cd7 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml @@ -10,10 +10,9 @@ # task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) # task_1: Online Domino training + export of the drafter checkpoint # -# NOTE: the inference side (vLLM / AR evaluation) is intentionally not wired up -# yet — the Domino correction head is not applied in pseudo_speculative_generate -# or in the serving stack. Add the vLLM smoke-test / MT-Bench AR-eval steps -# (see hf_online_dflash.yaml task_1/task_2) once the inference path lands. +# NOTE: serving (vLLM) is not wired up for Domino yet, so there is no smoke test. +# HF AR validation applies the correction head: to measure AR, append +# hf_online_dflash.yaml's task_2 with --steps 16 (Domino drafts the whole block). # # Reference: SpecForge PR #571 (z-lab) | drafter format: # huggingface.co/Huang2020/Qwen3-8B-Domino-b16