From a818df145feb991d87482edba103f3eff0ffe862 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 25 Oct 2026 07:12:52 +0000 Subject: [PATCH] perf(speculative): fuse the DFlash-family loss with the lm_head projection The DFlash family reads the draft's [rows, vocab] logits only through per-row statistics: KD cross-entropy against the teacher, label cross-entropy, the log-normalizer, argmax and top-k. New DFlashConfig.dflash_fused_loss (default off) computes them in row chunks fused with the frozen lm_head projection (plugins/dflash_fused_head.py). Each chunk is projected and reduced to its statistics, through torch.compile on CUDA, and every differentiable statistic writes its hidden-space Jacobian [rows, H] in the same pass. Backward is sum_t grad_t[r] * J_t[r], so any downstream combination is exact and no vocab-wide tensor outlives its chunk. The top-k logits are a differentiable output whose backward gathers the head rows one chunk at a time. Wired into HFDFlashModel._compute_loss (KD on the teacher's logits or its final hidden, label CE, D-PACE confidences, accuracy), the DFlash2 selector's candidates and the LiLiCorr lattice's candidates and normalizer. With the flag off every eager line runs as before. Convert refuses the flag when the class owning forward or _compute_loss does not declare support (DSpark, Domino), when lm_head is not a plain bias-free nn.Linear, and when lm_head is trainable. Loss path forward + backward at hidden 8192, vocab 202048, 4096 draft rows, one RTX PRO 6000, bf16 autocast, this tree's eager path vs fused: DFlash KD, hidden teacher 226 -> 184 ms, peak +12.3 -> +3.2 GiB DFlash KD, logits teacher 177 -> 130 ms, peak +12.3 -> +3.9 GiB DFlash CE 133 -> 122 ms, peak +7.7 -> +2.4 GiB DFlash2 CE+lambda + selector 147 -> 137 ms, peak +9.5 -> +2.4 GiB DFlash2 KD + selector 247 -> 201 ms, peak +12.3 -> +3.2 GiB LiLiCorr 242 -> 260 ms, peak +17.2 -> +5.1 GiB Loss rel <= 1.7e-5, accuracy and metrics equal, hidden-grad rel L2 3.5-4.4e-3. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/torch/speculative/config.py | 12 + .../torch/speculative/dflash/dflash_model.py | 1 + .../speculative/plugins/dflash_fused_head.py | 265 ++++++++++++++++++ .../torch/speculative/plugins/hf_dflash.py | 142 ++++++++-- .../torch/speculative/plugins/hf_dflash2.py | 29 +- .../torch/speculative/plugins/hf_lilicorr.py | 43 ++- .../speculative/plugins/test_hf_dflash.py | 48 ++++ .../plugins/test_dflash_fused_head.py | 88 ++++++ .../speculative/plugins/test_hf_dflash.py | 67 +++++ .../speculative/plugins/test_hf_dflash2.py | 28 ++ .../speculative/plugins/test_hf_lilicorr.py | 25 ++ 12 files changed, 718 insertions(+), 31 deletions(-) create mode 100644 modelopt/torch/speculative/plugins/dflash_fused_head.py create mode 100644 tests/unit/torch/speculative/plugins/test_dflash_fused_head.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f71f5c4c2b..b811dddfd6a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -30,6 +30,7 @@ Changelog - 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. - With ``training.dataloader_num_workers`` > 0, ``examples/speculative_decoding/main.py`` now yields each training batch as soon as a worker finishes it, so one slow sample, such as a slow hidden-state fetch in streaming training, no longer stalls every DDP rank. Sample order within the prefetch window is then no longer deterministic; set ``training.dataloader_in_order=true`` to restore it (context-parallel runs always keep it). - The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``. With PyTorch 2.9 or 2.10 its ``tf32: true`` makes Inductor fail on the first step; set ``training.tf32=false`` or ``dflash.dflash_use_torch_compile=false`` there. +- Add ``dflash_fused_loss`` to the DFlash config. DFlash, DFlash2 and LiLiCorr training then computes its loss in row chunks fused with the frozen ``lm_head`` projection instead of materializing the draft's full-vocabulary logits, which saves time and memory on large vocabularies. *Megatron Framework (M-LM / M-Bridge)* diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 5e1c3ae0794..99fb09148e2 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -162,6 +162,18 @@ class DFlashConfig(ModeloptBaseConfig): ), ) + dflash_fused_loss: bool = ModeloptField( + default=False, + description=( + "Compute the training loss from the draft's hidden states in row chunks, without " + "materializing its [rows, vocab] logits: each chunk is projected through lm_head " + "and reduced to the per-row statistics the loss reads, together with their " + "gradients. Supported by the dflash, dflash2 and lilicorr projector types with a " + "frozen, bias-free nn.Linear lm_head. On CUDA the per-chunk statistics run through " + "torch.compile." + ), + ) + dflash_swa_window_size: int | None = ModeloptField( default=None, description=( diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py index e4cb8753668..03c2c58caa7 100644 --- a/modelopt/torch/speculative/dflash/dflash_model.py +++ b/modelopt/torch/speculative/dflash/dflash_model.py @@ -53,4 +53,5 @@ def modify(self, config): self.dflash_attention_sink = config.dflash_attention_sink self.dflash_init_checkpoint = config.dflash_init_checkpoint self.dflash_use_flex_attention = config.dflash_use_flex_attention + self.dflash_fused_loss = config.dflash_fused_loss self.dflash_export_rope_scaling = config.dflash_export_rope_scaling diff --git a/modelopt/torch/speculative/plugins/dflash_fused_head.py b/modelopt/torch/speculative/plugins/dflash_fused_head.py new file mode 100644 index 00000000000..505c82669f9 --- /dev/null +++ b/modelopt/torch/speculative/plugins/dflash_fused_head.py @@ -0,0 +1,265 @@ +# 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. + +"""Fused, row-chunked lm_head statistics for DFlash-family training. + +The draft's vocab-wide logits are read only through per-row statistics: distillation +cross-entropy and total variation against a teacher, label cross-entropy, logsumexp, argmax +and top-k. Each chunk of rows is projected and reduced to those statistics, and the +hidden-space gradient of every differentiable one is written in the same pass, so no +``[rows, vocab]`` tensor outlives its chunk or is recomputed in backward. +""" + +import functools +from collections.abc import Iterable +from typing import NamedTuple + +import torch +import torch.nn.functional as F + +__all__ = ["TERMS", "HeadStats", "fused_head_stats"] + +# Per-row statistics that can carry a gradient back to the hidden states. +TERMS = ("kd_ce", "tv", "hard_ce", "lse") + + +class HeadStats(NamedTuple): + """Per-row statistics of the head's logits; a statistic that was not requested is ``None``.""" + + kd_ce: torch.Tensor | None # -sum(p * log q) against the teacher distribution p + tv: torch.Tensor | None # 0.5 * |p - q|_1: one minus the acceptance rate of q against p + hard_ce: torch.Tensor | None # -log q[target] + lse: torch.Tensor | None # logsumexp of the logits + top1: torch.Tensor # argmax + top_ids: torch.Tensor | None # top-k ids, descending + top_logits: torch.Tensor | None # their logits, differentiable + + +def _transform(raw, scale, cap): + """``cap * tanh(scale * raw / cap)`` computed in fp32 and returned in ``raw``'s dtype.""" + if scale is None and cap is None: + return raw + out = raw.float() + if scale is not None: + out = out * scale + if cap is not None: + out = cap * torch.tanh(out / cap) + return out.to(raw.dtype) + + +def _chunk_stats(raw, teacher, targets, teacher_is_logits, scale, cap, want, grads, top_k): + """Statistics of one chunk of rows, and each differentiable one's gradient w.r.t. ``raw``.""" + logits = _transform(raw, scale, cap) + z = logits.float() + lse = torch.logsumexp(z, dim=-1, keepdim=True) + log_q = z - lse + q = torch.exp(log_q) + vals, dz = {}, {} + if teacher is not None: + zt = (teacher if teacher_is_logits else _transform(teacher, scale, cap)).float() + p = torch.softmax(zt, dim=-1) + if "kd_ce" in want: + vals["kd_ce"] = -(p * log_q).sum(dim=-1) + if "kd_ce" in grads: + dz["kd_ce"] = q - p + if "tv" in want: + vals["tv"] = 0.5 * (q - p).abs().sum(dim=-1) + if "tv" in grads: + sign = torch.sign(q - p) + dz["tv"] = 0.5 * q * (sign - (sign * q).sum(dim=-1, keepdim=True)) + if "hard_ce" in want: + vals["hard_ce"] = -log_q.gather(-1, targets.unsqueeze(-1)).squeeze(-1) + if "hard_ce" in grads: + label = torch.arange(z.shape[-1], device=z.device) == targets.unsqueeze(-1) + dz["hard_ce"] = q - label.float() + if "lse" in want: + vals["lse"] = lse.squeeze(-1) + if "lse" in grads: + dz["lse"] = q + + # d logits / d raw through the transform; its dtype casts pass gradients through unchanged. + factor = 1.0 if scale is None else scale + if cap is not None: + th = torch.tanh(raw.float() * (factor / cap)) + factor = factor * (1.0 - th * th) + top_ids = top_logits = top_factor = None + if top_k: + top_logits, top_ids = logits.topk(top_k, dim=-1) + if isinstance(factor, torch.Tensor): + top_factor = factor.gather(-1, top_ids) + elif scale is not None: + top_factor = torch.full(top_ids.shape, scale, device=raw.device) + jac_inputs = tuple((dz[t] * factor).to(raw.dtype) for t in grads) + stats = tuple(vals.get(t) for t in TERMS) + return jac_inputs, stats, logits.argmax(dim=-1), top_ids, top_logits, top_factor + + +@functools.cache +def _compiled_chunk_stats(): + return torch.compile(_chunk_stats) + + +def _cat(parts): + return None if parts[0] is None else torch.cat(parts) + + +class _FusedHead(torch.autograd.Function): + @staticmethod + def forward( + ctx, + hidden, + weight, + teacher, + targets, + teacher_is_logits, + scale, + cap, + want, + grads, + top_k, + chunk, + ): + stats_fn = _compiled_chunk_stats() if hidden.device.type == "cuda" else _chunk_stats + # The head runs in its own dtype and the statistics in fp32, whatever autocast says. + with torch.autocast(hidden.device.type, enabled=False): + h = hidden.to(weight.dtype) + jac = {t: torch.empty_like(h) for t in grads} + parts = [] + for s in range(0, h.shape[0], chunk): + e = min(s + chunk, h.shape[0]) + rows = None + if teacher is not None: + rows = teacher[s:e] + if not teacher_is_logits: + rows = F.linear(rows.to(weight.dtype), weight) + jac_inputs, *outs = stats_fn( + F.linear(h[s:e], weight), + rows, + None if targets is None else targets[s:e], + teacher_is_logits, + scale, + cap, + want, + grads, + top_k, + ) + for t, dz in zip(grads, jac_inputs): + torch.mm(dz, weight, out=jac[t][s:e]) + parts.append(outs) + stats = [_cat([part[0][i] for part in parts]) for i in range(len(TERMS))] + top1, top_ids, top_logits, top_factor = ( + _cat([part[i] for part in parts]) for i in range(1, 5) + ) + ctx.save_for_backward(*(jac[t] for t in grads), top_ids, top_factor, weight) + ctx.set_materialize_grads(False) # statistics the loss never reads cost nothing in backward + ctx.grads, ctx.chunk, ctx.hidden_dtype = grads, chunk, hidden.dtype + ctx.mark_non_differentiable( + *(v for t, v in zip(TERMS, stats) if v is not None and t not in grads), + top1, + *(() if top_ids is None else (top_ids,)), + ) + return (*stats, top1, top_ids, top_logits) + + @staticmethod + def backward(ctx, *grad_outs): + *jacs, top_ids, top_factor, weight = ctx.saved_tensors + grad = None + for t, jac in zip(ctx.grads, jacs): + g = grad_outs[TERMS.index(t)] + if g is not None: + part = g.float().unsqueeze(-1) * jac.float() + grad = part if grad is None else grad + part + g_top = grad_outs[-1] + if g_top is not None: + g_top = g_top.float() if top_factor is None else g_top.float() * top_factor + # sum_k g[r, k] * weight[id[r, k]], gathering the head rows one chunk at a time. + part = torch.cat( + [ + torch.bmm( + g_top[s : s + ctx.chunk].unsqueeze(1).to(weight.dtype), + weight[top_ids[s : s + ctx.chunk]], + ).squeeze(1) + for s in range(0, g_top.shape[0], ctx.chunk) + ] + ).float() + grad = part if grad is None else grad + part + return (None if grad is None else grad.to(ctx.hidden_dtype), *([None] * 10)) + + +def fused_head_stats( + hidden: torch.Tensor, + weight: torch.Tensor, + *, + teacher_hidden: torch.Tensor | None = None, + teacher_logits: torch.Tensor | None = None, + targets: torch.Tensor | None = None, + want: Iterable[str] = (), + grads: Iterable[str] = (), + top_k: int = 0, + logit_scale: float | None = None, + soft_cap: float | None = None, + chunk: int = 2048, +) -> HeadStats: + """Per-row statistics of the head's logits ``hidden @ weight.T``, never materialized whole. + + On CUDA the per-chunk statistics run through ``torch.compile``. + + Args: + hidden: ``[R, H]`` rows. The ``grads`` statistics and ``top_logits`` backpropagate into + them, and nothing into ``weight``, which must be frozen. + weight: ``[V, H]`` head weight. + teacher_hidden: ``[R, H]`` teacher rows to project through the same head. ``kd_ce`` + and ``tv`` need this or ``teacher_logits``. + teacher_logits: ``[R, V]`` teacher logits, used as given. + targets: ``[R]`` label ids, needed for ``hard_ce``. + want: Statistics to return, named as in :data:`TERMS`; ``grads`` are returned too. + grads: Statistics to differentiate w.r.t. ``hidden``. + top_k: Also return each row's ``top_k`` ids and their logits. + logit_scale: Multiplier the head applies to its projection, if any. + soft_cap: ``c`` of a ``c * tanh(x / c)`` cap the head applies after ``logit_scale``. + chunk: Rows per projection; bounds the vocab-wide temporaries. + """ + want, grads = set(want), set(grads) + if unknown := (want | grads) - set(TERMS): + raise ValueError(f"Unknown statistics {sorted(unknown)}; choose from {TERMS}.") + if weight.requires_grad: + raise ValueError("fused_head_stats computes no gradient for the head weight; freeze it.") + if teacher_hidden is not None and teacher_logits is not None: + raise ValueError("Pass teacher_hidden or teacher_logits, not both.") + if top_k > weight.shape[0]: + raise ValueError(f"top_k={top_k} exceeds the vocabulary size {weight.shape[0]}.") + teacher = teacher_hidden if teacher_logits is None else teacher_logits + want |= grads + if not (torch.is_grad_enabled() and hidden.requires_grad): + grads = set() + if teacher is None and want & {"kd_ce", "tv"}: + raise ValueError("kd_ce and tv need teacher_hidden or teacher_logits.") + if targets is None and "hard_ce" in want: + raise ValueError("hard_ce needs targets.") + return HeadStats( + *_FusedHead.apply( + hidden, + weight, + teacher, + targets, + teacher_logits is not None, + logit_scale, + soft_cap, + tuple(t for t in TERMS if t in want), + tuple(t for t in TERMS if t in grads), + top_k, + chunk, + ) + ) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 1e20d5b3099..6c468d96bec 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -85,6 +85,7 @@ from ..dflash.conversion import DFlashDMRegistry from ..dflash.dflash_model import DFlashModel from .dflash_flex_attention import build_draft_block_mask +from .dflash_fused_head import HeadStats, fused_head_stats from .modeling_dflash import ( # noqa: F401 DFlashAttention, DFlashBaseModelOutput, @@ -133,6 +134,11 @@ def _gather_rows(x, positions): return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1)) +def _prediction_rows(label_positions): + """The base rows whose next-token logits predict the tokens at ``label_positions``.""" + return (label_positions - 1).clamp(min=0) + + def _block_positions(anchor_positions, block_size, offset, seq_len): """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. @@ -210,10 +216,15 @@ def _block_position_weights(mask, scheme, first, *, decay_factor, dpace_alpha, t return mask +def _masked_accuracy(preds, target_ids, mask): + """Fraction of ``preds`` equal to ``target_ids`` where ``mask`` is set (0-dim).""" + correct = (preds == target_ids) & (mask > 0.5) + return correct.sum().float() / (mask.sum() + 1e-6) + + def _masked_top1(logits, target_ids, mask): """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" - correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) - return correct.sum().float() / (mask.sum() + 1e-6) + return _masked_accuracy(logits.argmax(dim=-1), target_ids, mask) def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal): @@ -313,13 +324,15 @@ class _DFlashLossTerms(NamedTuple): ``ce_per_token`` is ``None`` on the KD path, which never forms a per-position cross-entropy. ``weights`` carries the position weighting (decay or D-PACE) and normalizes by ``weight_sum``; ``supervised_mask`` is the unweighted mask the - reported accuracy uses. + reported accuracy uses. ``head_stats`` is the fused pass's statistics of the draft's + logits under ``dflash_fused_loss``, and ``None`` otherwise. """ ce_per_token: torch.Tensor | None weights: torch.Tensor weight_sum: torch.Tensor supervised_mask: torch.Tensor + head_stats: HeadStats | None = None @DFlashDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) @@ -330,6 +343,10 @@ class HFDFlashModel(DFlashModel): # the anchor's own position onward, so slot 0 (the given anchor) is never supervised. _label_shift = 0 + # Set in the body of each class whose `forward` / `_compute_loss` read the draft's logits + # only through the fused pass. Checked on the overriding class, so it is not inherited. + _supports_fused_loss = True + @property def _base_model(self): return self.get_submodule(self.base_model_path) @@ -540,6 +557,8 @@ def modify(self, config): param.requires_grad = False self._find_base_model_parts() + if self.dflash_fused_loss: + self._check_fused_loss() # Factory hook: subclasses (e.g. Domino) override to build an augmented # draft module while reusing all of DFlash's modify() setup. @@ -580,6 +599,35 @@ def _build_draft_module(self, dflash_config): """Build the draft module. Subclasses override to use an augmented module.""" return DFlashModule(dflash_config) + def _check_fused_loss(self): + """Refuse ``dflash_fused_loss`` where it would be ignored or would drop a gradient.""" + for method in ("forward", "_compute_loss"): + owner = next(cls for cls in type(self).__mro__ if method in vars(cls)) + if not vars(owner).get("_supports_fused_loss", False): + raise ValueError( + f"dflash_fused_loss is not supported by {owner.__name__}: its {method} " + "reads the draft's logits directly. Unset dflash_fused_loss." + ) + head = self._base_model_lm_head + # The fused pass reads lm_head.weight directly: whatever a bias, a quantizer or a + # hook adds in lm_head's forward would be silently skipped. + plain_forward = ( + type(head).forward is torch.nn.Linear.forward + and "forward" not in vars(head) + and not (head._forward_hooks or head._forward_pre_hooks) + ) + if not plain_forward or head.bias is not None: + raise ValueError( + "dflash_fused_loss projects with lm_head.weight alone, so it needs a bias-free " + f"nn.Linear lm_head with no forward hooks; got {type(head).__name__}. Unset " + "dflash_fused_loss." + ) + if head.weight.requires_grad: + raise ValueError( + "dflash_fused_loss computes no gradient for lm_head, which is trainable here; " + "it needs dflash_freeze_base_model=True." + ) + def _base_device(self): """Device the frozen base model sits on. ``meta`` during a from_pretrained restore.""" if self.dflash_offline: @@ -922,13 +970,24 @@ def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None): projection to just those vocab entries. ``label_positions`` is keyword-only, so every call says which positions it passes: handing it the logit rows reads one row early. """ - positions = (label_positions - 1).clamp(min=0) + positions = _prediction_rows(label_positions) if base_outputs.logits is not None: if token_ids is None: return _gather_rows(base_outputs.logits, positions) batch = torch.arange(positions.shape[0], device=positions.device) batch = batch.view(-1, *[1] * positions.dim()) return base_outputs.logits[batch, positions.unsqueeze(-1), token_ids] + lm_head = self._base_model_lm_head + rows = self._teacher_hidden(base_outputs, label_positions=label_positions) + 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] + + @torch.no_grad() + def _teacher_hidden(self, base_outputs, *, label_positions): + """The base model's final hidden that predicts the tokens at ``label_positions``: [B, ..., H].""" if base_outputs.base_hidden is None: raise ValueError( "This objective needs the base model's distribution, but the batch carries " @@ -937,13 +996,9 @@ def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None): lm_head = self._base_model_lm_head # A producer can store the hidden states in a wider dtype than the target's weights. # Cast before the final norm too, which online training runs in the target's dtype. + positions = _prediction_rows(label_positions) 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] + return _apply_base_final_norm(rows, base_outputs.base_hidden_prenorm, self._base_model_norm) def _compute_loss( self, @@ -960,20 +1015,20 @@ def _compute_loss( """Compute weighted cross-entropy (or KD) loss and accuracy. Args: - logits: Draft model output [B, N*block_size, vocab]. + logits: Draft model output [B, N*block_size, vocab], or ``None`` under + ``dflash_fused_loss``, which reads ``draft_hidden`` instead. input_ids: Original input token IDs [B, seq_len]. anchor_positions: Anchor positions per block [B, N]. block_keep_mask: Valid block mask [B, N]. loss_mask: Token-level loss mask [B, seq_len]. draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``. - Unused here; passed for variants whose head consumes them. base_outputs: The step's ``DFlashBaseModelOutput``: the KD teacher under ``dflash_self_logit_distillation``, and read by variants. return_terms: Also return the unreduced pieces behind the loss, so a variant can recompose the block objective from a different divergence without rebuilding the target alignment and position weighting. TODO: promote this into a shared divergence seam when the DFlash-family - loss code is refactored; DFlash2 is the only consumer today. + loss code is refactored; DFlash2 and LiLiCorr are its consumers today. Returns: (loss, accuracy) tuple. @@ -984,7 +1039,6 @@ def _compute_loss( ) binary_eval_mask = targets.mask.view(-1) - flat_logits = logits.view(-1, logits.size(-1)) flat_targets = targets.token_ids.view(-1) kd = self.dflash_self_logit_distillation @@ -993,12 +1047,25 @@ def _compute_loss( "dflash_self_logit_distillation distills from base_outputs, but none was " "passed; an override of _compute_loss must forward it." ) + # Without logits (dflash_fused_loss), every statistic of them read below comes from + # one fused pass over the draft's hidden states. + stats = None + if logits is None: + stats = self._fused_head_stats( + draft_hidden, base_outputs, targets.positions, flat_targets + ) + else: + flat_logits = logits.view(-1, logits.size(-1)) # Non-KD loss is per-token cross-entropy; compute it once (grad enabled) so the # D-PACE confidences below can reuse it instead of a second CE pass. The KD path # optimizes KL, so its confidences need a dedicated no_grad CE pass. loss_per_token = None if not kd: - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") + loss_per_token = ( + stats.hard_ce + if stats is not None + else F.cross_entropy(flat_logits, flat_targets, reduction="none") + ) weight_mask = _block_position_weights( targets.mask, @@ -1009,6 +1076,8 @@ def _compute_loss( target_ce=lambda: ( loss_per_token.detach() if loss_per_token is not None + else stats.hard_ce + if stats is not None else F.cross_entropy(flat_logits, flat_targets, reduction="none") ).view(targets.mask.shape), ) @@ -1017,7 +1086,9 @@ def _compute_loss( valid_count = flat_weights.sum() + 1e-6 if valid_count > 1.0: - if kd: + if kd and stats is not None: + loss = (stats.kd_ce * flat_weights).sum() / valid_count + elif kd: flat_teacher = self._teacher_logits(base_outputs, label_positions=targets.positions) flat_teacher = flat_teacher.reshape(flat_logits.shape) target_soft = torch.softmax(flat_teacher, dim=-1) @@ -1028,9 +1099,12 @@ def _compute_loss( loss = (loss_per_token * flat_weights).sum() / valid_count with torch.no_grad(): - accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item() + preds = stats.top1 if stats is not None else flat_logits.argmax(dim=-1) + accuracy = _masked_accuracy(preds, flat_targets, binary_eval_mask).item() else: - loss = self._draft_zero_loss(flat_logits.device, through=(flat_logits,)) + # Under dflash_fused_loss the draft's output is its hidden states: no logits exist. + output = draft_hidden if stats is not None else flat_logits + loss = self._draft_zero_loss(output.device, through=(output,)) accuracy = 0.0 if return_terms: @@ -1039,10 +1113,40 @@ def _compute_loss( weights=flat_weights, weight_sum=valid_count, supervised_mask=binary_eval_mask, + head_stats=stats, ) return loss, accuracy, terms return loss, accuracy + def _fused_head_extras(self): + """``(statistics with gradient, top_k)`` a variant's own terms read from the fused pass.""" + return set(), 0 + + def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_targets): + """Every statistic of the draft's logits the loss reads, from one fused pass.""" + grads, top_k = self._fused_head_extras() + teacher = {} + if self.dflash_self_logit_distillation: + grads = {*grads, "kd_ce"} + if base_outputs.logits is not None: + rows = self._teacher_logits(base_outputs, label_positions=label_positions) + teacher["teacher_logits"] = rows.flatten(0, -2) + else: + rows = self._teacher_hidden(base_outputs, label_positions=label_positions) + teacher["teacher_hidden"] = rows.flatten(0, -2) + else: + grads = {*grads, "hard_ce"} + return fused_head_stats( + draft_hidden.flatten(0, -2), + self._base_model_lm_head.weight, + targets=flat_targets, + # D-PACE weighs positions by the draft's confidence in its label. + want={"hard_ce"} if self.dflash_loss_objective == "dpace" else (), + grads=grads, + top_k=top_k, + **teacher, + ) + def forward( self, input_ids=None, @@ -1197,7 +1301,7 @@ def forward( ) # 6. Compute loss and accuracy - logits = self._base_model_lm_head(hidden) + logits = None if self.dflash_fused_loss else self._base_model_lm_head(hidden) loss, accuracy = self._compute_loss( logits, input_ids, diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 88c1e3e38f7..73ba8e1775e 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -55,6 +55,8 @@ class HFDFlash2Model(HFDFlashModel): to it when ``dflash_architecture_config.projector_type == "dflash2"``. """ + _supports_fused_loss = True + def _build_draft_module(self, dflash_config): """Build the DFlash2 draft module (DFlash backbone + convolutions + selector).""" return DFlash2Module(dflash_config) @@ -108,13 +110,27 @@ def get_exporter(self): return DFlash2Exporter(self) + def _fused_head_extras(self): + """The selector scores the backbone's top-k candidates.""" + selector = self.dflash_module.candidate_selector + return set(), selector.top_k if self.dflash_selector_loss_alpha > 0 else 0 + def _selector_loss( - self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum + self, + logits, + target_ids, + hidden, + predecessor_ids, + supervised_mask, + loss_weights, + weight_sum, + head_stats=None, ): """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]``. + logits: Backbone logits per block position ``[B, N, block_size, V]``, or ``None`` + under ``dflash_fused_loss``, whose pass took their top-k into ``head_stats``. 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]``. @@ -131,7 +147,11 @@ def _selector_loss( selector = self.dflash_module.candidate_selector top_k = selector.top_k - unary_logits, candidate_ids = logits.topk(top_k, dim=-1) + if logits is None: + unary_logits = head_stats.top_logits.view(*target_ids.shape, top_k) + candidate_ids = head_stats.top_ids.view(*target_ids.shape, top_k) + else: + unary_logits, candidate_ids = logits.topk(top_k, dim=-1) # Train on the strict top-k, the candidate set serving actually builds. A gold # token the backbone did not propose is a backbone recall failure, not a @@ -260,13 +280,14 @@ def _compute_loss( # The backbone's own mask and position weights: a block position's choice only # matters if every earlier one was accepted, which decay / D-PACE model. selector_loss, selector_accuracy, selector_coverage = self._selector_loss( - logits.reshape(bsz, n_blocks, block_size, -1), + None if logits is None else logits.reshape(bsz, n_blocks, block_size, -1), target_ids, draft_hidden.reshape(bsz, n_blocks, block_size, -1), predecessor_ids, terms.supervised_mask.view_as(target_ids), terms.weights.view_as(target_ids), terms.weight_sum, + head_stats=terms.head_stats, ) self._selector_metrics = { "selector_accuracy": selector_accuracy, diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index b38234d78eb..4d7b604a84a 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -102,6 +102,8 @@ class HFLiLiCorrModel(HFDFlashModel): to it when ``dflash_architecture_config.projector_type == "lilicorr"``. """ + _supports_fused_loss = True + def _build_draft_module(self, dflash_config): """Build the LiLiCorr draft module (DFlash backbone + lattice reranker).""" return LiLiCorrModule(dflash_config) @@ -220,13 +222,15 @@ def _compute_loss( "committed target row." ) - origin_loss, origin_accuracy = super()._compute_loss( + origin_loss, origin_accuracy, terms = super()._compute_loss( logits, input_ids, anchor_positions, block_keep_mask, loss_mask, + draft_hidden=draft_hidden, base_outputs=base_outputs, + return_terms=True, ) # The backbone's own supervision, before its position weighting, so both losses # train the same positions. @@ -240,6 +244,7 @@ def _compute_loss( ) lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss( logits=logits, + head_stats=terms.head_stats, target_ids=targets.token_ids, # Block position 0 is the anchor -- the one row whose input embedding is a real # verified token -- so the lattice covers positions 1..bs-1. @@ -262,6 +267,10 @@ def _compute_loss( self._lilicorr_metrics = metrics return origin_loss + lilicorr_loss, accuracy + def _fused_head_extras(self): + """The lattice scores the backbone's top-k candidates against its log-normalizer.""" + return {"lse"}, self.dflash_module.lilicorr.candidate_topk + def _compute_lilicorr_loss( self, *, @@ -272,9 +281,13 @@ def _compute_lilicorr_loss( draft_hidden, target_hidden, base_outputs, + head_stats=None, ): """Score the candidate lattice and evaluate the LiLiCorr objective. + ``logits`` is ``None`` under ``dflash_fused_loss``, whose pass took the candidates, + their logits and the log-normalizer into ``head_stats``. + Returns ``(loss, accuracy, metrics)``. ``loss`` already carries the absolute term weights, ``accuracy`` is the greedy path-prefix fraction, and ``metrics`` holds Python floats read back in a single device sync. @@ -284,15 +297,22 @@ def _compute_lilicorr_loss( bsz, n_blocks, block_size = target_ids.shape num_slots = block_size - 1 topk = head.candidate_topk - vocab = logits.shape[-1] + vocab = self._base_model_lm_head.weight.shape[0] if logits is None else logits.shape[-1] batch_blocks = bsz * n_blocks if topk > vocab: raise ValueError(f"lilicorr_candidate_topk={topk} exceeds the vocabulary size {vocab}.") - logits_slots = logits.view(bsz, n_blocks, block_size, vocab)[:, :, 1:, :].reshape( - batch_blocks, num_slots, vocab - ) + def slots(per_position): + """Block positions 1..block_size-1 of a ``[B * N * block_size, ...]`` tensor.""" + return per_position.view(bsz, n_blocks, block_size, -1)[:, :, 1:].reshape( + batch_blocks, num_slots, -1 + ) + + if logits is not None: + logits_slots = logits.view(bsz, n_blocks, block_size, vocab)[:, :, 1:, :].reshape( + batch_blocks, num_slots, vocab + ) ground_truth = target_ids[:, :, 1:] gt_flat = ground_truth.reshape(batch_blocks, num_slots) # A block is trained only when every one of its slots is a supervised position, @@ -304,7 +324,10 @@ def _compute_lilicorr_loss( # is the rank order `rank_embedding` is indexed by. Taken on the logits # rather than on a log_softmax of them: the map is strictly monotone per row, # so the ids are identical, and it avoids a vocab-wide float32 buffer. - _, candidate_ids = torch.topk(logits_slots.detach(), k=topk, dim=-1) + if logits is None: + candidate_ids = slots(head_stats.top_ids) + else: + _, candidate_ids = torch.topk(logits_slots.detach(), k=topk, dim=-1) gt_hits = candidate_ids == gt_flat.unsqueeze(-1) gt_in_slot = gt_hits.any(dim=-1) gt_indices = gt_hits.long().argmax(dim=-1) @@ -314,8 +337,12 @@ def _compute_lilicorr_loss( # re-attaches the autograd graph, which is how the lattice terms reach the # drafter body and not only the head. logsumexp rather than a full log_softmax # so backward needs no second vocab-sized buffer. - candidate_logits = torch.gather(logits_slots, 2, candidate_ids).float() - log_denominator = torch.logsumexp(logits_slots.float(), dim=-1, keepdim=True) + if logits is None: + candidate_logits = slots(head_stats.top_logits).float() + log_denominator = slots(head_stats.lse.unsqueeze(-1)) + else: + candidate_logits = torch.gather(logits_slots, 2, candidate_ids).float() + log_denominator = torch.logsumexp(logits_slots.float(), dim=-1, keepdim=True) candidate_log_probs = candidate_logits - log_denominator draft_slots = draft_hidden.view(bsz, n_blocks, block_size, -1)[:, :, 1:, :] diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 5f2e8adc045..37b4cf70921 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -395,3 +395,51 @@ def test_backward_produces_finite_grads(self): ] assert grads, "no draft gradients were produced" assert all(torch.isfinite(g).all() for g in grads) + + +class TestDFlashFusedLossGPU: + """On CUDA the fused loss runs its statistics through torch.compile; they must still match.""" + + DFLASH2_ARCH = { + "projector_type": "dflash2", + "conv_kernel_size": 2, + "conv_group_size": 4, + "selector_rank": 8, + "selector_top_k": 5, + } + + @staticmethod + def _step(fused, overrides, arch): + model = get_tiny_llama(num_hidden_layers=4) + config = get_dflash_config() + config.update(dflash_fused_loss=fused, **overrides) + config["dflash_architecture_config"].update(arch) + mtsp.convert(model, [("dflash", config)]) + model = model.float().cuda().train() # fp32, so the two paths differ only by rounding + torch.manual_seed(0) + input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN), device="cuda") + torch.manual_seed(1) + out = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids)) + out.loss.backward() + params = model.dflash_module.named_parameters() + return out, {name: p.grad for name, p in params if p.grad is not None} + + @pytest.mark.parametrize( + ("overrides", "arch"), + [ + ({}, {}), # KD on the base's logits, D-PACE weighting + ( + {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "lambda"}, + DFLASH2_ARCH, + ), + ], + ids=["dflash-kd", "dflash2-selector"], + ) + def test_matches_the_materialized_logits(self, overrides, arch): + eager, eager_grads = self._step(False, overrides, arch) + fused, fused_grads = self._step(True, overrides, arch) + torch.testing.assert_close(fused.loss, eager.loss) + assert fused.train_acc == eager.train_acc + assert fused_grads.keys() == eager_grads.keys() + for name, grad in eager_grads.items(): + torch.testing.assert_close(fused_grads[name], grad, msg=name) diff --git a/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py b/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py new file mode 100644 index 00000000000..59d5d247143 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_dflash_fused_head.py @@ -0,0 +1,88 @@ +# 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 tests for the fused lm_head statistics behind ``dflash_fused_loss``.""" + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.torch.speculative.plugins.dflash_fused_head import TERMS, fused_head_stats + +ROWS, HIDDEN, VOCAB, TOP_K = 37, 16, 50, 4 +CHUNK = 8 # several chunks and a partial last one + + +def _transform(raw, scale, cap): + if scale is not None: + raw = raw * scale + return raw if cap is None else cap * torch.tanh(raw / cap) + + +@pytest.mark.parametrize("teacher_kind", ["hidden", "logits"]) +@pytest.mark.parametrize(("scale", "cap"), [(None, None), (0.7, None), (0.7, 2.0)]) +def test_matches_autograd_on_the_materialized_logits(teacher_kind, scale, cap): + """Every statistic, its gradient under arbitrary per-row upstream weights, and the top-k.""" + torch.manual_seed(0) + weight = torch.randn(VOCAB, HIDDEN) + hidden = torch.randn(ROWS, HIDDEN) + targets = torch.randint(0, VOCAB, (ROWS,)) + teacher_hidden = torch.randn(ROWS, HIDDEN) + teacher_logits = 3 * torch.randn(ROWS, VOCAB) # handed over as is, never transformed + coefs = {term: torch.randn(ROWS) for term in TERMS} + top_coefs = torch.randn(ROWS, TOP_K) + + ref_hidden = hidden.clone().requires_grad_(True) + logits = _transform(ref_hidden @ weight.T, scale, cap) + if teacher_kind == "hidden": + teacher = {"teacher_hidden": teacher_hidden} + log_p = _transform(teacher_hidden @ weight.T, scale, cap).log_softmax(-1) + else: + teacher = {"teacher_logits": teacher_logits} + log_p = teacher_logits.log_softmax(-1) + log_q = logits.log_softmax(-1) + ref = { + "kd_ce": -(log_p.exp() * log_q).sum(-1), + "tv": 0.5 * (log_q.exp() - log_p.exp()).abs().sum(-1), + "hard_ce": F.cross_entropy(logits, targets, reduction="none"), + "lse": logits.logsumexp(-1), + } + top_logits, top_ids = logits.topk(TOP_K, dim=-1) + ref_loss = sum((coefs[t] * ref[t]).sum() for t in TERMS) + (top_coefs * top_logits).sum() + ref_loss.backward() + + fused_hidden = hidden.clone().requires_grad_(True) + kwargs = dict( + targets=targets, + grads=TERMS, + top_k=TOP_K, + logit_scale=scale, + soft_cap=cap, + chunk=CHUNK, + **teacher, + ) + stats = fused_head_stats(fused_hidden, weight, **kwargs) + loss = sum((coefs[t] * getattr(stats, t)).sum() for t in TERMS) + (loss + (top_coefs * stats.top_logits).sum()).backward() + with torch.no_grad(): # values without the gradient work, e.g. an eval pass in train mode + no_grad_stats = fused_head_stats(fused_hidden, weight, **kwargs) + + for term in TERMS: + torch.testing.assert_close(getattr(stats, term), ref[term].detach(), msg=term) + torch.testing.assert_close(getattr(no_grad_stats, term), ref[term].detach(), msg=term) + assert torch.equal(stats.top1, logits.argmax(-1)) + assert torch.equal(stats.top_ids, top_ids) + torch.testing.assert_close(stats.top_logits, top_logits.detach()) + torch.testing.assert_close(fused_hidden.grad, ref_hidden.grad) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index d3bf4162cf7..0aae907c171 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -1727,3 +1727,70 @@ def test_offline_step_never_projects_the_full_sequence(self): ).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 + + +class TestDFlashFusedLoss: + """``dflash_fused_loss`` trains exactly what the loss on materialized logits trains.""" + + @staticmethod + def _step(fused, offline, **overrides): + """One training step's output and draft gradients, from fixed weights, data and anchors.""" + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=4) + if offline: + model.config.num_orig_hidden_layers = 4 + config = get_dflash_config(offline=offline) + config.update(dflash_fused_loss=fused, **overrides) + mtsp.convert(model, [("dflash", config)]) + model.float().train() # fp32, so the two paths differ only by rounding + batch = _dflash_batch(model.config.vocab_size) + if offline: + hidden = model.config.hidden_size + batch["base_model_outputs"] = { + "aux_hidden_states": torch.randn(2, SEQ_LEN, len(model.target_layer_ids) * hidden), + "base_model_hidden_states": torch.randn(2, SEQ_LEN, hidden), + } + torch.manual_seed(1) + out = model(**batch) + out.loss.backward() + params = model.dflash_module.named_parameters() + return out, {name: p.grad for name, p in params if p.grad is not None} + + @pytest.mark.parametrize( + ("offline", "overrides"), + [ + (False, {}), # KD on the base's logits, D-PACE weighting + (True, {"dflash_loss_objective": "decay", "dflash_loss_decay_factor": 4.0}), + (False, {"dflash_self_logit_distillation": False}), # cross-entropy, D-PACE + ], + ids=["kd-online", "kd-offline-decay", "ce"], + ) + def test_matches_the_materialized_logits(self, offline, overrides): + eager, eager_grads = self._step(False, offline, **overrides) + fused, fused_grads = self._step(True, offline, **overrides) + torch.testing.assert_close(fused.loss, eager.loss) + assert fused.train_acc == eager.train_acc + assert fused_grads.keys() == eager_grads.keys() + for name, grad in eager_grads.items(): + torch.testing.assert_close(fused_grads[name], grad, msg=name) + + @pytest.mark.parametrize( + ("overrides", "arch", "hook_lm_head"), + [ + ({}, {"projector_type": "dspark", "markov_rank": 4}, False), + ({}, {"projector_type": "domino", "gru_hidden_dim": 8, "emb_dim": 8}, False), + ({"dflash_freeze_base_model": False}, {}, False), + ({}, {}, True), + ], + ids=["dspark", "domino", "trainable-lm-head", "hooked-lm-head"], + ) + def test_unsupported_setups_are_rejected(self, overrides, arch, hook_lm_head): + """Each would make the fused pass skip a term or a gradient the eager loss has.""" + model = get_tiny_llama(num_hidden_layers=4) + if hook_lm_head: + model.lm_head.register_forward_hook(lambda module, args, output: output * 2) + config = get_dflash_config() + config.update(dflash_fused_loss=True, **overrides) + config["dflash_architecture_config"].update(arch) + with pytest.raises(ValueError, match="dflash_fused_loss"): + mtsp.convert(model, [("dflash", config)]) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index e1ea8c10dc0..67219179bbe 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -384,6 +384,34 @@ def test_overfits_a_single_batch(self): assert out["selector_metrics"]["selector_accuracy"].item() > 0.9 assert out["selector_metrics"]["selector_coverage"].item() > 0.9 + @pytest.mark.parametrize("lk_loss_type", ["ce", "lambda"]) + def test_fused_loss_matches_the_materialized_logits(self, lk_loss_type): + """The selector's top-k and its gradient come out of the fused pass unchanged.""" + results = [] + for fused in (False, True): + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash2_config() + config.update(dflash_fused_loss=fused, dflash_lk_loss_type=lk_loss_type) + mtsp.convert(model, [("dflash", config)]) + model.float().train() # fp32, so the two paths differ only by rounding + with torch.no_grad(): # a live selector, so its gradient reaches the backbone + torch.nn.init.normal_(model.dflash_module.candidate_selector.successor_codebook) + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + torch.manual_seed(1) + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + out.loss.backward() + params = model.dflash_module.named_parameters() + results.append((out, {name: p.grad for name, p in params if p.grad is not None})) + (eager, eager_grads), (fused, fused_grads) = results + torch.testing.assert_close(fused.loss, eager.loss) + assert fused.train_acc == eager.train_acc + for key, value in eager["selector_metrics"].items(): + torch.testing.assert_close(fused["selector_metrics"][key], value, msg=key) + assert fused_grads.keys() == eager_grads.keys() + for name, grad in eager_grads.items(): + torch.testing.assert_close(fused_grads[name], grad, msg=name) + class TestCandidateSelectorAlignment: """Pin the offset convention shared by the training objective and the lattice walk.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py index 0d132f9ce49..696b7b6ac36 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py +++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py @@ -428,6 +428,31 @@ def test_offline_forward_matches_online(self, prenorm): online_out.lilicorr_metrics[term], abs=1e-5 ) + def test_fused_loss_matches_the_materialized_logits(self): + """The candidates, their logits and the normalizer come out of the fused pass unchanged.""" + results = [] + for fused in (False, True): + torch.manual_seed(0) + # Half the vocabulary per slot: the lattice covers about half the blocks. + model = _converted( + variant="margin", topk=16, dflash_lilicorr_w_cal=0.5, dflash_fused_loss=fused + ) + model.float().train() # fp32, so the two paths differ only by rounding + batch = _make_batch(model.dflash_config.vocab_size) + torch.manual_seed(1) + out = model(**batch) + out.loss.backward() + params = model.dflash_module.named_parameters() + results.append((out, {name: p.grad for name, p in params if p.grad is not None})) + (eager, eager_grads), (fused, fused_grads) = results + torch.testing.assert_close(fused.loss, eager.loss) + assert fused.lilicorr_metrics == pytest.approx(eager.lilicorr_metrics, rel=1e-5, abs=1e-6) + assert fused_grads.keys() == eager_grads.keys() + for name, grad in eager_grads.items(): + # Scaled to each gradient: the normalizer's share of it is below a fixed atol. + atol = 1e-5 * grad.abs().max().item() + torch.testing.assert_close(fused_grads[name], grad, rtol=0, atol=atol, msg=name) + @pytest.mark.parametrize( ("kwargs", "degenerate"), [