diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 05e877cf7d58..0fde5f506e24 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -37,6 +37,7 @@ Changelog - Add support for the `SPEED-Bench v1.1 `_ dataset. - 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 and ``full_determinism`` runs always keep it). - The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``, and turns ``tf32`` off, which compilation needs on PyTorch 2.9 and 2.10. To train eagerly as before, start the run with ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``; a run resumed from a checkpoint keeps the checkpoint's DFlash settings. +- 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 5e1c3ae0794a..99fb09148e2e 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 e4cb87536685..03c2c58caa7a 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_head_stats.py b/modelopt/torch/speculative/plugins/dflash_head_stats.py new file mode 100644 index 000000000000..2832902399e9 --- /dev/null +++ b/modelopt/torch/speculative/plugins/dflash_head_stats.py @@ -0,0 +1,353 @@ +# 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. + +"""Per-row statistics of an lm_head's logits for DFlash-family training. + +A loss reads the draft's vocab-wide logits only through per-row statistics: distillation +cross-entropy against a teacher, label cross-entropy, logsumexp, argmax and top-k. A +:class:`HeadRequest` names the ones it reads, and either backend returns them as a +:class:`HeadStats`: + +* :func:`materialized_head_stats` reads them off logits the head already produced, through + autograd. +* :func:`fused_head_stats` projects row chunks through the head's weight and reduces each to + the statistics in the same pass, writing the hidden-space gradient of every differentiable + one, so no ``[rows, vocab]`` tensor outlives its chunk or is recomputed in backward. +""" + +import functools +from typing import NamedTuple + +import torch +import torch.nn.functional as F +from torch import nn + +__all__ = [ + "TERMS", + "HeadRequest", + "HeadStats", + "fused_head_stats", + "fused_projection", + "materialized_head_stats", +] + +# Per-row statistics that can carry a gradient back to the head's input. +TERMS = ("kd_ce", "hard_ce", "lse") + + +class HeadRequest(NamedTuple): + """The statistics of the head's logits a loss reads, honored alike by both backends. + + ``grads`` come back differentiable. ``values`` come back detached, for weights and + metrics: a loss term built on one trains nothing, so whatever the loss differentiates + belongs in ``grads``. ``top_k`` also returns each row's ``top_k`` ids and their logits, + differentiable. Each row's argmax is always returned. + """ + + grads: frozenset[str] = frozenset() + values: frozenset[str] = frozenset() + top_k: int = 0 + + def extend(self, *, grads=(), values=(), top_k=0) -> "HeadRequest": + """This request and more; one head pass serves a single ``top_k``.""" + if top_k and self.top_k and top_k != self.top_k: + raise ValueError(f"One head pass serves one top_k; got {self.top_k} and {top_k}.") + return HeadRequest(self.grads | set(grads), self.values | set(values), top_k or self.top_k) + + +class HeadStats(NamedTuple): + """Per-row statistics of the head's logits; a statistic not requested is ``None``.""" + + kd_ce: torch.Tensor | None # -sum(p * log q) against the teacher distribution p + hard_ce: torch.Tensor | None # -log q[target] + lse: torch.Tensor | None # logsumexp of the logits, in fp32 + top1: torch.Tensor # argmax + top_ids: torch.Tensor | None # top-k ids, descending + top_logits: torch.Tensor | None # their logits, differentiable + + +def _check_request(request, *, vocab, teacher, targets): + """Refuse a request the inputs cannot serve, the same way from either backend.""" + named = request.grads | request.values + if unknown := named - set(TERMS): + raise ValueError(f"Unknown statistics {sorted(unknown)}; choose from {TERMS}.") + if "kd_ce" in named and teacher is None: + raise ValueError("kd_ce needs the teacher's logits or hidden states.") + if "hard_ce" in named and targets is None: + raise ValueError("hard_ce needs targets.") + if request.top_k > vocab: + raise ValueError(f"top_k={request.top_k} exceeds the vocabulary size {vocab}.") + + +def materialized_head_stats( + logits: torch.Tensor, + request: HeadRequest, + *, + targets: torch.Tensor | None = None, + teacher_logits: torch.Tensor | None = None, +) -> HeadStats: + """The requested statistics of ``logits``, which the head already produced, through autograd. + + Args: + logits: ``[R, V]`` output of the head, whatever transform its forward applies. + request: The statistics to return. + targets: ``[R]`` label ids, needed for ``hard_ce``. + teacher_logits: ``[R, V]`` teacher logits, needed for ``kd_ce``. + """ + _check_request(request, vocab=logits.shape[-1], teacher=teacher_logits, targets=targets) + + def stat(term, compute): + if term in request.grads: + return compute() + if term in request.values: + with torch.no_grad(): + return compute() + return None + + kd_ce = stat( + "kd_ce", + lambda: ( + -(torch.softmax(teacher_logits, dim=-1) * torch.log_softmax(logits, dim=-1)).sum(dim=-1) + ), + ) + hard_ce = stat("hard_ce", lambda: F.cross_entropy(logits, targets, reduction="none")) + top_logits = top_ids = None + if request.top_k: + top_logits, top_ids = logits.topk(request.top_k, dim=-1) + lse = stat("lse", lambda: torch.logsumexp(logits.float(), dim=-1)) + with torch.no_grad(): + top1 = logits.argmax(dim=-1) + return HeadStats(kd_ce, hard_ce, lse, top1, top_ids, top_logits) + + +def fused_projection(head: nn.Module) -> tuple[torch.Tensor, float | None, float | None]: + """``(weight, logit_scale, soft_cap)`` with which :func:`fused_head_stats` reproduces ``head``. + + The fused pass reads ``head.weight`` alone, so it reproduces a bias-free ``nn.Linear`` + without forward hooks whose forward is either ``nn.Linear``'s own, or applies after the + projection exactly the transform its ``logit_transform()`` returns as + ``(logit_scale, soft_cap)``. Anything else, such as a bias, a quantizer or a hook, would be + skipped silently, so it is refused. + """ + hooked = bool(head._forward_hooks or head._forward_pre_hooks) or "forward" in vars(head) + if isinstance(head, nn.Linear) and head.bias is None and not hooked: + if type(head).forward is nn.Linear.forward: + return head.weight, None, None + if hasattr(type(head), "logit_transform"): + return (head.weight, *head.logit_transform()) + raise ValueError( + "The fused pass projects with the head's weight alone, so it needs a bias-free " + "nn.Linear head with no forward hooks, running nn.Linear's forward or one whose " + f"transform logit_transform() declares; got {type(head).__name__}." + ) + + +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 "kd_ce" in want: + zt = (teacher if teacher_is_logits else _transform(teacher, scale, cap)).float() + p = torch.softmax(zt, dim=-1) + vals["kd_ce"] = -(p * log_q).sum(dim=-1) + if "kd_ce" in grads: + dz["kd_ce"] = q - p + 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, + request: HeadRequest, + *, + targets: torch.Tensor | None = None, + teacher_hidden: torch.Tensor | None = None, + teacher_logits: torch.Tensor | None = None, + logit_scale: float | None = None, + soft_cap: float | None = None, + chunk: int = 2048, +) -> HeadStats: + """The requested 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 ``request.grads`` statistics and ``top_logits`` + backpropagate into them, and nothing into ``weight``, which must be frozen. + weight: ``[V, H]`` head weight. + request: The statistics to return. + targets: ``[R]`` label ids, needed for ``hard_ce``. + teacher_hidden: ``[R, H]`` teacher rows to project through the same head. ``kd_ce`` + needs this or ``teacher_logits``. + teacher_logits: ``[R, V]`` teacher logits, used as given. + 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. + """ + 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.") + teacher = teacher_hidden if teacher_logits is None else teacher_logits + _check_request(request, vocab=weight.shape[0], teacher=teacher, targets=targets) + want = request.grads | request.values + grads = request.grads if torch.is_grad_enabled() and hidden.requires_grad else frozenset() + 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), + request.top_k, + chunk, + ) + ) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 1e20d5b30999..805dc3ef0448 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -76,7 +76,6 @@ from typing import NamedTuple import torch -import torch.nn.functional as F from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -85,6 +84,13 @@ from ..dflash.conversion import DFlashDMRegistry from ..dflash.dflash_model import DFlashModel from .dflash_flex_attention import build_draft_block_mask +from .dflash_head_stats import ( + HeadRequest, + HeadStats, + fused_head_stats, + fused_projection, + materialized_head_stats, +) from .modeling_dflash import ( # noqa: F401 DFlashAttention, DFlashBaseModelOutput, @@ -133,6 +139,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 +221,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 +329,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`` holds every statistic of the draft's logits + the model's ``_head_request`` names, for a variant's own terms. """ ce_per_token: torch.Tensor | None weights: torch.Tensor weight_sum: torch.Tensor supervised_mask: torch.Tensor + head_stats: HeadStats @DFlashDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) @@ -330,6 +348,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` reads the draft's logits only through + # `_draft_head_stats`. Checked on the class defining `forward`, so it is not inherited. + _supports_fused_loss = True + @property def _base_model(self): return self.get_submodule(self.base_model_path) @@ -540,6 +562,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 +604,26 @@ 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.""" + owner = next(cls for cls in type(self).__mro__ if "forward" 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 forward builds the " + "draft's logits itself instead of reading them through _draft_head_stats. Unset " + "dflash_fused_loss." + ) + head = self._base_model_lm_head + try: + fused_projection(head) + except ValueError as err: + raise ValueError(f"dflash_fused_loss: {err} Unset dflash_fused_loss.") from None + 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 +966,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,43 +992,37 @@ 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, - logits, + draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, *, - draft_hidden=None, base_outputs=None, return_terms=False, ): """Compute weighted cross-entropy (or KD) loss and accuracy. Args: - logits: Draft model output [B, N*block_size, vocab]. + draft_hidden: Draft hidden states [B, N*block_size, H]. The loss reads the logits + they project to only through ``_draft_head_stats``. 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 +1033,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 +1041,8 @@ def _compute_loss( "dflash_self_logit_distillation distills from base_outputs, but none was " "passed; an override of _compute_loss must forward it." ) - # Non-KD loss is per-token cross-entropy; compute it once (grad enabled) so the - # D-PACE confidences below can reuse it instead of a second CE pass. The KD path - # optimizes KL, so its confidences need a dedicated no_grad CE pass. - loss_per_token = None - if not kd: - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") + stats = self._draft_head_stats(draft_hidden, base_outputs, targets) + loss_per_token = None if kd else stats.hard_ce weight_mask = _block_position_weights( targets.mask, @@ -1006,31 +1050,19 @@ def _compute_loss( _first_predicted_position(self._label_shift), decay_factor=self.dflash_loss_decay_factor, dpace_alpha=self.dflash_dpace_alpha, - target_ce=lambda: ( - loss_per_token.detach() - if loss_per_token is not None - else F.cross_entropy(flat_logits, flat_targets, reduction="none") - ).view(targets.mask.shape), + target_ce=lambda: stats.hard_ce.detach().view(targets.mask.shape), ) flat_weights = weight_mask.view(-1) valid_count = flat_weights.sum() + 1e-6 if valid_count > 1.0: - if kd: - flat_teacher = self._teacher_logits(base_outputs, label_positions=targets.positions) - flat_teacher = flat_teacher.reshape(flat_logits.shape) - target_soft = torch.softmax(flat_teacher, dim=-1) - draft_logsoft = torch.log_softmax(flat_logits, dim=-1) - kd_loss = -(target_soft * draft_logsoft).sum(dim=-1) - loss = (kd_loss * flat_weights).sum() / valid_count - else: - loss = (loss_per_token * flat_weights).sum() / valid_count - + per_token = stats.kd_ce if kd else loss_per_token + loss = (per_token * flat_weights).sum() / valid_count with torch.no_grad(): - accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item() + accuracy = _masked_accuracy(stats.top1, flat_targets, binary_eval_mask).item() else: - loss = self._draft_zero_loss(flat_logits.device, through=(flat_logits,)) + loss = self._draft_zero_loss(draft_hidden.device, through=(draft_hidden,)) accuracy = 0.0 if return_terms: @@ -1039,10 +1071,57 @@ 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 _head_request(self): + """The statistics of the draft's logits the loss reads; a variant extends it.""" + return HeadRequest( + grads=frozenset({"kd_ce" if self.dflash_self_logit_distillation else "hard_ce"}), + # D-PACE weighs positions by the draft's confidence in its label. + values=frozenset({"hard_ce"} if self.dflash_loss_objective == "dpace" else ()), + ) + + def _draft_head_stats(self, draft_hidden, base_outputs, targets): + """The statistics ``_head_request`` names of the logits ``draft_hidden`` projects to. + + ``dflash_fused_loss`` computes them in row chunks fused with the lm_head projection; + otherwise the logits are materialized whole. Either way they come back as one + ``HeadStats``, so no loss term depends on which. + """ + request = self._head_request() + flat_targets = targets.token_ids.view(-1) + kd = "kd_ce" in request.grads | request.values + if self.dflash_fused_loss: + teacher = {} + if kd and base_outputs.logits is not None: + rows = self._teacher_logits(base_outputs, label_positions=targets.positions) + teacher["teacher_logits"] = rows.flatten(0, -2) + elif kd: + rows = self._teacher_hidden(base_outputs, label_positions=targets.positions) + teacher["teacher_hidden"] = rows.flatten(0, -2) + weight, logit_scale, soft_cap = fused_projection(self._base_model_lm_head) + return fused_head_stats( + draft_hidden.flatten(0, -2), + weight, + request, + targets=flat_targets, + logit_scale=logit_scale, + soft_cap=soft_cap, + **teacher, + ) + logits = self._base_model_lm_head(draft_hidden) + flat_logits = logits.view(-1, logits.size(-1)) + teacher_logits = None + if kd: + teacher_logits = self._teacher_logits(base_outputs, label_positions=targets.positions) + teacher_logits = teacher_logits.reshape(flat_logits.shape) + return materialized_head_stats( + flat_logits, request, targets=flat_targets, teacher_logits=teacher_logits + ) + def forward( self, input_ids=None, @@ -1197,14 +1276,12 @@ def forward( ) # 6. Compute loss and accuracy - logits = self._base_model_lm_head(hidden) loss, accuracy = self._compute_loss( - logits, + hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, - draft_hidden=hidden, base_outputs=base_outputs, ) diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 88c1e3e38f72..9d012dca5872 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,30 @@ def get_exporter(self): return DFlash2Exporter(self) + def _head_request(self): + """The backbone loss's statistics, and the top-k candidates the selector scores.""" + request = super()._head_request() + if self.dflash_selector_loss_alpha <= 0: + return request + return request.extend(top_k=self.dflash_module.candidate_selector.top_k) + def _selector_loss( - self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum + self, + unary_logits, + candidate_ids, + target_ids, + hidden, + predecessor_ids, + supervised_mask, + loss_weights, + weight_sum, ): """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]``. + unary_logits: The backbone's logits of its top-k candidates per block position + ``[B, N, block_size, k]``, differentiable. + candidate_ids: Those candidates' ids, in descending logit order, same shape. 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,8 +150,6 @@ def _selector_loss( selector = self.dflash_module.candidate_selector top_k = selector.top_k - 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 # selector classification example: substituting it in would teach the selector @@ -208,13 +225,12 @@ def _lk_loss(self, terms): def _compute_loss( self, - logits, + draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, *, - draft_hidden=None, base_outputs=None, ): """Backbone DFlash loss plus the candidate-selector cross-entropy. @@ -225,18 +241,17 @@ def _compute_loss( the selector's own accuracy is logged separately. """ loss, accuracy, terms = super()._compute_loss( - logits, + draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, - draft_hidden=draft_hidden, base_outputs=base_outputs, return_terms=True, ) if self.dflash_lk_loss_type != "ce": loss = self._lk_loss(terms) - if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: + if self.dflash_selector_loss_alpha <= 0: return loss, accuracy bsz = input_ids.shape[0] @@ -259,10 +274,12 @@ 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. + per_position = (bsz, n_blocks, block_size, -1) selector_loss, selector_accuracy, selector_coverage = self._selector_loss( - logits.reshape(bsz, n_blocks, block_size, -1), + terms.head_stats.top_logits.view(per_position), + terms.head_stats.top_ids.view(per_position), target_ids, - draft_hidden.reshape(bsz, n_blocks, block_size, -1), + draft_hidden.reshape(per_position), predecessor_ids, terms.supervised_mask.view_as(target_ids), terms.weights.view_as(target_ids), diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index b38234d78ebe..df6d06ced4b4 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) @@ -192,13 +194,12 @@ def get_exporter(self): def _compute_loss( self, - logits, + draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, *, - draft_hidden=None, base_outputs=None, ): """Add the LiLiCorr lattice terms to the DFlash block loss. @@ -213,20 +214,20 @@ def _compute_loss( ``HFDFlashModel.forward`` already builds for them. """ target_hidden = base_outputs.target_hidden if base_outputs is not None else None - if draft_hidden is None or target_hidden is None: + if target_hidden is None: raise ValueError( - "LiLiCorr requires draft_hidden and target_hidden in _compute_loss: the " - "reranker reads the draft's per-slot hidden states and anchors on a " + "LiLiCorr requires target_hidden in _compute_loss: the reranker anchors on a " "committed target row." ) - origin_loss, origin_accuracy = super()._compute_loss( - logits, + origin_loss, origin_accuracy, terms = super()._compute_loss( + draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask, base_outputs=base_outputs, + return_terms=True, ) # The backbone's own supervision, before its position weighting, so both losses # train the same positions. @@ -238,8 +239,11 @@ def _compute_loss( self.dflash_block_size, self._label_shift, ) + stats = terms.head_stats lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss( - logits=logits, + top_ids=stats.top_ids, + top_logits=stats.top_logits, + lse=stats.lse, 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,10 +266,17 @@ def _compute_loss( self._lilicorr_metrics = metrics return origin_loss + lilicorr_loss, accuracy + def _head_request(self): + """The backbone loss's statistics, and the candidates the lattice scores.""" + candidates = self.dflash_module.lilicorr.candidate_topk + return super()._head_request().extend(grads={"lse"}, top_k=candidates) + def _compute_lilicorr_loss( self, *, - logits, + top_ids, + top_logits, + lse, target_ids, slot_mask, anchor_positions, @@ -275,6 +286,10 @@ def _compute_lilicorr_loss( ): """Score the candidate lattice and evaluate the LiLiCorr objective. + ``top_ids`` / ``top_logits`` are the backbone's top-k candidates and their + (differentiable) logits per block position, ``[B * N * block_size, k]`` in descending + order, and ``lse`` their rows' logsumexp, ``[B * N * block_size]``. + 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 +299,14 @@ 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] batch_blocks = bsz * n_blocks - if topk > vocab: - raise ValueError(f"lilicorr_candidate_topk={topk} exceeds the vocabulary size {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 + ) - 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, @@ -300,23 +314,15 @@ def _compute_lilicorr_loss( valid_block = slot_mask.all(dim=-1).reshape(batch_blocks, 1) with torch.no_grad(): - # topk already returns candidates in descending log-probability order, which - # 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) + # Descending logits: the rank order `rank_embedding` is indexed by. + candidate_ids = slots(top_ids) gt_hits = candidate_ids == gt_flat.unsqueeze(-1) gt_in_slot = gt_hits.any(dim=-1) gt_indices = gt_hits.long().argmax(dim=-1) - # Soft-value channel. The candidate *ids* are the hard top-k, so the forward - # values are unchanged; recomputing their log-probs from the LIVE logits - # 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) - candidate_log_probs = candidate_logits - log_denominator + # Differentiable log-probabilities of the hard top-k: how the lattice terms reach the + # drafter body and not only the head. + candidate_log_probs = slots(top_logits).float() - slots(lse.unsqueeze(-1)) draft_slots = draft_hidden.view(bsz, n_blocks, block_size, -1)[:, :, 1:, :] start_scores, pair_scores = module.score_lattice( diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 01ed2713068e..9b9569d1205f 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -428,3 +428,51 @@ def test_compiled_flex_matches_eager_flex(self): grads = [p.grad for p in compiled.dflash_module.parameters() if p.grad is not None] 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_head_stats.py b/tests/unit/torch/speculative/plugins/test_dflash_head_stats.py new file mode 100644 index 000000000000..140158fa0e48 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_dflash_head_stats.py @@ -0,0 +1,112 @@ +# 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 both backends of the lm_head statistics behind the DFlash-family loss.""" + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.torch.speculative.plugins.dflash_head_stats import ( + TERMS, + HeadRequest, + fused_head_stats, + materialized_head_stats, +) + +ROWS, HIDDEN, VOCAB, TOP_K = 37, 16, 50, 4 +CHUNK = 8 # several chunks and a partial last one +BACKENDS = ["fused", "materialized"] + + +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) + + +def _stats(backend, hidden, weight, request, scale=None, cap=None, **inputs): + """The statistics of the head ``_transform(hidden @ weight.T)``, from ``backend``.""" + if backend == "fused": + return fused_head_stats( + hidden, weight, request, logit_scale=scale, soft_cap=cap, chunk=CHUNK, **inputs + ) + if (teacher_hidden := inputs.pop("teacher_hidden", None)) is not None: + inputs["teacher_logits"] = _transform(teacher_hidden @ weight.T, scale, cap) + return materialized_head_stats(_transform(hidden @ weight.T, scale, cap), request, **inputs) + + +@pytest.mark.parametrize("backend", BACKENDS) +@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_logits(backend, 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), + "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() + + request = HeadRequest(grads=frozenset(TERMS), top_k=TOP_K) + test_hidden = hidden.clone().requires_grad_(True) + stats = _stats(backend, test_hidden, weight, request, scale, cap, targets=targets, **teacher) + 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 = _stats( + backend, test_hidden, weight, request, scale, cap, targets=targets, **teacher + ) + + 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(test_hidden.grad, ref_hidden.grad) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_values_carry_no_gradient_from_either_backend(backend): + """Only ``grads`` statistics are differentiable, so the backends cannot train differently.""" + torch.manual_seed(0) + hidden = torch.randn(ROWS, HIDDEN, requires_grad=True) + targets = torch.randint(0, VOCAB, (ROWS,)) + request = HeadRequest(grads=frozenset({"lse"}), values=frozenset({"hard_ce"})) + stats = _stats(backend, hidden, torch.randn(VOCAB, HIDDEN), request, targets=targets) + assert stats.lse.requires_grad + assert not stats.hard_ce.requires_grad + assert stats.kd_ce is None diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index d3bf4162cf79..a2aa88d5f52e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -253,15 +253,15 @@ class TestDPaceLossIntegration: """Exercise the _compute_loss block-weighting branches on CPU.""" @staticmethod - def _make_inputs(vocab=32, seq_len=SEQ_LEN, n_blocks=2): - """Synthetic CPU inputs for _compute_loss (no model forward needed).""" + def _make_inputs(model, seq_len=SEQ_LEN, n_blocks=2): + """Synthetic CPU inputs for _compute_loss (no draft forward needed).""" bsz = 1 - logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab) - input_ids = torch.randint(0, vocab, (bsz, seq_len)) + draft_hidden = _random_draft_hidden(model, bsz, n_blocks) + input_ids = torch.randint(0, model.config.vocab_size, (bsz, seq_len)) anchor_positions = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks] block_keep_mask = torch.ones(bsz, n_blocks) loss_mask = torch.ones(bsz, seq_len) - return logits, input_ids, anchor_positions, block_keep_mask, loss_mask + return draft_hidden, input_ids, anchor_positions, block_keep_mask, loss_mask def _converted_model(self, objective, **overrides): model = get_tiny_llama(num_hidden_layers=4) @@ -274,7 +274,7 @@ def _converted_model(self, objective, **overrides): def test_compute_loss_dpace_branch(self): """Default dpace objective produces a finite loss and valid accuracy.""" model = self._converted_model("dpace", dflash_self_logit_distillation=False) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = model._compute_loss(*self._make_inputs(model)) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 @@ -283,16 +283,17 @@ def test_compute_loss_decay_branch(self): model = self._converted_model( "decay", dflash_loss_decay_factor=4.0, dflash_self_logit_distillation=False ) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = model._compute_loss(*self._make_inputs(model)) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 def test_compute_loss_dpace_kd_branch(self): """dpace + KD: confidences use a dedicated no_grad CE pass.""" - vocab = 32 model = self._converted_model("dpace") - inputs = self._make_inputs(vocab=vocab) - teacher = DFlashBaseModelOutput(None, logits=torch.randn(1, SEQ_LEN, vocab)) + inputs = self._make_inputs(model) + teacher = DFlashBaseModelOutput( + None, logits=torch.randn(1, SEQ_LEN, model.config.vocab_size) + ) loss, acc = model._compute_loss(*inputs, base_outputs=teacher) assert torch.isfinite(loss).item() assert 0.0 <= acc <= 1.0 @@ -301,7 +302,7 @@ def test_kd_without_base_outputs_raises(self): """KD is decided by the config, so a missing teacher cannot fall back to CE.""" model = self._converted_model("dpace") with pytest.raises(ValueError, match="base_outputs"): - model._compute_loss(*self._make_inputs()) + model._compute_loss(*self._make_inputs(model)) class TestDFlashSaveRestore: @@ -971,6 +972,12 @@ def test_multi_turn_masks_only_assistant(self, tiny_tokenizer): assert "How are you?" not in decoded +def _random_draft_hidden(model, bsz, n_blocks): + """Draft hidden states for ``n_blocks`` blocks, in the dtype the lm_head projects.""" + dtype = model._base_model_lm_head.weight.dtype + return torch.randn(bsz, n_blocks * BLOCK_SIZE, model.config.hidden_size, dtype=dtype) + + def _dflash_batch(vocab_size, bsz=2, seq_len=SEQ_LEN): torch.manual_seed(0) input_ids = torch.randint(1, vocab_size, (bsz, seq_len)) @@ -1500,13 +1507,13 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self): torch.manual_seed(0) input_ids = torch.randint(0, vocab, (bsz, SEQ_LEN)) loss_mask = torch.ones(bsz, SEQ_LEN) - logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab) + draft_hidden = _random_draft_hidden(model, bsz, n_blocks) anchors = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks] keep = torch.ones(bsz, n_blocks) - base_loss, base_acc = model._compute_loss(logits, input_ids, anchors, keep, loss_mask) + base_loss, base_acc = model._compute_loss(draft_hidden, input_ids, anchors, keep, loss_mask) padded_loss, padded_acc = model._compute_loss( - torch.cat([logits, torch.randn(bsz, pad * BLOCK_SIZE, vocab)], dim=1), + torch.cat([draft_hidden, _random_draft_hidden(model, bsz, pad)], dim=1), input_ids, torch.cat([anchors, torch.zeros(bsz, pad, dtype=anchors.dtype)], dim=1), torch.cat([keep, torch.zeros(bsz, pad)], dim=1), @@ -1692,7 +1699,7 @@ def test_kd_loss_matches_full_logits(self): """KD read from the base hidden equals KD from the same logits built in full.""" model = self._model() hidden = self._hidden(model, bsz=1) - inputs = TestDPaceLossIntegration._make_inputs(vocab=model.config.vocab_size) + inputs = TestDPaceLossIntegration._make_inputs(model) full = DFlashBaseModelOutput(None, logits=model._base_model_lm_head(hidden)) loss_full, _ = model._compute_loss(*inputs, base_outputs=full) loss, _ = model._compute_loss( @@ -1727,3 +1734,102 @@ 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 _TransformedHead(torch.nn.Linear): + """An lm_head that scales and soft-caps its logits, and declares so for the fused pass.""" + + SCALE, CAP = 0.5, 3.0 + + def forward(self, x): + return self.CAP * torch.tanh(super().forward(x) * self.SCALE / self.CAP) + + def logit_transform(self): + return self.SCALE, self.CAP + + +class _UndeclaredHead(torch.nn.Linear): + """An lm_head whose forward does more than project, without declaring it.""" + + def forward(self, x): + return 2 * super().forward(x) + + +def _swap_lm_head(model, head_cls): + """Replace ``model.lm_head`` with a ``head_cls`` holding the same weight.""" + old = model.lm_head + model.lm_head = head_cls(old.in_features, old.out_features, bias=False, dtype=old.weight.dtype) + model.lm_head.weight.data.copy_(old.weight.data) + + +class TestDFlashFusedLoss: + """``dflash_fused_loss`` trains exactly what the loss on materialized logits trains.""" + + @staticmethod + def _step(fused, offline, transformed_head=False, **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 transformed_head: + _swap_lm_head(model, _TransformedHead) + 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", "transformed_head", "overrides"), + [ + (False, False, {}), # KD on the base's logits, D-PACE weighting + (True, False, {"dflash_loss_objective": "decay", "dflash_loss_decay_factor": 4.0}), + (False, False, {"dflash_self_logit_distillation": False}), # cross-entropy, D-PACE + (True, True, {}), # KD on the base's hidden, through a declared logit transform + ], + ids=["kd-online", "kd-offline-decay", "ce", "kd-offline-transformed-head"], + ) + def test_matches_the_materialized_logits(self, offline, transformed_head, overrides): + eager, eager_grads = self._step(False, offline, transformed_head, **overrides) + fused, fused_grads = self._step(True, offline, transformed_head, **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", "lm_head"), + [ + ({}, {"projector_type": "dspark", "markov_rank": 4}, None), + ({}, {"projector_type": "domino", "gru_hidden_dim": 8, "emb_dim": 8}, None), + ({"dflash_freeze_base_model": False}, {}, None), + ({}, {}, "hooked"), + ({}, {}, "undeclared"), + ], + ids=["dspark", "domino", "trainable-lm-head", "hooked-lm-head", "undeclared-lm-head"], + ) + def test_unsupported_setups_are_rejected(self, overrides, arch, 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 lm_head == "hooked": + model.lm_head.register_forward_hook(lambda module, args, output: output * 2) + elif lm_head == "undeclared": + _swap_lm_head(model, _UndeclaredHead) + 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 e1ea8c10dc06..6ef7a97f8b91 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.""" @@ -410,7 +438,7 @@ def _selector_terms(self, selector, target_ids, candidate_ids, loss_weights=None logits.scatter_(-1, candidate_ids[..., slot : slot + 1], float(k - slot)) return HFDFlash2Model._selector_loss( model, - logits, + *logits.topk(k, dim=-1), target_ids, torch.randn( b, length, 8, generator=torch.Generator().manual_seed(1), dtype=torch.double diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py index 0d132f9ce49d..696b7b6ac363 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"), [