LLM模型添加自定义Token代码示例：为Llama 3.2模型添加思考与回答标记

本文将介绍如何为大型语言模型(LLM)添加自定义token并进行训练，使模型能够有效地利用这些新增token。以Llama 3.2模型为基础，实现了类似DeepSeek R1中think和answer标记功能的扩展方法，通过监督微调使模型学习使用这些标记进行推理过程与答案输出的区分。![]

本文聚焦于如何通过监督微调和标记示例训练模型使用新token，这类似于DeepSeek在其主要训练迭代前的"冷启动"训练阶段，不涉及RLHF或GRPO等强化学习训练方法。

环境配置

本文可以在A100 GPU的Google Colab环境中运行，但任何具备足够内存的GPU环境均可适用。我们将使用Llama-3.2-1B-instruct作为基础模型，这需要接受其服务条款并在环境中完成HuggingFace身份验证。理论上，本方法应与HuggingFace库中的大多数模型兼容。

硬件需求：约32GB GPU内存，Colab环境下运行时间约3小时。通过调整训练部分的超参数，可以适应较低GPU内存环境的需求，相关参数将在后文中详细说明。

依赖包安装

首先，安装所需的Python库：

 !pip install --upgrade transformers bitsandbytes peft accelerate datasets trl

定义实验使用的模型，使用1B参数量的Llama 3.2模型进行实验。该技术同样适用于更大规模的模型，但可能需要更长的训练时间。

 model_id = "meta-llama/Llama-3.2-1B-Instruct"

向Tokenizer添加自定义Token

首先加载并准备模型的tokenizer，同时定义必要的padding token和相关参数。

 from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 定义padding token和相关参数
# 这些是训练器后续所需的配置
tokenizer.pad_token = "<|finetune_right_pad_id|>"
tokenizer.pad_token_id = 128004
 tokenizer.padding_side = 'right'

在添加新token前，先检查tokenizer如何处理我们计划用作自定义token的文本字符串，以便进行后续比较。我们将添加用于表示LLM输出中思考(think)和回答(answer)部分的token，总共4个token。

 tokenizer("<think></think><answer></answer")

输出结果：

{'input_ids': [128000, 14023, 771, 1500, 27963, 1822, 9399, 1500, 9399], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1]}

可以看到，默认情况下tokenizer使用了8个token来表示这些文本（不包括初始的begin text token [128000]）。现在使用

add_tokens

方法添加自定义token：

 tokenizer.add_tokens("<think>")
 tokenizer.add_tokens("</think>")
 tokenizer.add_tokens("<answer>")
 tokenizer.add_tokens("</answer>")

验证新token的编码效果：

 tokenizer("<think></think><answer></answer>")

输出结果：

{'input_ids': [128000, 128256, 128257, 128258, 128259], 'attention_mask': [1, 1, 1, 1, 1]}

可以观察到，tokenizer现在仅使用4个新token对相同文本进行编码。进一步验证解码过程：

 tokenizer.decode([128256]),tokenizer.decode([128257]),tokenizer.decode([128258]),tokenizer.decode([128259])

输出结果：

(' <think>', '</think>', '<answer>', '</answer>')

验证成功，tokenizer已正确添加并处理新token的编码与解码。

加载和调整模型

虽然tokenizer已准备完毕，但模型尚未适配新token。如果直接传入新token，模型会因嵌入层缺少对应权重而报错。需要扩展模型以容纳新token，这可通过HuggingFace提供的内置函数实现，该函数会调整模型的token嵌入层大小，同时保留现有token权重。

 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 
 # 以全精度加载模型，不进行量化处理
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

调整模型大小以匹配扩展后的tokenizer：

 # 记录调整前的嵌入层和语言模型头部大小
embedding_size = model.get_input_embeddings().weight.shape
print(f"Embedding layer size before resize: {embedding_size}")
lm_head_size = model.lm_head.weight.shape
print(f"LM head size before resize: {lm_head_size}")
print("-"*10)

# 调整token嵌入层大小以适应扩展后的tokenizer
# 此操作保留现有token的训练权重，仅为新token添加权重
model.resize_token_embeddings(len(tokenizer))

# 验证调整后的大小
embedding_size = model.get_input_embeddings().weight.shape
print(f"Embedding layer size after resize: {embedding_size}")
lm_head_size = model.lm_head.weight.shape
 print(f"LM head size after resize: {lm_head_size}")

输出结果：

_Embedding layer size before resize: torch.Size([128256, 2048])

LM head size before resize: torch.Size([128256, 2048])

