resize_token_embeddings

torchtune.modules.embedding_utils.resize_token_embeddings(model: Module, num_embeddings: int) → None

Resizes the token embeddings and the final output projection layer of a TransformerDecoder model. The default init strategy is taking the mean of all the embeddings, new embeddings will be instantiated to this value.

This function modifies the model in-place.

The primary purpose is to adjust the vocabulary size of a pre-trained model. This is useful when fine-tuning a model on a dataset with a different vocabulary or when adding special tokens.

Example

>>> model = setup_model(...)
>>> tokenizer = setup_tokenizer(...)
>>> resize_token_embedding(model, tokenizer.num_embeddings)

Parameters:

• model (nn.Module) – The transformer model to modify. The model is expected to have tok_embeddings (an nn.Embedding layer) and output (e.g., nn.Linear or TiedLinear) attributes.

• num_embeddings (int) – The desired number of embeddings in the resized embedding layer and output projection layer.

Returns:

The function modifies the model in-place.

Return type:

None

Raises:

AssertionError – When trying to resize a model with FusionEmbedding.