本文约3600字，建议阅读7分钟
本文将介绍如何为大型语言模型(LLM)添加自定义token并进行训练，使模型能够有效地利用这些新增token。
本文将介绍如何为大型语言模型(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 AutoTokenizertokenizer = AutoTokenizer.from_pretrained(model_id)
# 定义padding token和相关参数# 这些是训练器后续所需的配置tokenizer.pad_token = ""tokenizer.pad_token_id = 128004 tokenizer.padding_side = 'right'
在添加新token前，先检查tokenizer如何处理我们计划用作自定义token的文本字符串，以便进行后续比较。我们将添加用于表示LLM输出中思考(think)和回答(answer)部分的token，总共4个token。
输出结果：
{'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("") tokenizer.add_tokens("") tokenizer.add_tokens("") tokenizer.add_tokens("")
验证新token的编码效果：
输出结果：
{'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])
输出结果：
验证成功，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.shapeprint(f"Embedding layer size before resize: {embedding_size}")lm_head_size = model.lm_head.weight.shapeprint(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.shapeprint(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|>user
Hello! assistant
Hello! How can I assist you today?
模型运行正常。接下来分析模型对新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_prob = get_token_probability(model, tokens, think_id) answer_id = tokenizer.convert_tokens_to_ids("") answer_prob = get_token_probability(model, tokens, answer_id) print(f"Probability of : {think_prob:.6f}") print(f"Probability of : {answer_prob:.6f}") print_think_answer_probabilibites_on_test()
Probability of
: 0.000000 Probability of
: 0.000000
如预期，当前模型对新token的预测概率接近零。若不进行特殊处理，模型需要更长的训练时间才能提高这些权重。为加速学习过程，我们可以将模型已有的某些高概率token的权重克隆到新token上，为学习提供更好的起点：
import torch.nn as nn# 获取模型的输入嵌入层embedding_layer = model.get_input_embeddings()
# 选择参考token：使用start_header_id tokenreference_token_id = tokenizer.convert_tokens_to_ids("")# 将参考token的嵌入权重复制到新tokenfor token in ["", "", "", ""]: 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
hink>: 0.199994 Probability of
: 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 = "%s%s"%(reasoning, answer) messages = [ {"role": "user", "content": question}, {"role":"assistant", "content": assistant_response} ] text = tokenizer.apply_chat_template(messages, tokenize=False) return {"text": text}
# 并行处理训练集和测试集import multiprocessingdata_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 })
