vocab_swap

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.

Functions

swap_draft_vocabulary

Rebuild draft's embeddings against the base tokenizer's vocabulary.

swap_draft_vocabulary(draft, draft_tokenizer, base_tokenizer, base_vocab_size)

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.

Parameters:

base_vocab_size (int) – 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.

Return type:

dict