英文大语言模型中文适配之继续预训练

内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。

1. 为什么需要进行继续预训练？

在构建中文领域的 tokenization 之后，新增的中文词汇尚未得到训练，因此需要在指令微调之前进行继续预训练（Continued Pre-training）。

预训练的核心思想：根据上一个字预测下一个字是什么（即因果语言建模，CLM）。

本文以 IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese 模型及其自带 tokenizer 为例进行说明。

2. 如何对继续预训练数据进行预处理？

使用的数据为斗破苍穹小说数据，位于 data 目录下，分别为 corpus.txt 和 test_corpus.txt，每一行为一句或多句话。

2.1 数据预处理流程

数据预处理位于 test_dataset.py 中，主要包含以下步骤：

import os
import logging
import datasets
import transformers
from pprint import pprint
from itertools import chain
from datasets import load_dataset, concatenate_datasets
from transformers.testing_utils import CaptureLogger
from transformers import AutoTokenizer, LlamaTokenizer

tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
logger = logging.getLogger(__name__)

lm_datasets = []
files = ["data/test_corpus.txt"]
data_cache_dir = "./cache_data"
preprocessing_num_workers = 1

tokenizer = AutoTokenizer.from_pretrained("IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese")

2.2 Tokenize 函数

def tokenize_function(examples):
    with CaptureLogger(tok_logger) as cl:
        output = tokenizer(examples["text"])
        if "Token indices sequence length is longer than the" in cl.out:
            tok_logger.warning(
                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input "
                "will be chunked into smaller bits before being passed to the model."
            )
        return output

2.3 文本拼接与分块

block_size = 128

def group_texts(examples):
    # 拼接所有文本
    concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])

    # 丢弃不足 block_size 的余数部分
    if total_length >= block_size:
        total_length = (total_length // block_size) * block_size

    # 按 block_size 分块
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result

2.4 数据加载与缓存

for idx, file in enumerate(files):
    filename = ''.join(file.split(".")[:-1])
    cache_path = os.path.join(data_cache_dir, filename)
    os.makedirs(cache_path, exist_ok=True)
    try:
        processed_dataset = datasets.load_from_disk(cache_path, keep_in_memory=False)
        print(f'training datasets-{filename} has been loaded from disk')
    except Exception:
        cache_dir = os.path.join(data_cache_dir, filename + "_text")
        os.makedirs(cache_dir, exist_ok=True)
        raw_dataset = load_dataset("text", data_files=data_file, cache_dir=cache_dir, keep_in_memory=False)

        # 直接进行 tokenize，只需在句子开头加上 bos_token
        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns="text",
            load_from_cache_file=True,
            keep_in_memory=False,
            desc="Running tokenizer on dataset",
        )

        grouped_datasets = tokenized_dataset.map(
            group_texts,
            batched=True,
            num_proc=preprocessing_num_workers,
            load_from_cache_file=True,
            keep_in_memory=False,
            desc=f"Grouping texts in chunks of {block_size}",
        )
        processed_dataset = grouped_datasets
        processed_dataset.save_to_disk(cache_path)

    if idx == 0:
        lm_datasets = processed_dataset['train']
    else:
        lm_datasets = concatenate_datasets([lm_datasets, processed_dataset['train']])

lm_datasets = lm_datasets.train_test_split(test_size=0.1)

2.5 处理结果示例

字段	说明
input_ids	token 化后的输入 ID 序列，前后添加 21134（bos）和 21133（eos）
token_type_ids	token 类型标识，全为 0
attention_mask	注意力掩码，全为 1
labels	标签，与 input_ids 一致

关键点：先使用 tokenizer() 得到输入，然后将所有文本的 input_ids、attention_mask、token_type_ids 各自拼接起来，再按 block_size 分块得到最终输入。

3. 如何构建模型？

在 test_model.py 中可以初步使用预训练模型查看效果：

from transformers import BertTokenizer, GPT2LMHeadModel, AutoModelForCausalLM

hf_model_path = 'IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese'
tokenizer = BertTokenizer.from_pretrained(hf_model_path)
model = AutoModelForCausalLM.from_pretrained(hf_model_path)

