diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2907f626226..3a12f4e5b7c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -18,6 +18,8 @@ Changelog *Speculative Decoding* - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint. +- Add ``dflash_loss_terms`` to compose the DFlash-family block objective from weighted loss terms (``ce``, ``kd``, ``tvd``, ``tv`` and ``lambda``), e.g. ``{ce: 0.1, tvd: 0.9}``, for any ``projector_type``, online or offline. Left unset, each variant keeps its own objective. +- Add ``dflash_teacher_softcap`` to cap the distillation teacher with the base model's ``final_logit_softcapping`` (Gemma family), the way the model caps its own logits. *Megatron Framework (M-LM / M-Bridge)* @@ -37,6 +39,8 @@ 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, LiLiCorr, DSpark and Domino 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``, ``lilicorr_metrics``, ``dspark_metrics`` and ``domino_metrics``. DSpark's ``ce_loss``/``l1_loss``/``base_accuracy`` are now ``loss_ce``/``loss_tvd``/``backbone_accuracy``, and Domino's ``final_loss``/``base_loss``/``base_accuracy`` are now ``loss_final``/``loss_backbone``/``backbone_accuracy``. +- Online DFlash, DFlash2 and LiLiCorr training now distills text and multimodal batches alike from the base model's ``lm_head`` applied to its final hidden state, as offline training and DSpark already did. The teacher therefore leaves out the logit scaling some bases apply at the ``lm_head`` (``logits_scaling``, ``logit_scale``, ``lm_head_multiplier``: Granite, Cohere, Falcon-H1 and others), which conversion now warns about. +- On a base that caps its logits with ``final_logit_softcapping`` (Gemma family, including the Gemma 3n and Gemma 4 VLMs), online DFlash, DFlash2 and LiLiCorr training now needs ``dflash_teacher_softcap: true`` to keep a capped teacher. A run resumed from an HF-format checkpoint saved before this release restores that checkpoint's config, in which the switch is off, and so continues uncapped; a run resumed from an FSDP-sharded checkpoint takes the switch from the recipe. **Deprecations** @@ -55,6 +59,8 @@ Changelog - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target. - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit. - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. +- Fix online DSpark and Domino training on multimodal batches, which ran the base model's language model on the raw token ids and so dropped the image and video features; they now run the base model's multimodal forward, as DFlash's do. Like DFlash's, their Gemma 3 and PaliGemma image batches reach that forward without ``token_type_ids``, so on transformers 5.5 it raises and on 5.12 and later it attends causally over the image. +- Fix online DFlash-family training on Gemma 4 image and video batches, which failed in the vision tower because their ``image_position_ids`` and ``video_position_ids`` were not passed on. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/speculative_decoding/doc/dflash.md b/examples/speculative_decoding/doc/dflash.md index c468c7652be..ffafdf4a0ae 100644 --- a/examples/speculative_decoding/doc/dflash.md +++ b/examples/speculative_decoding/doc/dflash.md @@ -165,6 +165,8 @@ See [`modelopt_recipes/general/speculative_decoding/dflash.yaml`](../../../model | `dflash.dflash_loss_objective` | `dpace` | Position weighting: `decay` (static) or `dpace` (dynamic, see below) | | `dflash.dflash_dpace_alpha` | 0.5 | D-PACE smoothing factor in (0, 1]; only used when objective is `dpace` | | `dflash.dflash_self_logit_distillation` | true | Use target model logits as soft labels (vs hard CE) | +| `dflash.dflash_loss_terms` | unset | Block objective as `{term: weight}` over `ce`, `kd`, `tvd`, `tv` and `lambda`, e.g. `{ce: 0.1, tvd: 0.9}`; unset keeps the variant's own | +| `dflash.dflash_teacher_softcap` | false | Cap the teacher with the base's `final_logit_softcapping` (Gemma family) | | `dflash.dflash_mask_token_id` | auto | Token ID for masked positions (see note below) | | `dflash.dflash_architecture_config.num_hidden_layers` | 5 | Draft decoder layers | | `training.answer_only_loss` | false | Mask loss on non-assistant tokens | diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 1a53b21e4d4..86e4bd77689 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -16,9 +16,9 @@ """Configurations for speculative decoding modes.""" from copy import deepcopy -from typing import Literal +from typing import Annotated, Literal -from pydantic import model_validator +from pydantic import Field, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField @@ -75,6 +75,17 @@ def _get_dflash_default_config(): }, } +# The terms a DFlash-family block objective is built from; plugins/dflash_losses.py implements them. +DFlashLossTerm = Literal["ce", "kd", "tvd", "tv", "lambda"] + +# The per-projector_type selectors of the block objective that dflash_loss_terms replaces. +_LEGACY_OBJECTIVE_FIELDS = ( + "dflash_self_logit_distillation", + "dflash_lk_loss_type", + "dflash_ce_loss_alpha", + "dflash_l1_loss_alpha", +) + class DFlashConfig(ModeloptBaseConfig): """DFlash config for block-wise parallel speculative decoding.""" @@ -98,10 +109,48 @@ class DFlashConfig(ModeloptBaseConfig): default=True, description="Whether to freeze base model during DFlash module training." ) + dflash_loss_terms: ( + dict[DFlashLossTerm, Annotated[float, Field(ge=0.0, allow_inf_nan=False)]] | None + ) = ModeloptField( + default=None, + description=( + "The block objective as ``{term: weight}``, summed in the order given, e.g. " + "``{ce: 0.1, tvd: 0.9}``. Terms: ``ce`` (cross-entropy on the labels), ``kd`` " + "(cross-entropy against the base model's distribution), ``tvd`` (total variation " + "to the base model's distribution), ``tv`` (``1 - q(gold)``, total variation to " + "the label) and ``lambda`` (``ce`` and ``tv`` blended by acceptance, see " + "dflash_lk_ce_scale). Any term works with any projector_type, online and offline; " + "``kd`` and ``tvd`` read the base model's distribution, which offline batches must " + "then carry. Weights are >= 0 with at least one > 0; a term of weight 0 is " + "reported without being trained on. None (default) keeps each projector_type's " + "objective: ``kd`` or ``ce`` per dflash_self_logit_distillation for DFlash and " + "LiLiCorr, the same or dflash_lk_loss_type's term for DFlash2, " + "dflash_ce_loss_alpha * ``ce`` + dflash_l1_loss_alpha * ``tvd`` for DSpark, and " + "``ce`` for Domino. Setting it requires those selectors to keep their defaults. " + "The terms a variant adds itself (the DFlash2 selector, the LiLiCorr lattice, the " + "DSpark confidence head) and Domino's blend of the objective on its corrected and " + "backbone logits keep their own settings." + ), + ) + dflash_self_logit_distillation: bool = ModeloptField( default=True, description="Whether to use logit distillation from base model." ) + dflash_teacher_softcap: bool = ModeloptField( + default=False, + description=( + "Cap the teacher's logits the way the base model caps its own, " + "``cap * tanh(logits / cap)`` with ``cap`` the base config's " + "``final_logit_softcapping`` (Gemma family). The teacher is the base's ``lm_head`` " + "applied to its final hidden state for every batch, text or multimodal, online or " + "offline, which leaves that cap out unless this is set. Other scaling a base applies " + "at its ``lm_head``, such as ``logits_scaling`` or ``logit_scale``, stays out; " + "conversion warns about it. Requires a base config that sets " + "``final_logit_softcapping``. The draft's own logits are never capped." + ), + ) + dflash_loss_decay_factor: float = ModeloptField( default=0.0, description="Gamma for exponential loss decay weighting (paper Eq.4). " @@ -148,7 +197,7 @@ class DFlashConfig(ModeloptBaseConfig): dflash_use_torch_compile: bool = ModeloptField( default=False, description=( - "Whether to torch.compile the draft decoder stack and the DSpark TVD chunk in " + "Whether to torch.compile the draft decoder stack and the TVD term's chunk in " "training. Compiles once, on the first step, and needs a working Inductor " "toolchain." ), @@ -424,8 +473,9 @@ class DFlashConfig(ModeloptBaseConfig): default=1.0, ge=0.0, description=( - "DFlash2 only: scale of the CE share in the dflash_lk_loss_type='lambda' " - "blend. 1.0 starts the run as pure CE. Ignored for other loss types." + "Scale of the CE share in the 'lambda' blend (dflash_lk_loss_type='lambda', or the " + "'lambda' term of dflash_loss_terms). 1.0 starts the run as pure CE. Ignored " + "without that term." ), ) @@ -433,9 +483,8 @@ class DFlashConfig(ModeloptBaseConfig): default=1.0, ge=0.0, description=( - "DFlash2 only: how fast the CE share decays as acceptance rises in the " - "dflash_lk_loss_type='lambda' blend. 0 pins the blend at dflash_lk_ce_scale. " - "Ignored for other loss types." + "How fast the CE share decays as acceptance rises in the 'lambda' blend. 0 pins " + "the blend at dflash_lk_ce_scale. Ignored without that term." ), ) @@ -466,6 +515,28 @@ def _check_dpace_alpha(self) -> "DFlashConfig": ) return self + @model_validator(mode="after") + def _check_loss_terms(self) -> "DFlashConfig": + if self.dflash_loss_terms is None: + return self + if not any(self.dflash_loss_terms.values()): + raise ValueError( + f"dflash_loss_terms needs a term of weight > 0, got {self.dflash_loss_terms}." + ) + # Against the defaults, not model_fields_set: main.py converts from model_dump() and + # modelopt_state restores from it, which marks every field as set. + moved = [ + name + for name in _LEGACY_OBJECTIVE_FIELDS + if getattr(self, name) != type(self).model_fields[name].default + ] + if moved: + raise ValueError( + f"dflash_loss_terms sets the block objective, so {moved} must keep their " + "defaults; express the objective in dflash_loss_terms alone." + ) + return self + class MedusaConfig(ModeloptBaseConfig): """Medusa config.""" diff --git a/modelopt/torch/speculative/plugins/dflash_losses.py b/modelopt/torch/speculative/plugins/dflash_losses.py index 9b71fb1cf42..d18a1aef602 100644 --- a/modelopt/torch/speculative/plugins/dflash_losses.py +++ b/modelopt/torch/speculative/plugins/dflash_losses.py @@ -42,9 +42,9 @@ (: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. -The divergences a block objective is built from live in :data:`DIVERGENCES`. An objective -is an ordered list of ``(name, weight, term)`` entries, summed left to right by -:func:`objective_loss`. +The divergences a block objective is built from live in :data:`DIVERGENCES`, under the +names ``dflash_loss_terms`` takes. An objective is an ordered list of +``(name, weight, term)`` entries, summed left to right by :func:`objective_loss`. """ import functools @@ -54,6 +54,8 @@ import torch.nn.functional as F import torch.utils.checkpoint +from ..config import DFlashLossTerm + __all__ = [ "DIVERGENCES", "DraftLogits", @@ -387,7 +389,7 @@ def __call__(self, ctx, source): # The closed set of terms an objective is built from, by the name its metric reports. -DIVERGENCES = { +DIVERGENCES: dict[DFlashLossTerm, type[_Divergence]] = { "ce": _CrossEntropy, "kd": _SoftCrossEntropy, "tvd": _TeacherTotalVariation, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a0c2cc34bb1..3fcf336e00d 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -82,6 +82,8 @@ from transformers.trainer_pt_utils import LabelSmoother from transformers.utils import ModelOutput +from modelopt.torch.utils import warn_rank_0 + from ..dflash.conversion import DFlashDMRegistry from ..dflash.dflash_model import DFlashModel from .dflash_flex_attention import build_draft_block_mask @@ -105,6 +107,7 @@ _BASE_MODEL_PATHS, _EMBED_TOKENS_PATHS, _LM_HEAD_PATHS, + _LOGIT_SCALING_FIELDS, find_final_norm_path, ) from .modeling_final_norm import _apply_base_final_norm @@ -120,6 +123,8 @@ "pixel_values_videos", "image_grid_thw", "video_grid_thw", + "image_position_ids", + "video_position_ids", "mm_token_type_ids", "image_sizes", "images", @@ -226,9 +231,12 @@ def _base_model_norm(self): path = getattr(self, "base_model_norm_path", None) return self.get_submodule(path) if path else None - @property - def _needs_base_logits(self) -> bool: - """Whether the loss reads the target's logits; offline training then reconstructs them.""" + def _reads_teacher(self, config) -> bool: + """Whether a loss term reads the base distribution; offline batches must then carry it. + + Becomes ``_needs_base_logits`` at convert time. A variant that adds a term reading the + teacher declares it here, from ``config``. + """ return any(term.needs_teacher for _, _, term in self._objective) @property @@ -280,6 +288,14 @@ def modify(self, config): self.dflash_fp32_master_weights = config.dflash_fp32_master_weights base_config = self._base_llm_config + base_softcap = getattr(base_config, "final_logit_softcapping", None) + if config.dflash_teacher_softcap and base_softcap is None: + raise ValueError( + "dflash_teacher_softcap=True, but the base config sets no " + "final_logit_softcapping to cap the teacher's logits with." + ) + self._teacher_softcap = base_softcap if config.dflash_teacher_softcap else None + # Use Qwen3Config (not generic PretrainedConfig) so rope_parameters is # auto-populated from rope_theta. DFlash draft uses Qwen3 components. self.dflash_config = _Qwen3Config(**config.dflash_architecture_config) @@ -460,6 +476,18 @@ def modify(self, config): self.dflash_module.compile_body() self._objective = self._build_objective(config) + self._needs_base_logits = self._reads_teacher(config) + scaling = { + name: getattr(base_config, name) + for name in _LOGIT_SCALING_FIELDS + if getattr(base_config, name, None) not in (None, 1.0) + } + if self._needs_base_logits and scaling: + warn_rank_0( + f"The base config sets {', '.join(f'{k}={v}' for k, v in scaling.items())}, " + "which the distillation teacher does not apply: the teacher is lm_head's raw " + "projection of the base model's final hidden state." + ) # Compiling fuses the TVD chunk's six vocab-wide elementwise passes. self._tvd_chunk_fn = ( torch.compile(_tvd_chunk, dynamic=False, fullgraph=True) @@ -472,8 +500,19 @@ def _build_draft_module(self, dflash_config): return DFlashModule(dflash_config) def _build_objective(self, config): - """The block objective, ``[(name, weight, term)]``: KD from the base model, or CE.""" - return [objective_term("kd" if config.dflash_self_logit_distillation else "ce", 1.0)] + """The block objective, ``[(name, weight, term)]``, from ``dflash_loss_terms`` if set.""" + terms = config.dflash_loss_terms + if terms is None: + terms = self._default_loss_terms(config) + blend = {"scale": config.dflash_lk_ce_scale, "decay": config.dflash_lk_ce_decay} + return [ + objective_term(name, weight, **(blend if name == "lambda" else {})) + for name, weight in terms.items() + ] + + def _default_loss_terms(self, config): + """The variant's ``{term: weight}`` without ``dflash_loss_terms``: KD from the base, or CE.""" + return {"kd" if config.dflash_self_logit_distillation else "ce": 1.0} def _base_device(self): """Device the frozen base model sits on. ``meta`` during a from_pretrained restore.""" @@ -817,7 +856,8 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None): Next-token logits: the ones that predict the token at position ``i`` sit at ``i - 1``. Only those rows of the base hidden go through the final norm and lm_head, so no full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the - projection to just those vocab entries. + projection to just those vocab entries. ``dflash_teacher_softcap`` caps what is + projected; logits handed over in ``base_outputs`` are the model's own and only gathered. """ positions = (label_positions - 1).clamp(min=0) if base_outputs.logits is not None: @@ -837,10 +877,15 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None): rows = _gather_rows(base_outputs.base_hidden, positions).to(lm_head.weight.dtype) rows = _apply_base_final_norm(rows, base_outputs.base_hidden_prenorm, self._base_model_norm) if token_ids is None: - return lm_head(rows) - logits = torch.einsum("...h,...kh->...k", rows, lm_head.weight[token_ids]) - bias = getattr(lm_head, "bias", None) - return logits if bias is None else logits + bias[token_ids] + logits = lm_head(rows) + else: + logits = torch.einsum("...h,...kh->...k", rows, lm_head.weight[token_ids]) + bias = getattr(lm_head, "bias", None) + if bias is not None: + logits = logits + bias[token_ids] + cap = self._teacher_softcap + # In the base model's op order, so the cap rounds as the base's own does. + return logits if cap is None else torch.tanh(logits / cap) * cap def _online_base_outputs( self, @@ -853,13 +898,32 @@ def _online_base_outputs( cache_position=None, **kwargs, ): - """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. + """Run the frozen base model: the target hidden states, and its final hidden state. + + The final hidden state is the inner model's output for every batch, and + ``_teacher_logits`` projects it only at the rows the loss reads. A text batch runs the + inner model once. A multimodal batch needs the top-level forward to insert its + image/video features, and a hook on the inner model reads its output during it. + """ base_forward_kwargs = _multimodal_forward_kwargs(kwargs) - with torch.no_grad(): - if base_forward_kwargs: + if not base_forward_kwargs: + with torch.no_grad(): + base_out = self._base_model( + input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True + ) + return DFlashBaseModelOutput( + target_hidden=self._target_hidden(base_out.hidden_states), + base_hidden=base_out.last_hidden_state, + ) + + # Read where a text batch reads it, not at lm_head's input, which a CausalLM may rescale + # first (Inkling divides it by a multiplier). + base_hidden = [] + hook = self._base_model.register_forward_hook( + lambda _module, _args, output: base_hidden.append(output[0]) + ) + try: + with torch.no_grad(): raw_outputs = super().forward( input_ids=input_ids, attention_mask=attention_mask, @@ -873,20 +937,22 @@ def _online_base_outputs( return_dict=True, **base_forward_kwargs, ) - else: - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - + finally: + hook.remove() 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." ) + if len(base_hidden) != 1 or base_hidden[0].shape[:2] != input_ids.shape: + raise RuntimeError( + "The teacher reads the base model's final hidden state off its inner model, so the " + "multimodal forward must run that once over the whole sequence; it ran it on " + f"{[tuple(h.shape) for h in base_hidden]} for input_ids {tuple(input_ids.shape)}." + ) return DFlashBaseModelOutput( - target_hidden=self._target_hidden(raw_outputs.hidden_states), logits=raw_outputs.logits + target_hidden=self._target_hidden(raw_outputs.hidden_states), + base_hidden=base_hidden[0], ) def _draft_logits(self, batch): @@ -974,8 +1040,8 @@ def forward( - 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``. + A variant changes the step only through its hooks: ``_draft_logits`` (its head), + ``_position_weighting`` and ``_compute_loss``. """ if not self.training: if self.dflash_offline: diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index f716f4eaad1..8d5b047da1d 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -18,7 +18,7 @@ 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:`_build_objective` and :meth:`_compute_loss`. +:meth:`_default_loss_terms` and :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 @@ -39,7 +39,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .dflash_losses import LossOutput, _block_tokens, objective_term +from .dflash_losses import LossOutput, _block_tokens from .hf_dflash import HFDFlashModel from .modeling_dflash2 import DFlash2Module @@ -145,17 +145,11 @@ 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 _build_objective(self, config): - """DFlash's objective, or the label TV / CE-TV blend ``dflash_lk_loss_type`` selects.""" - if config.dflash_lk_loss_type == "tv": - return [objective_term("tv", 1.0)] - if config.dflash_lk_loss_type == "lambda": - return [ - objective_term( - "lambda", 1.0, scale=config.dflash_lk_ce_scale, decay=config.dflash_lk_ce_decay - ) - ] - return super()._build_objective(config) + def _default_loss_terms(self, config): + """DFlash's terms, or the label TV / CE-TV blend ``dflash_lk_loss_type`` selects.""" + if config.dflash_lk_loss_type == "ce": + return super()._default_loss_terms(config) + return {config.dflash_lk_loss_type: 1.0} def _compute_loss(self, ctx): """Backbone DFlash loss plus the candidate-selector cross-entropy. diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 0a4bc9fdbd4..b33734728f4 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -51,7 +51,8 @@ loss = (1 - lambda_base) * final_loss + lambda_base * base_loss -where ``final_loss`` is CE on the corrected logits and ``base_loss`` is CE on the +where ``final_loss`` is the block objective (CE unless ``dflash_loss_terms`` sets +another) on the corrected logits and ``base_loss`` the same objective on the backbone-only logits. ``lambda_base`` decays linearly from ``lambda_base_start`` to 0 over ``lambda_base_decay_ratio`` of training (curriculum: learn a good parallel backbone first, then the causal correction). The schedule is driven by @@ -65,9 +66,8 @@ from transformers.trainer_callback import TrainerCallback from ..dflash.conversion import DominoDMRegistry -from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_loss, objective_term +from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_loss from .hf_dflash import HFDFlashModel -from .modeling_dflash import DFlashBaseModelOutput from .modeling_domino import DominoModule logger = logging.getLogger(__name__) @@ -134,25 +134,15 @@ def get_exporter(self): def _current_lambda_base(self) -> float: return float(getattr(self, "_lambda_base", self.dflash_lambda_base_start)) - def _build_objective(self, config): + def _default_loss_terms(self, config): """CE on the labels, for the corrected logits and for the backbone's.""" - return [objective_term("ce", 1.0)] + return {"ce": 1.0} @property def _position_weighting(self): # dflash_loss_objective is not wired for Domino yet: it always decays. return "decay" - def _online_base_outputs(self, input_ids, attention_mask, **kwargs): - """The target hidden states; Domino does not use target-logit KD.""" - with torch.no_grad(): - base_out = self._base_model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - return DFlashBaseModelOutput(target_hidden=self._target_hidden(base_out.hidden_states)) - def _draft_logits(self, batch): """The backbone's logits with the GRU causal correction added to each block's suffix.""" backbone = super()._draft_logits(batch).backbone @@ -189,9 +179,9 @@ def _compute_loss(self, ctx): backbone_loss, _ = objective_loss(self._objective, ctx, "backbone") loss = (1.0 - lambda_base) * out.loss + lambda_base * backbone_loss metrics = { + **out.metrics, "loss_final": out.loss.detach(), "loss_backbone": backbone_loss.detach(), - "backbone_accuracy": out.metrics["backbone_accuracy"], # float64 holds the scheduled value exactly. "lambda_base": torch.full((), lambda_base, dtype=torch.float64, device=out.loss.device), } diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 69b5fdc3f1f..4f249dfded5 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -57,6 +57,7 @@ where ``TVD`` is the total-variation distance between the corrected draft distribution and the target (base-model) distribution at the aligned position, and the confidence target is ``c* = 1 - 0.5 * TVD`` (the analytical accept rate). +``dflash_loss_terms`` replaces the first two terms with any composition. """ import logging @@ -66,9 +67,8 @@ from transformers import PreTrainedModel from ..dflash.conversion import DSparkDMRegistry -from .dflash_losses import DraftLogits, LossOutput, _block_tokens, objective_term +from .dflash_losses import DraftLogits, LossOutput, _block_tokens from .hf_dflash import HFDFlashModel -from .modeling_dflash import DFlashBaseModelOutput from .modeling_dspark import DSparkModule logger = logging.getLogger(__name__) @@ -112,32 +112,19 @@ def get_exporter(self): return DSparkExporter(self) - def _build_objective(self, config): + def _default_loss_terms(self, config): """CE plus TVD to the base model; the DeepSpec defaults are TVD-dominant.""" - return [ - objective_term("ce", config.dflash_ce_loss_alpha), - objective_term("tvd", config.dflash_l1_loss_alpha), - ] + return {"ce": config.dflash_ce_loss_alpha, "tvd": config.dflash_l1_loss_alpha} + + def _reads_teacher(self, config) -> bool: + # The confidence head's target, the accept rate 1 - TVD / 2, reads the base distribution. + return super()._reads_teacher(config) or config.dflash_confidence_head_alpha > 0 @property def _position_weighting(self): # dflash_loss_objective is not wired for DSpark yet: it always decays. return "decay" - def _online_base_outputs(self, input_ids, attention_mask, **kwargs): - """The target hidden states, and the base's final hidden for the loss to project.""" - # The inner model, not the CausalLM: the loss projects only the rows it reads. - with torch.no_grad(): - base_out = self._base_model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - return DFlashBaseModelOutput( - target_hidden=self._target_hidden(base_out.hidden_states), - base_hidden=base_out.last_hidden_state, - ) - def _head_inputs(self, batch): """Each block position's teacher-forced previous token, and the draft's hidden states. diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index be3100e0278..c4216a66236 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -175,13 +175,12 @@ def modify(self, config): self.dflash_lilicorr_margin, ) - @property - def _needs_base_logits(self) -> bool: + def _reads_teacher(self, config) -> bool: """The distractor penalty and the calibration term both read the target's logits.""" return ( - super()._needs_base_logits - or self.dflash_lilicorr_w_pen > 0.0 - or self.dflash_lilicorr_w_cal > 0.0 + super()._reads_teacher(config) + or config.dflash_lilicorr_w_pen > 0.0 + or config.dflash_lilicorr_w_cal > 0.0 ) def get_exporter(self): diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 45870d8b6c7..836234fc052 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -118,10 +118,9 @@ def _get_sink_attention_fn(): class DFlashBaseModelOutput: """What a DFlash-family draft takes from the base model in a training step. - ``logits`` holds the full base logits when the producer already has them (DFlash's - online CausalLM forward, or an offline batch carrying ``base_model_logits``); - otherwise ``base_hidden`` stays unprojected, and ``HFDFlashModel._teacher_logits`` - projects only the rows a loss reads. + ``logits`` holds the full base logits when an offline batch carries them + (``base_model_logits``); otherwise ``base_hidden`` stays unprojected, and + ``HFDFlashModel._teacher_logits`` projects only the rows a loss reads. """ target_hidden: torch.Tensor # concatenated hidden states from target layers [B, seq, N*H] diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 85a3fe3a9f4..8354ce45eff 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -92,6 +92,15 @@ def find_final_norm_path(model) -> str | None: "language_model.backbone", ] _VLM_CONFIG_ATTRS = ["text_config", "llm_config"] +# Base config fields by which a CausalLM scales its logits at lm_head (1.0: no scaling). The +# DFlash teacher is lm_head's raw projection and leaves them out, so conversion warns about them. +_LOGIT_SCALING_FIELDS = ( + "logits_scaling", # Granite family, HyperCLOVA X, MiniCPM3 + "logit_scale", # Cohere family + "lm_head_multiplier", # Falcon-H1 + "output_multiplier", # MuseGlimmer + "logits_mup_width_multiplier", # Inkling +) _SAFETENSORS_INDEX_FILENAME = "model.safetensors.index.json" # Single-file safetensors names to try, in order. Mistral native checkpoints # use ``consolidated.safetensors`` instead of the HF-standard ``model.safetensors``. @@ -118,11 +127,14 @@ def __init__( rms_norm_eps=1e-6, rope_theta=None, final_norm_type=None, + final_logit_softcapping=None, **kwargs, ): """Initialize FakeBaseConfig with minimal model configuration parameters.""" super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) self.rms_norm_eps = rms_norm_eps + # The base's logit cap (Gemma family), for dflash_teacher_softcap. + self.final_logit_softcapping = final_logit_softcapping # Which self-implemented final-norm class FakeBaseModel builds, or None to build no norm # (model whose final-norm type we don't know). See _FINAL_NORM_CLASSES / # _FINAL_NORM_TYPE_BY_MODEL_TYPE. Persisted so a reloaded config rebuilds the same norm. @@ -242,6 +254,13 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM final_norm_type=_select_final_norm_type( getattr(base_cfg, "model_type", None), base_cfg ), + final_logit_softcapping=getattr(base_cfg, "final_logit_softcapping", None), + # So that DFlash's conversion warns about them on a fake base too. + **{ + name: getattr(base_cfg, name) + for name in _LOGIT_SCALING_FIELDS + if getattr(base_cfg, name, None) is not None + }, ) model = cls(config) # Load lm_head, embed_tokens, and (for known models) the final norm into the model. diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e7237b11b8f..659359be418 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/modelopt_recipes/general/speculative_decoding/domino.yaml @@ -62,8 +62,6 @@ dflash: dflash_block_size: 16 dflash_num_anchors: 256 dflash_use_torch_compile: false - # Domino does not use target-logit KD; it trains its own base/final CE losses. - 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. diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index db69f3c4b9e..fc0cf3e3788 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml @@ -65,9 +65,6 @@ dflash: dflash_block_size: 16 dflash_num_anchors: 256 dflash_use_torch_compile: false - # DSpark computes the target distribution internally for its TVD/confidence - # terms; the DFlash KD path is unused. - 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. 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 d96780e05d2..263537a40ed 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 @@ -81,7 +81,6 @@ dflash: # --- training knobs --- dflash_num_anchors: 64 dflash_use_torch_compile: false - dflash_self_logit_distillation: false dflash_loss_objective: dpace dflash_ce_loss_alpha: 0.1 dflash_l1_loss_alpha: 0.9 diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md index 90d7f457369..9450ad868de 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md @@ -87,6 +87,7 @@ in `examples/speculative_decoding/README.md#dflash-block-diffusion-for-speculati | `dflash.dflash_num_anchors` | 512 | Random anchor positions sampled per sequence | | `dflash.dflash_loss_decay_factor` | 4.0 | Exponential decay gamma; 0 disables | | `dflash.dflash_self_logit_distillation` | true | Logit distillation from the target | +| `dflash.dflash_loss_terms` | unset | `{term: weight}` over `ce`/`kd`/`tvd`/`tv`/`lambda`, any variant; unset keeps the variant's objective. Leave the variant's own loss selectors at their defaults with it | | `dflash.dflash_architecture_config.num_hidden_layers` | 5 | Draft decoder layers — must equal the dump's draft depth (`--num-draft-layers` on vLLM; hardcoded 5 on HF / TRT-LLM) | | `dflash.dflash_mask_token_id` | auto | See *Per-model adjustments* | | `dflash.dflash_swa_window_size` | unset | Sliding-window attention for the draft; must be >= `dflash_block_size` | @@ -108,6 +109,7 @@ Export is automatic: after training, rank 0 exports every `checkpoint-` to | `trust_remote_code` MoE with an older transformers pin | Set `OVERRIDE_TRANSFORMERS` in the task environment (MiniMax-M2.7 needs 4.57.1). Set `ACCELERATE_CONFIG` when the model needs FSDP2 via accelerate config rather than transformers-native `ParallelismConfig`. | | Very large MoE base | Use the offline variant with `model.use_fake_base_for_offline=true`; plain DDP suffices, so no FSDP2 patches. Set `MIXED_PRECISION: "no"` with `training.bf16=false` if the model requires it. | | Draft trained at short context, served long | Set `dflash.dflash_export_rope_scaling` (YaRN); factor = target context / `training_seq_len`. | +| Base with `final_logit_softcapping` (Gemma family) and a loss term reading the target's distribution (`kd`, `tvd`; DFlash, DSpark and LiLiCorr by default) | Set `dflash.dflash_teacher_softcap=true` to distill from the capped distribution, text and image batches alike. | | Multi-node | Set `NUM_NODES` in the environment; `HEAD_NODE_IP` is auto-detected from Slurm. | ## Success markers diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 8494f6f4b76..f6181b398f4 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -49,7 +49,7 @@ Everything in `dflash.md` applies. Domino adds: | `dflash.dflash_lambda_base_start` | 1.0 | Curriculum start weight on the base loss | | `dflash.dflash_lambda_base_decay_ratio` | 1.0 | Fraction of training over which `lambda_base` decays to 0 | -`dflash_self_logit_distillation` is **false** — Domino trains its own base/final CE +Domino ignores `dflash_self_logit_distillation` — it trains its own base/final CE losses rather than distilling target logits. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256, `num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04, `dflash_loss_decay_factor` 7.0. diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index bb2c6992ac5..d0bad8ee183 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -88,7 +88,7 @@ Everything in `dflash.md` applies. DSpark adds: Total loss is `ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE`. -`dflash_self_logit_distillation` is **false** for DSpark — it computes the target +DSpark ignores `dflash_self_logit_distillation` — it computes the target distribution internally for the TVD and confidence terms, so the DFlash KD path is unused. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256, `num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04. diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index a02128e5e39..f3943f01f2c 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -59,6 +59,7 @@ ) from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list from modelopt.torch.quantization.mode import CalibrateModeRegistry, get_modelike_from_algo_cfg +from modelopt.torch.speculative.config import DFlashConfig # --------------------------------------------------------------------------- # Static YAML fixtures @@ -329,13 +330,13 @@ def test_import_resolution_honors_huggingface_alias(): _PTQ_SCHEMA = "modelopt.recipe.config.ModelOptPTQRecipe" -def _all_shipped_ptq_recipe_paths(): - """Every shipped PTQ recipe, discovered from disk rather than a hardcoded list. +def _all_shipped_recipe_paths(recipe_type): + """Every shipped recipe of ``recipe_type``, discovered from disk rather than a hardcoded list. A recipe says it is PTQ with ``metadata.recipe_type``, with a ``# modelopt-schema:`` comment naming :class:`ModelOptPTQRecipe`, or -- for a checkpoint alias, which states neither -- by delegating to a recipe that does. All three are picked up, so every - shipped recipe is swept by the tests below. + shipped recipe is swept by the tests below. Other recipe types state theirs the same way. """ root = files("modelopt_recipes") paths = [] @@ -351,7 +352,7 @@ def _all_shipped_ptq_recipe_paths(): # Ask the loader's own resolver rather than re-deriving the rules here: a # top-level ``$import`` does not imply PTQ (an alias of a speculative recipe # delegates the same way), and this cannot drift from what load_recipe does. - if _peek_recipe_type(path) == RecipeType.PTQ: + if _peek_recipe_type(path) == recipe_type: paths.append(str(rel.with_suffix(""))) return paths @@ -359,7 +360,7 @@ def _all_shipped_ptq_recipe_paths(): # Discovered from disk (not hardcoded) so the smoke tests cover every shipped PTQ # recipe — general/, model_type//, and models/// — and # never drift as recipes are added, moved, or removed. -_BUILTIN_PTQ_RECIPES = _all_shipped_ptq_recipe_paths() +_BUILTIN_PTQ_RECIPES = _all_shipped_recipe_paths(RecipeType.PTQ) def test_ptq_recipes_are_discovered(): @@ -1059,6 +1060,25 @@ def test_load_recipe_dflash_builtin(): assert recipe.training.training_seq_len == 4096 +def test_dspark_and_domino_recipes_take_loss_terms_alone(): + """Neither variant reads ``dflash_self_logit_distillation``, so none of their recipes pins it. + + ``dflash_loss_terms`` needs the old loss selectors at their defaults; a pin there would have + to be undone before a user could set the terms. + """ + checked = [] + for path in _all_shipped_recipe_paths(RecipeType.SPECULATIVE_DFLASH): + projector_type = load_recipe(path).dflash.dflash_architecture_config.get("projector_type") + if projector_type not in ("dspark", "domino"): + continue + recipe = load_recipe(path, overrides=["dflash.dflash_loss_terms={ce: 0.5, tvd: 0.5}"]) + # main.py converts from a model_dump(), which marks every field as set. + dumped = DFlashConfig(**recipe.dflash.model_dump()) + assert dumped.dflash_loss_terms == {"ce": 0.5, "tvd": 0.5}, path + checked.append(path) + assert len(checked) >= 3, f"too few DSpark/Domino recipes discovered: {checked}" + + def test_load_recipe_dflash_missing_section_raises(tmp_path): """load_recipe raises ValueError when 'dflash' is absent for a SPECULATIVE_DFLASH recipe.""" bad = tmp_path / "bad.yml" diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 44cc19ed90c..1d3ae5d8af7 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -25,6 +25,7 @@ import transformers from modelopt.torch.speculative.plugins.modeling_fakebase import ( + _LOGIT_SCALING_FIELDS, FakeBaseConfig, FakeBaseModel, find_final_norm_path, @@ -71,6 +72,20 @@ def test_fakebase_local_happy_path(fake_checkpoint): assert model.embed_tokens.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE]) +@pytest.mark.parametrize("softcap", [None, 30.0]) +def test_fakebase_carries_the_logit_softcap(fake_checkpoint, fake_config, softcap): + """dflash_teacher_softcap reads the cap off the base config, a fake one included.""" + fake_config.final_logit_softcapping = softcap + assert FakeBaseModel.from_source(str(fake_checkpoint)).config.final_logit_softcapping == softcap + + +@pytest.mark.parametrize("field", _LOGIT_SCALING_FIELDS) +def test_fakebase_carries_the_logit_scaling(fake_checkpoint, fake_config, field): + """DFlash's conversion warns about logit scaling off the base config, a fake one included.""" + setattr(fake_config, field, 8.0) + assert getattr(FakeBaseModel.from_source(str(fake_checkpoint)).config, field) == 8.0 + + def test_fakebase_missing_index_raises(tmp_path, fake_config): with pytest.raises(FileNotFoundError, match="safetensors"): FakeBaseModel.from_source(str(tmp_path)) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ec86aa115c7..107af93409e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -18,10 +18,13 @@ GPU-dependent tests (training forward, module forward) are in tests/gpu/. """ +import copy import json import logging import os +import warnings from types import SimpleNamespace +from typing import get_args from unittest.mock import MagicMock, patch import pytest @@ -30,17 +33,26 @@ from _test_utils.torch.speculative.dflash import get_dflash_config from _test_utils.torch.transformers_models import ( get_tiny_llama, + get_tiny_qwen3vl, tf_modelopt_state_and_output_tester, ) from torch.nn.attention.flex_attention import create_mask -from transformers import AutoModelForCausalLM +from transformers import ( + AutoModelForCausalLM, + Gemma4Config, + Gemma4ForConditionalGeneration, + GraniteConfig, + GraniteForCausalLM, +) from transformers.modeling_layers import GradientCheckpointingLayer import modelopt.torch.opt as mto 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.config import DFlashConfig, DFlashLossTerm from modelopt.torch.speculative.plugins.dflash_losses import ( + DIVERGENCES, DraftLogits, _block_position_weights, _block_targets, @@ -64,7 +76,8 @@ MasterWeightAdamW, VerifyMasterWeightsCallback, ) -from modelopt.torch.speculative.utils import AcceptanceRateValidation +from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel +from modelopt.torch.speculative.utils import AcceptanceRateValidation, load_vlm_or_llm from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator BLOCK_SIZE = 4 @@ -359,6 +372,189 @@ def test_each_term_reduces_as_before(self, name): assert torch.equal(loss, legacy) +# What each variant's conversion requires beyond get_dflash_config(), sized for the tiny +# model: its dflash_architecture_config fields, then its other config fields. +_VARIANTS = { + "dflash": ({}, {}), + "dflash2": ( + {"conv_kernel_size": 2, "conv_group_size": 4, "selector_rank": 8, "selector_top_k": 5}, + {}, + ), + "lilicorr": ( + { + "lilicorr_candidate_topk": 4, + "lilicorr_num_layers": 1, + "lilicorr_num_heads": 4, + "lilicorr_mlp_ratio": 2.0, + "lilicorr_factor_dim": 8, + "lilicorr_vector_eps": 1e-4, + "lilicorr_logit_scale": 8.0, + }, + { + "dflash_lilicorr_w_ce": 0.25, + "dflash_lilicorr_w_margin": 0.0, + "dflash_lilicorr_w_pen": 0.25, + }, + ), + "dspark": ({"markov_rank": 16}, {}), + # lambda_base below 1, so that Domino's corrected logits train too. + "domino": ({"emb_dim": 16, "gru_hidden_dim": 32}, {"dflash_lambda_base_start": 0.5}), +} + + +def _variant_model(projector_type, offline=False, arch=None, **overrides): + """A tiny model converted to ``projector_type``, with ``arch`` and ``overrides`` applied.""" + arch_fields, fields = _VARIANTS[projector_type] + config = get_dflash_config(offline=offline) + config["dflash_architecture_config"].update( + projector_type=projector_type, **arch_fields, **(arch or {}) + ) + config.update(fields, **overrides) + model = get_tiny_llama(num_hidden_layers=4) + model.config.num_orig_hidden_layers = 4 + mtsp.convert(model, [("dflash", config)]) + return model + + +def _objective_spec(model): + """The objective as data: each term's name, weight, class and bound parameters.""" + return [(name, weight, type(term), vars(term)) for name, weight, term in model._objective] + + +class TestLossTerms: + """``dflash_loss_terms`` composes the block objective of any variant from the registry.""" + + def test_every_configurable_term_is_implemented(self): + assert set(get_args(DFlashLossTerm)) == set(DIVERGENCES) + + @pytest.mark.parametrize( + ("terms", "match"), + [ + ({"ce": 1.0, "l1": 1.0}, "'ce', 'kd', 'tvd', 'tv' or 'lambda'"), + ({"ce": -0.5}, "greater than or equal to 0"), + ({"ce": 0.0}, "weight > 0"), + ], + ids=["unknown", "negative", "all-zero"], + ) + def test_bad_terms_are_rejected(self, terms, match): + with pytest.raises(ValueError, match=match): + DFlashConfig(dflash_loss_terms=terms) + + @pytest.mark.parametrize( + ("selector", "value"), + [ + ("dflash_self_logit_distillation", False), + ("dflash_lk_loss_type", "tv"), + ("dflash_ce_loss_alpha", 0.5), + ("dflash_l1_loss_alpha", 0.5), + ], + ) + def test_terms_and_a_legacy_selector_are_rejected(self, selector, value): + """One source for the objective.""" + with pytest.raises(ValueError, match=selector): + DFlashConfig(dflash_loss_terms={"ce": 1.0}, **{selector: value}) + + @pytest.mark.parametrize( + ("projector_type", "legacy", "terms"), + [ + ("dflash", {}, {"kd": 1.0}), + ("dflash", {"dflash_self_logit_distillation": False}, {"ce": 1.0}), + ( + "dflash2", + {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "tv"}, + {"tv": 1.0}, + ), + ( + "dflash2", + {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "lambda"}, + {"lambda": 1.0}, + ), + ("lilicorr", {}, {"kd": 1.0}), + ("dspark", {"dflash_ce_loss_alpha": 0.3}, {"ce": 0.3, "tvd": 0.9}), + ("domino", {}, {"ce": 1.0}), + ], + ids=[ + "dflash-kd", + "dflash-ce", + "dflash2-tv", + "dflash2-lambda", + "lilicorr", + "dspark", + "domino", + ], + ) + def test_unset_terms_keep_each_variants_objective(self, projector_type, legacy, terms): + """Unset, each variant builds its own objective, which the same terms reproduce.""" + blend = {"dflash_lk_ce_scale": 0.8, "dflash_lk_ce_decay": 2.0} + model = _variant_model(projector_type, **legacy, **blend) + assert [(name, weight) for name, weight, _ in model._objective] == list(terms.items()) + explicit = _variant_model(projector_type, dflash_loss_terms=terms, **blend) + assert _objective_spec(explicit) == _objective_spec(model) + + def test_terms_survive_the_config_round_trips(self): + """main.py converts from a model_dump() and a restore reads one back: all fields set.""" + config = DFlashConfig(**get_dflash_config(), dflash_loss_terms={"tvd": 0.5, "ce": 0.5}) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config.model_dump())]) + restored = mto.restore_from_modelopt_state( + get_tiny_llama(num_hidden_layers=4), mto.modelopt_state(model) + ) + assert [name for name, _, _ in restored._objective] == ["tvd", "ce"] + assert _objective_spec(restored) == _objective_spec(model) + + @pytest.mark.parametrize( + ("projector_type", "terms"), + [ + ("dflash", {"ce": 0.5, "tvd": 0.5}), + ("dflash2", {"tv": 0.5, "kd": 0.5}), + ("lilicorr", {"tvd": 1.0}), + ("dspark", {"kd": 1.0}), + ("domino", {"kd": 0.5, "lambda": 0.5, "tvd": 0.0}), + ], + ids=["dflash", "dflash2", "lilicorr", "dspark", "domino"], + ) + def test_a_composed_objective_trains_every_variant(self, projector_type, terms): + """Online, the teacher every variant now has; every draft parameter gets a gradient. + + Every term is reported, one of weight 0 too. + """ + model = _variant_model(projector_type, dflash_loss_terms=terms).float() + model.train() + out = model(**_dflash_batch(model.config.vocab_size)) + assert torch.isfinite(out.loss) + assert {f"loss_{name}" for name in terms} <= out.draft_metrics.keys() + out.loss.backward() + ungraded = [ + name + for name, p in model.dflash_module.named_parameters() + if p.requires_grad and (p.grad is None or not torch.isfinite(p.grad).all()) + ] + assert not ungraded, ungraded + + @pytest.mark.parametrize( + ("projector_type", "terms", "confidence", "needs"), + [ + ("dspark", {"ce": 1.0}, 0.0, False), + ("dspark", {"ce": 1.0}, 0.5, True), + ("domino", None, 0.0, False), + ("domino", {"kd": 1.0}, 0.0, True), + ], + ids=["dspark-ce", "dspark-ce-confidence", "domino", "domino-kd"], + ) + def test_offline_needs_the_teacher_only_for_terms_that_read_it( + self, projector_type, terms, confidence, needs + ): + """Objective terms and variant terms declare it; nothing else does.""" + model = _variant_model( + projector_type, + offline=True, + arch={"use_confidence_head": bool(confidence)}, + dflash_loss_terms=terms, + dflash_confidence_head_alpha=confidence, + ) + assert model._needs_base_logits is needs + + class TestDFlashSaveRestore: """Test DFlash model save and restore.""" @@ -1622,9 +1818,14 @@ class TestTeacherLogits: """ @staticmethod - def _model(): + def _model(softcap=None): + """A converted tiny model; ``softcap`` sets the base's cap and turns the switch on.""" model = get_tiny_llama(num_hidden_layers=4) - mtsp.convert(model, [("dflash", get_dflash_config())]) + config = get_dflash_config() + if softcap is not None: + model.config.final_logit_softcapping = softcap + config["dflash_teacher_softcap"] = True + mtsp.convert(model, [("dflash", config)]) return model @staticmethod @@ -1667,8 +1868,28 @@ def test_token_ids_pick_entries_of_the_full_rows(self): model._teacher_logits(outputs, labels, token_ids), rows.gather(-1, token_ids) ) - def test_handed_over_logits_are_gathered_not_recomputed(self): - model = self._model() + @pytest.mark.parametrize("narrow", [False, True], ids=["full", "token_ids"]) + def test_softcap_caps_the_projected_teacher(self, narrow): + cap = 0.5 + model = self._model(softcap=cap).float() + hidden = 10 * self._hidden(model) + labels = self._labels() + rows = self._rows(model._base_model_lm_head(hidden), labels - 1) + assert rows.abs().max() > cap, "fixture does not reach the cap" + want = cap * torch.tanh(rows / cap) + token_ids = None + if narrow: + token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5)) + want = want.gather(-1, token_ids) + got = model._teacher_logits( + DFlashBaseModelOutput(None, base_hidden=hidden), labels, token_ids + ) + torch.testing.assert_close(got, want) + + @pytest.mark.parametrize("softcap", [None, 0.5], ids=["uncapped", "capped"]) + def test_handed_over_logits_are_gathered_not_recomputed(self, softcap): + """Handed-over logits are the model's own: gathered as they are, never capped again.""" + model = self._model(softcap) logits = torch.randn(2, SEQ_LEN, model.config.vocab_size) outputs = DFlashBaseModelOutput(None, logits=logits) labels = self._labels() @@ -1714,11 +1935,12 @@ def test_kd_loss_matches_full_logits(self): ) torch.testing.assert_close(loss, loss_full) - def test_offline_step_never_projects_the_full_sequence(self): - """Offline/streaming with KD on, lm_head only ever sees draft rows and teacher rows.""" + @pytest.mark.parametrize("offline", [False, True], ids=["online", "offline"]) + def test_step_never_projects_the_full_sequence(self, offline): + """With KD on, lm_head only ever sees draft rows and teacher rows.""" model = get_tiny_llama(num_hidden_layers=4) model.config.num_orig_hidden_layers = 4 - mtsp.convert(model, [("dflash", get_dflash_config(offline=True))]) + mtsp.convert(model, [("dflash", get_dflash_config(offline=offline))]) assert model.dflash_self_logit_distillation model.train() seen = [] @@ -1727,17 +1949,382 @@ def test_offline_step_never_projects_the_full_sequence(self): ) bsz, hidden = 2, model.config.hidden_size dtype = next(model.dflash_module.parameters()).dtype - base_model_outputs = { - "aux_hidden_states": torch.randn( - bsz, SEQ_LEN, len(model.target_layer_ids) * hidden, dtype=dtype - ), - "base_model_hidden_states": torch.randn(bsz, SEQ_LEN, hidden, dtype=dtype), - } + kwargs = {} + if offline: + kwargs["base_model_outputs"] = { + "aux_hidden_states": torch.randn( + bsz, SEQ_LEN, len(model.target_layer_ids) * hidden, dtype=dtype + ), + "base_model_hidden_states": torch.randn(bsz, SEQ_LEN, hidden, dtype=dtype), + } input_ids = torch.randint(1, model.config.vocab_size, (bsz, SEQ_LEN)) model( - input_ids=input_ids, - attention_mask=torch.ones_like(input_ids), - base_model_outputs=base_model_outputs, + input_ids=input_ids, attention_mask=torch.ones_like(input_ids), **kwargs ).loss.backward() assert len(seen) == 2, seen # the draft's logits, then the KD teacher rows assert (bsz, SEQ_LEN, hidden) not in seen, seen + + def test_online_teacher_is_the_causal_lm_distribution(self): + """The rows projected from the inner model's final hidden are the CausalLM's logits.""" + model = self._model().float() + labels = self._labels() + input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN)) + attention_mask = torch.ones_like(input_ids) + outputs = model._online_base_outputs(input_ids, attention_mask) + model.eval() + with torch.no_grad(): + full = model(input_ids=input_ids, attention_mask=attention_mask).logits + torch.testing.assert_close( + model._teacher_logits(outputs, labels), self._rows(full, labels - 1) + ) + + +# Qwen3-VL's vision placeholder ids, moved inside get_tiny_qwen3vl's 32-token vocab. +_VISION_START, _IMAGE = 28, 30 + + +def _tiny_vlm(arch=None, softcap=None): + """A tiny Qwen3-VL converted online, in fp32; ``arch`` updates the draft's architecture. + + ``softcap`` names a cap in its text config, which Qwen3-VL's own forward never applies, and + turns ``dflash_teacher_softcap`` on. + """ + model = get_tiny_qwen3vl(num_hidden_layers=4) + model.config.vision_start_token_id, model.config.image_token_id = _VISION_START, _IMAGE + config = get_dflash_config() + config["dflash_architecture_config"].update(arch or {}) + if softcap is not None: + model.config.text_config.final_logit_softcapping = softcap + config["dflash_teacher_softcap"] = True + mtsp.convert(model, [("dflash", config)]) + return model.float() + + +def _image_batch(): + """Two sequences, each holding a 2x2-patch image at positions 4-7, and its pixels.""" + g = torch.Generator().manual_seed(0) + input_ids = torch.randint(1, _VISION_START, (2, SEQ_LEN), generator=g) + input_ids[:, 3] = _VISION_START + input_ids[:, 4:8] = _IMAGE + return { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + # One row per 4x4 patch of 3 channels: get_tiny_qwen3vl merges no patches. + "pixel_values": torch.randn(2 * 4, 3 * 4 * 4, generator=g), + "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 2]]), + "mm_token_type_ids": (input_ids == _IMAGE).long(), + } + + +def _tiny_gemma4_and_batch(modality): + """A tiny Gemma 4 VLM converted online, in fp32, and a batch with an image or a video at 3-6.""" + torch.manual_seed(0) + image_token, video_token = 62, 63 + text_config = { + "vocab_size": 64, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "max_position_embeddings": 64, + "sliding_window": 8, + "vocab_size_per_layer_input": 64, + "hidden_size_per_layer_input": 8, + } + vision_config = { + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "head_dim": 8, + "patch_size": 4, + "position_embedding_size": 64, + } + model = Gemma4ForConditionalGeneration( + Gemma4Config( + text_config=text_config, + vision_config=vision_config, + image_token_id=image_token, + # The other special tokens, inside the tiny vocabulary too. + boi_token_id=60, + eoi_token_id=61, + video_token_id=video_token, + audio_token_id=57, + boa_token_id=59, + eoa_token_index=58, + ) + ) + mtsp.convert(model, [("dflash", get_dflash_config())]) + g = torch.Generator().manual_seed(1) + input_ids = torch.randint(1, 56, (2, SEQ_LEN), generator=g) + # 6x6 patches of 4x4 pixels, pooled to 2x2 tokens; a video has one such frame. + input_ids[:, 3:7] = image_token if modality == "image" else video_token + ys, xs = torch.meshgrid(torch.arange(6), torch.arange(6), indexing="ij") + pixels = torch.randn(2, 6 * 6, 3 * 4 * 4, generator=g) + positions = torch.stack([xs.flatten(), ys.flatten()], -1).expand(2, -1, -1) + if modality == "image": + visual = {"pixel_values": pixels, "image_position_ids": positions} + else: + visual = {"pixel_values_videos": pixels[:, None], "video_position_ids": positions[:, None]} + batch = { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + **visual, + # Gemma 4's processor marks image tokens 1 and video tokens 2. + "mm_token_type_ids": (input_ids == image_token).long() + 2 * (input_ids == video_token), + } + return model.float(), batch + + +class TestMultimodalTeacher: + """An image batch gets a text batch's teacher: lm_head projecting the inner model's output. + + The top-level forward inserts the image, and a hook reads the inner model's output during it. + """ + + @staticmethod + def _labels(): + torch.manual_seed(0) + return torch.randint(1, SEQ_LEN, (2, 3, BLOCK_SIZE)) + + @staticmethod + def _own_rows(model, batch, positions): + """Rows of the model's own logits, which for Qwen3-VL are lm_head's output as it is.""" + model.eval() + with torch.no_grad(): + logits = model(**batch).logits + return logits[torch.arange(positions.shape[0]).view(-1, 1, 1), positions] + + def test_teacher_projects_the_final_hidden_state_over_the_whole_sequence(self): + model = _tiny_vlm() + batch = _image_batch() + lm_head_inputs = [] + model._base_model_lm_head.register_forward_pre_hook( + lambda _module, args: lm_head_inputs.append(args[0]) + ) + outputs = model._online_base_outputs(**batch) + assert outputs.logits is None + # Qwen3-VL feeds the inner model's output to lm_head as it is. + assert torch.equal(outputs.base_hidden, lm_head_inputs[0]) + labels = self._labels() + torch.testing.assert_close( + model._teacher_logits(outputs, labels), self._own_rows(model, batch, labels - 1) + ) + + def test_text_and_image_batches_read_the_same_final_hidden_state(self): + """Up to the image, also when the CausalLM rescales lm_head's input, as Inkling does.""" + model = _tiny_vlm() + model._base_model_lm_head.register_forward_pre_hook(lambda _module, args: (args[0] / 24,)) + batch = _image_batch() + image = model._online_base_outputs(**batch).base_hidden + text = model._online_base_outputs(batch["input_ids"], batch["attention_mask"]).base_hidden + torch.testing.assert_close(image[:, :4], text[:, :4]) + + def test_softcap_caps_the_image_batch_teacher(self): + cap = 0.1 + model = _tiny_vlm(softcap=cap) + batch = _image_batch() + outputs = model._online_base_outputs(**batch) + labels = self._labels() + rows = self._own_rows(model, batch, labels - 1) + assert rows.abs().max() > cap, "fixture does not reach the cap" + torch.testing.assert_close( + model._teacher_logits(outputs, labels), cap * torch.tanh(rows / cap) + ) + + @pytest.mark.parametrize( + "arch", + [ + {}, + {"projector_type": "dspark", "markov_rank": 16}, + {"projector_type": "domino", "emb_dim": 16, "gru_hidden_dim": 32}, + ], + ids=["dflash", "dspark", "domino"], + ) + def test_every_variant_trains_on_the_image(self, arch): + """Another image in the same text changes the base's rows from the image on, not before.""" + model = _tiny_vlm(arch) + batch = _image_batch() + other = {**batch, "pixel_values": torch.randn_like(batch["pixel_values"])} + first, second = (model._online_base_outputs(**b) for b in (batch, other)) + for name in ("target_hidden", "base_hidden"): + a, b = getattr(first, name), getattr(second, name) + torch.testing.assert_close(a[:, :4], b[:, :4], msg=name) + assert not torch.allclose(a[:, 4:], b[:, 4:]), name + model.train() + out = model(**batch) + assert torch.isfinite(out.loss) + out.loss.backward() + + @pytest.mark.parametrize("modality", ["image", "video"]) + def test_a_gemma4_image_or_video_reaches_the_base(self, modality): + """Gemma 4's vision tower reads the patch positions its processor sends with the pixels.""" + model, batch = _tiny_gemma4_and_batch(modality) + pixels = "pixel_values" if modality == "image" else "pixel_values_videos" + other = {**batch, pixels: torch.randn_like(batch[pixels])} + first, second = (model._online_base_outputs(**b).base_hidden for b in (batch, other)) + torch.testing.assert_close(first[:, :3], second[:, :3]) + assert not torch.allclose(first[:, 3:], second[:, 3:]) + + def test_the_inner_model_is_hooked_only_inside_an_image_batch_forward(self): + """Never for a text batch; an image batch's hook goes even when its forward fails.""" + model = _tiny_vlm() + base = model._base_model + hooks_seen = [] + base.register_forward_pre_hook( + lambda module, _args: hooks_seen.append(len(module._forward_hooks)) + ) + batch = _image_batch() + model._online_base_outputs(batch["input_ids"], batch["attention_mask"]) + model._online_base_outputs(**batch) + assert hooks_seen == [0, 1] + assert not base._forward_hooks + with pytest.raises(RuntimeError): # pixels that do not fit the vision tower + model._online_base_outputs(**{**batch, "pixel_values": torch.randn(3, 5)}) + assert not base._forward_hooks + + def test_logits_to_keep_does_not_narrow_the_teacher(self, monkeypatch): + """It narrows lm_head's input, not the inner model's output that the teacher reads.""" + model = _tiny_vlm() + batch = _image_batch() + want = model._online_base_outputs(**batch).base_hidden + monkeypatch.setattr( + hf_dflash, + "_MULTIMODAL_FORWARD_KWARGS", + hf_dflash._MULTIMODAL_FORWARD_KWARGS | {"logits_to_keep"}, + ) + assert torch.equal(model._online_base_outputs(**batch, logits_to_keep=1).base_hidden, want) + + @pytest.mark.parametrize("run", ["twice", "short"]) + def test_the_inner_model_must_run_once_over_the_whole_sequence(self, run): + """Otherwise the forward gives the teacher no one final hidden state for every row.""" + model = _tiny_vlm() + batch = _image_batch() + base = model._base_model + + def run_it_again(_module, _args): + base(input_ids=batch["input_ids"]) + + def drop_the_last_row(_module, _args, output): + output.last_hidden_state = output.last_hidden_state[:, :-1] + + if run == "twice": + model._base_model_lm_head.register_forward_pre_hook(run_it_again) + else: + base.register_forward_hook(drop_the_last_row) + hooks = dict(base._forward_hooks) + with pytest.raises(RuntimeError, match="once over the whole sequence"): + model._online_base_outputs(**batch) + assert base._forward_hooks == hooks + + +class TestTeacherSoftcapAtConvert: + """``dflash_teacher_softcap`` alone decides the cap, checked against the base at build time.""" + + @staticmethod + def _base(multimodal=True, cap=30.0): + """A tiny base capping its logits at ``cap``; a nested text config marks a multimodal one.""" + model = get_tiny_llama(num_hidden_layers=4) + model.config.num_orig_hidden_layers = 4 + text_config = model.config + if multimodal: + text_config = model.config.text_config = copy.deepcopy(model.config) + text_config.final_logit_softcapping = cap + return model + + def test_a_base_without_a_cap_is_rejected(self): + config = get_dflash_config() + config["dflash_teacher_softcap"] = True + with pytest.raises(ValueError, match="final_logit_softcapping"): + mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) + + @pytest.mark.parametrize("switch", [False, True], ids=["switch-off", "switch-on"]) + @pytest.mark.parametrize("multimodal", [False, True], ids=["text-only", "multimodal"]) + def test_a_capping_base_converts_with_the_cap_the_switch_sets(self, multimodal, switch): + """Online, with KD reading the teacher: a multimodal base gets the same rule.""" + model = self._base(multimodal) + config = get_dflash_config() + config["dflash_teacher_softcap"] = switch + mtsp.convert(model, [("dflash", config)]) + assert model.dflash_self_logit_distillation + assert model._teacher_softcap == (30.0 if switch else None) + + @pytest.mark.parametrize( + ("saved", "cap"), + [(True, 30.0), (False, None), (None, None)], + ids=["on", "off", "saved-before-the-switch"], + ) + def test_restore_reads_the_switch_from_the_saved_config(self, saved, cap): + """A checkpoint saved before the switch existed restores with it off.""" + model = self._base() + config = get_dflash_config() + config["dflash_teacher_softcap"] = bool(saved) + mtsp.convert(model, [("dflash", config)]) + state = mto.modelopt_state(model) + if saved is None: + del state["modelopt_state_dict"][0][1]["config"]["dflash_teacher_softcap"] + restored = mto.restore_from_modelopt_state(self._base(), state) + assert restored._teacher_softcap == cap + + +class TestLogitScalingWarning: + """Conversion names the logit scaling that the teacher, lm_head's raw projection, leaves out. + + Once, on rank 0, and only when a loss term reads the teacher. + """ + + @staticmethod + def _granite(logits_scaling): + """A tiny Granite, whose CausalLM divides its logits by ``logits_scaling``.""" + config = GraniteConfig( + vocab_size=32, + hidden_size=32, + intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + logits_scaling=logits_scaling, + ) + return GraniteForCausalLM(config) + + @staticmethod + def _warnings(model, kd=True, offline=None): + config = get_dflash_config(offline=offline) + config["dflash_self_logit_distillation"] = kd + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + mtsp.convert(model, [("dflash", config)]) + return [str(w.message) for w in caught if "teacher does not apply" in str(w.message)] + + def test_a_scaling_base_warns_once(self): + (message,) = self._warnings(self._granite(8.0)) + assert "logits_scaling=8.0" in message + + def test_an_offline_fake_base_warns_too(self, tmp_path): + """Offline VLMs always train on a fake base, which carries the scaling fields over.""" + self._granite(8.0).save_pretrained(tmp_path) + model = load_vlm_or_llm(str(tmp_path), use_fake_base=True, use_offline_training=True) + assert isinstance(model, FakeBaseModel) + (message,) = self._warnings(model, offline=True) + assert "logits_scaling=8.0" in message + + @pytest.mark.parametrize( + ("base", "kd"), + [("llama", True), ("granite-unscaled", True), ("granite", False)], + ids=["llama", "identity-scaling", "teacher-unread"], + ) + def test_no_warning_otherwise(self, base, kd): + model = { + "llama": lambda: get_tiny_llama(num_hidden_layers=4), + "granite-unscaled": lambda: self._granite(1.0), + "granite": lambda: self._granite(8.0), + }[base]() + assert not self._warnings(model, kd) + + def test_only_rank_0_warns(self, monkeypatch): + monkeypatch.setattr("modelopt.torch.utils.distributed.is_master", lambda group=None: False) + assert not self._warnings(self._granite(8.0)) 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 7adf3400cd7..9153dcccacd 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml @@ -65,7 +65,6 @@ pipeline: - 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_lambda_base_start=1.0 - dflash.dflash_lambda_base_decay_ratio=1.0 environment: