diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index f03ec0efa23..142d058ebb5 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -59,6 +59,7 @@ ) from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.speculative.config import ExternalDraftConfig +from modelopt.torch.speculative.external.vocab_swap import swap_draft_vocabulary from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel from modelopt.torch.speculative.plugins.hf_domino import DominoLambdaCallback from modelopt.torch.speculative.plugins.hf_external import SPARSE_CAPABLE_LOSSES @@ -232,12 +233,38 @@ def setup_external_draft( trust_remote_code=recipe.model.trust_remote_code, ) external_cfg: dict = recipe.external.model_dump() - # The draft's own tokenizer; the base's is used only for its lm_head. + # The draft's own tokenizer, needed before conversion if its vocabulary has to be + # re-indexed onto the base's. tokenizer = transformers.AutoTokenizer.from_pretrained( recipe.draft_model_name_or_path, model_max_length=training_args.training_seq_len, trust_remote_code=recipe.model.trust_remote_code, ) + if base_tokenizer is None: + base_tokenizer = transformers.AutoTokenizer.from_pretrained( + recipe.model.model_name_or_path, + model_max_length=training_args.training_seq_len, + trust_remote_code=recipe.model.trust_remote_code, + ) + # Gate on the flag, not on a size difference: different tokenizers are routinely padded + # to the same width. Must precede convert so the saved ModelOpt state describes the + # post-swap architecture. + if external_cfg["external_vocab_swap"]: + # The checkpoint's saved config says whether its rows are already base-indexed; + # swapping those again would scramble them while every check below still passes. + saved = dict(modelopt_state(draft)["modelopt_state_dict"]).get("external") + if saved is not None and saved["config"].get("external_vocab_swap"): + print_rank_0("[external] draft is already swapped; keeping its embeddings as saved") + else: + stats = swap_draft_vocabulary(draft, tokenizer, base_tokenizer, base_vocab_size) + print_rank_0( + f"[external] vocabulary swapped onto the base: " + f"{stats['matched']}/{stats['total']} rows carried over, " + f"{stats['frequent_coverage']:.1%} of the first 50k ids" + ) + # The draft's rows are indexed by the base vocabulary from here on, so train and + # validate it against the base tokenizer. + tokenizer = base_tokenizer # Sparse data carries only the teacher's truncated policy. Objectives that need full # teacher logits would silently train against a zero-padded distribution, so gate on # the list the plugin exports. @@ -270,11 +297,6 @@ def setup_external_draft( # with no acceptance. Equal sizes are not enough -- different tokenizers are # routinely padded to the same width -- so also check that the two agree on the ids. draft.validate_against_base(base_vocab_size) - if base_tokenizer is None: - base_tokenizer = transformers.AutoTokenizer.from_pretrained( - recipe.model.model_name_or_path, - trust_remote_code=recipe.model.trust_remote_code, - ) draft.validate_tokenizer_against_base(tokenizer, base_tokenizer) # The norm is not always at base_model.model.norm -- FakeBaseModel keeps it at .norm # -- so search the same paths the other offline modes use. It stays optional; diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index d6b42f4e1b0..2ce3a75f305 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -490,6 +490,17 @@ class ExternalDraftConfig(ModeloptBaseConfig): ), ) + external_vocab_swap: bool = ModeloptField( + default=False, + description=( + "Re-index the draft onto the base model's vocabulary when the two tokenizers " + "differ, instead of raising. Rebuilds the draft's embeddings, carrying over " + "rows for tokens both tokenizers spell identically. The replacement matrix is " + "``base_vocab_size * hidden_size`` parameters, so this is a large change to a " + "small draft; validate with an acceptance measurement, not a loss curve." + ), + ) + external_tvd_alpha: float = ModeloptField( default=0.9, ge=0.0, diff --git a/modelopt/torch/speculative/external/__init__.py b/modelopt/torch/speculative/external/__init__.py index d7379c45d76..2f4e1ed728b 100644 --- a/modelopt/torch/speculative/external/__init__.py +++ b/modelopt/torch/speculative/external/__init__.py @@ -19,3 +19,4 @@ from .conversion import * from .external_model import * from .sparse_data import * +from .vocab_swap import * diff --git a/modelopt/torch/speculative/external/vocab_swap.py b/modelopt/torch/speculative/external/vocab_swap.py new file mode 100644 index 00000000000..97203f145f6 --- /dev/null +++ b/modelopt/torch/speculative/external/vocab_swap.py @@ -0,0 +1,183 @@ +# 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. + +"""Re-index a pretrained draft onto the base model's vocabulary. + +Speculative decoding requires the draft to propose tokens the target can verify, +so the two must share a vocabulary. A draft trained with a different tokenizer +can still be used once its input embeddings and lm_head are rebuilt against the +base's vocabulary: rows for tokens both tokenizers spell identically are carried +over, and the rest are initialised from the draft's own embedding statistics. + +The replacement matrix is ``base_vocab_size * hidden_size`` parameters, a large +fraction of a small draft, and the swap costs acceptance. Prefer a +vocabulary-matched draft where one exists, and judge the result by an acceptance +measurement rather than the loss curve. +""" + +import torch + +__all__ = ["swap_draft_vocabulary"] + + +def _retarget_special_tokens(draft, base_tokenizer): + """Point the draft's special-token ids at the base tokenizer's ids for the same tokens. + + Only ever overwrites with an id that resolved. Writing ``None`` through would be the + damaging case: a base with no ``pad_token`` is ordinary, and ``eos_token_id = None`` + on the exported draft means it never signals end-of-sequence, so every consumer + generates to ``max_new_tokens``. + """ + unk_id = getattr(base_tokenizer, "unk_token_id", None) + for attr in ("bos_token", "eos_token", "pad_token", "unk_token"): + token = getattr(base_tokenizer, attr, None) + if not isinstance(token, str): + continue + # convert_tokens_to_ids, not get_vocab: for several fast tokenizers get_vocab() + # omits added tokens, and chat-template markers such as <|eot_id|> are added + # tokens -- looking those up in the raw vocab yields None and erases a real id. + new_id = base_tokenizer.convert_tokens_to_ids(token) + if new_id is None or (attr != "unk_token" and unk_id is not None and new_id == unk_id): + continue + for target in (draft.config, getattr(draft, "generation_config", None)): + if target is not None and hasattr(target, f"{attr}_id"): + setattr(target, f"{attr}_id", new_id) + + +def _fill_from_statistics(old_weight, rows: int, seed: int = 0): + """Rows drawn from ``old_weight``'s per-column statistics, in fp32 and seeded. + + ``seed`` must differ between the embeddings and the head. Both draws request the + same shape, so a shared seed hands them an identical normal stream and every + unmatched row gets ``head_i . emb_i`` positive in expectation -- a systematically + raised self-logit on exactly the ids the draft has learned nothing about. Fixing + the seed is what buys cross-rank determinism; sharing it is not. + """ + stats = old_weight.float() + mean, std = stats.mean(dim=0), stats.std(dim=0).clamp_min(1e-6) + generator = torch.Generator(device=old_weight.device).manual_seed(seed) + draw = torch.normal( + mean.expand(rows, stats.shape[1]), + std.expand(rows, stats.shape[1]), + generator=generator, + ) + return draw.to(old_weight.dtype) + + +def swap_draft_vocabulary(draft, draft_tokenizer, base_tokenizer, base_vocab_size: int) -> dict: + """Rebuild ``draft``'s embeddings against the base tokenizer's vocabulary. + + Must run *before* ``mtsp.convert`` so the converted module and the saved + ModelOpt state describe the post-swap architecture; doing it afterwards leaves + the checkpoint disagreeing with its config on ``vocab_size``. + + Args: + base_vocab_size: taken from the base model's config rather than + ``len(base_tokenizer)``, which disagrees whenever the config pads the + embedding table. + + Returns: + ``matched``, ``total``, and the fraction of the base's most frequent ids that + carried over. These describe the swap, not how well the swapped draft will + accept. + """ + old_embed = draft.get_input_embeddings() + old_weight = old_embed.weight.data + + base_vocab = base_tokenizer.get_vocab() + draft_vocab = draft_tokenizer.get_vocab() + + # Fill unmatched rows from the draft's own distribution rather than zeros: a + # zero row is a uniform logit over the hidden dimension and trains slowly. + # The statistics are accumulated in fp32: a bf16 sum over tens of thousands of + # rows saturates, and the std (a sum of squares) worst of all. The draw is seeded + # so every rank fills the unmatched rows identically. + new_weight = _fill_from_statistics(old_weight, base_vocab_size) + + # The output head is a different learned matrix from the input embeddings when the + # two are untied, so it has to be rebuilt from its own rows -- copying embeddings + # over it would discard the pretrained head entirely. + old_head = draft.get_output_embeddings() + untied = old_head is not None and not draft.config.tie_word_embeddings + if untied: + # Built from the head's own row width, not the embeddings': they coincide on + # standard drafts, but assuming it turns a resized head into a confusing + # shape error at the copy below rather than a named failure here. + old_head_weight = old_head.weight.data + new_head_weight = _fill_from_statistics(old_head_weight, base_vocab_size, seed=1) + # A head with a bias has one scalar per token. resize keeps those at their old + # ids while the rows move, so a matched token would inherit another token's + # bias. Llama and Qwen heads carry none, which is why this is easy to miss. + old_head_bias = None if old_head.bias is None else old_head.bias.data + if old_head_bias is not None: + new_head_bias = _fill_from_statistics( + old_head_bias.unsqueeze(-1), base_vocab_size, seed=2 + ).squeeze(-1) + + # Collect the pairs first, then move the rows in two gathers. A per-token tensor + # assignment costs a kernel launch each, which at a 128k base vocabulary is 128k of + # them before training starts. + pairs = [ + (base_id, draft_id) + for token, base_id in base_vocab.items() + if base_id < base_vocab_size + and (draft_id := draft_vocab.get(token)) is not None + and draft_id < old_weight.shape[0] + ] + matched = len(pairs) + frequent_matched = sum(1 for base_id, _ in pairs if base_id < 50000) + if pairs: + device = old_weight.device + base_ids = torch.tensor([b for b, _ in pairs], dtype=torch.long, device=device) + draft_ids = torch.tensor([d for _, d in pairs], dtype=torch.long, device=device) + new_weight.index_copy_(0, base_ids, old_weight.index_select(0, draft_ids)) + if untied: + # The head can be narrower than the embeddings on a draft whose head was + # resized on its own, so it keeps its own bound rather than reusing theirs. + keep = draft_ids < old_head_weight.shape[0] + new_head_weight.index_copy_( + 0, base_ids[keep], old_head_weight.index_select(0, draft_ids[keep]) + ) + if old_head_bias is not None: + new_head_bias.index_copy_( + 0, base_ids[keep], old_head_bias.index_select(0, draft_ids[keep]) + ) + + # resize_token_embeddings rewires config.vocab_size and re-ties the lm_head; + # doing that by hand silently leaves a tied head pointing at the old matrix. + # mean_resizing=False because every row it would compute is overwritten below, + # and its multivariate-normal fit dominates the cost of the swap. + draft.resize_token_embeddings(base_vocab_size, mean_resizing=False) + draft.get_input_embeddings().weight.data.copy_(new_weight) + if untied: + draft.get_output_embeddings().weight.data.copy_(new_head_weight) + if old_head_bias is not None: + draft.get_output_embeddings().bias.data.copy_(new_head_bias) + + # The special-token ids still index the draft's old vocabulary. The checkpoint is + # saved with the base tokenizer, so leaving them would pair base tokenizer files + # with draft ids and every consumer would stop on the wrong token. + _retarget_special_tokens(draft, base_tokenizer) + # resize keeps the old padding_idx when it is still in range, so it would point at + # whatever the base spells at the draft's former pad id. + embedding = draft.get_input_embeddings() + if getattr(embedding, "padding_idx", None) is not None: + embedding.padding_idx = draft.config.pad_token_id + + return { + "matched": matched, + "total": base_vocab_size, + "frequent_coverage": frequent_matched / min(50000, base_vocab_size), + } diff --git a/modelopt/torch/speculative/plugins/hf_external.py b/modelopt/torch/speculative/plugins/hf_external.py index 77f27a4027f..4ae91987e8f 100644 --- a/modelopt/torch/speculative/plugins/hf_external.py +++ b/modelopt/torch/speculative/plugins/hf_external.py @@ -235,7 +235,9 @@ def validate_against_base(self, base_vocab_size: int) -> None: if draft_vocab_size != int(base_vocab_size): raise ValueError( "External draft and base model must share a vocabulary: draft " - f"vocab_size={draft_vocab_size}, base vocab_size={int(base_vocab_size)}." + f"vocab_size={draft_vocab_size}, base vocab_size={int(base_vocab_size)}. " + "Set external_vocab_swap=true to re-index the draft onto the base " + "vocabulary, which rebuilds its embeddings and lm_head." ) def validate_tokenizer_against_base(self, draft_tokenizer, base_tokenizer): diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index c73ead5659f..f0fb2c6113d 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -70,8 +70,8 @@ def _tiny_causal_lm(vocab_size=VOCAB, hidden=HIDDEN): class _FakeTok: """Enough of a tokenizer for the vocabulary checks; both sides share one vocab.""" - def __init__(self): - self._vocab = {chr(i): i for i in range(VOCAB)} + def __init__(self, vocab=None): + self._vocab = {chr(i): i for i in range(VOCAB)} if vocab is None else dict(vocab) def get_vocab(self): return dict(self._vocab) @@ -245,6 +245,30 @@ def _from_pretrained(cls, *a, **k): assert calls == ["draft"], f"base tokenizer re-loaded: {calls}" +# None: a fresh draft; otherwise the swap flag the checkpoint was trained with +@pytest.mark.parametrize("saved_swap", [None, False, True]) +def test_vocab_swap_skips_only_an_already_swapped_draft(main_mod, monkeypatch, saved_swap): + """Swapping a base-indexed draft again scrambles its trained rows; equal widths hide it.""" + draft_tok = _FakeTok({chr(i): VOCAB - 1 - i for i in range(VOCAB)}) # reversed ids + monkeypatch.setattr( + main_mod.transformers.AutoTokenizer, + "from_pretrained", + classmethod(lambda cls, name, *a, **k: draft_tok if name == "draft" else _FakeTok()), + ) + recipe = _recipe(offline_data_path="/tmp/hidden") + recipe.external.external_vocab_swap = True + draft = _tiny_causal_lm() + if saved_swap is not None: + mtsp.convert(draft, [("external", {"external_vocab_swap": saved_swap})]) + before = draft.get_input_embeddings().weight.data.clone() + + main_mod.setup_external_draft( + recipe, _training_args(), base_model=_tiny_causal_lm(hidden=2 * HIDDEN), draft=draft + ) + after = draft.get_input_embeddings().weight.data + assert torch.equal(after, before if saved_swap else before.flip(0)) + + def test_resuming_a_different_speculative_mode_is_refused(main_mod, stub_tokenizers, tmp_path): """A checkpoint from another speculative mode must fail with a clear message. diff --git a/tests/unit/torch/speculative/test_external_draft_vocab_swap.py b/tests/unit/torch/speculative/test_external_draft_vocab_swap.py new file mode 100644 index 00000000000..13c0d9671ca --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft_vocab_swap.py @@ -0,0 +1,149 @@ +# 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. + +"""Tests for re-indexing an external draft onto the base model's vocabulary.""" + +import pytest +import torch + +from modelopt.torch.speculative.external.vocab_swap import swap_draft_vocabulary + +transformers = pytest.importorskip("transformers") + +HIDDEN = 32 + + +class _FakeTok: + def __init__(self, vocab, specials=None, added=None): + self._v = vocab + # added tokens are deliberately absent from get_vocab(), as they are for several + # fast tokenizers -- that is the case that used to erase a real id + self._added = added or {} + for attr, tok in (specials or {}).items(): + setattr(self, attr, tok) + + def get_vocab(self): + return self._v + + def convert_tokens_to_ids(self, token): + if token in self._added: + return self._added[token] + return self._v.get(token) + + +def _causal_lm(hidden=HIDDEN, vocab_size=64): + cfg = transformers.LlamaConfig( + vocab_size=vocab_size, + hidden_size=hidden, + intermediate_size=2 * hidden, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=128, + ) + return transformers.LlamaForCausalLM(cfg) + + +def test_vocab_swap_carries_over_matching_tokens_and_reports_coverage(): + """Rows for tokens both tokenizers spell the same must survive the swap, and the + frequent-id coverage is the number worth looking at.""" + draft = _causal_lm(hidden=HIDDEN, vocab_size=8) + before = draft.get_input_embeddings().weight.data.clone() + stats = swap_draft_vocabulary( + draft, + _FakeTok({"a": 0, "b": 1, "c": 2, "zz": 7}), + _FakeTok({"c": 0, "a": 1, "q": 2, "b": 3}), + 4, + ) + + assert stats["matched"] == 3 # a, b, c; 'q' has no draft row + assert stats["frequent_coverage"] == 0.75 + assert draft.config.vocab_size == 4 + after = draft.get_input_embeddings().weight.data + assert torch.allclose(after[0], before[2]) # base 'c' <- draft id 2 + assert torch.allclose(after[1], before[0]) # base 'a' <- draft id 0 + assert torch.allclose(after[3], before[1]) # base 'b' <- draft id 1 + assert not torch.allclose(after[2], before[0]) # 'q' freshly initialised + + +def test_untied_lm_head_is_rebuilt_from_its_own_rows(): + """An untied output head must carry over *head* rows, not input-embedding rows.""" + draft = _causal_lm(vocab_size=8) + draft.config.tie_word_embeddings = False + # Make the head plainly distinct from the embeddings so a mix-up is visible. + draft.get_output_embeddings().weight.data.normal_(mean=5.0, std=0.01) + draft.lm_head.bias = torch.nn.Parameter(torch.arange(8.0)) # Llama has none; others do + head_before = draft.get_output_embeddings().weight.data.clone() + bias_before = draft.get_output_embeddings().bias.data.clone() + embed_before = draft.get_input_embeddings().weight.data.clone() + + swap_draft_vocabulary(draft, _FakeTok({"a": 0, "b": 1}), _FakeTok({"b": 0, "a": 1}), 4) + + head_after = draft.get_output_embeddings().weight.data + assert torch.allclose(head_after[0], head_before[1]) # base 'b' <- draft head row 1 + assert torch.allclose(head_after[1], head_before[0]) # base 'a' <- draft head row 0 + assert not torch.allclose(head_after[0], embed_before[1]), ( + "the output head must not be overwritten with input-embedding rows" + ) + # the bias is one scalar per token, so it has to move with the row it belongs to + bias_after = draft.get_output_embeddings().bias.data + assert bias_after[0] == bias_before[1] and bias_after[1] == bias_before[0], ( + "a matched token kept another token's bias" + ) + + +def test_unmatched_embedding_and_head_rows_are_independent(): + """Embedding and head fills must not share a normal stream.""" + draft = _causal_lm(vocab_size=4) # LlamaConfig defaults to an untied head + # identical statistics, so a shared seed would draw identical rows + draft.get_output_embeddings().weight.data.copy_(draft.get_input_embeddings().weight.data) + swap_draft_vocabulary(draft, _FakeTok({"x": 0}), _FakeTok({"x": 0, "q": 1, "r": 2}), 4) + + emb = draft.get_input_embeddings().weight.data + head = draft.get_output_embeddings().weight.data + assert not torch.allclose(emb[1:], head[1:]), "head and embedding fills share a stream" + + +def test_unmatched_rows_are_reproducible_across_calls(): + """The draw is seeded, so every rank fills unmatched rows identically.""" + a = _causal_lm(vocab_size=4) + b = _causal_lm(vocab_size=4) + b.get_input_embeddings().weight.data.copy_(a.get_input_embeddings().weight.data) + tok_d, tok_b = _FakeTok({"x": 0}), _FakeTok({"x": 0, "q": 1}) + swap_draft_vocabulary(a, tok_d, tok_b, 4) + swap_draft_vocabulary(b, tok_d, tok_b, 4) + assert torch.allclose( + a.get_input_embeddings().weight.data, b.get_input_embeddings().weight.data + ) + + +def test_special_tokens_are_never_overwritten_with_none(): + """A base without a pad_token, or whose eos is an added token, must not erase the draft's. + + ``eos_token_id = None`` on the exported draft means it never signals end-of-sequence, + so every consumer generates to ``max_new_tokens``. + """ + draft = _causal_lm(vocab_size=4) + draft.config.eos_token_id = 3 + draft.config.pad_token_id = 2 + base = _FakeTok( + {"a": 0, "b": 1}, + # eos is an added token; pad is set but unknown to the base + specials={"eos_token": "<|eot_id|>", "pad_token": ""}, + added={"<|eot_id|>": 1}, + ) + swap_draft_vocabulary(draft, _FakeTok({"a": 0}), base, 4) + assert draft.config.eos_token_id == 1, "an added-token eos must resolve, not erase" + assert draft.config.pad_token_id == 2, "an unresolved pad_token must leave the draft's"