def generate_word_level(input_text, n_return=5, max_length=128, top_p=0.9):
    inputs = tokenizer(input_text, return_tensors='pt', add_special_tokens=False).to(model.device)
    gen = model.generate(
        inputs=inputs['input_ids'],
        max_length=max_length,
        do_sample=True,
        top_p=top_p,
        eos_token_id=21133,
        pad_token_id=0,
        num_return_sequences=n_return
    )
    sentences = tokenizer.batch_decode(gen)
    for idx, sentence in enumerate(sentences):
        print(f'sentence {idx}: {sentence}')
        print('*' * 20)
    return gen

outputs = generate_word_level('西湖的景色', n_return=5, max_length=128)

继续预训练注意事项

注意事项	说明
词表大小重设	如果使用自定义 tokenizer，需重新设置嵌入层和 lm_head 层的词表数目
LoRA 微调参数	需设置额外保存的参数：transformer.wte、lm_head
权重保存	原始 chinese-llama-alpaca 保存 LoRA 参数有问题，需修改为只保存一份
预训练前准备	使用 test_pretrained_model.py 时也需先重设 vocab_size

4. 如何进行继续预训练？

使用 torchrun 启动训练，结合 DeepSpeed ZeRO-2 降低显存占用：

torchrun --nnodes 1 --nproc_per_node 1 run_clm_pt_with_peft.py \
  --deepspeed ds_zero2_no_offload.json \
  --model_name_or_path IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese \
  --tokenizer_name_or_path IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese \
  --dataset_dir data \
  --data_cache_dir temp_data_cache_dir \
  --validation_split_percentage 0.001 \
  --per_device_train_batch_size 32 \
  --per_device_eval_batch_size 16 \
  --do_train --seed $RANDOM \
  --fp16 \
  --max_steps 2500 \
  --lr_scheduler_type cosine \
  --learning_rate 2e-4 \
  --warmup_ratio 0.05 \
  --weight_decay 0.01 \
  --logging_strategy steps --logging_steps 10 \
  --save_strategy steps --save_total_limit 3 --save_steps 50 \
  --gradient_accumulation_steps 1 \
  --preprocessing_num_workers 8 \
  --block_size 512 \
  --output_dir output_dir --overwrite_output_dir \
  --ddp_timeout 30000 --logging_first_step True \
  --lora_rank 8 --lora_alpha 32 \
  --trainable c_attn \
  --modules_to_save transformer.wte,lm_head \
  --lora_dropout 0.05 \
  --torch_dtype float16 \
  --gradient_checkpointing \
  --ddp_find_unused_parameters False

使用 DeepSpeed ZeRO 后占用的显存会更小。

5. 如何使用继续预训练后的模型？

在 test_pretrained_model.py 中加载 LoRA 微调后的模型：

import os
import torch
from transformers import BertTokenizer, GPT2LMHeadModel, AutoModelForCausalLM
from peft import PeftModel

hf_model_path = 'IDEA-CCNL/Wenzhong2.0-GPT2-110M-BertTokenizer-chinese'
tokenizer = BertTokenizer.from_pretrained(hf_model_path)
model = AutoModelForCausalLM.from_pretrained(hf_model_path)

# 重设词表大小
model_vocab_size = model.get_output_embeddings().weight.size(0)
model.resize_token_embeddings(len(tokenizer))

# 加载 LoRA 权重
model = PeftModel.from_pretrained(model, os.path.join("output_dir", "adapter_model"), torch_dtype=torch.float32)
model.cuda()
model.eval()

def generate_word_level(input_text, n_return=5, max_length=128, top_p=0.9):
    inputs = tokenizer(input_text, return_tensors='pt', add_special_tokens=False).to(model.device)
    gen = model.generate(
        inputs=inputs['input_ids'],
        max_length=max_length,
        do_sample=True,
        top_p=top_p,
        eos_token_id=21133,
        pad_token_id=0,
        num_return_sequences=n_return
    )
    sentences = tokenizer.batch_decode(gen)
    for idx, sentence in enumerate(sentences):
        print(f'sentence {idx}: {sentence}')
        print('*' * 20)
    return gen

训练效果对比

• 未经过继续预训练：生成内容偏离小说风格，出现人物身份混乱（如将"萧炎"描述为1964年出生的真人）
• 经过继续预训练：生成内容符合小说语境，能够续写符合角色性格的对话和动作

模型确实得到了有效的训练，能够生成与训练数据风格一致的文本。