Implementation:Hiyouga LLaMA Factory Embedding Resize

Knowledge Sources	Hiyouga_LLaMA_Factory
Domains	Tokenization, Model Architecture
Last Updated	2026-02-06 19:00 GMT

Overview

Resizes model embedding and output layers when the vocabulary changes, supporting multiple initialization strategies including noisy mean and description-based semantic initialization.

Description

This module handles the critical task of expanding a pretrained model's embedding layers when new tokens are added to the vocabulary. resize_embedding_layer pads the new vocabulary size to a multiple of 64 for hardware efficiency, then resizes both input and output embedding matrices. Three initialization strategies are available for the new token embeddings: noise_init (default) computes the mean of existing embeddings plus Gaussian noise, desc_init initializes each new token by averaging the embeddings of its textual description tokens, and desc_init_w_noise combines semantic initialization with added noise. The module is fully compatible with DeepSpeed ZeRO-3 through the use of GatheredParameters context managers, and it properly handles tied and untied weight configurations.

Usage

Use resize_embedding_layer after adding new special tokens to the tokenizer and before training begins. This is called automatically during model patching when custom tokens are detected. Provide a new_special_tokens_config dictionary mapping token strings to their descriptions for semantic initialization.

Code Reference

Source Location

• Repository: Hiyouga_LLaMA_Factory
• File: src/llamafactory/model/model_utils/embedding.py
• Lines: 1-218

Signature

def _noisy_mean_initialization(
    embed_weight: "torch.Tensor",
    num_new_tokens: int,
) -> None:
    ...

def _description_based_initialization(
    embed_weight: "torch.Tensor",
    num_new_tokens: int,
    descriptions: dict[str, str],
    tokenizer: "PreTrainedTokenizer",
    model: "PreTrainedModel",
    add_noise: bool = False,
) -> None:
    ...

def _initialize_embeddings(
    embed_weight: "torch.Tensor",
    num_new_tokens: int,
    init_method: str,
    new_special_tokens_config: Optional[dict],
    tokenizer: "PreTrainedTokenizer",
    model: "PreTrainedModel",
) -> None:
    ...

def resize_embedding_layer(
    model: "PreTrainedModel",
    tokenizer: "PreTrainedTokenizer",
    new_special_tokens_config: Optional[dict] = None,
    init_special_tokens: str = "noise_init",
) -> None:
    ...

Import

from llamafactory.model.model_utils.embedding import resize_embedding_layer

I/O Contract

Inputs

Name	Type	Required	Description
model	PreTrainedModel	Yes	The model whose embeddings will be resized
tokenizer	PreTrainedTokenizer	Yes	The tokenizer (used to determine target vocabulary size)
new_special_tokens_config	dict or None	No	Dictionary mapping new token strings to their textual descriptions for semantic initialization
init_special_tokens	str	No (default: "noise_init")	Initialization method: noise_init, desc_init, or desc_init_w_noise

Outputs

Name	Type	Description
(side effect)	None	Resizes input/output embedding matrices in-place and updates model.config.vocab_size

Usage Examples

from llamafactory.model.model_utils.embedding import resize_embedding_layer

# Basic resize with noisy mean initialization
resize_embedding_layer(model, tokenizer)

# Resize with semantic description-based initialization
descriptions = {
    "<|START_OF_SVG|>": "Marks the beginning of an SVG document",
    "<|END_OF_SVG|>": "Marks the end of an SVG document",
}
resize_embedding_layer(
    model,
    tokenizer,
    new_special_tokens_config=descriptions,
    init_special_tokens="desc_init",
)

# Resize with semantic initialization plus noise
resize_embedding_layer(
    model,
    tokenizer,
    new_special_tokens_config=descriptions,
    init_special_tokens="desc_init_w_noise",
)

Related Pages

• Hiyouga_LLaMA_Factory_Model_Loader - Model loader that triggers embedding resize during model patching
• Hiyouga_LLaMA_Factory_MoE_Config - MoE configuration that also requires DeepSpeed ZeRO-3 compatibility