How to Expand Llama3 Embedding and LM Head Layers for Vocabulary Extension

To expand Llama3 embedding and LM head layers for a larger vocabulary, use the expand_embedding_and_lmhead.py script to resize the weight matrices, initialize new tokens by averaging their compositional sub-token embeddings, and export a compatible PyTorch checkpoint.

Expanding the vocabulary of LLaMA 3 to support additional languages like Chinese requires resizing the model's embedding and language modeling head matrices without destroying pretrained English knowledge. The crazyboym/llama3-chinese-chat repository provides a specialized script that automates this process while preserving the semantic structure of the original embedding space. This guide explains the exact steps to expand Llama3 embedding and LM head layers using the official tooling and source code.

Understanding the Embedding and LM Head Architecture

LLaMA 3 stores token representations in two critical weight matrices that must remain synchronized during vocabulary expansion.

The Token Embedding Layer

The embedding layer (model.embed_tokens.weight) holds a matrix of shape (vocab_size, d_model) that converts integer token IDs into dense vectors fed to the first Transformer block. Each row represents a unique token's vector representation in the model's hidden dimension.

The Language Model Head

The LM head (lm_head.weight) uses an identical shape (vocab_size, d_model) to project final hidden states back to vocabulary logits during generation. In LLaMA 3, this matrix is typically tied to the embedding weights, meaning both matrices must expand simultaneously to accommodate new tokens.

The Vocabulary Expansion Process

The tools/expand_embedding_and_lmhead.py script automates the resizing and initialization through a five-step pipeline.

Loading Model Shards and Weights

The script first consolidates sharded Safetensors checkpoints into a single dictionary. It reads all shards (model-{index:05d}-of-{total:05d}.safetensors) sequentially and extracts the two target matrices (model.embed_tokens.weight and lm_head.weight) from the merged state dict.

According to the source code in tools/expand_embedding_and_lmhead.py, lines 69-83 handle the shard loading and tensor extraction, supporting any number of input shards via the num_shards parameter.

Allocating Expanded Matrices

After determining the new vocabulary size from the target tokenizer, the script allocates zero-filled tensors of shape (new_vocab_size, d_model) for both layers. It then copies the existing English embedding and LM-head rows into the corresponding positions of the new matrices, leaving rows for new tokens initialized to zero.

This allocation occurs in lines 84-90 of the expansion script, preserving the original dtype (float16 or bfloat16) to avoid precision mismatches.

Initializing New Tokens via Averaging

For each new token ID, the init_embeddings_average function (lines 11-35) performs compositional initialization:

1. Decodes the new token text using the new tokenizer
2. Re-tokenizes that text using the old tokenizer to obtain constituent sub-token IDs
3. Averages the embedding vectors (and LM-head vectors) of those original sub-tokens
4. Writes the averaged vector into the new token's position

This approach ensures that a Chinese token composed of multiple English sub-pieces receives a semantically reasonable starting point derived from its components, rather than random initialization. The function is invoked at lines 92-99 after matrix allocation.

Handling Unknown Tokens

If a new token cannot be represented by the old tokenizer's vocabulary, the script falls back to the <unk> token ID (0) and prints a warning (lines 28-31), ensuring the process continues without failure.

Running the Expansion Script

Execute the vocabulary expansion from your terminal using the following command structure:

# Install dependencies

pip install torch safetensors transformers fire matplotlib

# Run expansion

python tools/expand_embedding_and_lmhead.py \
    --old_tokenizer /path/to/llama3-original/tokenizer \
    --new_tokenizer /path/to/chinese-extended/tokenizer \
    --num_shards 4 \
    --old_model /path/to/llama3/model/shards \
    --new_model /path/to/output/expanded_model \
    --save_embedding_plots true

The --num_shards argument must match the number of Safetensors files in your original model directory. The --save_embedding_plots option triggers the draw function (lines 37-54) to generate PNG visualizations of the first 128 dimensions, allowing visual verification that new rows contain non-zero values.

Loading and Validating the Expanded Model

After expansion, load the resulting pytorch_model.bin using standard Transformers workflows:

from transformers import AutoTokenizer, AutoModelForCausalLM

# Load new tokenizer and expanded model

tokenizer = AutoTokenizer.from_pretrained("/path/to/new/tokenizer")
model = AutoModelForCausalLM.from_pretrained(
    "/path/to/output/expanded_model",
    torch_dtype="auto",
    device_map="auto"
)

# Test Chinese generation

prompt = "今天天气怎么样？"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Remember to manually update config.json in the output directory to reflect the new vocab_size before loading the model.

Critical Implementation Details

Tokenizer Alignment Requirements

The old and new tokenizers must share the same base vocabulary. The script relies on old_tokenizer(text)["input_ids"] to map new token strings back to existing IDs. If the base tokenizers differ fundamentally (e.g., different pre-tokenization rules), the averaging logic will produce nonsensical embeddings.

Checkpoint Format Conversion

The input uses Safetensors sharding, but the output is a single monolithic pytorch_model.bin file. If you require sharded outputs for distribution, use the companion tools/merge_weight.py utility to re-split the expanded checkpoint.

Data Type Preservation

New tensors inherit the dtype of the original weights. If your base model uses bfloat16, the expanded matrices maintain bfloat16 precision throughout the process.

Summary

• File location: The expansion logic resides in tools/expand_embedding_and_lmhead.py within the crazyboym/llama3-chinese-chat repository.
• Initialization method: New tokens are initialized by averaging embeddings of their compositional sub-tokens via the init_embeddings_average function.
• Matrix targets: The script modifies model.embed_tokens.weight and lm_head.weight, both shaped (vocab_size, d_model).
• Output format: Produces a single pytorch_model.bin checkpoint from sharded Safetensors inputs.
• Post-processing: You must manually update config.json to match the new vocabulary size before inference.

Frequently Asked Questions

Why initialize new embeddings by averaging existing tokens?

Averaging provides a semantically grounded starting point that respects the original embedding space. When a new Chinese token decomposes into multiple English sub-tokens (e.g., byte-pair encoding pieces), their average represents a centroid in vector space that approximates the new token's meaning while preserving the model's existing knowledge.

Can I use this script with sharded Safetensors checkpoints?

Yes. The num_shards parameter accepts any integer value, and the script sequentially loads all model-XXXXX-of-XXXXX.safetensors files into a unified dictionary before processing. The output consolidates these into a single PyTorch checkpoint.

Do I need to modify config.json after running the expansion?

Yes. The script saves the expanded weight matrices but does not automatically update the configuration file. You must manually edit config.json in the output directory to set "vocab_size": <new_size> to match your expanded tokenizer.

What happens if a new token cannot be tokenized by the old tokenizer?

The script falls back to the unknown token ID (0) and initializes that row from the <unk> token's embedding. It prints a warning to stderr indicating which tokens triggered this fallback, allowing you to audit the coverage of your vocabulary expansion.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}