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

> Learn how to expand Llama3 embedding and LM head layers for larger vocabularies. Use the provided script to resize weights, initialize new tokens, and export a PyTorch checkpoint.

- Repository: [Xinlu Lai/llama3-chinese-chat](https://github.com/crazyboym/llama3-chinese-chat)
- Tags: tutorial
- Published: 2026-02-28

---

**To expand Llama3 embedding and LM head layers for a larger vocabulary, use the [`expand_embedding_and_lmhead.py`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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:

```bash

# 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:

```python
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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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`](https://github.com/crazyboym/llama3-chinese-chat/blob/main/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.