diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85b5577494c..ea4814557da 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -116,6 +116,8 @@ repos: modelopt/torch/speculative/plugins/modeling_domino.py| modelopt/torch/speculative/plugins/hf_dflash.py| modelopt/torch/speculative/plugins/modeling_dflash.py| + modelopt/torch/speculative/plugins/hf_dflash2.py| + modelopt/torch/speculative/plugins/modeling_dflash2.py| modelopt/torch/speculative/plugins/hf_dspark.py| modelopt/torch/speculative/plugins/modeling_dspark.py| modelopt/torch/speculative/plugins/hf_medusa.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4580a045f39..6253d17f37e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,10 @@ Changelog - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. +*Speculative Decoding* + +- Add the DFlash2 draft variant, selected with ``dflash_architecture_config.projector_type="dflash2"``. It keeps DFlash's one-pass parallel backbone and adds a grouped dynamic convolution around every attention/MLP sublayer (sized by ``conv_kernel_size`` / ``conv_group_size``) plus a low-rank candidate selector (``selector_rank`` / ``selector_top_k``) that scores transitions between adjacent block positions' top-k candidates. The selector's training term is weighted by ``dflash_selector_loss_alpha`` (default 1.0). Exported checkpoints declare ``DFlash2DraftModel`` and match the SGLang/vLLM DFlash2 loaders. + *Misc* - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 255b1d9ab04..5b8b81016db 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -31,11 +31,15 @@ def _get_rope_theta(config, default=None): - """Get RoPE theta from either legacy or Transformers 5 config fields.""" - rope_theta = getattr(config, "rope_theta", None) - if rope_theta is not None: - return rope_theta - + """Get RoPE theta from either legacy or Transformers 5 config fields. + + ``rope_parameters`` is checked FIRST. A config can carry both fields with + different values: Transformers 5 stores the real base under + ``rope_parameters`` while the class default (10000.0 for Qwen3) may still be + visible as a top-level ``rope_theta``. Reading ``rope_theta`` first silently + exports a draft whose RoPE base is 100x off the target's, which breaks + serving because DFlash injects the target's KV into every draft layer. + """ # Transformers 5 stores this under rope_parameters (and exposes the same # data through rope_scaling for backwards compatibility). for attr in ("rope_parameters", "rope_scaling"): @@ -43,6 +47,10 @@ def _get_rope_theta(config, default=None): if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None: return rope_config["rope_theta"] + rope_theta = getattr(config, "rope_theta", None) + if rope_theta is not None: + return rope_theta + return default @@ -533,3 +541,45 @@ def _export_config(self): } ) return config + + +class DFlash2Exporter(DFlashExporter): + """Draft model exporter for DFlash2 (DFlash backbone + convolutions + selector). + + Same z-lab-compatible format as DFlash, plus the DFlash2 weights + (``layers.*.attention_conv.*`` / ``layers.*.mlp_conv.*`` / + ``candidate_selector.*``, already captured by the inherited ``dflash_module.`` + stripping) and the config fields the SGLang/vLLM ``DFlash2DraftModel`` loader + needs to rebuild them (``conv_kernel_size``, ``conv_group_size``, + ``selector_rank``, ``selector_top_k``). + + The architecture name is what selects the DFlash2 serving path: a checkpoint + declaring ``DFlashDraftModel`` loads as a plain DFlash draft and would silently + ignore the convolutions and the selector. + """ + + def _export_config(self): + """Extend the DFlash config with the DFlash2 architecture fields.""" + config = super()._export_config() + draft_config = self.model.dflash_config + + config["architectures"] = ["DFlash2DraftModel"] + # Present because HFDFlash2Model.modify validates them at convert time. + config["dflash_config"].update( + { + "projector_type": getattr(draft_config, "projector_type", "dflash2"), + "conv_kernel_size": draft_config.conv_kernel_size, + "conv_group_size": draft_config.conv_group_size, + "selector_rank": draft_config.selector_rank, + "selector_top_k": draft_config.selector_top_k, + # The published DFlash2 checkpoints carry block_size inside + # dflash_config; the DFlash loader reads it from the top level. + # Emit both so either contract resolves to the same value. + "block_size": config["block_size"], + } + ) + # Published DFlash2 checkpoints state causality explicitly rather than + # leaving it to be inferred from layer_types. Only set it when the SWA + # block above has not already written a `causal` entry. + config.setdefault("is_causal", config["dflash_config"].get("causal", False)) + return config diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 4df57d3035f..bd199652392 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -228,6 +228,17 @@ class DFlashConfig(ModeloptBaseConfig): ), ) + dflash_selector_loss_alpha: float = ModeloptField( + default=1.0, + ge=0.0, + description=( + "DFlash2 only: weight of the candidate-selector cross-entropy term, added to " + "the backbone loss. The selector re-ranks the backbone's top-k candidates per " + "block position; 0 trains the backbone and convolutions only. " + "Ignored unless dflash_architecture_config.projector_type == 'dflash2'." + ), + ) + @model_validator(mode="after") def _check_dpace_alpha(self) -> "DFlashConfig": # Validate at construction regardless of the active objective, so a bad alpha diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 0516ab7181d..a2fc92bd0d5 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/modelopt/torch/speculative/dflash/conversion.py @@ -35,6 +35,12 @@ # ``dflash_architecture_config.projector_type == "dspark"`` and kept in its own # registry so its wrapper (HFDSparkModel) does not overwrite HFDFlashModel. DSparkDMRegistry = _DMRegistryCls(prefix="DSpark") +# DFlash2 also reuses the dflash mode/config/recipe, converting the base model to a +# DFlash backbone whose sublayers are wrapped in grouped dynamic convolutions, plus a +# low-rank candidate selector. Selected via +# ``dflash_architecture_config.projector_type == "dflash2"`` and kept in its own +# registry so its wrapper (HFDFlash2Model) does not overwrite HFDFlashModel. +DFlash2DMRegistry = _DMRegistryCls(prefix="DFlash2") def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertReturnType: @@ -53,12 +59,14 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe registry = DominoDMRegistry elif projector_type == "dspark": registry = DSparkDMRegistry + elif projector_type == "dflash2": + registry = DFlash2DMRegistry elif projector_type in (None, "dflash"): registry = DFlashDMRegistry else: raise ValueError( f"Unsupported dflash_architecture_config.projector_type: {projector_type!r}. " - "Expected 'dflash' (default), 'domino' or 'dspark'." + "Expected 'dflash' (default), 'domino', 'dspark' or 'dflash2'." ) original_cls = type(model) diff --git a/modelopt/torch/speculative/plugins/__init__.py b/modelopt/torch/speculative/plugins/__init__.py index ea789388ad4..87f66947bd2 100644 --- a/modelopt/torch/speculative/plugins/__init__.py +++ b/modelopt/torch/speculative/plugins/__init__.py @@ -31,6 +31,7 @@ with import_plugin("transformers"): from .hf_dflash import * + from .hf_dflash2 import * from .hf_domino import * from .hf_dspark import * from .hf_eagle import * diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index e0d63bde136..0b146a39064 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -419,10 +419,21 @@ def modify(self, config): # overwrite any user value and warn. (rope_scaling is intentionally NOT inherited: # DFlash uses standard Qwen3 RotaryEmbedding; the long-context YaRN scaling is # added only at export via dflash_export_rope_scaling.) + # A config can carry BOTH a top-level rope_theta and a rope_parameters dict + # with different values: Transformers 5 keeps the real base in + # rope_parameters while the class default (10000.0 for Qwen3) stays visible + # as rope_theta. rope_parameters wins, otherwise the draft trains against a + # RoPE base 100x off the target's. + base_rope_params = getattr(base_config, "rope_parameters", None) + if not isinstance(base_rope_params, dict): + base_rope_params = {} for attr in ("rope_theta", "rope_type", "rope_interleaved"): - if not hasattr(base_config, attr): + if attr in base_rope_params: + base_val = base_rope_params[attr] + elif hasattr(base_config, attr): + base_val = getattr(base_config, attr) + else: continue - base_val = getattr(base_config, attr) user_val = getattr(self.dflash_config, attr, None) if user_val is not None and user_val != base_val: logger.warning( @@ -434,6 +445,12 @@ def modify(self, config): base_val, ) setattr(self.dflash_config, attr, base_val) + # Qwen3Config populates rope_parameters at construction, so a later + # setattr on the flat field alone would leave the dict — which is what + # the rotary module reads — holding the stale value. + draft_rope_params = getattr(self.dflash_config, "rope_parameters", None) + if isinstance(draft_rope_params, dict) and attr in draft_rope_params: + draft_rope_params[attr] = base_val self.dflash_config.head_dim = getattr( self.dflash_config, @@ -632,7 +649,14 @@ def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device): return attn_mask def _compute_loss( - self, logits, input_ids, anchor_positions, block_keep_mask, loss_mask, base_logits=None + self, + logits, + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + base_logits=None, + draft_hidden=None, ): """Compute weighted cross-entropy (or KD) loss and accuracy. @@ -643,6 +667,8 @@ def _compute_loss( block_keep_mask: Valid block mask [B, N]. loss_mask: Token-level loss mask [B, seq_len]. base_logits: Base model logits for KD loss [B, seq_len, vocab], or None for CE. + draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``. + Unused here; DFlash2 needs them for its candidate-selector term. Returns: (loss, accuracy) tuple. @@ -921,6 +947,7 @@ def forward( block_keep_mask, loss_mask, base_outputs.logits if self.dflash_self_logit_distillation else None, + draft_hidden=hidden, ) return ModelOutput( diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py new file mode 100644 index 00000000000..fb4ede5c70a --- /dev/null +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""HF DFlash2 model wrapper — DFlash training plus the candidate-selector objective. + +DFlash2 differs from DFlash only in the draft module (grouped dynamic convolutions +around every sublayer, plus a candidate selector) and in one extra loss term, so +this wrapper reuses ``HFDFlashModel``'s forward wholesale and overrides just +:meth:`_compute_loss`. + +The convolutions need no supervision of their own: they sit inside the backbone +and are trained by the backbone loss. The selector does, because at serving time +it — not an independent argmax — picks the drafted token at each block position. + +Selector supervision (following the SGLang/SpecForge reference): + +- Take the backbone's top-k candidates per block position. +- Score each candidate against its *teacher-forced* predecessor token, so the + positions train in parallel exactly as the backbone does. +- When the gold token is missing from the top-k, substitute it into the last + candidate slot. Without this the selector sees no positive class on the hard + positions and never learns those edges. +""" + +import torch +import torch.nn.functional as F +from transformers import PreTrainedModel + +from ..dflash.conversion import DFlash2DMRegistry +from .hf_dflash import HFDFlashModel +from .modeling_dflash2 import DFlash2Module + +__all__ = ["HFDFlash2Model"] + + +@DFlash2DMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFDFlash2Model(HFDFlashModel): + """DFlash model with DFlash2's sublayer convolutions and candidate selector. + + Registered in ``DFlash2DMRegistry`` so that ``convert_to_dflash_model`` routes + to it when ``dflash_architecture_config.projector_type == "dflash2"``. + """ + + def _build_draft_module(self, dflash_config): + """Build the DFlash2 draft module (DFlash backbone + convolutions + selector).""" + return DFlash2Module(dflash_config) + + def modify(self, config): + """Initialize the DFlash2 draft module and read the selector loss weight.""" + arch_config = config.dflash_architecture_config + missing = [ + name + for name in ("conv_kernel_size", "conv_group_size", "selector_rank", "selector_top_k") + if arch_config.get(name) is None + ] + if missing: + raise ValueError( + f"DFlash2 (projector_type='dflash2') requires {missing} in " + "dflash_architecture_config (convolution taps/group size and the " + "candidate selector's rank/top-k)." + ) + super().modify(config) + self.dflash_selector_loss_alpha = getattr(config, "dflash_selector_loss_alpha", 1.0) + + def get_exporter(self): + """Get the exporter for the DFlash2 draft model.""" + from modelopt.torch.export.plugins.hf_spec_export import DFlash2Exporter + + return DFlash2Exporter(self) + + def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask): + """Cross-entropy over the selector's candidate set, and its top-1 accuracy. + + Args: + logits: Backbone logits per block position ``[B, N, block_size, V]``. + target_ids: Gold token ids ``[B, N, block_size]``. + hidden: Backbone hidden states ``[B, N, block_size, H]``. + predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``. + weight_mask: Per-position loss weights ``[B, N, block_size]``. + + Returns: + ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised + positions whose gold token was already in the backbone's top-k, i.e. how + often the selector is choosing rather than being handed the answer. + """ + selector = self.dflash_module.candidate_selector + top_k = selector.top_k + + unary_logits, candidate_ids = logits.topk(top_k, dim=-1) + + # Where the gold token is absent from the top-k, overwrite the last (lowest + # scoring) slot with it, so every supervised position has a correct class. + gold_in_topk = (candidate_ids == target_ids.unsqueeze(-1)).any(dim=-1) + gold_slot = torch.where( + gold_in_topk, + (candidate_ids == target_ids.unsqueeze(-1)).float().argmax(dim=-1), + torch.full_like(target_ids, top_k - 1), + ) + gold_unary = logits.gather(-1, target_ids.unsqueeze(-1)) + candidate_ids = candidate_ids.scatter(-1, gold_slot.unsqueeze(-1), target_ids.unsqueeze(-1)) + unary_logits = unary_logits.scatter(-1, gold_slot.unsqueeze(-1), gold_unary) + + selector_logits = selector.score_candidates( + candidate_ids, unary_logits, hidden, predecessor_ids + ) + + flat_weights = weight_mask.reshape(-1) + denominator = flat_weights.sum() + 1e-6 + per_token = F.cross_entropy( + selector_logits.float().reshape(-1, top_k), + gold_slot.reshape(-1), + reduction="none", + ) + loss = (per_token * flat_weights).sum() / denominator + + with torch.no_grad(): + chosen = selector_logits.argmax(dim=-1).reshape(-1) + accuracy = ( + (chosen == gold_slot.reshape(-1)).float() * flat_weights + ).sum() / denominator + coverage = (gold_in_topk.reshape(-1).float() * flat_weights).sum() / denominator + return loss, accuracy.item(), coverage.item() + + def _compute_loss( + self, + logits, + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + base_logits=None, + draft_hidden=None, + ): + """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. + """ + loss, accuracy = super()._compute_loss( + logits, input_ids, anchor_positions, block_keep_mask, loss_mask, base_logits + ) + if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: + return loss, accuracy + + bsz, seq_len = input_ids.shape + block_size = self.dflash_block_size + n_blocks = anchor_positions.shape[1] + device = input_ids.device + + offsets = torch.arange(block_size, device=device).view(1, 1, -1) + label_indices = anchor_positions.unsqueeze(-1) + offsets + valid_label = label_indices < seq_len + safe_label_indices = label_indices.clamp(max=seq_len - 1) + expanded_ids = input_ids.unsqueeze(1).expand(-1, n_blocks, -1) + target_ids = torch.gather(expanded_ids, 2, safe_label_indices) + + # Same supervision mask as the backbone loss: valid block, in bounds, not the + # anchor slot, and inside the answer span. Position weighting (decay/D-PACE) is + # deliberately not applied — it shapes *where* the backbone spends capacity, + # while the selector should learn every position's transition equally. + weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() + weight_mask = weight_mask * valid_label.float() + weight_mask = weight_mask * (offsets > 0).float() + weight_mask = weight_mask * torch.gather( + loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + ) + + # Teacher-forced predecessor of block position k is the real token at anchor+k-1; + # position 0's predecessor is the anchor itself, matching the serving-side walk + # which starts from the last verified token. + predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0)) + + selector_loss, selector_accuracy, selector_coverage = self._selector_loss( + logits.reshape(bsz, n_blocks, block_size, -1), + target_ids, + draft_hidden.reshape(bsz, n_blocks, block_size, -1), + predecessor_ids, + weight_mask, + ) + self._selector_metrics = { + "selector_accuracy": selector_accuracy, + "selector_coverage": selector_coverage, + } + return loss + self.dflash_selector_loss_alpha * selector_loss, accuracy diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 6463cb4109d..6bcff38d6c6 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -221,6 +221,27 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m return self.o_proj(attn_output) +class _IdentitySublayerWrapper(nn.Module): + """No-op sublayer wrapper: the default around each attention/MLP sublayer. + + ``DFlashDecoderLayer`` calls ``prepare()`` before a sublayer and ``finish()`` + after it, so a variant can transform the sublayer's input and output without + the layer's forward growing a branch. This default does nothing and holds no + parameters, so it neither appears in ``state_dict()`` nor changes the numerics + of a plain DFlash (or Domino/DSpark) draft. + + DFlash2 substitutes ``DFlashGroupedConv`` here (see ``modeling_dflash2.py``). + """ + + def prepare(self, hidden_states): + """Return the sublayer input unchanged, with no state to carry to ``finish``.""" + return hidden_states, None + + def finish(self, hidden_states, state): + """Return the sublayer output unchanged.""" + return hidden_states + + class DFlashDecoderLayer(nn.Module): """Draft decoder layer with KV injection.""" @@ -231,19 +252,26 @@ def __init__(self, config, layer_idx): self.mlp = _MLP_CLS(config) self.input_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + # Sublayer wrappers; no-ops unless a variant replaces them (DFlash2). + self.attention_conv = _IdentitySublayerWrapper() + self.mlp_conv = _IdentitySublayerWrapper() def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None): """Forward pass with residual connections.""" residual = hidden_states hidden_states = self.input_layernorm(hidden_states) + hidden_states, conv_state = self.attention_conv.prepare(hidden_states) hidden_states = self.self_attn( hidden_states, target_hidden, position_embeddings, attention_mask ) + hidden_states = self.attention_conv.finish(hidden_states, conv_state) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, conv_state = self.mlp_conv.prepare(hidden_states) hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp_conv.finish(hidden_states, conv_state) hidden_states = residual + hidden_states return hidden_states diff --git a/modelopt/torch/speculative/plugins/modeling_dflash2.py b/modelopt/torch/speculative/plugins/modeling_dflash2.py new file mode 100644 index 00000000000..309bb4c54e2 --- /dev/null +++ b/modelopt/torch/speculative/plugins/modeling_dflash2.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Adapted from https://github.com/sgl-project/SpecForge/pull/772 +# Copyright (c) 2025 sgl-project +# +# 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. + +"""DFlash2 draft module — DFlash backbone plus local convolution and candidate selection. + +DFlash2 (Inco AI / Z Lab, https://inco.ai/blog/dflash2/) keeps DFlash's one-pass +parallel backbone and adds two small components that address the two ways a +purely parallel draft loses acceptance: + +- :class:`DFlashGroupedConv` — a grouped *dynamic* depthwise convolution wrapped + around every attention and MLP sublayer. Each block position mixes in its + predecessors inside the block, which injects the intra-block sequential + dependency the parallel backbone lacks (mitigating suffix acceptance decay) + without a second backbone pass. Taps do not cross the block boundary. + +- :class:`CandidateSelector` — a low-rank transition scorer. Instead of an + independent argmax per block position, the drafter keeps the target head's + top-k candidates per position and scores adjacent transitions, so serving can + walk one coherent path through the block. + +Where Domino uses a GRU and DSpark a Markov transition bias, DFlash2 spends its +extra capacity on these two pieces: both are cheap (a few percent of draft +parameters, ~1% of serving step latency in the reference measurements). + +This module owns the parameters only; the training wrapper (``HFDFlash2Model`` +in ``hf_dflash2.py``) orchestrates the forward and the selector loss. Module and +parameter names (``attention_conv`` / ``mlp_conv`` / ``base_kernel`` / +``kernel_projection`` / ``candidate_selector`` / ``predecessor_codebook`` / +``successor_codebook`` / ``hidden_projection``) match the SGLang and vLLM +``DFlash2DraftModel`` loaders so an exported checkpoint is served directly. +""" + +import torch +import torch.nn.functional as F +from torch import nn + +from .modeling_dflash import DFlashModule + +__all__ = ["CandidateSelector", "DFlash2Module", "DFlashGroupedConv"] + + +class DFlashGroupedConv(nn.Module): + """Grouped dynamic depthwise convolution over positions within a proposal block. + + Wraps one sublayer: :meth:`prepare` convolves the sublayer input and emits the + dynamic kernel for the output side, :meth:`finish` convolves the sublayer + output. One projection of the sublayer input produces both sides' kernel + deltas. + + ``base_kernel`` starts as an identity (tap 0 weight 1, later taps 0), so a + freshly built DFlash2 draft computes exactly what its DFlash backbone would. + That makes the convolution a stable extension rather than a perturbation, and + lets a DFlash checkpoint warm-start a DFlash2 run. + """ + + def __init__(self, hidden_size: int, block_size: int, taps: int, group_size: int): + """Build the identity-initialized base kernel and the dynamic-kernel projection.""" + super().__init__() + if taps < 1: + raise ValueError(f"DFlash2 conv_kernel_size must be >= 1, got {taps}.") + if taps > block_size: + raise ValueError( + f"DFlash2 conv_kernel_size ({taps}) must not exceed " + f"dflash_block_size ({block_size})." + ) + if group_size < 1 or hidden_size % group_size: + raise ValueError( + f"DFlash2 conv_group_size ({group_size}) must be >= 1 and divide " + f"hidden_size ({hidden_size})." + ) + + self.block_size = int(block_size) + self.taps = int(taps) + self.group_size = int(group_size) + self.num_groups = int(hidden_size) // self.group_size + + # [input/output side, tap, channel]; identity at tap 0. Layout matches the + # SGLang/vLLM DFlash2 weight loader. + base_kernel = torch.zeros(2, self.taps, int(hidden_size)) + base_kernel[:, 0] = 1.0 + self.base_kernel = nn.Parameter(base_kernel) + self.kernel_projection = nn.Linear( + int(hidden_size), 2 * self.taps * self.num_groups, bias=False + ) + + def _convolve(self, hidden_states, delta, side: int): + """Apply the depthwise convolution for one side, with taps clipped at block starts.""" + bsz, seq_len, hidden_size = hidden_states.shape + if seq_len % self.block_size: + raise ValueError( + f"DFlash2 convolution needs a sequence length divisible by " + f"block_size ({self.block_size}), got {seq_len}." + ) + + n_blocks = seq_len // self.block_size + blocks = hidden_states.reshape( + bsz, n_blocks, self.block_size, self.num_groups, self.group_size + ) + dynamic = delta.reshape(bsz, n_blocks, self.block_size, self.taps, self.num_groups) + base = self.base_kernel[side].reshape(1, 1, 1, self.taps, self.num_groups, self.group_size) + # Per-position, per-group coefficients: static base plus the dynamic delta. + coefficients = base + dynamic.unsqueeze(-1) + + output = coefficients[:, :, :, 0] * blocks + for tap in range(1, self.taps): + # Shift within the block only: position k reads k-tap, and the first + # `tap` positions of each block read zeros rather than the previous block. + shifted = F.pad(blocks[:, :, : self.block_size - tap], (0, 0, 0, 0, tap, 0)) + output = output + coefficients[:, :, :, tap] * shifted + return output.reshape(bsz, seq_len, hidden_size) + + def prepare(self, hidden_states): + """Convolve the sublayer input; return it with the output side's dynamic kernel.""" + coefficients = self.kernel_projection(hidden_states).reshape( + *hidden_states.shape[:-1], 2, self.taps, self.num_groups + ) + return self._convolve(hidden_states, coefficients[..., 0, :, :], side=0), coefficients[ + ..., 1, :, : + ] + + def finish(self, hidden_states, state): + """Convolve the sublayer output using the kernel produced by :meth:`prepare`.""" + return self._convolve(hidden_states, state, side=1) + + +class CandidateSelector(nn.Module): + """Low-rank scorer for transitions between adjacent block positions' candidates. + + Scores an edge from a predecessor token ``p`` to a candidate token ``c`` at a + block position with hidden state ``h`` as:: + + edge(p -> c) = + unary_logit[c] + + i.e. a bilinear form between the two token codebooks, gated by the context. + Training scores each position's candidate set independently under teacher + forcing (:meth:`score_candidates`); serving walks the resulting lattice. + """ + + def __init__(self, hidden_size: int, vocab_size: int, rank: int, top_k: int, std: float): + """Build the predecessor/successor codebooks and the context projection.""" + super().__init__() + if rank < 1: + raise ValueError(f"DFlash2 selector_rank must be >= 1, got {rank}.") + if not 1 <= top_k <= vocab_size: + raise ValueError( + f"DFlash2 selector_top_k must be in [1, vocab_size={vocab_size}], got {top_k}." + ) + self.top_k = int(top_k) + self.rank = int(rank) + self.predecessor_codebook = nn.Parameter(torch.empty(int(vocab_size), int(rank))) + self.successor_codebook = nn.Parameter(torch.empty(int(vocab_size), int(rank))) + self.hidden_projection = nn.Linear(int(hidden_size), int(rank), bias=False) + nn.init.normal_(self.predecessor_codebook, std=std) + nn.init.normal_(self.successor_codebook, std=std) + + def score_candidates(self, candidate_ids, unary_logits, hidden_states, predecessor_ids): + """Add the predecessor transition score to a candidate set's unary logits. + + Args: + candidate_ids: Candidate token ids ``[..., K]``. + unary_logits: Backbone logits for those candidates ``[..., K]``. + hidden_states: Backbone hidden at this position ``[..., H]``. + predecessor_ids: Teacher-forced predecessor token ids ``[...]``. + + Returns: + Selector logits over the candidate set ``[..., K]``. + """ + predecessor = self.predecessor_codebook[predecessor_ids] + successor = self.successor_codebook[candidate_ids] + context = predecessor * self.hidden_projection(hidden_states) + transition = torch.einsum("...r,...kr->...k", context.to(successor.dtype), successor) + return unary_logits + transition + + @torch.no_grad() + def greedy_path(self, candidate_ids, unary_logits, hidden_states, anchor_token_ids): + """Walk the candidate lattice greedily, mirroring the serving-side path walk. + + Args: + candidate_ids: ``[B, L, K]`` candidate ids per block position. + unary_logits: ``[B, L, K]`` backbone logits for those candidates. + hidden_states: ``[B, L, H]`` backbone hidden per block position. + anchor_token_ids: ``[B]`` the verified token preceding the block. + + Returns: + Selected token ids ``[B, L]``. + """ + predecessor_ids = anchor_token_ids + path = [] + for position in range(candidate_ids.shape[1]): + scores = self.score_candidates( + candidate_ids[:, position], + unary_logits[:, position], + hidden_states[:, position], + predecessor_ids, + ) + selected = scores.argmax(dim=-1, keepdim=True) + predecessor_ids = candidate_ids[:, position].gather(1, selected)[:, 0] + path.append(predecessor_ids) + return torch.stack(path, dim=1) + + +class DFlash2Module(DFlashModule): + """DFlash draft backbone with per-sublayer convolutions and a candidate selector.""" + + def __init__(self, config): + """Initialize the DFlash backbone, then attach the convolutions and the selector.""" + super().__init__(config) + + self.projector_type = getattr(config, "projector_type", "dflash2") + + def required_int(name: str) -> int: + """Read an int architecture field, rejecting missing values and bools.""" + value = getattr(config, name, None) + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError( + f"DFlash2 (projector_type='dflash2') requires an integer " + f"'{name}' in dflash_architecture_config, got {value!r}." + ) + return value + + taps = required_int("conv_kernel_size") + group_size = required_int("conv_group_size") + rank = required_int("selector_rank") + top_k = required_int("selector_top_k") + + std = getattr(config, "initializer_range", 0.02) + + # Replace each layer's no-op sublayer wrappers with real convolutions. The + # backbone layer forward already calls prepare()/finish() around attention + # and the MLP, so nothing else in the layer changes. + for layer in self.layers: + for wrapper_name in ("attention_conv", "mlp_conv"): + setattr( + layer, + wrapper_name, + DFlashGroupedConv( + hidden_size=config.hidden_size, + block_size=self.block_size, + taps=taps, + group_size=group_size, + ), + ) + + self.candidate_selector = CandidateSelector( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + rank=rank, + top_k=top_k, + std=std, + ) + + # DFlashModule.__init__ already ran _init_weights before these modules + # existed, so initialize the new Linear layers explicitly. base_kernel and + # the codebooks keep the init set in their own constructors. + self._init_head_weights(std) + + def _init_head_weights(self, std: float): + """Initialize the convolution and selector Linear layers (matching HF _init_weights).""" + linears = [self.candidate_selector.hidden_projection] + for layer in self.layers: + linears += [layer.attention_conv.kernel_projection, layer.mlp_conv.kernel_projection] + for module in linears: + nn.init.normal_(module.weight, mean=0.0, std=std) + if module.bias is not None: + nn.init.zeros_(module.bias) diff --git a/modelopt_recipes/general/speculative_decoding/dflash2.yaml b/modelopt_recipes/general/speculative_decoding/dflash2.yaml new file mode 100644 index 00000000000..86555f357f9 --- /dev/null +++ b/modelopt_recipes/general/speculative_decoding/dflash2.yaml @@ -0,0 +1,98 @@ +# DFlash2 speculative-decoding training recipe. +# +# DFlash2 (https://inco.ai/blog/dflash2/) reuses the DFlash mode/pipeline and adds +# two components, selected via dflash_architecture_config.projector_type=dflash2: +# - a grouped dynamic depthwise convolution around every attention/MLP sublayer, +# giving each block position a view of its predecessors inside the block; +# - a low-rank candidate selector that scores transitions between adjacent block +# positions' top-k candidates, so serving walks one coherent path. +# The selector is trained by an extra cross-entropy term weighted by +# dflash_selector_loss_alpha. Online training is the default path (data.mode=online). +# Override fields via an OmegaConf dotlist. + +metadata: + recipe_type: speculative_dflash + description: DFlash2 training recipe (DFlash backbone + sublayer conv + candidate selector). + +# maps to ModelArguments (main.py) +model: + model_name_or_path: + trust_remote_code: false + use_fake_base_for_offline: false + +# maps to DataArguments (main.py) +data: + mode: online + data_path: + offline_data_path: + # Jinja chat template with {% generation %} tags for answer_only_loss. + chat_template: + +# maps to TrainingArguments (main.py) +training: + # --- commonly modified --- + output_dir: + num_train_epochs: 6 + per_device_train_batch_size: 1 + learning_rate: 6.0e-4 + warmup_ratio: 0.04 + training_seq_len: 3072 + logging_steps: 50 + save_steps: 2000 + 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 + + # --- rarely modified --- + do_eval: false + lr_scheduler_type: linear + save_strategy: steps + weight_decay: 0.0 + max_grad_norm: 1.0 + dataloader_drop_last: true + bf16: true + tf32: true + remove_unused_columns: false + # Safe default: the selector params are unused when + # dflash_selector_loss_alpha == 0, which would otherwise trip DDP. + ddp_find_unused_parameters: true + ddp_timeout: 1800 + report_to: tensorboard + +# maps to DFlashConfig (modelopt/torch/speculative/config.py). +dflash: + dflash_block_size: 16 + dflash_num_anchors: 256 + dflash_use_torch_compile: false + dflash_self_logit_distillation: false + # gamma for exponential loss decay (block_size=16 -> 7). + dflash_loss_decay_factor: 7.0 + # Qwen3 has no native mask token; 151669 is an unused id used by the reference. + dflash_mask_token_id: 151669 + # Weight of the candidate-selector cross-entropy term (0 disables it and trains + # the backbone + convolutions only). + dflash_selector_loss_alpha: 1.0 + dflash_architecture_config: + num_hidden_layers: 5 + # Draft attention/MLP dims — set explicitly (the draft is an independent + # Qwen3 model and does NOT inherit these from the base). GQA: 8 KV heads. + num_attention_heads: 32 + num_key_value_heads: 8 + head_dim: 128 + intermediate_size: 12288 + projector_type: dflash2 + # Grouped dynamic depthwise convolution. conv_kernel_size is the number of taps + # (2 = each position also sees its predecessor); it must not exceed the block + # size. conv_group_size must divide hidden_size. + conv_kernel_size: 2 + conv_group_size: 16 + # Candidate selector: rank of the transition codebooks, and how many of the + # backbone top-k candidates per position it re-ranks. + selector_rank: 256 + selector_top_k: 16 diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py new file mode 100644 index 00000000000..b2ad7528665 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CPU unit tests for the DFlash2 speculative decoding plugin. + +DFlash2 reuses the DFlash mode/pipeline and adds grouped dynamic convolutions +around every attention/MLP sublayer plus a low-rank candidate selector. These +tests cover conversion routing, the convolution's two structural invariants +(identity at initialization, no leakage across the block boundary), the selector +training objective, and the export format against the SGLang/vLLM +``DFlash2DraftModel`` layout (``attention_conv.*`` / ``mlp_conv.*`` / +``candidate_selector.*``). +""" + +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 + +import modelopt.torch.speculative as mtsp +from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG +from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel +from modelopt.torch.speculative.plugins.hf_dflash2 import HFDFlash2Model +from modelopt.torch.speculative.plugins.modeling_dflash import ( + DFlashModule, + _IdentitySublayerWrapper, +) +from modelopt.torch.speculative.plugins.modeling_dflash2 import DFlash2Module, DFlashGroupedConv + +BLOCK_SIZE = 4 +NUM_DRAFT_LAYERS = 2 +SEQ_LEN = 16 # must be a multiple of BLOCK_SIZE +CONV_KERNEL_SIZE = 2 +CONV_GROUP_SIZE = 4 +SELECTOR_RANK = 8 +SELECTOR_TOP_K = 5 + +ARCH_FIELDS = ["conv_kernel_size", "conv_group_size", "selector_rank", "selector_top_k"] + + +def _get_dflash2_config(selector_loss_alpha=1.0, block_size=BLOCK_SIZE, **arch_overrides): + """Create a DFlash2 config for testing (dflash mode + projector_type=dflash2).""" + config = deepcopy(DFLASH_DEFAULT_CFG["config"]) + config["dflash_block_size"] = block_size + config["dflash_use_torch_compile"] = False + config["dflash_mask_token_id"] = 0 # token 0 as mask for the tiny model + config["dflash_self_logit_distillation"] = False + config["dflash_selector_loss_alpha"] = selector_loss_alpha + config["dflash_architecture_config"] = { + "num_hidden_layers": NUM_DRAFT_LAYERS, + "projector_type": "dflash2", + "conv_kernel_size": CONV_KERNEL_SIZE, + "conv_group_size": CONV_GROUP_SIZE, + "selector_rank": SELECTOR_RANK, + "selector_top_k": SELECTOR_TOP_K, + **arch_overrides, + } + return config + + +def _make_batch(vocab_size): + torch.manual_seed(0) + input_ids = torch.randint(1, vocab_size, (2, SEQ_LEN)) + return input_ids, torch.ones_like(input_ids), input_ids.clone() + + +class TestDFlash2Convert: + """Test DFlash2 conversion routing and module construction.""" + + def test_convert_creates_dflash2_model(self): + """projector_type=dflash2 routes to HFDFlash2Model (a HFDFlashModel subclass).""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + assert isinstance(model, HFDFlash2Model) + assert isinstance(model, HFDFlashModel) + assert isinstance(model.dflash_module, DFlash2Module) + + def test_every_sublayer_wrapped_in_a_convolution(self): + """Both sublayer wrappers on every draft layer become real convolutions.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + layers = model.dflash_module.layers + assert len(layers) == NUM_DRAFT_LAYERS + for layer in layers: + for conv in (layer.attention_conv, layer.mlp_conv): + assert isinstance(conv, DFlashGroupedConv) + assert conv.taps == CONV_KERNEL_SIZE + assert conv.group_size == CONV_GROUP_SIZE + + def test_selector_shapes(self): + """The candidate selector's codebooks and projection are sized from the config.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + selector = model.dflash_module.candidate_selector + vocab = model.dflash_config.vocab_size + assert selector.top_k == SELECTOR_TOP_K + assert selector.predecessor_codebook.shape == (vocab, SELECTOR_RANK) + assert selector.successor_codebook.shape == (vocab, SELECTOR_RANK) + assert selector.hidden_projection.out_features == SELECTOR_RANK + assert selector.hidden_projection.bias is None + + def test_new_params_trainable(self): + """The convolution and selector parameters are trainable.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + new = [ + (n, p) + for n, p in model.named_parameters() + if "_conv." in n or "candidate_selector" in n + ] + assert len(new) >= 2 * NUM_DRAFT_LAYERS * 2 + 3 + assert all(p.requires_grad for _, p in new) + + @pytest.mark.parametrize("field", ARCH_FIELDS) + def test_missing_architecture_field_raises(self, field): + """projector_type=dflash2 without a required architecture field is an error.""" + config = _get_dflash2_config() + del config["dflash_architecture_config"][field] + model = get_tiny_llama(num_hidden_layers=4) + with pytest.raises(ValueError, match=field): + mtsp.convert(model, [("dflash", config)]) + + def test_conv_kernel_larger_than_block_raises(self): + """A convolution tap count exceeding the block size is an error.""" + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash2_config(conv_kernel_size=BLOCK_SIZE + 1) + with pytest.raises(ValueError, match="conv_kernel_size"): + mtsp.convert(model, [("dflash", config)]) + + def test_conv_group_size_must_divide_hidden(self): + """A conv_group_size that does not divide hidden_size is an error.""" + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash2_config(conv_group_size=model.config.hidden_size - 1) + with pytest.raises(ValueError, match="conv_group_size"): + mtsp.convert(model, [("dflash", config)]) + + def test_dflash_mode_still_creates_plain_dflash(self): + """Without projector_type=dflash2, conversion still yields a plain DFlash model.""" + config = deepcopy(DFLASH_DEFAULT_CFG["config"]) + config["dflash_mask_token_id"] = 0 + config["dflash_architecture_config"] = {"num_hidden_layers": NUM_DRAFT_LAYERS} + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + assert isinstance(model, HFDFlashModel) + assert not isinstance(model, HFDFlash2Model) + assert type(model.dflash_module) is DFlashModule + # The sublayer seam stays a parameterless no-op for a plain DFlash draft. + for layer in model.dflash_module.layers: + assert isinstance(layer.attention_conv, _IdentitySublayerWrapper) + assert isinstance(layer.mlp_conv, _IdentitySublayerWrapper) + assert not any("_conv." in n for n, _ in model.dflash_module.named_parameters()) + + +class TestDFlashGroupedConv: + """Test the convolution's structural invariants directly.""" + + def _conv(self, hidden_size=32, taps=2): + torch.manual_seed(0) + return DFlashGroupedConv( + hidden_size=hidden_size, block_size=BLOCK_SIZE, taps=taps, group_size=CONV_GROUP_SIZE + ).double() + + def test_identity_at_initialization(self): + """With the dynamic kernel zeroed, the identity base kernel is a no-op. + + This is what makes enabling DFlash2 a stable extension of a DFlash backbone + rather than a perturbation of it. + """ + conv = self._conv() + with torch.no_grad(): + conv.kernel_projection.weight.zero_() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + out = conv.finish(*conv.prepare(x)) + assert torch.allclose(out, x, atol=1e-12) + + def test_taps_do_not_cross_the_block_boundary(self): + """Perturbing the last position of a block leaves later blocks untouched.""" + conv = self._conv() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + baseline = conv.finish(*conv.prepare(x)) + + perturbed_input = x.clone() + perturbed_input[:, BLOCK_SIZE - 1] += 5.0 + perturbed = conv.finish(*conv.prepare(perturbed_input)) + + assert torch.allclose(baseline[:, BLOCK_SIZE:], perturbed[:, BLOCK_SIZE:], atol=1e-12) + assert not torch.allclose(baseline[:, :BLOCK_SIZE], perturbed[:, :BLOCK_SIZE]) + + def test_intra_block_dependency_is_backward_only(self): + """A position influences its successors inside the block, never its predecessors. + + This is the point of the convolution: it injects the sequential dependency the + parallel backbone lacks, without letting a position see the future. + """ + conv = self._conv() + x = torch.randn(2, SEQ_LEN, 32, dtype=torch.double) + baseline = conv.finish(*conv.prepare(x)) + + perturbed_input = x.clone() + perturbed_input[:, 1] += 5.0 + perturbed = conv.finish(*conv.prepare(perturbed_input)) + + assert torch.allclose(baseline[:, 0], perturbed[:, 0], atol=1e-12) + assert not torch.allclose(baseline[:, 2], perturbed[:, 2]) + + def test_sequence_length_must_be_block_aligned(self): + """A sequence length not divisible by the block size is an error.""" + conv = self._conv() + with pytest.raises(ValueError, match="block_size"): + conv.prepare(torch.randn(1, BLOCK_SIZE + 1, 32, dtype=torch.double)) + + +class TestDFlash2Forward: + """Test the DFlash2 training forward (online path on CPU).""" + + def test_forward_grads_reach_conv_and_selector(self): + """Backward fills gradients on the convolutions, the selector and the backbone.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + assert out.loss.requires_grad + assert out.loss.dim() == 0 + out.loss.backward() + + module = model.dflash_module + for grad in ( + module.candidate_selector.predecessor_codebook.grad, + module.candidate_selector.successor_codebook.grad, + module.layers[0].attention_conv.base_kernel.grad, + module.layers[0].mlp_conv.kernel_projection.weight.grad, + module.fc.weight.grad, + ): + assert grad is not None and torch.isfinite(grad).all() + assert grad.abs().sum() > 0 + + def test_selector_metrics_reported(self): + """The forward records selector accuracy and top-k coverage.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + metrics = model._selector_metrics + for key in ("selector_accuracy", "selector_coverage"): + assert 0.0 <= metrics[key] <= 1.0 + + def test_selector_alpha_zero_disables_the_term(self): + """alpha=0 trains the backbone and convolutions only; the selector gets no grad.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config(selector_loss_alpha=0.0))]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + out.loss.backward() + + module = model.dflash_module + codebook_grad = module.candidate_selector.predecessor_codebook.grad + assert codebook_grad is None or codebook_grad.abs().sum() == 0 + # The convolutions still train: they live inside the backbone. + conv_grad = module.layers[0].attention_conv.kernel_projection.weight.grad + assert conv_grad is not None and conv_grad.abs().sum() > 0 + + def test_selector_loss_increases_total_loss(self): + """The selector term adds to the backbone loss rather than replacing it.""" + input_ids, attention_mask, labels = _make_batch(32) + + losses = {} + for alpha in (0.0, 1.0): + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config(selector_loss_alpha=alpha))]) + model.train() + torch.manual_seed(0) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + losses[alpha] = float(out.loss.detach()) + assert losses[1.0] > losses[0.0] + + def test_overfits_a_single_batch(self): + """A few steps on one batch drive backbone and selector accuracy up. + + Guards the target/predecessor alignment: a misaligned selector objective still + produces a finite decreasing loss, but its accuracy does not reach 1. + """ + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + model.train() + + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-3) + for _ in range(60): + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + optimizer.zero_grad() + out.loss.backward() + optimizer.step() + + assert out.train_acc[0][0] > 0.9 + assert model._selector_metrics["selector_accuracy"] > 0.9 + + +class TestDFlash2Export: + """Test the DFlash2 export format (weights + config).""" + + def _export(self, tmp_path): + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash2_config())]) + export_dir = tmp_path / "exported" + model.get_exporter().export(export_dir) + return export_dir + + def test_export_weight_keys_match_reference(self, tmp_path): + """Exported weights carry the DFlash2 tensors under reference names, no prefix.""" + sd = load_file(str(self._export(tmp_path) / "model.safetensors")) + for key in sd: + assert "dflash_module." not in key + assert "rotary_emb" not in key + + assert "candidate_selector.predecessor_codebook" in sd + assert "candidate_selector.successor_codebook" in sd + assert "candidate_selector.hidden_projection.weight" in sd + for layer_idx in range(NUM_DRAFT_LAYERS): + for wrapper in ("attention_conv", "mlp_conv"): + assert f"layers.{layer_idx}.{wrapper}.base_kernel" in sd + assert f"layers.{layer_idx}.{wrapper}.kernel_projection.weight" in sd + + def test_export_config_declares_dflash2_architecture(self, tmp_path): + """config.json selects the DFlash2 serving path and carries its fields. + + The architecture name matters: a checkpoint declaring ``DFlashDraftModel`` + loads as a plain DFlash draft and silently ignores these weights. + """ + with open(self._export(tmp_path) / "config.json") as f: + cfg = json.load(f) + + assert cfg["architectures"] == ["DFlash2DraftModel"] + dflash_config = cfg["dflash_config"] + assert dflash_config["projector_type"] == "dflash2" + assert dflash_config["conv_kernel_size"] == CONV_KERNEL_SIZE + assert dflash_config["conv_group_size"] == CONV_GROUP_SIZE + assert dflash_config["selector_rank"] == SELECTOR_RANK + assert dflash_config["selector_top_k"] == SELECTOR_TOP_K + assert "mask_token_id" in dflash_config + assert "target_layer_ids" in dflash_config diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml new file mode 100644 index 00000000000..e300e0665b4 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -0,0 +1,107 @@ +# DFlash2 online speculative decoding training for Qwen3-8B. +# +# DFlash2 = the DFlash draft backbone plus two additions: +# * a grouped dynamic depthwise convolution wrapped around every attention and +# MLP sublayer (taps clipped at block boundaries, identity-initialized so a +# fresh DFlash2 draft computes exactly what its DFlash backbone would), and +# * a low-rank candidate selector that scores transitions between adjacent +# block positions' top-k candidates, so serving walks one coherent path. +# See the dflash2.yaml recipe and +# modelopt/torch/speculative/plugins/{modeling,hf}_dflash2.py. +# +# 2-step pipeline: +# task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) +# task_1: Online DFlash2 training + export of the drafter checkpoint +# +# As configured this is a short convergence check (max_steps=2000), matching the +# other Qwen3-8B online examples so it finishes on one node. To reproduce the +# published Qwen3-8B DFlash2 curve instead, see "Full run" below. +# +# Reference: inco.ai/blog/dflash2 | vLLM PR #52816 (serving support) +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml --yes +# +# Full run (the strict A/B against DFlash / DSpark / Domino, 3 epochs = ~92K +# steps on the 1.96M-conversation Spec-Decoding-Dataset-v1, 8 nodes x 8 H100): +# - point data.data_path at that corpus instead of task_0's output +# - training.num_train_epochs=3 and drop training.max_steps +# - training.save_steps=4000 +# - slurm_config.nodes=8 (global batch stays 64 = nodes x gpus x bs x accum) +# Every other knob below is already the A/B setting. + +job_name: Qwen3-8B_DFlash2_online +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + + # Step 1: Build input conversations. example_data_config.yaml enables only the + # daring-anteater source (train: 50000) — multi-turn SFT with real assistant + # completions. --full-conversations keeps those completions so answer_only_loss + # has assistant spans to mask. make_dataset.sh writes /scratchspace/data/train.jsonl. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Online DFlash2 training (the script exports the drafter at the end). + # Consumes the conversations built in task_0 (shared via /scratchspace). + task_1: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash2.yaml + - model.model_name_or_path=<> + - data.data_path=/scratchspace/data/train.jsonl + - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja + - training.output_dir=/scratchspace/dflash2_bs16 + - training.per_device_train_batch_size=1 + - training.num_train_epochs=1 + - training.max_steps=2000 + - training.training_seq_len=4096 + - training.learning_rate=6.0e-4 + - training.warmup_ratio=0.04 + - training.warmup_steps=0 + - training.lr_scheduler_type=linear + - training.save_steps=5000 + - training.logging_steps=100 + - training.disable_tqdm=true + - training.answer_only_loss=true + # Draft backbone — identical to the DFlash / DSpark / Domino arms so the + # only difference between them is the correction head. + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_self_logit_distillation=false + - dflash.dflash_architecture_config.num_hidden_layers=5 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=12288 + # DFlash2 knobs (also set in the recipe; repeated here for visibility). + # A draft dim NOT set explicitly falls back to the Qwen3Config default, not + # to the base model's — hence the five dims above are always spelled out. + - dflash.dflash_architecture_config.projector_type=dflash2 + - dflash.dflash_architecture_config.conv_kernel_size=2 + - dflash.dflash_architecture_config.conv_group_size=16 + - dflash.dflash_architecture_config.selector_rank=256 + - dflash.dflash_architecture_config.selector_top_k=16 + - dflash.dflash_selector_loss_alpha=1.0 + # Sliding-window draft attention, matching the published DFlash2 drafters. + - dflash.dflash_swa_window_size=2048 + environment: + - MAX_FINAL_LOSS: "5.0" + - MIN_FINAL_ACC: "0.15" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8