LLM微调实战：为HuggingFace模型添加自定义Token的完整指南

1. 为什么需要自定义Token？

在大型语言模型(LLM)的实际应用中，我们经常遇到原始词表无法满足特定需求的情况。以下是几种典型场景：

• 领域专业术语：医学、法律等专业领域的术语可能被拆分为多个子词，影响模型理解
• 特殊功能标记：对话系统中的角色标识符（如<user>、<bot>）或推理标记（如<think>、<answer>）
• 多语言支持：为中文、日文等非拉丁语系添加更高效的token表示
• API调用标识：函数调用场景中的特殊指令标记（如<call_api:search>）

# 典型问题示例：专业术语被错误拆分
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
print(tokenizer.tokenize("troponin"))  # 输出: ['tro', '##pon', '##in']

2. 核心操作流程

2.1 Tokenizer扩展

添加新token的第一步是扩展tokenizer的词表。HuggingFace提供了两种主要方法：

from transformers import AutoTokenizer

# 加载原始tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 方法1：添加普通token
new_tokens = ["<think>", "</think>", "<answer>", "</answer>"]
num_added = tokenizer.add_tokens(new_tokens)
print(f"Added {num_added} new tokens")

# 方法2：添加特殊token
special_tokens_dict = {"additional_special_tokens": ["<USER>", "<BOT>"]}
tokenizer.add_special_tokens(special_tokens_dict)

注意：add_special_tokens会确保这些token在分词时不会被拆分，适合用于功能标记

2.2 模型适配

仅扩展tokenizer是不够的，还需要调整模型的embedding层：

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 记录原始embedding大小
original_size = model.get_input_embeddings().weight.shape[0]
print(f"Original embedding size: {original_size}")

# 调整embedding层大小
model.resize_token_embeddings(len(tokenizer))

# 验证新size
new_size = model.get_input_embeddings().weight.shape[0]
print(f"New embedding size: {new_size}")

3. 权重初始化策略

新添加token的embedding默认是随机初始化的，这可能导致训练初期不稳定。以下是几种优化方案：

3.1 均值初始化

# 计算现有token embedding的均值
existing_embeddings = model.get_input_embeddings().weight[:original_size]
mean_embedding = torch.mean(existing_embeddings, dim=0)

# 将新token初始化为均值
with torch.no_grad():
    new_embeddings = model.get_input_embeddings().weight[original_size:]
    new_embeddings.copy_(mean_embedding)

3.2 相似token复制

# 选择语义相关的参考token
reference_token = "explanation"
reference_id = tokenizer.convert_tokens_to_ids(reference_token)
reference_embedding = model.get_input_embeddings().weight[reference_id]

# 复制到新token
with torch.no_grad():
    for token in ["<think>", "</think>"]:
        token_id = tokenizer.convert_tokens_to_ids(token)
        model.get_input_embeddings().weight[token_id] = reference_embedding.clone()

3.3 描述文本平均

对于有明确语义的新token，可以用描述文本的平均embedding初始化：

def init_with_description(token, description):
    desc_ids = tokenizer.encode(description, add_special_tokens=False)
    desc_embeddings = model.get_input_embeddings().weight[desc_ids]
    mean_embedding = torch.mean(desc_embeddings, dim=0)
    
    token_id = tokenizer.convert_tokens_to_ids(token)
    with torch.no_grad():
        model.get_input_embeddings().weight[token_id] = mean_embedding

init_with_description("<answer>", "final response output")
init_with_description("</answer>", "end of response marker")

4. 完整实战案例

下面我们以实现一个带有思考链(Chain-of-Thought)标记的对话模型为例：

4.1 环境准备

# 安装必要库
pip install transformers torch datasets sentencepiece

4.2 添加推理标记

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 添加CoT相关token
cot_tokens = ["<think>", "</think>", "<answer>", "</answer>"]
tokenizer.add_tokens(cot_tokens)

# 调整模型并初始化权重
model.resize_token_embeddings(len(tokenizer))
init_with_description("<think>", "internal reasoning process")
init_with_description("</think>", "end of reasoning")

4.3 数据准备与微调

使用包含思考过程的数据集进行微调：

from datasets import load_dataset
from transformers import TrainingArguments, Trainer

dataset = load_dataset("SkunkworksAI/reasoning-0.01")

def format_with_cot(example):
    return {
        "text": f"User: {example['instruction']}\nAssistant: <think>{example['reasoning']}</think><answer>{example['output']}</answer>"
    }

dataset = dataset.map(format_with_cot)

training_args = TrainingArguments(
    output_dir="./cot_finetuned",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    num_train_epochs=3,
    logging_steps=100,
    save_strategy="epoch"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer
)

trainer.train()

4.4 效果验证

def generate_with_cot(prompt):
    input_text = f"User: {prompt}\nAssistant: <think>"
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
    
    # 首先生成思考过程
    think_outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        stopping_criteria=[lambda tokens: tokenizer.decode(tokens[0][-8:]).endswith("</think>")]
    )
    think_text = tokenizer.decode(think_outputs[0])
    
    # 然后生成回答
    answer_inputs = tokenizer(think_text + "<answer>", return_tensors="pt").to(model.device)
    final_outputs = model.generate(
        **answer_inputs,
        max_new_tokens=200,
        stopping_criteria=[lambda tokens: tokenizer.decode(tokens[0][-9:]).endswith("</answer>")]
    )
    
    return tokenizer.decode(final_outputs[0])

print(generate_with_cot("解释量子隧穿效应"))

5. 高级技巧与优化

5.1 词表扩展对比

方法	优点	缺点	适用场景
添加新token	精确控制新增词汇	增加模型大小	少量关键术语
训练子词	自动学习词表	需要大量数据	大规模领域适应
合并词表	保留原词表特性	实现复杂	多语言场景

5.2 显存优化策略

当添加大量token时，可采用以下方法减少显存占用：

# 8-bit量化
from transformers import BitsAndBytesConfig
quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0
)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quant_config
)

# 梯度检查点
model.gradient_checkpointing_enable()

5.3 混合精度训练

training_args = TrainingArguments(
    fp16=True,  # 或bf16=True
    gradient_accumulation_steps=4,
    ...
)

6. 常见问题解决

问题1：添加token后模型输出乱码

解决方案：

1. 检查是否调用了resize_token_embeddings
2. 验证新token的embedding是否合理初始化
3. 确保微调数据包含足够的新token示例

问题2：训练过程中loss波动大

解决方案：

# 降低学习率
training_args.learning_rate = 1e-6

# 使用warmup
training_args.warmup_steps = 500

问题3：长文本中新token识别错误

解决方案：

# 调整tokenizer的模型最大长度
tokenizer.model_max_length = 4096
model.config.max_position_embeddings = 4096

在实际项目中，我们发现合理初始化新token的embedding可以将模型收敛速度提升30-50%。特别是在few-shot场景下，采用语义相近token的embedding作为初始化，效果明显优于随机初始化。