Embedding layer size after resize: torch.Size([128260, 2048]) LM head size after resize: torch.Size([128260, 2048])_

执行简单测试，确认模型在调整大小后仍能正常运行：

 messages = [{"role": "user", "content": "Hello!"}]
 tokens = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt")
 tokens = tokens.to(model.device)
 outputs = model.generate(tokens, max_new_tokens=100)
 decoded_outputs = tokenizer.decode(outputs[0])
 print(decoded_outputs)

部分输出内容：

< |eot_id|><|start_header_id|>user<|end_header_id|>

Hello! <|eot_id|><|start_header_id|>assistant<|end_header_id|>

Hello! How can I assist you today? <|eot_id|>

模型运行正常。接下来分析模型对新token的预测概率：

 import torch

# 辅助函数：计算模型对特定token的预测概率
def get_token_probability(model, input_tokens, target_token):
    with torch.no_grad():
        outputs = model(input_tokens)
    # 获取模型输出的logits
    logits = outputs.logits[:, -1, :]
    # 计算softmax概率
    probs = torch.softmax(logits, dim=-1)
    token_prob = probs[0, target_token]
    return token_prob

# 测试函数：分析模型对think和answer token的预测概率
def print_think_answer_probabilibites_on_test():
    question = "Why is the sky blue?"
    messages = [{"role": "user", "content": question}]
    tokens = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt")
    tokens = tokens.to(model.device)

    think_id = tokenizer.convert_tokens_to_ids("<think>")
    think_prob = get_token_probability(model, tokens, think_id)

    answer_id = tokenizer.convert_tokens_to_ids("<answer>")
    answer_prob = get_token_probability(model, tokens, answer_id)
    print(f"Probability of <think>: {think_prob:.6f}")
    print(f"Probability of <answer>: {answer_prob:.6f}")

 print_think_answer_probabilibites_on_test()

输出结果：

Probability of <think>: 0.000000 Probability of <answer>: 0.000000

如预期，当前模型对新token的预测概率接近零。若不进行特殊处理，模型需要更长的训练时间才能提高这些权重。为加速学习过程，我们可以将模型已有的某些高概率token的权重克隆到新token上，为学习提供更好的起点：

 import torch.nn as nn

# 获取模型的输入嵌入层
embedding_layer = model.get_input_embeddings()

# 选择参考token：使用start_header_id token
reference_token_id = tokenizer.convert_tokens_to_ids("<|start_header_id|>")

# 将参考token的嵌入权重复制到新token
for token in ["<think>", "</think>", "<answer>", "</answer>"]:
    token_id = tokenizer.convert_tokens_to_ids(token)
    embedding_layer.weight.data[token_id] = embedding_layer.weight.data[reference_token_id].clone()

# 再次测试新token的概率
 print_think_answer_probabilibites_on_test()

输出结果：

Probability of <think>: 0.199994 Probability of <answer>: 0.199994

现在新token的预测概率明显提高，为后续训练创造了更好的条件。

准备训练数据

模型和tokenizer扩展完成后，需要准备训练数据以教导模型如何使用新token。这里使用SkunkworksAI/reasoning-0.01数据集，这是一个包含推理过程和最终答案的数据集，适合用于训练模型区分思考过程和回答内容。

 from datasets import load_dataset

# 加载数据集，选择前10000个样本，按9:1比例划分训练集和测试集
data_set = load_dataset("SkunkworksAI/reasoning-0.01", split='train[:10000]').train_test_split(test_size=.1)

# 数据处理函数：将样本格式化为包含think和answer标记的对话格式
def create_sample_conversation(row):
    reasoning = row['reasoning']
    question = row['instruction']
    answer = row['output']
    assistant_response = "<think>%s</think><answer>%s</answer>"%(reasoning, answer)
    messages = [
        {"role": "user", "content": question},
        {"role":"assistant", "content": assistant_response}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return {"text": text}

# 并行处理训练集和测试集
import multiprocessing
data_set['train'] = data_set['train'].map(
    create_sample_conversation,
    num_proc= multiprocessing.cpu_count(),
    load_from_cache_file=False
)
data_set['test'] = data_set['test'].map(
    create_sample_conversation,
    num_proc= multiprocessing.cpu_count(),
    load_from_cache_file=False
)

# 显示数据集信息
print(data_set['train'])
 print(data_set['test'])

输出结果：

Dataset({ features: ['instruction', 'reasoning', 'output', 'reasoning_chains', 'text'], num_rows: 9000 }) Dataset({ features: ['instruction', 'reasoning', 'output', 'reasoning_chains', 'text'], num_rows: 1000 })

配置训练器

准备训练数据后，需要配置训练器。由于使用的是相对较小的模型，可以直接在单个GPU上训练完整模型，而无需使用参数高效微调方法。

 from trl import SFTConfig, SFTTrainer

# 以下参数可根据GPU内存进行调整
# 减小批量大小可降低内存需求，但仍需足够空间存储模型和梯度
samples_per_training_step = 32
batch_size = 4
gradient_accumulation_steps = int(samples_per_training_step/batch_size)

# 训练配置
training_arguments = SFTConfig(
    output_dir="./training_outputs",
    eval_strategy="steps",
    do_eval=True,
    optim="adamw_8bit",
    per_device_train_batch_size=batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    per_device_eval_batch_size=batch_size,
    log_level="debug",
    save_strategy="epoch",
    logging_steps=50,
    learning_rate=1e-5,
    eval_steps=50,
    num_train_epochs=2,
    warmup_ratio=0.1,
    lr_scheduler_type="linear",
    dataset_text_field="text",
    max_seq_length=1024,
    report_to='none'
)

# 初始化训练器
trainer = SFTTrainer(
    model=model,
    train_dataset=data_set['train'],
    eval_dataset=data_set['test'],
    processing_class=tokenizer,
    args=training_arguments
 )

执行模型训练

配置完成后，启动训练过程：

 trainer.train()

训练将需要数小时完成，共计约562个训练步骤（数据的2个epoch）。以下是训练过程中的损失函数变化：

训练完成后，检查GPU内存占用：

 # 测量训练期间的最大GPU内存使用量
 import torch
 
 max_memory = torch.cuda.max_memory_allocated() / (1024 ** 3)  # 将字节转换为GB
 print(f"Max GPU memory used: {max_memory:.2f} GB")

输出结果：

Max GPU memory used: 31.03 GB

模型评估

训练完成后，进行简单的手动测试，验证模型是否已学会使用新添加的token：

 question = "Write a Python script to check if two string variables are anagrams or not."
 messages = [{"role": "user", "content": question}]
 tokens = trainer.tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt")
 tokens = tokens.to(model.device)
 outputs = trainer.model.generate(tokens, max_new_tokens=1024)
 new_tokens = outputs[0]
 decoded_outputs = tokenizer.decode(new_tokens)
 print(decoded_outputs)

测试结果显示模型成功学会了使用think和answer token，能够分别将思考过程和最终答案包含在相应标记中。

保存和加载训练后的模型

训练完成的模型需要保存以便后续使用：

 # 保存模型和tokenizer
 final_model_path = "./model/final_model"
 final_model = trainer.model
 final_tokenizer = trainer.tokenizer
 
 final_model.save_pretrained(final_model_path)
 final_tokenizer.save_pretrained(final_model_path)

验证保存的模型能否正确加载并使用：

 # 加载保存的模型和tokenizer
loaded_model = AutoModelForCausalLM.from_pretrained(
    final_model_path,
    device_map="auto"
)
loaded_tokenizer = AutoTokenizer.from_pretrained(final_model_path)

# 验证模型结构
print(f"Embedding layer size after resize: {loaded_model.get_input_embeddings().weight.shape}")
 print(f"LM head size after resize: {loaded_model.lm_head.weight.shape}")

输出结果：

Embedding layer size after resize: torch.Size([128260, 2048]) LM head size after resize: torch.Size([128260, 2048])

对加载的模型进行测试：

 # 使用加载的模型进行推理测试
question = "Write a Python script to check if two string variables are anagrams or not."
messages = [{"role": "user", "content": question}]
tokens = loaded_tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt")
tokens = tokens.to(loaded_model.device)
outputs = loaded_model.generate(tokens, max_new_tokens=1024)
new_tokens = outputs[0][tokens.shape[-1]:]
decoded_outputs = loaded_tokenizer.decode(new_tokens)
 print(decoded_outputs)

测试结果表明，模型成功保存并正确加载，能够按预期使用自定义的think和answer标记。

总结

本文详细介绍了为HuggingFace模型添加自定义token并训练模型使用这些token的完整流程。这一技术可用于多种场景，如区分模型的思考过程与最终答案、增强模型的推理表达能力等。通过对tokenizer进行扩展、调整模型结构、准备专门的训练数据和进行有效的微调，我们成功实现了自定义token的集成与应用，为大语言模型的个性化定制提供了有效方法。

代码：https://avoid.overfit.cn/post/808effdaed4240f2a908c4323c832af8