大语言模型词表裁剪方法与实践

大语言模型词表裁剪方法与实践

前言

在多语言大语言模型（LLM）的应用场景中，原始模型的词表（Vocabulary）往往非常庞大，包含了全球多种语言的字符和子词。然而在实际的下游任务中，我们可能只需要支持特定的语言，例如仅中文和英文。此时，对词表进行裁剪是一个有效的优化手段。

通过裁剪词表，我们可以显著减少模型的参数量，降低显存占用，同时保留模型在目标语言上的性能表现。本文将基于 Bloom 模型为例，详细介绍如何进行词表裁剪的操作流程、代码实现及注意事项。

词表裁剪原理

大语言模型的词表大小直接决定了嵌入层（Embedding Layer）和语言模型头（LM Head）的维度。假设原词表大小为 $V_{old}$，隐藏层维度为 $H$，则嵌入层的参数量为 $V_{old} \times H$，LM Head 的参数量为 $H \times V_{new}$（通常 $V_{new} = V_{old}$）。

当我们将词表裁剪至 $V_{new}$ 时，我们需要从原模型中提取对应于新词表中 token 的权重向量，并构建新的模型结构。核心逻辑如下：

1. Token 映射：建立新词表 Token ID 到原词表 Token ID 的映射关系。
2. 权重提取：根据映射关系，从原 Embedding 矩阵和 LM Head 矩阵中复制对应的行向量。
3. 配置更新：更新模型配置中的 vocab_size 字段。
4. 一致性验证：确保裁剪后的模型在相同输入下生成的文本与原模型一致。

准备工作

首先，我们需要准备两个关键组件：

1. 原模型：包含完整多语言词表的预训练模型（如 bigscience/bloom-560m）。
2. 新词表：经过筛选的、仅包含目标语言 Token 的 tokenizer 文件（通常是 tokenizer.json 或 vocab.txt）。

注意：新词表必须是原词表的子集，即所有新词表中的 Token 都必须存在于原词表中，否则会导致索引越界错误。

代码实现

以下是一个完整的词表裁剪类实现，基于 Hugging Face Transformers 库。

1. 基础裁剪类

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm

class VocabularyPruner(object):
    def check(self, old_model_name_or_path, new_model_name_or_path, text):
        """
        检查模型裁剪后，生成结果是否一致
        """
        max_length = 20

        # 使用老模型对文本编码
        old_model = AutoModelForCausalLM.from_pretrained(old_model_name_or_path)
        old_tokenizer = AutoTokenizer.from_pretrained(old_model_name_or_path)
        old_input_ids = old_tokenizer(text, return_tensors=).input_ids
        old_output = old_model.generate(old_input_ids, max_length=max_length)
        old_output_text = old_tokenizer.batch_decode(old_output)
        (.(old_output_text))

        
        new_model = AutoModelForCausalLM.from_pretrained(new_model_name_or_path)
        new_tokenizer = AutoTokenizer.from_pretrained(new_model_name_or_path)
        new_input_ids = new_tokenizer(text, return_tensors=).input_ids
        new_output = new_model.generate(new_input_ids, max_length=max_length)
        new_output_text = new_tokenizer.batch_decode(new_output)
        (.(new_output_text))

         old_output_text == new_output_text:
            ()
        :
            ()

     ():
         NotImplementedError

     ():
        
          os.path.exists(save_path):
            os.makedirs(save_path)
        
        
        new_tokenizer = AutoTokenizer.from_pretrained(new_tokenizer_name_or_path)
        
        old_tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

        
        old_vocab = old_tokenizer.vocab
        new_vocab = new_tokenizer.vocab
         token  tqdm(new_vocab.keys()):
             token   old_vocab:
                 Exception(.(token))
        ()

        
        new2old_token_id = {}
         token, token_id  tqdm(new_vocab.items()):
            old_token_id = old_vocab[token]
            new2old_token_id[token_id] = old_token_id

        
        model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=)
        
        old_params = (p.numel()  p  model.parameters())
        ( % (old_params / ))

        
        vocab_size = (new_tokenizer)
        hidden_size = model.config.hidden_size

        new_embeds = torch.nn.Embedding(vocab_size, hidden_size, dtype=model.dtype)
        new_lm_head = torch.nn.Linear(in_features=hidden_size, out_features=vocab_size, bias=, dtype=model.dtype)
        
        
        .update_embeddings(model, new2old_token_id, new_embeds, new_lm_head)

        model.config.__dict__[] = vocab_size
         new_name_or_path   :
            model.config.__dict__[] = new_name_or_path

        
        new_params = (p.numel()  p  model.parameters())
        ( % (new_params / ))

        (.(((new_tokenizer) / (old_tokenizer), )*))
        (.((new_params / old_params, )*))
        
        model.save_pretrained(save_path)
        new_tokenizer.save_pretrained(save_path)