How to Train Custom Tokens with LoRA (Fixing Untrained Embeddings)

Why the new tokens break

When we adapt a model to a new domain, it is common to extend the vocabulary with domain-specific tokens, such as special markers like <|start_header_id|> or custom XML tags. In practice this is two steps: we add the tokens with tokenizer.add_tokens() and grow the embedding table to match with model.resize_token_embeddings(len(tokenizer)).

The resize creates a new row for each added token, initialized at random. That row only becomes meaningful if training updates it. With LoRA, however, the entire base model is frozen, and only the adapters on the modules you target receive gradients. A typical configuration targets the attention projections (q_proj, v_proj, and so on), and sometimes the MLP layers, but never the embedding table. Subsequently, the new rows keep their random initialization, and the model produces nonsense whenever one of these tokens appears.

The fix for untied models

To get gradients into the new rows, we need to tell PEFT to train the input embeddings and the output head in full, alongside the LoRA adapters. This is what the modules_to_save argument is for.

from peft import LoraConfig, get_peft_model

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    # train the input embeddings and output head in full
    modules_to_save=["embed_tokens", "lm_head"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, peft_config)

This works as long as the embeddings and the head are two separate tensors. Before relying on it, it is worth checking that this is actually the case.

When the weights are tied

Many causal language models tie their weights: the input embedding (embed_tokens) and the output projection (lm_head) are in fact the same tensor, shared to reduce the parameter count. Gemma does this, as do the smaller Llama 3.2 models and most small models. It is easy to check:

print(model.config.tie_word_embeddings)  # True means they are shared

On a tied model, the configuration from the previous section becomes a trap. Passing both embed_tokens and lm_head to modules_to_save makes PEFT create a separate trainable copy of each, which breaks the tie. Training still runs, but the two tensors are now independent, and once the adapter is merged or the model is reloaded the result is broken generations. Passing only embed_tokens is not a fix either: PEFT copies the embeddings into a trainable module while lm_head continues to point at the original, frozen tensor, so the output side never learns the new tokens. Either way, the tokens remain wrong at generation time.

Fortunately, there are three reliable options for tied models:

1. Train only the new tokens with trainable_token_indices. PEFT learns a delta on just the rows you specify, and for models that follow the standard Transformers tying convention it keeps the tied head in sync automatically. This is the cheapest option and usually the right one.
2. Keep the tie while training in full by passing ensure_weight_tying=True in the LoraConfig (available in recent PEFT versions). You can then list embed_tokens and lm_head in modules_to_save, and PEFT will keep them sharing weights rather than splitting them.
3. Apply LoRA to the embeddings by adding the embedding layer to target_modules instead of saving it in full. This is useful when the embedding table is large, as in Gemma, where full fine-tuning would inflate the trainable parameter count considerably.

from peft import LoraConfig, get_peft_model

# the ids of the tokens you added to the tokenizer
new_token_ids = tokenizer.convert_tokens_to_ids(["<custom_1>", "<custom_2>"])

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    # train just the new rows; the tied head is handled for you
    trainable_token_indices=new_token_ids,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, peft_config)

Whichever option you choose, it is worth verifying the outcome directly. Read the rows corresponding to the new token ids from both the input embeddings and the output head before and after a few training steps, and confirm that they have actually changed. On a tied model, the important check is that the head moved as well. Put another way, training without errors is not the same as training the right weights, and on tied models that gap is exactly where the silent failure hides.