diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 51f13f98c3e..859d960d0b2 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/modelopt/torch/speculative/dflash/conversion.py @@ -48,6 +48,16 @@ # registry so its wrapper (HFLiLiCorrModel) does not overwrite HFDFlashModel. LiLiCorrDMRegistry = _DMRegistryCls(prefix="LiLiCorr") +# dflash_architecture_config.projector_type -> the registry holding that variant's wrapper. +_REGISTRY_BY_PROJECTOR_TYPE = { + None: DFlashDMRegistry, + "dflash": DFlashDMRegistry, + "domino": DominoDMRegistry, + "dspark": DSparkDMRegistry, + "dflash2": DFlash2DMRegistry, + "lilicorr": LiLiCorrDMRegistry, +} + def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertReturnType: """Convert the model to a DFlash (or Domino) model as per `config`.""" @@ -59,22 +69,13 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe custom_config = config.dflash_architecture_config config.dflash_architecture_config = {**default_dflash_config, **custom_config} - # Route to the Domino registry when the architecture asks for the Domino head. projector_type = config.dflash_architecture_config.get("projector_type") - if projector_type == "domino": - registry = DominoDMRegistry - elif projector_type == "dspark": - registry = DSparkDMRegistry - elif projector_type == "dflash2": - registry = DFlash2DMRegistry - elif projector_type == "lilicorr": - registry = LiLiCorrDMRegistry - elif projector_type in (None, "dflash"): - registry = DFlashDMRegistry - else: + registry = _REGISTRY_BY_PROJECTOR_TYPE.get(projector_type) + if registry is None: + supported = [name for name in _REGISTRY_BY_PROJECTOR_TYPE if name is not None] raise ValueError( f"Unsupported dflash_architecture_config.projector_type: {projector_type!r}. " - "Expected 'dflash' (default), 'domino', 'dspark', 'dflash2' or 'lilicorr'." + f"Expected one of {supported} ('dflash' is the default)." ) registry.register_subclass(type(model), "base_model_class") diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index c1c3389e428..5d3fc293efc 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -97,7 +97,7 @@ _LM_HEAD_PATHS, find_final_norm_path, ) -from .modeling_final_norm import _maybe_apply_base_final_norm +from .modeling_final_norm import _apply_base_final_norm logger = logging.getLogger(__name__) @@ -133,6 +133,89 @@ def _gather_rows(x, positions): return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1)) +def _block_positions(anchor_positions, block_size, offset, seq_len): + """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. + + Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end. + """ + k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1) + unclamped = anchor_positions.unsqueeze(-1) + k + offset + return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len + + +def _block_gather(x, positions): + """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs].""" + return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions) + + +def _block_tokens(input_ids, anchor_positions, block_size, offset): + """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs].""" + positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1]) + return _block_gather(input_ids, positions) + + +class _BlockTargets(NamedTuple): + """What each block position is trained on: the token at ``anchor + k + label_shift``.""" + + token_ids: torch.Tensor # [B, N, bs] + positions: torch.Tensor # [B, N, bs] their sequence indices, clamped onto the sequence + mask: torch.Tensor # [B, N, bs] 1.0 where the position is supervised + + +def _first_predicted_position(shift): + """The first block position whose label comes after the anchor, for a label shift.""" + return max(0, 1 - shift) + + +def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift): + """Labels of every block position and where they are supervised. + + A position is supervised when its block is kept and its label lies inside the sequence + and the loss mask, and is not the anchor itself, which the draft is given as input. + """ + positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1]) + k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1) + mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() + mask = mask * inside.float() + mask = mask * (k >= _first_predicted_position(shift)).float() + mask = mask * _block_gather(loss_mask, positions) + return _BlockTargets(_block_gather(input_ids, positions), positions, mask) + + +def _block_position_weights(mask, scheme, first, *, decay_factor, dpace_alpha, target_ce=None): + """``mask`` weighted per block position, by D-PACE or by exponential decay. + + ``first`` is the first predicted block position. ``target_ce`` returns the draft's + cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. The two floats are + keyword-only because swapping them would still give a plausible weight curve. + """ + block_size = mask.shape[-1] + if scheme == "dpace" and block_size > 1: + if target_ce is None: + raise ValueError( + "D-PACE weighs each block position by the draft's confidence on its label, " + "so it needs target_ce." + ) + with torch.no_grad(): + # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9). + confidences = torch.exp(-target_ce()[..., first:].float()) + dpace = torch.ones_like(mask) + dpace[..., first:] = _dpace_position_weights( + confidences, dpace_alpha, valid_mask=mask[..., first:] + ) + return mask * dpace + if decay_factor > 0: + k = torch.arange(block_size, device=mask.device).view(1, 1, -1) + return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor) + return mask + + +def _masked_top1(logits, target_ids, mask): + """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" + correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) + return correct.sum().float() / (mask.sum() + 1e-6) + + def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal): """The draft's attention visibility rule, as a FlexAttention ``mask_mod``. @@ -162,6 +245,19 @@ def mask_mod(b, h, q_idx, kv_idx): return mask_mod +def _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device): + """``mask_mod`` evaluated on broadcast index grids, as SDPA's additive [B, 1, Q, KV] mask.""" + visible = mask_mod( + torch.arange(bsz, device=device).view(-1, 1, 1, 1), + None, + torch.arange(q_len, device=device).view(1, 1, -1, 1), + torch.arange(kv_len, device=device).view(1, 1, 1, -1), + ) + attn_mask = torch.zeros(bsz, 1, q_len, kv_len, device=device, dtype=dtype) + attn_mask.masked_fill_(~visible, torch.finfo(dtype).min) + return attn_mask + + def _dpace_position_weights( confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None ) -> torch.Tensor: @@ -230,6 +326,10 @@ class _DFlashLossTerms(NamedTuple): class HFDFlashModel(DFlashModel): """DFlash Model for HuggingFace transformers.""" + # Block position k is trained on the token at anchor + k + _label_shift: DFlash predicts + # the anchor's own position onward, so slot 0 (the given anchor) is never supervised. + _label_shift = 0 + @property def _base_model(self): return self.get_submodule(self.base_model_path) @@ -293,7 +393,7 @@ def modify(self, config): """Initialize DFlash draft module.""" super().modify(config) - self.dflash_fp32_master_weights = getattr(config, "dflash_fp32_master_weights", False) + self.dflash_fp32_master_weights = config.dflash_fp32_master_weights base_config = self._base_llm_config # Use Qwen3Config (not generic PretrainedConfig) so rope_parameters is @@ -418,7 +518,7 @@ def modify(self, config): "contains duplicates; each draft layer needs a distinct capture layer " "(the streaming producer captures each base layer at most once)." ) - # forward() indexes hidden_states[lid + 1], where a negative id silently wraps. + # _target_hidden indexes hidden_states[lid + 1], where a negative id silently wraps. if min(user_target_layer_ids) < 0: raise ValueError( f"dflash_architecture_config.target_layer_ids {user_target_layer_ids} " @@ -702,56 +802,124 @@ def _build_draft_attention_mask( head_dims={layer.self_attn.head_dim for layer in self.dflash_module.layers}, ) - visible = mask_mod( - torch.arange(bsz, device=device).view(-1, 1, 1, 1), - None, - torch.arange(q_len, device=device).view(1, 1, -1, 1), - torch.arange(kv_len, device=device).view(1, 1, 1, -1), - ) # [B, 1, Q, KV] - attn_mask = torch.zeros(bsz, 1, q_len, kv_len, device=device, dtype=dtype) - attn_mask.masked_fill_(~visible, torch.finfo(dtype).min) - return attn_mask + return _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device) def _build_generate_swa_mask(self, ctx_len, bsz, dtype, device): """Generation-time mask [B, 1, block_size, ctx_len + block_size], or None. Returns None only when there is nothing to mask: full attention over the context - *and* bidirectional blocks (KV cache with no mask). With sliding-window attention, - each block query only sees context within ``dflash_swa_window_size`` tokens before - its real position (ctx_len + position-in-block), matching training and vLLM - inference; block kv is left un-windowed. With ``dflash_draft_attention="causal"`` - the block is additionally lower-triangular, mirroring - :meth:`_build_draft_attention_mask` so generation matches training. + *and* bidirectional blocks (KV cache with no mask). Otherwise it is the training rule, + :func:`_draft_mask_mod`, for the one block being drafted, anchored at ``ctx_len``: + the window counts from each query's real position, and under + ``dflash_draft_attention="causal"`` the block is lower-triangular. """ window = self.dflash_swa_window_size causal = self.dflash_draft_attention == "causal" if window is None and not causal: return None block_size = self.dflash_block_size - kv_len = ctx_len + block_size - kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1) - q_pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1, 1) - q_real_pos = ctx_len + q_pos_in_block - is_ctx = kv_idx < ctx_len - # Context kv kept iff within the window (when windowing); block kv always visible - # unless the block is causal, in which case only positions <= the query's. - keep_ctx = is_ctx if window is None else (is_ctx & (kv_idx > q_real_pos - window)) - keep_block = ~is_ctx - if causal: - keep_block = keep_block & ((kv_idx - ctx_len) <= q_pos_in_block) - keep = keep_ctx | keep_block - attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype) - attn_mask.masked_fill_(~keep, torch.finfo(dtype).min) - return attn_mask + anchors = torch.full((bsz, 1), ctx_len, dtype=torch.long, device=device) + keep = torch.ones(bsz, 1, dtype=torch.bool, device=device) + mask_mod = _draft_mask_mod(ctx_len, anchors, keep, block_size, window, causal) + return _dense_additive_mask(mask_mod, bsz, block_size, ctx_len + block_size, dtype, device) - @torch.no_grad() - def _teacher_logits(self, base_outputs, positions, token_ids=None): - """Base-model logits at ``positions`` ([B, ...] sequence indices) -> [B, ..., vocab]. + def _target_hidden(self, hidden_states, device=None): + """The base hidden states the draft's ``fc`` reads, concatenated: [B, seq, n * H]. - Only the requested rows of the base hidden go through the final norm and lm_head, so - no full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the - projection to just those vocab entries. + ``hidden_states[0]`` is the embedding output, so layer ``lid``'s output is at ``lid + 1``. + ``device`` gathers the layers first, for a base model sharded across GPUs. """ + selected = [hidden_states[lid + 1] for lid in self.target_layer_ids] + if device is not None: + selected = [h.to(device) for h in selected] + return torch.cat(selected, dim=-1) + + def _generation_base(self, input_ids): + """The base model's greedy next token after ``input_ids`` ([B, 1]), and its hidden states.""" + if self.dflash_offline: + raise RuntimeError( + "An offline DFlash-family model cannot run AR validation / " + "pseudo_speculative_generate — base model layers were deleted during offline " + "conversion to save memory. Reload the full base model before running AR validation." + ) + # Call the base model's inner model directly (avoids DynamicModule dispatch) + model_output = self._base_model(input_ids=input_ids, output_hidden_states=True) + base_logits = self._base_model_lm_head(model_output.last_hidden_state) + base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device) + return base_token, model_output.hidden_states + + def _draft_block(self, base_token, hidden_states): + """Draft hidden states [B, block_size, H] for the block that ``base_token`` anchors. + + Re-runs no base layer: the context is ``hidden_states`` from :meth:`_generation_base`. + """ + # target_layer_ids spans early and late layers, so under device_map="auto" these come + # off different GPUs and the cat fails. Everything the draft consumes is gathered + # onto the draft's own device -- the last base layer's, per modify() -- which is + # also not input_ids.device once the base model is sharded. All no-ops single-device. + device = self._base_device() + target_hidden = self._target_hidden(hidden_states, device=device) + + block_size = self.dflash_block_size + bsz = base_token.shape[0] + + # Block: first token is base_token (anchor), rest are mask + block_ids = torch.full( + (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device + ) + block_ids[:, 0] = base_token.squeeze(-1).to(device) + # The embedding table is on the first shard; its output follows it, not block_ids. + noise_embedding = self._base_model_embeddings(block_ids).to(device) + + # The training layout for one block anchored right after the context. + ctx_len = target_hidden.shape[1] + anchors = torch.full((bsz, 1), ctx_len, dtype=torch.long, device=device) + pos_ids = self._build_position_ids(ctx_len, anchors, device) + + attn_mask = self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device) + + return self.dflash_module( + noise_embedding=noise_embedding, + target_hidden=target_hidden, + position_ids=pos_ids, + attention_mask=attn_mask, + ) + + def _draft_zero_loss(self, device, through=()): + """A zero loss with a backward path to the draft, for a step with nothing to train. + + DDP expects a gradient for each trainable parameter on every rank; a rank returning a + constant loss would leave the others waiting on its reduction. Before the draft has + run, the loss reads one element of every trainable draft parameter. Once it has run, + pass every one of its outputs as ``through`` (``None`` for an absent head): the loss + then reaches only the parameters behind them, so an output left out leaves its + parameters unreduced. Reading the parameters directly instead would put this rank's + collectives out of step when FSDP2 shards the draft as its own unit. + """ + if through: + outputs = [out for out in through if out is not None] + if not outputs: + raise ValueError( + "through= holds no draft output; the parameters cannot stand in for one " + "once the draft has run." + ) + return sum(out.sum() * 0.0 for out in outputs) + return sum( + (p.reshape(-1)[0] * 0.0 for p in self.dflash_module.parameters() if p.requires_grad), + torch.zeros((), device=device), + ) + + @torch.no_grad() + def _teacher_logits(self, base_outputs, *, label_positions, token_ids=None): + """The base model's logits for the tokens at ``label_positions`` ([B, ...]) -> [B, ..., vocab]. + + Next-token logits: the ones that predict the token at position ``i`` sit at ``i - 1``. + Only those rows of the base hidden go through the final norm and lm_head, so no + full-sequence logits are built from it; ``token_ids`` ([B, ..., k]) narrows the + projection to just those vocab entries. ``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) if base_outputs.logits is not None: if token_ids is None: return _gather_rows(base_outputs.logits, positions) @@ -767,11 +935,7 @@ def _teacher_logits(self, base_outputs, positions, token_ids=None): # 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. rows = _gather_rows(base_outputs.base_hidden, positions).to(lm_head.weight.dtype) - rows = _maybe_apply_base_final_norm( - rows, - {"base_hidden_prenorm": base_outputs.base_hidden_prenorm}, - self._base_model_norm, - ) + 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]) @@ -811,35 +975,14 @@ def _compute_loss( Returns: (loss, accuracy) tuple. """ - bsz, seq_len = input_ids.shape block_size = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - device = input_ids.device - - label_offsets = torch.arange(0, block_size, device=device).view(1, 1, -1) - label_indices = anchor_positions.unsqueeze(-1) + label_offsets - valid_label = label_indices < seq_len - safe_label_indices = label_indices.clamp(max=seq_len - 1) - - target_ids = torch.gather( - input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + targets = _block_targets( + input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift ) - - # Weight mask: valid block * in bounds * exclude anchor (pos 0) * loss_mask - weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() - weight_mask = weight_mask * valid_label.float() - pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1) - weight_mask = weight_mask * (pos_in_block > 0).float() - - orig_loss_mask = torch.gather( - loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices - ) - weight_mask = weight_mask * orig_loss_mask - - binary_eval_mask = weight_mask.view(-1) + binary_eval_mask = targets.mask.view(-1) flat_logits = logits.view(-1, logits.size(-1)) - flat_targets = target_ids.view(-1) + flat_targets = targets.token_ids.view(-1) kd = self.dflash_self_logit_distillation if kd and base_outputs is None: @@ -854,37 +997,25 @@ def _compute_loss( if not kd: loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - # Block-position loss weighting: dynamic D-PACE weights or static exponential decay. - if self.dflash_loss_objective == "dpace" and block_size > 1: - # Draft confidence q_i = exp(-CE) on the target-selected token, over the - # predicted positions (slot 0 is the given anchor, already masked above). - # Weights are detached (paper Eq.9), so this adds the documented ~2.3% - # training overhead without altering the cross-entropy gradient. - with torch.no_grad(): - conf_ce = ( - loss_per_token.detach() - if loss_per_token is not None - else F.cross_entropy(flat_logits, flat_targets, reduction="none") - ).view(bsz, n_blocks, block_size) - confidences = torch.exp(-conf_ce[..., 1:].float()) - dpace = torch.ones_like(weight_mask) - dpace[..., 1:] = _dpace_position_weights( - confidences, self.dflash_dpace_alpha, valid_mask=weight_mask[..., 1:] - ) - weight_mask = weight_mask * dpace - elif self.dflash_loss_decay_factor > 0: - k = torch.arange(block_size, device=device).view(1, 1, -1) - decay = torch.exp(-(k - 1).clamp(min=0).float() / self.dflash_loss_decay_factor) - weight_mask = weight_mask * decay + weight_mask = _block_position_weights( + targets.mask, + self.dflash_loss_objective, + _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), + ) flat_weights = weight_mask.view(-1) valid_count = flat_weights.sum() + 1e-6 if valid_count > 1.0: if kd: - # KD loss: teacher logits for token anchor+k are at position anchor+k-1 - teacher_indices = (safe_label_indices - 1).clamp(min=0) - flat_teacher = self._teacher_logits(base_outputs, teacher_indices) + 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) @@ -894,12 +1025,9 @@ def _compute_loss( loss = (loss_per_token * flat_weights).sum() / valid_count with torch.no_grad(): - preds = flat_logits.argmax(dim=-1) - correct = (preds == flat_targets) & (binary_eval_mask > 0.5) - accuracy = correct.sum().float() / (binary_eval_mask.sum() + 1e-6) - accuracy = accuracy.item() + accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item() else: - loss = flat_logits.sum() * 0.0 + loss = self._draft_zero_loss(flat_logits.device, through=(flat_logits,)) accuracy = 0.0 if return_terms: @@ -1011,9 +1139,7 @@ def forward( "The base model did not return hidden states required for DFlash training. " "Ensure its top-level multimodal forward supports output_hidden_states=True." ) - offset = 1 - selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids] - target_hidden = torch.cat(selected, dim=-1) # [B, seq, num_layers * H] + target_hidden = self._target_hidden(raw_outputs.hidden_states) base_outputs = DFlashBaseModelOutput( target_hidden=target_hidden, logits=raw_outputs.logits ) @@ -1039,17 +1165,10 @@ def forward( n_blocks = anchor_positions.shape[1] if n_blocks == 0 or not block_keep_mask.any(): - # Keep all trainable draft parameters in the graph so DDP can reduce a rank - # that receives an all-masked answer-only batch. - dummy = sum( - ( - parameter.reshape(-1)[0] * 0.0 - for parameter in self.dflash_module.parameters() - if parameter.requires_grad - ), - torch.zeros((), device=device), + # An all-masked answer-only batch: nothing to train, but DDP still reduces. + return ModelOutput( + loss=self._draft_zero_loss(device), logits=base_outputs.logits, train_acc=[[0.0]] ) - return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]]) # 4. Build draft inputs noise_embedding = self._build_noise_embedding( @@ -1120,7 +1239,7 @@ def pseudo_speculative_generate(self, input_ids, steps=1): - ``noise_embedding``: embedded block tokens - ``target_hidden``: extracted hidden states from step 2 - ``position_ids``: context + block positions - - ``attention_mask=None``: no mask at inference (all positions attend freely) + - ``attention_mask``: the training rule for this one block (None when it masks nothing) The draft model's KV injection concatenates projected target_hidden as K/V with the block's own K/V, enabling the draft to "see" the target's context. @@ -1144,68 +1263,10 @@ def pseudo_speculative_generate(self, input_ids, steps=1): base_token: Next token from base model [B, 1]. draft_tokens: Draft tokens [B, min(steps, block_size-1)] or None if steps < 1. """ - if self.dflash_offline: - raise RuntimeError( - "DFlash offline model cannot run AR validation / pseudo_speculative_generate — " - "base model layers were deleted during offline conversion to save memory. " - "Reload the full base model before running AR validation." - ) - # Call the base model's inner model directly (avoids DynamicModule dispatch) - model_output = self._base_model( - input_ids=input_ids, - output_hidden_states=True, - ) - # Compute logits via lm_head - base_logits = self._base_model_lm_head(model_output.last_hidden_state) - # Build output with hidden_states - base_outputs = ModelOutput( - logits=base_logits, - hidden_states=model_output.hidden_states, - ) - base_logits = base_outputs.logits - base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device) - + base_token, hidden_states = self._generation_base(input_ids) if steps < 1: return base_token, None - - # Extract target hidden states (raw, before FC projection) - hid_offset = 1 - selected = [base_outputs.hidden_states[lid + hid_offset] for lid in self.target_layer_ids] - # target_layer_ids spans early and late layers, so under device_map="auto" these come - # off different GPUs and the cat below fails. Everything the draft consumes is gathered - # onto the draft's own device -- the last base layer's, per modify() -- which is - # also not input_ids.device once the base model is sharded. All no-ops single-device. - device = self._base_device() - target_hidden = torch.cat([h.to(device) for h in selected], dim=-1) - - block_size = self.dflash_block_size - bsz = input_ids.shape[0] - - # Block: first token is base_token (anchor), rest are mask - block_ids = torch.full( - (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device - ) - block_ids[:, 0] = base_token.squeeze(-1).to(device) - # The embedding table is on the first shard; its output follows it, not block_ids. - noise_embedding = self._base_model_embeddings(block_ids).to(device) - - # Position IDs: training uses [0..L-1, 0..L-1] where noise positions - # mirror context positions. At inference, block predicts tokens at - # seq_len..seq_len+B-1, so noise positions continue from ctx_len. - ctx_len = target_hidden.shape[1] - ctx_positions = torch.arange(ctx_len, device=device) - block_positions = torch.arange(ctx_len, ctx_len + block_size, device=device) - pos_ids = torch.cat([ctx_positions, block_positions]).unsqueeze(0).expand(bsz, -1) - - attn_mask = self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device) - - # Draft forward - draft_hidden = self.dflash_module( - noise_embedding=noise_embedding, - target_hidden=target_hidden, - position_ids=pos_ids, - attention_mask=attn_mask, - ) + draft_hidden = self._draft_block(base_token, hidden_states) # Logits on positions 1..block_size-1 (skip anchor at position 0) draft_logits = self._base_model_lm_head(draft_hidden[:, 1:, :]) @@ -1213,5 +1274,5 @@ def pseudo_speculative_generate(self, input_ids, steps=1): # Return up to `steps` tokens. base_token already follows input_ids; draft_tokens is # produced on the draft's device, and callers cat both onto the running sequence. - num_tokens = min(steps, block_size - 1) + num_tokens = min(steps, self.dflash_block_size - 1) return base_token, draft_tokens[:, :num_tokens].to(input_ids.device) diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 6060923b546..6e4171c9210 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -39,7 +39,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .hf_dflash import HFDFlashModel +from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens from .modeling_dflash2 import DFlash2Module __all__ = ["HFDFlash2Model"] @@ -72,10 +72,10 @@ def modify(self, config): "candidate selector's rank/top-k)." ) super().modify(config) - self.dflash_selector_loss_alpha = getattr(config, "dflash_selector_loss_alpha", 1.0) - self.dflash_lk_loss_type = getattr(config, "dflash_lk_loss_type", "ce") - self.dflash_lk_ce_scale = getattr(config, "dflash_lk_ce_scale", 1.0) - self.dflash_lk_ce_decay = getattr(config, "dflash_lk_ce_decay", 1.0) + self.dflash_selector_loss_alpha = config.dflash_selector_loss_alpha + self.dflash_lk_loss_type = config.dflash_lk_loss_type + self.dflash_lk_ce_scale = config.dflash_lk_ce_scale + self.dflash_lk_ce_decay = config.dflash_lk_ce_decay if self.dflash_lk_loss_type != "ce" and self.dflash_self_logit_distillation: raise ValueError( f"dflash_lk_loss_type={self.dflash_lk_loss_type!r} needs the draft's " @@ -231,45 +231,35 @@ def _compute_loss( if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: return loss, accuracy - bsz, seq_len = input_ids.shape + bsz = input_ids.shape[0] block_size = self.dflash_block_size n_blocks = anchor_positions.shape[1] - device = input_ids.device - - offsets = torch.arange(block_size, device=device).view(1, 1, -1) - label_indices = anchor_positions.unsqueeze(-1) + offsets - valid_label = label_indices < seq_len - safe_label_indices = label_indices.clamp(max=seq_len - 1) - expanded_ids = input_ids.unsqueeze(1).expand(-1, n_blocks, -1) - target_ids = torch.gather(expanded_ids, 2, safe_label_indices) - - # Same supervision mask as the backbone loss: valid block, in bounds, not the - # anchor slot, and inside the answer span. Position weighting (decay/D-PACE) is - # deliberately not applied — it shapes *where* the backbone spends capacity, - # while the selector should learn every position's transition equally. - weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() - weight_mask = weight_mask * valid_label.float() - weight_mask = weight_mask * (offsets > 0).float() - weight_mask = weight_mask * torch.gather( - loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + + # The backbone's supervision mask, without its position weighting (decay/D-PACE): + # that shapes *where* the backbone spends capacity, while the selector should learn + # every position's transition equally. + targets = _block_targets( + input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift ) # Teacher-forced predecessor of block position k is the real token at anchor+k-1. - # Only offsets 1..block_size-1 are supervised (weight_mask zeroes slot 0), so the + # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the # first supervised position, k=1, has the anchor's own token as its predecessor. # Slot 0's entry here resolves to anchor-1 and never contributes. # # This is the train/serve contract: CandidateSelector.greedy_path seeds its walk # from the anchor token, so ITS position 0 corresponds to block offset 1. A caller # that hands greedy_path the full 0..block_size-1 candidate set is off by one. - predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0)) + predecessor_ids = _block_tokens( + input_ids, anchor_positions, block_size, self._label_shift - 1 + ) selector_loss, selector_accuracy, selector_coverage = self._selector_loss( logits.reshape(bsz, n_blocks, block_size, -1), - target_ids, + targets.token_ids, draft_hidden.reshape(bsz, n_blocks, block_size, -1), predecessor_ids, - weight_mask, + targets.mask, ) self._selector_metrics = { "selector_accuracy": selector_accuracy, diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 45d4e1be6c7..efc22a47c46 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -68,7 +68,14 @@ from transformers.utils import ModelOutput from ..dflash.conversion import DominoDMRegistry -from .hf_dflash import HFDFlashModel +from .hf_dflash import ( + HFDFlashModel, + _block_position_weights, + _block_targets, + _block_tokens, + _first_predicted_position, + _masked_top1, +) from .modeling_dflash import DFlashBaseModelOutput from .modeling_domino import DominoModule @@ -101,6 +108,8 @@ class HFDominoModel(HFDFlashModel): route to it when ``dflash_architecture_config.projector_type == "domino"``. """ + _label_shift = 1 + def _build_draft_module(self, dflash_config): """Build the Domino draft module (DFlash backbone + GRU correction head).""" return DominoModule(dflash_config) @@ -121,8 +130,8 @@ def modify(self, config): # (DFlashConfig carries the two fields); updated each step by # DominoLambdaCallback. Defaults keep a single forward (e.g. unit tests) # well-defined without a scheduler. - self.dflash_lambda_base_start = getattr(config, "dflash_lambda_base_start", 1.0) - self.dflash_lambda_base_decay_ratio = getattr(config, "dflash_lambda_base_decay_ratio", 1.0) + self.dflash_lambda_base_start = config.dflash_lambda_base_start + self.dflash_lambda_base_decay_ratio = config.dflash_lambda_base_decay_ratio self._lambda_base = self.dflash_lambda_base_start if not getattr(self.dflash_module, "shift_label", True): raise NotImplementedError( @@ -151,18 +160,15 @@ def _apply_domino_head(self, hidden, base_logits, input_ids, anchor_positions, n Returns: Corrected logits [B, N*block_size, vocab]. """ - bsz, seq_len = input_ids.shape + bsz = input_ids.shape[0] bs = self.dflash_block_size - device = input_ids.device suffix_start = self.dflash_module.pure_draft_prefix_len hidden4d = hidden.reshape(bsz, n_blocks, bs, hidden.size(-1)) base4d = base_logits.reshape(bsz, n_blocks, bs, -1) # Teacher-forced previous tokens: the real token at anchor+j for j in [0, bs). - prev_offsets = torch.arange(bs, device=device).view(1, 1, -1) - prev_idx = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp(max=seq_len - 1) - prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx) + prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1) block_emb = self._base_model_embeddings(prev_ids) # [B, N, bs, H] gru_in = block_emb.reshape(bsz * n_blocks, bs, block_emb.size(-1)) @@ -187,40 +193,27 @@ def _compute_domino_loss( Uses next-token (shift_label) alignment: position k predicts the token at anchor+k+1, and position 0 is *not* excluded (unlike base DFlash). """ - bsz, seq_len = input_ids.shape - bs = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - device = input_ids.device - - # shift_label=True: label for block position k is the token at anchor+k+1. - label_offsets = torch.arange(1, 1 + bs, device=device).view(1, 1, -1) - label_indices = anchor_positions.unsqueeze(-1) + label_offsets - valid_label = label_indices < seq_len - safe_label_indices = label_indices.clamp(max=seq_len - 1) - - target_ids = torch.gather( - input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + targets = _block_targets( + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + self.dflash_block_size, + self._label_shift, ) - - # Weight mask: valid block * in bounds * loss_mask. No pos-0 exclusion. - weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, bs).float() - weight_mask = weight_mask * valid_label.float() - orig_loss_mask = torch.gather( - loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + binary_eval_mask = targets.mask.view(-1) + # dflash_loss_objective is not wired for Domino yet: it always decays. + weight_mask = _block_position_weights( + targets.mask, + "decay", + _first_predicted_position(self._label_shift), + decay_factor=self.dflash_loss_decay_factor, + dpace_alpha=self.dflash_dpace_alpha, ) - weight_mask = weight_mask * orig_loss_mask - - binary_eval_mask = weight_mask.view(-1) - - # Loss decay: exp(-k/gamma) so the first prediction (k=0) gets weight 1.0. - if self.dflash_loss_decay_factor > 0: - k = torch.arange(bs, device=device).view(1, 1, -1) - decay = torch.exp(-k.clamp(min=0).float() / self.dflash_loss_decay_factor) - weight_mask = weight_mask * decay flat_final = final_logits.reshape(-1, final_logits.size(-1)) flat_base = base_logits.reshape(-1, base_logits.size(-1)) - flat_targets = target_ids.reshape(-1) + flat_targets = targets.token_ids.reshape(-1) flat_weights = weight_mask.reshape(-1) valid_count = flat_weights.sum() + 1e-6 @@ -236,15 +229,10 @@ def _compute_domino_loss( loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss with torch.no_grad(): - eval_count = binary_eval_mask.sum() + 1e-6 - final_correct = (flat_final.argmax(dim=-1) == flat_targets) & ( - binary_eval_mask > 0.5 - ) - base_correct = (flat_base.argmax(dim=-1) == flat_targets) & (binary_eval_mask > 0.5) - accuracy = (final_correct.sum().float() / eval_count).item() - base_accuracy = (base_correct.sum().float() / eval_count).item() + accuracy = _masked_top1(flat_final, flat_targets, binary_eval_mask).item() + base_accuracy = _masked_top1(flat_base, flat_targets, binary_eval_mask).item() else: - loss = flat_final.sum() * 0.0 + loss = self._draft_zero_loss(flat_final.device, through=(flat_final,)) final_loss = loss base_loss = loss accuracy = 0.0 @@ -323,9 +311,7 @@ def forward( attention_mask=attention_mask, output_hidden_states=True, ) - offset = 1 - selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids] - target_hidden = torch.cat(selected, dim=-1) # [B, seq, num_layers * H] + target_hidden = self._target_hidden(base_out.hidden_states) # 2. Build loss mask (same convention as DFlash). if labels is not None: @@ -344,9 +330,7 @@ def forward( n_blocks = anchor_positions.shape[1] if n_blocks == 0 or not block_keep_mask.any(): - # Zero loss that still flows through all draft params for DDP sync. - dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0 - return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]]) + return ModelOutput(loss=self._draft_zero_loss(device), logits=None, train_acc=[[0.0]]) # 4. Build draft inputs (inherited helpers). noise_embedding = self._build_noise_embedding( diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 7c81f4cc478..bf0c5d004f5 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -69,7 +69,14 @@ from transformers.utils import ModelOutput from ..dflash.conversion import DSparkDMRegistry -from .hf_dflash import HFDFlashModel +from .hf_dflash import ( + HFDFlashModel, + _block_position_weights, + _block_targets, + _block_tokens, + _first_predicted_position, + _masked_top1, +) from .modeling_dflash import DFlashBaseModelOutput from .modeling_dspark import DSparkModule @@ -117,6 +124,8 @@ class HFDSparkModel(HFDFlashModel): to it when ``dflash_architecture_config.projector_type == "dspark"``. """ + _label_shift = 1 + def _build_draft_module(self, dflash_config): """Build the DSpark draft module (DFlash backbone + Markov/confidence head).""" return DSparkModule(dflash_config) @@ -132,9 +141,9 @@ def modify(self, config): super().modify(config) # Three-term loss weights (DSpark only). Defaults follow the DeepSpec recipe # (L1/TVD-dominant) so a config that only sets the head still trains sensibly. - self.dflash_ce_loss_alpha = getattr(config, "dflash_ce_loss_alpha", 0.1) - self.dflash_l1_loss_alpha = getattr(config, "dflash_l1_loss_alpha", 0.9) - self.dflash_confidence_head_alpha = getattr(config, "dflash_confidence_head_alpha", 0.0) + self.dflash_ce_loss_alpha = config.dflash_ce_loss_alpha + self.dflash_l1_loss_alpha = config.dflash_l1_loss_alpha + self.dflash_confidence_head_alpha = config.dflash_confidence_head_alpha if self.dflash_confidence_head_alpha > 0 and not self.dflash_module.use_confidence_head: raise ValueError( "dflash_confidence_head_alpha > 0 but the confidence head was not built; " @@ -158,18 +167,15 @@ def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_position Returns ``(final_logits [B, N, bs, V], confidence_logits [B, N, bs] | None)``. """ - bsz, seq_len = input_ids.shape + bsz = input_ids.shape[0] bs = self.dflash_block_size - device = input_ids.device hidden4d = hidden.reshape(bsz, n_blocks, bs, hidden.size(-1)) base4d = backbone_logits.reshape(bsz, n_blocks, bs, -1) # Teacher-forced previous token for block position k: the real token at # anchor+k (so position 0's predecessor is the anchor itself). - prev_offsets = torch.arange(bs, device=device).view(1, 1, -1) - prev_idx = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp(max=seq_len - 1) - prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx) + prev_ids = _block_tokens(input_ids, anchor_positions, bs, self._label_shift - 1) bias = self.dflash_module.compute_markov_bias(prev_ids, hidden4d) final4d = base4d + bias @@ -197,57 +203,41 @@ def _compute_dspark_loss( next-token distribution at position anchor+k (= label index - 1), read from ``base_outputs``. """ - bsz, seq_len = input_ids.shape - bs = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - device = input_ids.device vocab = final_logits.size(-1) - - # shift_label=True: label for block position k is the token at anchor+k+1. - label_offsets = torch.arange(1, 1 + bs, device=device).view(1, 1, -1) - label_indices = anchor_positions.unsqueeze(-1) + label_offsets - valid_label = label_indices < seq_len - safe_label_indices = label_indices.clamp(max=seq_len - 1) - - target_ids = torch.gather( - input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + targets = _block_targets( + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + self.dflash_block_size, + self._label_shift, ) - - # Weight mask: valid block * in bounds * loss_mask (no pos-0 exclusion). - weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, bs).float() - weight_mask = weight_mask * valid_label.float() - orig_loss_mask = torch.gather( - loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices + binary_eval_mask = targets.mask.view(-1) + # dflash_loss_objective is not wired for DSpark yet: it always decays. + weight_mask = _block_position_weights( + targets.mask, + "decay", + _first_predicted_position(self._label_shift), + decay_factor=self.dflash_loss_decay_factor, + dpace_alpha=self.dflash_dpace_alpha, ) - weight_mask = weight_mask * orig_loss_mask - - binary_eval_mask = weight_mask.view(-1) - - # Exponential position decay (exp(-k/gamma); position 0 gets weight 1). - if self.dflash_loss_decay_factor > 0: - k = torch.arange(bs, device=device).view(1, 1, -1) - decay = torch.exp(-k.float() / self.dflash_loss_decay_factor) - weight_mask = weight_mask * decay flat_final = final_logits.reshape(-1, vocab) flat_base = backbone_logits.reshape(-1, vocab) - flat_targets = target_ids.reshape(-1) + flat_targets = targets.token_ids.reshape(-1) flat_weights = weight_mask.reshape(-1) valid_count = flat_weights.sum() + 1e-6 if valid_count <= 1.0: - # Touch every draft parameter, as forward()'s early return does, so DDP with - # find_unused_parameters=False still sees the confidence head's gradient. - loss = ( - flat_final.sum() * 0.0 + sum(p.sum() for p in self.dflash_module.parameters()) * 0.0 - ) + # The confidence head is not behind final_logits. + loss = self._draft_zero_loss(flat_final.device, through=(flat_final, confidence_logits)) metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0} return loss, 0.0, metrics - # Aligned target distribution: base-model logits that predict token anchor+k+1 - # sit at position anchor+k (= label index - 1). - teacher_indices = (safe_label_indices - 1).clamp(min=0) - flat_teacher = self._teacher_logits(base_outputs, teacher_indices).reshape(-1, vocab) + # Aligned target distribution: the base model's own prediction of each label. + flat_teacher = self._teacher_logits( + base_outputs, label_positions=targets.positions + ).reshape(-1, vocab) # Term 1: cross-entropy on the corrected (final) logits. ce_per_token = F.cross_entropy(flat_final, flat_targets, reduction="none") @@ -278,12 +268,8 @@ def _compute_dspark_loss( ) with torch.no_grad(): - eval_count = binary_eval_mask.sum() + 1e-6 - keep = binary_eval_mask > 0.5 - acc = ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count - base_acc = ( - (flat_base.argmax(dim=-1) == flat_targets) & keep - ).sum().float() / eval_count + acc = _masked_top1(flat_final, flat_targets, binary_eval_mask) + base_acc = _masked_top1(flat_base, flat_targets, binary_eval_mask) # One device sync for all five scalars instead of one per .item(). acc_v, base_acc_v, ce_v, l1_v, conf_v = torch.stack( [acc, base_acc, ce_loss.detach(), l1_loss.detach(), confidence_loss.detach()] @@ -354,10 +340,8 @@ def forward( attention_mask=attention_mask, output_hidden_states=True, ) - offset = 1 - selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids] base_outputs = DFlashBaseModelOutput( - target_hidden=torch.cat(selected, dim=-1), # [B, seq, num_layers * H] + target_hidden=self._target_hidden(base_out.hidden_states), base_hidden=base_out.last_hidden_state, ) target_hidden = base_outputs.target_hidden @@ -379,9 +363,7 @@ def forward( n_blocks = anchor_positions.shape[1] if n_blocks == 0 or not block_keep_mask.any(): - # Zero loss that still flows through all draft params for DDP sync. - dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0 - return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]]) + return ModelOutput(loss=self._draft_zero_loss(device), logits=None, train_acc=[[0.0]]) # 4. Build draft inputs (inherited helpers). noise_embedding = self._build_noise_embedding( @@ -436,53 +418,17 @@ def pseudo_speculative_generate(self, input_ids, steps=1): (the anchor for k=0). Without this override DSpark would fall back to the DFlash backbone-only generate, under-reporting acceptance length. """ - if self.dflash_offline: - raise RuntimeError( - "DSpark offline model cannot run AR validation / pseudo_speculative_generate — " - "base model layers were deleted during offline conversion. Reload the full " - "base model before running AR validation." - ) - model_output = self._base_model(input_ids=input_ids, output_hidden_states=True) - base_logits = self._base_model_lm_head(model_output.last_hidden_state) - base_token = base_logits[:, -1:, :].argmax(dim=-1).to(input_ids.device) - + base_token, hidden_states = self._generation_base(input_ids) if steps < 1: return base_token, None - - hid_offset = 1 - selected = [model_output.hidden_states[lid + hid_offset] for lid in self.target_layer_ids] - # Same sharding hazard as HFDFlashModel.pseudo_speculative_generate: gather onto the - # draft's device before the cat rather than assuming one GPU holds every target layer. - device = self._base_device() - target_hidden = torch.cat([h.to(device) for h in selected], dim=-1) - - block_size = self.dflash_block_size - bsz = input_ids.shape[0] - - # Block input: anchor at position 0, mask tokens elsewhere (parallel backbone). - block_ids = torch.full( - (bsz, block_size), self.mask_token_id, dtype=torch.long, device=device - ) - block_ids[:, 0] = base_token.squeeze(-1).to(device) - noise_embedding = self._base_model_embeddings(block_ids).to(device) - - ctx_len = target_hidden.shape[1] - ctx_positions = torch.arange(ctx_len, device=device) - block_positions = torch.arange(ctx_len, ctx_len + block_size, device=device) - pos_ids = torch.cat([ctx_positions, block_positions]).unsqueeze(0).expand(bsz, -1) - - draft_hidden = self.dflash_module( - noise_embedding=noise_embedding, - target_hidden=target_hidden, - position_ids=pos_ids, - attention_mask=self._build_generate_swa_mask(ctx_len, bsz, target_hidden.dtype, device), - ) + draft_hidden = self._draft_block(base_token, hidden_states) backbone_logits = self._base_model_lm_head(draft_hidden) # [B, block_size, V] # Autoregressive Markov sampling over the block. m = self.dflash_module - num_tokens = min(steps, block_size) - prev_token = base_token.squeeze(-1).to(device) # anchor precedes block position 0 + num_tokens = min(steps, self.dflash_block_size) + # anchor precedes block position 0 + prev_token = base_token.squeeze(-1).to(draft_hidden.device) state = None draft_tokens = [] for k in range(num_tokens): diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index 42f6a8bbed5..b38234d78eb 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -77,7 +77,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import LiLiCorrDMRegistry -from .hf_dflash import HFDFlashModel +from .hf_dflash import HFDFlashModel, _block_positions, _block_targets from .modeling_lilicorr import LiLiCorrModule logger = logging.getLogger(__name__) @@ -126,7 +126,7 @@ def modify(self, config): ) super().modify(config) - weights = {name: float(getattr(config, name, -1.0)) for name in _TERM_WEIGHT_FIELDS} + weights = {name: float(getattr(config, name)) for name in _TERM_WEIGHT_FIELDS} unset = [name for name, value in weights.items() if value < 0.0] if unset: raise ValueError( @@ -137,7 +137,7 @@ def modify(self, config): self.dflash_lilicorr_w_ce = weights["dflash_lilicorr_w_ce"] self.dflash_lilicorr_w_margin = weights["dflash_lilicorr_w_margin"] self.dflash_lilicorr_w_pen = weights["dflash_lilicorr_w_pen"] - self.dflash_lilicorr_w_cal = float(getattr(config, "dflash_lilicorr_w_cal", 0.0)) + self.dflash_lilicorr_w_cal = float(config.dflash_lilicorr_w_cal) if not any(weights.values()) and self.dflash_lilicorr_w_cal == 0.0: raise ValueError( "All LiLiCorr objective weights are 0, so the reranker would never " @@ -145,7 +145,7 @@ def modify(self, config): "least one of dflash_lilicorr_w_ce / _w_margin / _w_pen / _w_cal above 0." ) - self.dflash_lilicorr_margin = float(getattr(config, "dflash_lilicorr_margin", -1.0)) + self.dflash_lilicorr_margin = float(config.dflash_lilicorr_margin) if self.dflash_lilicorr_w_margin > 0.0 and self.dflash_lilicorr_margin <= 0.0: raise ValueError( "dflash_lilicorr_w_margin > 0 requires a positive dflash_lilicorr_margin " @@ -190,46 +190,6 @@ def get_exporter(self): return LiLiCorrExporter(self) - def _block_targets(self, input_ids, anchor_positions, block_keep_mask, loss_mask): - """Return ``(target_ids [B, N, bs], slot_mask [B, N, bs-1])`` for the lattice. - - ``target_ids[b, n, j]`` is the ground-truth **vocabulary id** at absolute position - ``anchor[b, n] + j``, covering the whole block including position 0 (the anchor - token, which the caller slices off). Distinct from ``gt_indices`` in - :meth:`_compute_lilicorr_loss`, which is a position *within* a slot's k candidates. - - ``slot_mask[b, n, s]`` says whether predicted slot ``s`` — block position - ``s + 1`` — is supervised: its block is kept, its position is in bounds, and - ``loss_mask`` allows it. - - Recomputed rather than threaded out of :meth:`HFDFlashModel._compute_loss`: both - are exact functions of the anchors and the masks, so recomputing them keeps the - shared method's signature to the tensors a variant cannot derive on its own. - The slot mask is the parent's *pre-weighting* eval mask, so the two losses - supervise the same positions regardless of the block-position weighting. - """ - bsz, seq_len = input_ids.shape - block_size = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - device = input_ids.device - - label_offsets = torch.arange(0, block_size, device=device).view(1, 1, -1) - label_indices = anchor_positions.unsqueeze(-1) + label_offsets - valid_label = label_indices < seq_len - safe_label_indices = label_indices.clamp(max=seq_len - 1) - - target_ids = torch.gather( - input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices - ) - keep = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() - keep = keep * valid_label.float() - keep = keep * torch.gather( - loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices - ) - # 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. - return target_ids, keep[:, :, 1:] > 0.5 - def _compute_loss( self, logits, @@ -268,13 +228,22 @@ def _compute_loss( loss_mask, base_outputs=base_outputs, ) - target_ids, slot_mask = self._block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask + # The backbone's own supervision, before its position weighting, so both losses + # train the same positions. + targets = _block_targets( + input_ids, + anchor_positions, + block_keep_mask, + loss_mask, + self.dflash_block_size, + self._label_shift, ) lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss( logits=logits, - target_ids=target_ids, - slot_mask=slot_mask, + target_ids=targets.token_ids, + # Block position 0 is the anchor -- the one row whose input embedding is a real + # verified token -- so the lattice covers positions 1..bs-1. + slot_mask=targets.mask[:, :, 1:] > 0.5, anchor_positions=anchor_positions, draft_hidden=draft_hidden, target_hidden=target_hidden, @@ -545,14 +514,13 @@ def _candidate_target_logits( batch_blocks, num_slots, topk = candidate_ids.shape bsz, n_blocks = anchor_positions.shape seq_len = base_outputs.target_hidden.shape[1] - # The target's next-token logits that predict the token at anchor+1+s sit at - # position anchor+s. Only the [blocks, slots, k] candidate logits are computed, - # never whole vocabulary rows. - slot_offsets = torch.arange(num_slots, device=candidate_ids.device) - positions = (anchor_positions.unsqueeze(-1) + slot_offsets).clamp(min=0, max=seq_len - 1) + # Slot s is block position s + 1, labelled like the backbone's. Only the + # [blocks, slots, k] candidate logits are computed, never whole vocabulary rows. + labels, _ = _block_positions(anchor_positions, num_slots + 1, self._label_shift, seq_len) + labels = labels[..., 1:] candidate_target_logits = self._teacher_logits( base_outputs, - positions, + label_positions=labels, token_ids=candidate_ids.reshape(bsz, n_blocks, num_slots, topk), ) return candidate_target_logits.reshape(batch_blocks, num_slots, topk).float() diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index d35badb5aa8..28b1ff5a0db 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -123,21 +123,20 @@ def _select_final_norm_type(model_type: str | None, base_cfg=None) -> str | None return _FINAL_NORM_TYPE_BY_MODEL_TYPE.get(model_type or "") -def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm): - """Re-apply the base model's final norm to ``hidden`` before lm_head, per the producer. +def _apply_base_final_norm(hidden, prenorm, base_model_norm): + """Re-apply the base model's final norm to ``hidden`` before lm_head, if it was captured before it. - When the offline/streaming producer captured a *pre-final-norm* hidden it sets - ``base_hidden_prenorm=True`` in ``base_model_outputs``; the consumer must re-apply the base - final norm to reconstruct correct logits. Shared by the DFlash and EAGLE offline forwards. + When the offline/streaming producer captured a *pre-final-norm* hidden (``prenorm``), the + consumer must re-apply the base final norm to reconstruct correct logits. - - ``base_hidden_prenorm`` falsy (post-norm capture, e.g. HF/TRT-LLM): return ``hidden`` as-is. - - ``base_hidden_prenorm`` true and a norm is available: return the normed hidden. - - ``base_hidden_prenorm`` true but ``base_model_norm is None``: raise — we cannot reconstruct - correct logits, and silently feeding an un-normed hidden into lm_head would corrupt the + - ``prenorm`` false (post-norm capture, e.g. HF/TRT-LLM): return ``hidden`` as-is. + - ``prenorm`` true and a norm is available: return the normed hidden. + - ``prenorm`` true but ``base_model_norm is None``: raise — we cannot reconstruct correct + logits, and silently feeding an un-normed hidden into lm_head would corrupt the distillation target. The base ``model_type`` is not enabled in ``_FINAL_NORM_TYPE_BY_MODEL_TYPE``. """ - if not base_model_outputs.get("base_hidden_prenorm", False): + if not prenorm: return hidden if base_model_norm is None: raise RuntimeError( @@ -149,3 +148,13 @@ def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm): "matching norm class in _FINAL_NORM_CLASSES if it needs a new norm flavor." ) return base_model_norm(hidden) + + +def _maybe_apply_base_final_norm(hidden, base_model_outputs, base_model_norm): + """:func:`_apply_base_final_norm`, reading ``prenorm`` from the producer's outputs dict. + + The producer sets ``base_hidden_prenorm=True`` in ``base_model_outputs`` for a pre-norm + capture; older dumps omit it and are post-norm. + """ + prenorm = bool(base_model_outputs.get("base_hidden_prenorm", False)) + return _apply_base_final_norm(hidden, prenorm, base_model_norm) diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md index b0b72b7f45c..90d7f457369 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md @@ -160,7 +160,7 @@ DFlash-specific: | --- | --- | --- | | `seq_len (N) must be divisible by block_size (B)` | `training.training_seq_len` not a multiple of `dflash_block_size` | Adjust either value, or pad | | `DFlash offline model cannot run eval/inference forward` | Offline conversion deletes base-model layers to save memory | Don't run eval on the offline model; reload the full base first | -| `DFlash offline model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs | +| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Same cause, hit via AR validation | Keep `training.estimate_ar=false` and `training.ar_validate_steps=0` in offline runs | | `dflash_swa_window_size (N) must be >= dflash_block_size (B)` | Config validation | Raise the window or lower the block size | | `The base model did not return hidden states required for DFlash training` | Base model's top-level forward ignores `output_hidden_states=True` | Usually a multimodal wrapper — needs a model-side fix | | `ERROR: DRAFT_CKPT_DIR=... contains no exported-checkpoint-* directory` | Upstream training produced no draft | Fix training; do not chase the smoke test | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index 18482ac7613..f857d193d3f 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -152,7 +152,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `DSpark requires markov_rank > 0, got N` | Non-positive value | Set a positive rank | | `Unsupported markov_head_type: '...'. Expected 'vanilla', 'gated' or 'rnn'` | Typo or unsupported head | Use one of the three | | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 | -| `DSpark offline model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export | +| `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export | | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) | | Acceptance length capped despite clean training (streaming) | `EAGLE_CAPTURE_IDS` final layer wrong, or the vLLM aux-capture fix (vllm#46788) missing | Correct the ids; use a container with the fix | | Trainer init crash on a serve container | tensorboard absent | `training.report_to=none` | diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ff8931149e5..55d2978a7e2 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -44,9 +44,14 @@ DFlashBaseModelOutput, DFlashModule, HFDFlashModel, + _block_targets, _dpace_position_weights, build_target_layer_ids, ) +from modelopt.torch.speculative.plugins.hf_dflash2 import HFDFlash2Model +from modelopt.torch.speculative.plugins.hf_domino import HFDominoModel +from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel +from modelopt.torch.speculative.plugins.hf_lilicorr import HFLiLiCorrModel from modelopt.torch.speculative.plugins.master_weight_adamw import ( MasterWeightAdamW, VerifyMasterWeightsCallback, @@ -562,6 +567,49 @@ def test_padded_query_rows_index_in_bounds(self): assert visible.shape == (padded_q.numel(), kv.numel()) +def _legacy_generate_swa_mask(model, ctx_len, bsz, dtype, device): + """Verbatim copy of the hand-written generation mask, kept as the semantic reference.""" + window = model.dflash_swa_window_size + causal = model.dflash_draft_attention == "causal" + if window is None and not causal: + return None + block_size = model.dflash_block_size + kv_len = ctx_len + block_size + kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, -1) + q_pos_in_block = torch.arange(block_size, device=device).view(1, 1, -1, 1) + q_real_pos = ctx_len + q_pos_in_block + is_ctx = kv_idx < ctx_len + keep_ctx = is_ctx if window is None else (is_ctx & (kv_idx > q_real_pos - window)) + keep_block = ~is_ctx + if causal: + keep_block = keep_block & ((kv_idx - ctx_len) <= q_pos_in_block) + keep = keep_ctx | keep_block + attn_mask = torch.zeros(bsz, 1, block_size, kv_len, device=device, dtype=dtype) + attn_mask.masked_fill_(~keep, torch.finfo(dtype).min) + return attn_mask + + +class TestGenerateMaskRule: + """Generation evaluates the training rule for its one block, anchored at ``ctx_len``.""" + + @pytest.mark.parametrize("window", [None, BLOCK_SIZE, 2 * BLOCK_SIZE]) + @pytest.mark.parametrize("attention", ["bidirectional", "causal"]) + def test_matches_the_hand_written_mask(self, window, attention): + model = get_tiny_llama(num_hidden_layers=4) + config = get_dflash_config(block_size=BLOCK_SIZE) + config["dflash_swa_window_size"] = window + config["dflash_draft_attention"] = attention + mtsp.convert(model, [("dflash", config)]) + for ctx_len in (1, 5, 16): + args = (ctx_len, 2, torch.float32, torch.device("cpu")) + got = model._build_generate_swa_mask(*args) + want = _legacy_generate_swa_mask(model, *args) + if want is None: + assert got is None, f"ctx_len={ctx_len}" + else: + assert torch.equal(got, want), f"ctx_len={ctx_len}" + + class TestDFlashSwaMask: """Test all-layer non-causal sliding-window attention mask (MiMo-style).""" @@ -1445,8 +1493,92 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self): assert padded_acc == pytest.approx(base_acc) +# Anchors 1 and 4 of an 8-token sequence whose loss_mask is 0 only at position 3: +# (label positions, supervision mask) per block position. +_LABELS_AT_ANCHOR = ([[1, 2, 3, 4], [4, 5, 6, 7]], [[0, 1, 0, 1], [0, 1, 1, 1]]) +_LABELS_AFTER_ANCHOR = ([[2, 3, 4, 5], [5, 6, 7, 7]], [[1, 0, 1, 1], [1, 1, 1, 0]]) + + +@pytest.mark.parametrize( + ("model_cls", "expected"), + [ + (HFDFlashModel, _LABELS_AT_ANCHOR), + (HFDFlash2Model, _LABELS_AT_ANCHOR), + (HFLiLiCorrModel, _LABELS_AT_ANCHOR), + (HFDSparkModel, _LABELS_AFTER_ANCHOR), + (HFDominoModel, _LABELS_AFTER_ANCHOR), + ], + ids=["dflash", "dflash2", "lilicorr", "dspark", "domino"], +) +def test_block_targets_follow_the_label_shift(model_cls, expected): + """Each block position's label and supervision; a wrong shift trains silently worse.""" + positions, mask = expected + input_ids = 100 + torch.arange(8).unsqueeze(0) # token id = 100 + its position + loss_mask = torch.ones(1, 8) + loss_mask[0, 3] = 0.0 + targets = _block_targets( + input_ids, + torch.tensor([[1, 4]]), + torch.tensor([[True, True]]), + loss_mask, + 4, + model_cls._label_shift, + ) + assert torch.equal(targets.positions, torch.tensor([positions])) + assert torch.equal(targets.token_ids, 100 + torch.tensor([positions])) + assert torch.equal(targets.mask, torch.tensor([mask], dtype=torch.float32)) + + +# DSpark's zero losses, whose confidence head adds a second output, are tested in +# test_hf_dspark.py::TestDdpGradientCoverage. +_ZERO_LOSS_HEADS = { + "dflash": {}, + "domino": {"projector_type": "domino", "gru_hidden_dim": 32, "emb_dim": 16}, +} + + +@pytest.mark.parametrize("draft_runs", [False, True], ids=["no_block", "zero_weight"]) +@pytest.mark.parametrize("head", list(_ZERO_LOSS_HEADS)) +def test_a_step_with_nothing_to_train_reaches_every_draft_parameter(head, draft_runs): + """DDP needs each trainable draft parameter's gradient on every rank, even at zero loss. + + Once the draft has run, the loss must reach it through its output, which an FSDP2 unit + around the draft needs. + """ + model = get_tiny_llama(num_hidden_layers=4) + config = get_dflash_config() + config["dflash_architecture_config"].update(_ZERO_LOSS_HEADS[head]) + mtsp.convert(model, [("dflash", config)]) + model.train() + outputs = [] + + def keep_output(module, args, output): + output.retain_grad() + outputs.append(output) + + model.dflash_module.register_forward_hook(keep_output) + input_ids = torch.randint(1, model.config.vocab_size, (2, SEQ_LEN)) + loss_mask = torch.zeros(2, SEQ_LEN) + if draft_runs: + # Anchors are drawn from the loss mask, so each row keeps one block, anchored at 2 + # or 9, and none of its labels is in the mask. + loss_mask[:, [2, 9]] = 1.0 + out = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), loss_mask=loss_mask) + assert out.loss.item() == 0.0 + out.loss.backward() + ungraded = [ + n for n, p in model.dflash_module.named_parameters() if p.requires_grad and p.grad is None + ] + assert not ungraded, f"no gradient for {ungraded}" + assert len(outputs) == int(draft_runs) + assert all(output.grad is not None for output in outputs) + + class TestTeacherLogits: - """The base distribution is projected only at the rows, and vocab entries, asked for.""" + """The base distribution is projected only at the rows, and vocab entries, asked for. + + It is the next-token distribution, so label ``i`` reads row ``i - 1``. + """ @staticmethod def _model(): @@ -1455,9 +1587,9 @@ def _model(): return model @staticmethod - def _positions(bsz=2): + def _labels(bsz=2): torch.manual_seed(0) - return torch.randint(0, SEQ_LEN, (bsz, 3, BLOCK_SIZE)) + return torch.randint(1, SEQ_LEN, (bsz, 3, BLOCK_SIZE)) @staticmethod def _rows(full, positions): @@ -1471,45 +1603,56 @@ def _hidden(model, bsz=2): def test_matches_the_full_sequence_projection(self): model = self._model() hidden = self._hidden(model) - positions = self._positions() - got = model._teacher_logits(DFlashBaseModelOutput(None, base_hidden=hidden), positions) - want = self._rows(model._base_model_lm_head(hidden), positions) + labels = self._labels() + got = model._teacher_logits( + DFlashBaseModelOutput(None, base_hidden=hidden), label_positions=labels + ) + want = self._rows(model._base_model_lm_head(hidden), labels - 1) torch.testing.assert_close(got, want) def test_prenorm_hidden_gets_the_base_final_norm(self): model = self._model() hidden = self._hidden(model) - positions = self._positions() + labels = self._labels() outputs = DFlashBaseModelOutput(None, base_hidden=hidden, base_hidden_prenorm=True) - want = self._rows(model._base_model_lm_head(model._base_model_norm(hidden)), positions) - torch.testing.assert_close(model._teacher_logits(outputs, positions), want) + want = self._rows(model._base_model_lm_head(model._base_model_norm(hidden)), labels - 1) + torch.testing.assert_close(model._teacher_logits(outputs, label_positions=labels), want) def test_token_ids_pick_entries_of_the_full_rows(self): model = self._model() outputs = DFlashBaseModelOutput(None, base_hidden=self._hidden(model)) - positions = self._positions() - token_ids = torch.randint(0, model.config.vocab_size, (*positions.shape, 5)) - rows = model._teacher_logits(outputs, positions) + labels = self._labels() + token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5)) + rows = model._teacher_logits(outputs, label_positions=labels) torch.testing.assert_close( - model._teacher_logits(outputs, positions, token_ids), rows.gather(-1, token_ids) + model._teacher_logits(outputs, label_positions=labels, token_ids=token_ids), + rows.gather(-1, token_ids), ) def test_handed_over_logits_are_gathered_not_recomputed(self): model = self._model() logits = torch.randn(2, SEQ_LEN, model.config.vocab_size) outputs = DFlashBaseModelOutput(None, logits=logits) - positions = self._positions() - token_ids = torch.randint(0, model.config.vocab_size, (*positions.shape, 5)) - want = self._rows(logits, positions) - assert torch.equal(model._teacher_logits(outputs, positions), want) + labels = self._labels() + token_ids = torch.randint(0, model.config.vocab_size, (*labels.shape, 5)) + want = self._rows(logits, labels - 1) + assert torch.equal(model._teacher_logits(outputs, label_positions=labels), want) assert torch.equal( - model._teacher_logits(outputs, positions, token_ids), want.gather(-1, token_ids) + model._teacher_logits(outputs, label_positions=labels, token_ids=token_ids), + want.gather(-1, token_ids), ) + def test_label_zero_does_not_wrap_to_the_last_row(self): + model = self._model() + logits = torch.randn(1, SEQ_LEN, model.config.vocab_size) + outputs = DFlashBaseModelOutput(None, logits=logits) + got = model._teacher_logits(outputs, label_positions=torch.tensor([[0]])) + assert torch.equal(got[0, 0], logits[0, 0]) + def test_missing_base_distribution_raises(self): model = self._model() with pytest.raises(ValueError, match="base_model_hidden_states"): - model._teacher_logits(DFlashBaseModelOutput(None), self._positions()) + model._teacher_logits(DFlashBaseModelOutput(None), label_positions=self._labels()) def test_offline_dict_keeps_the_hidden_unprojected(self): d = { diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 2dabbe37d57..c9000273988 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -1061,6 +1061,8 @@ def test_loss_branch_with_zero_total_weight(self): final_logits, confidence_logits = model._apply_markov_head( hidden, backbone_logits, input_ids, anchors, n_blocks ) + final_logits.retain_grad() + confidence_logits.retain_grad() loss, _, _ = model._compute_dspark_loss( backbone_logits, final_logits, @@ -1073,3 +1075,5 @@ def test_loss_branch_with_zero_total_weight(self): ) loss.backward() assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}" + # Through the draft's outputs, which an FSDP2 unit around the draft needs. + assert final_logits.grad is not None and confidence_logits.grad is not None