diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a0c2cc34bb1..b686478232c 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -1156,3 +1156,39 @@ def pseudo_speculative_generate(self, input_ids, steps=1): # produced on the draft's device, and callers cat both onto the running sequence. num_tokens = min(steps, self.dflash_block_size - first) return base_token, draft_tokens[:, :num_tokens].to(input_ids.device) + + def _decode_left_to_right(self, input_ids, steps): + """``pseudo_speculative_generate`` for a sequential head, drafting the block left to right. + + Block position ``k`` drafts the token after the anchor and the ``k`` tokens drafted before + it (``_label_shift = 1``): the argmax of its backbone logits plus the head's correction. + """ + base_token, hidden_states = self._generation_base(input_ids) + if steps < 1: + return base_token, None + draft_hidden = self._draft_block(base_token, hidden_states) + device = draft_hidden.device + # Under device_map="auto" a tied lm_head sits with the embedding table on the first GPU, + # not the draft's, and its output stays there. + backbone_logits = self._base_model_lm_head(draft_hidden).to(device) + + prev_token = base_token.squeeze(-1).to(device) + state = None + draft_tokens = [] + for k in range(min(steps, self.dflash_block_size)): + correction, state = self._correction_step(k, prev_token, draft_hidden[:, k, :], state) + logits = backbone_logits[:, k, :] + if correction is not None: + logits = logits + correction + prev_token = logits.argmax(dim=-1) + draft_tokens.append(prev_token) + return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device) + + def _correction_step(self, k, prev_token, hidden, state): + """A sequential head's correction to block position ``k``'s logits, and its new ``state``. + + ``prev_token`` [B] is the token decoded at ``k - 1`` (the anchor for ``k = 0``), ``hidden`` + [B, H] the draft's hidden state at ``k``, and ``state`` the head's carry, None at ``k = 0``. + The correction is [B, vocab], or None to leave the backbone's logits as they are. + """ + raise NotImplementedError diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 0a4bc9fdbd4..661810967a3 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -205,29 +205,17 @@ def pseudo_speculative_generate(self, input_ids, steps=1): before it. From ``pure_draft_prefix_len`` on, its backbone logits get the correction from the GRU state over exactly those tokens. """ - base_token, hidden_states = self._generation_base(input_ids) - if steps < 1: - return base_token, None - draft_hidden = self._draft_block(base_token, hidden_states) - device = draft_hidden.device - # Under device_map="auto" the embedding table and a tied lm_head sit on the first GPU, - # not the draft's, and their outputs stay there. - backbone_logits = self._base_model_lm_head(draft_hidden).to(device) # [B, block_size, V] + return self._decode_left_to_right(input_ids, steps) + def _correction_step(self, k, prev_token, hidden, state): + """The GRU's correction to block position ``k``; none before ``pure_draft_prefix_len``.""" m = self.dflash_module - prev_token = base_token.squeeze(-1).to(device) - gru_state = None - draft_tokens = [] - for k in range(min(steps, self.dflash_block_size)): - prev_emb = self._base_model_embeddings(prev_token).to(device) - gru_out, gru_state = m.prefix_gru(prev_emb.unsqueeze(1), gru_state) - logits = backbone_logits[:, k, :] - if k >= m.pure_draft_prefix_len: - correction = m.embed_proj(torch.cat([draft_hidden[:, k, :], gru_out[:, 0, :]], -1)) - logits = logits + correction - prev_token = logits.argmax(dim=-1) - draft_tokens.append(prev_token) - return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device) + # Under device_map="auto" the embedding table sits on the first GPU, not the draft's. + prev_emb = self._base_model_embeddings(prev_token).to(hidden.device) + gru_out, state = m.prefix_gru(prev_emb.unsqueeze(1), state) + if k < m.pure_draft_prefix_len: + return None, state + return m.embed_proj(torch.cat([hidden, gru_out[:, 0, :]], -1)), state class DominoLambdaCallback(TrainerCallback): diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 69b5fdc3f1f..e93ea5b09d4 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -193,23 +193,8 @@ 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. """ - base_token, hidden_states = self._generation_base(input_ids) - if steps < 1: - return base_token, None - draft_hidden = self._draft_block(base_token, hidden_states) - # Under device_map="auto" a tied lm_head sits with the embedding table, not the draft. - backbone_logits = self._base_model_lm_head(draft_hidden).to(draft_hidden.device) - - # Autoregressive Markov sampling over the block. - m = self.dflash_module - 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): - bias, state = m.markov_step(prev_token, draft_hidden[:, k, :], state) - tok = (backbone_logits[:, k, :] + bias).argmax(dim=-1) - draft_tokens.append(tok) - prev_token = tok - return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device) + return self._decode_left_to_right(input_ids, steps) + + def _correction_step(self, k, prev_token, hidden, state): + """Block position ``k``'s Markov bias; only the rnn head carries ``state``.""" + return self.dflash_module.markov_step(prev_token, hidden, state) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 2e2c541a090..b60d754c89b 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -269,6 +269,40 @@ def spy(*args, **kwargs): assert draft_tokens.shape == (1, 3) +class TestDSparkGenerate: + """AR validation drafts with the head that is trained and exported.""" + + @pytest.mark.parametrize("head_type", HEAD_TYPES) + def test_drafts_are_the_training_head_argmax(self, head_type): + """Greedy drafts equal the training head's argmax when it is teacher-forced on them.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dspark_config(head_type=head_type))]) + model.float().eval() + torch.manual_seed(0) + with torch.no_grad(): + # Large enough to overturn the backbone's argmax, so a draft that skips the bias or + # feeds the head the wrong prefix, hidden state or recurrent state disagrees. + for name, p in model.dflash_module.named_parameters(): + if name.startswith(("markov_", "gate_proj.", "joint_proj.")): + p.normal_() + input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN)) + + base_token, drafts = model.pseudo_speculative_generate(input_ids, steps=BLOCK_SIZE) + + with torch.no_grad(): + _, hidden_states = model._generation_base(input_ids) + batch = DraftBatch( + torch.cat([base_token, drafts], dim=1), # the anchor, then the drafts + loss_mask=None, + anchors=torch.zeros(2, 1, dtype=torch.long), + keep=None, + base_outputs=None, + hidden=model._draft_block(base_token, hidden_states), + ) + final_logits = model._draft_logits(batch).final # [B, 1 block, bs, V] + assert torch.equal(final_logits[:, 0].argmax(dim=-1), drafts) + + class TestDSparkExporter: """Test vanilla serving exports and the native gated/recurrent formats."""