如何为HuggingFace模型添加新Token并仅训练其Embedding？

不用修改模型类就能适配多款LLM，核心思路是通过Tokenizer扩展新Token后，精准锁定新Token对应的Embedding参数，仅将这部分参数纳入优化器训练，其余参数全程冻结，具体步骤如下：

1. 扩展Tokenizer与模型Embedding层 #

• 先记录原始词表大小：old_vocab_size = len(tokenizer)
• 添加新Token：new_tokens = ["token1", "token2"]，执行tokenizer.add_tokens(new_tokens)
• 调整模型Embedding层大小以适配新Token：model.resize_token_embeddings(len(tokenizer))

  注：resize操作会保留原有Embedding参数，新增Token的Embedding会随机初始化

2. 冻结全局参数，仅开放新Token的Embedding训练 #

• 先冻结模型所有参数：

  for param in model.parameters():
      param.requires_grad = False
  

• 获取输入Embedding层，开放新Token对应参数的梯度更新：

  embedding_layer = model.get_input_embeddings()
  # 新Token的索引范围是[old_vocab_size, len(tokenizer)-1]
  embedding_layer.weight[old_vocab_size:].requires_grad = True
  

• 部分模型的输出头（如lm_head）与输入Embedding不共享权重，需同步开放对应新Token的参数训练：

  if hasattr(model, 'lm_head') and not model.config.tie_word_embeddings:
      lm_head = model.lm_head
      lm_head.weight[old_vocab_size:].requires_grad = True
  

3. 构建仅包含目标参数的优化器 #

• 把需要训练的参数单独传给优化器，避免冗余计算：

  trainable_params = []
  trainable_params.append(embedding_layer.weight[old_vocab_size:])
  if hasattr(model, 'lm_head') and not model.config.tie_word_embeddings:
      trainable_params.append(model.lm_head.weight[old_vocab_size:])
  
  optimizer = torch.optim.AdamW(trainable_params, lr=5e-4)
  

4. 验证训练状态 #

• 训练过程中可通过以下方式确认仅目标参数在更新：

  for name, param in model.named_parameters():
      if param.requires_grad:
          print(f"可训练参数：{name}，梯度是否存在：{param.grad is not None}")
  

内容的提问来源于stack exchange，提问作者Bipolo