中文GPT2预训练实战

GPT2是一个很好的长文本生成模型，但官方版本并没有开源中文预训练好的模型。因此，最近用开源的中文新闻，wiki，评论等从头训练了一个中文GPT2用于文本生成任务。

预训练使用的是HuggingFace的transformers库，这库是个好东西，把当前主流的transfomer-based模型都封装了一遍，使用起来方便很多。但由于不同模型的结构、参数等等细节不同，封装成统一的interface还是有难度，因此此库上也有一些折衷，也并不像想像中那么好使。就pretrain和fine-tune来说，都是训练一个language model，理论上调用它examples/run_language_modeling.py即可。不过真训练起来发现还是有些坑：

• 中文Tokenizer用什么？
• 大规模语料爆内存
• fp16混合精度训练？

最后一个问题坑有点深，nvidia的apex库用起来并不是那么顺滑，尝试了多次并未成功。此外听朋友说即使fp16也仅能提升约20%的训练速度，看起来并不太promising。另外inference时也要使用apex，增加了额外依赖，麻烦。

这里就说说前两个问题。

中文Tokenizer

GPT2有对应的GPT2Tokenizer，用的是BPE分词，但是并没有中文分词器。理论上来说BPE是语言无关的，但对于中日韩文来说，还是字级别character-level的分词更合理些，直接用GPT2Tokenizer不太好。因此，要么自己train一个中文tokenizer，要么换tokenizer。

训练中文GPT2Tokenizer(失败)

根据语言自己训练一个tokenizer想法很直接，看了眼GPT2Tokenizer的源码，其实就是个BPETokenizer，于是直接用HuggingFace的tokenizer库就可以训练。这个库的底层是用Rust写的，可以最大程度地并行处理。

训练代码：

import os
from tokenizers import ByteLevelBPETokenizer
from transformers import GPT2Tokenizer
from transformers import CONFIG_MAPPING, AutoConfig, GPT2Config

model_type = 'gpt2'
model_path = './modeltest'
vocab_size = 50000

# training tokenizer
tokenizer = ByteLevelBPETokenizer()
tokenizer.train(files=['data/traint.txt'], min_frequency=5, vocab_size=vocab_size, special_tokens=["<|endoftext|>"])
tokenizer.save(model_path)

tokenizer = GPT2Tokenizer(os.path.join(model_path, 'vocab.json'), os.path.join(model_path, 'merges.txt'), errors='replace')
tokenizer.model_max_length = 1024
tokenizer.save_pretrained(model_path)

config = GPT2Config.from_pretrained(model_type)
config.bos_token_id = tokenizer.bos_token_id
config.eos_token_id = tokenizer.eos_token_id
config.vocab_size = vocab_size
config.save_pretrained(model_path)

tokenizer = GPT2Tokenizer.from_pretrained(model_path)
config = AutoConfig.from_pretrained(model_path)

text = '中国手机'
print (tokenizer.tokenize(text))
print (tokenizer.convert_tokens_to_ids(tokenizer.tokenize(text)))

注意ModelConfig和Tokenizer本身是耦合的，需要里面的bos_token_id, eos_token_id及vocab_size，因此需要加载默认的GPT2Config，改好这些配置后存储到自己的model文件夹中。

运行结果：

[00:00:00] Reading files                            █████████████████████████████████████████████████████████████████████████                 100
[00:00:00] Tokenize words                           █████████████████████████████████████████████████████████████████████████ 576      /      576
[00:00:00] Count pairs                              █████████████████████████████████████████████████████████████████████████ 576      /      576
[00:00:00] Compute merges                           █████████████████████████████████████████████████████████████████████████ 696      /      696

['ä¸Ń', 'åĽ½', 'æīĭ', 'æľ', 'º']
[435, 456, 484, 261, 119]

看起来不错？那就太naive了。数据量太小，换了个小数据集再测，发现tokenizer直接崩掉：

[00:00:00] Reading files                            █████████████████████████████████████████████████████████████████████████                 100
[00:00:00] Tokenize words                           █████████████████████████████████████████████████████████████████████████ 774      /      774
[00:00:00] Count pairs                              █████████████████████████████████████████████████████████████████████████ 774      /      774
thread 'thread '<unnamed>' panicked at 'called `Option::unwrap()` on a `None` value', /__w/tokenizers/tokenizers/tokenizers/src/models/bpe/trainer.rs:373:34
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
<unnamed>' panicked at 'called `Option::unwrap()` on a `None` value', /__w/tokenizers/tokenizers/tokenizers/src/models/bpe/trainer.rs:373:34
fatal runtime error: failed to initiate panic, error 5
Aborted (core dumped)

Debug了好长时间，发现原来是因为数据集中有一个~字符……这只是一个测试，当然可以通过处理数据集的方式来解决问题，但鬼知道还有哪些特殊字符会引起tokenizer core dump。而且因为底层是Rust，这玩意儿也不懂，无法通过改源码的方式Fix……

如果看下生成的merges.txt和vocab.json，会发现里面都是bytes，并非可见字符，也是BPE分词应有的效果，但考虑到在中文上的适用性，还是算了。

换BertTokenizer

看到大部分中文预训练模型都是直接用的Google开源的BertTokenizer，我们也直接用吧：

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
tokenizer.save_pretrained(model_path)

但用的时候要注意设置bos_token_id, eos_token_id及vocab_size到GPT2Config中:

config = CONFIG_MAPPING['gpt2']()
config.bos_token_id = tokenizer.cls_token_id
config.eos_token_id = tokenizer.sep_token_id
config.vocab_size = tokenizer.vocab_size

BertTokenizerFast

用大规模语料训练时发现，预处理数据太慢了，htop一看只用了一个CPU核，不慌，BertTokenizerFast可以上场了，它们都是基于tokenizer库直接开发的，可以用上所有的CPU核，分词速度与核数按比例提升：

from transformers import BertTokenizerFast

tokenizer = BertTokenizerFast.from_pretrained("bert-base-chinese")

大规模语料爆内存

训练语料大概10几G，在creating feature时就崩了且没有任何有用的提示：

transformers.data.datasets.language_modeling -   Creating features from dataset file at /data/transformers/data
./pretrain.sh: line 13:  2692 Killed                  CUDA_VISIBLE_DEVICES=0 python run_pretrain.py
--output_dir output --model_type gpt2 --model_name_or_path $MODEL_PATH --do_train --train_data_file
$TRAIN_FILE --per_gpu_train_batch_size 16 --per_gpu_eval_batch_size 4 --seed 1024

htop观察了一下，是内存爆了out of memory，注意是内存不是显存，如果是显存问题可以通过减少batch_size，但内存……

流式Dataset

原始数据用的是LineByLineTextDataset，看了一下 LineByLineTextDataset实现：

class LineByLineTextDataset(Dataset):
    """
    This will be superseded by a framework-agnostic approach
    soon.
    """

    def __init__(self, tokenizer: PreTrainedTokenizer, file_path: str, block_size: int, local_rank=-1):
        assert os.path.isfile(file_path)
        # Here, we do not cache the features, operating under the assumption
        # that we will soon use fast multithreaded tokenizers from the
        # `tokenizers` repo everywhere =)
        logger.info("Creating features from dataset file at %s", file_path)

        with open(file_path, encoding="utf-8") as f:
            lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        batch_encoding = tokenizer.batch_encode_plus(lines, add_special_tokens=True, max_length=block_size)
        self.examples = batch_encoding["input_ids"]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i) -> torch.Tensor:
        return torch.tensor(self.examples[i], dtype=torch.long)

注意到它是把所有数据都load到内存中，然后存储batch_encoding["input_ids"]，再看下batch_encoding的代码，里面好多额外的数据结构，内存撑不住了。最好改成流式读取和处理的方式，下面实现了一个StreamingLineByLineTextDataset，简单起见，只在class中存原始数据，而在__getitem__方法中将其转换成tensor:

import torch
from torch.utils.data.dataset import Dataset

class StreamingLineByLineTextDataset(Dataset):
    def __init__(self, tokenizer: PreTrainedTokenizer, file_path: str, block_size: int, local_rank=-1):
        assert os.path.isfile(file_path)

        self.tokenizer = tokenizer
        self.block_size = block_size

        logger.info("Creating features from dataset file at %s", file_path)

        with open(file_path, encoding="utf-8") as f:
            self.lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        logger.info("Loaded lines: %s", len(self.lines))

    def __len__(self):
        return len(self.lines)

    def __getitem__(self, i) -> torch.Tensor:
        batch_encoding = self.tokenizer.batch_encode_plus([self.lines[i]], add_special_tokens = True, max_length = self.block_size)
        input_ids = batch_encoding["input_ids"]
        return torch.tensor(input_ids[0], dtype=torch.long)

小结

以上简单记录了下训练中文GPT2模型中遇到的一些坑和解决方案，期待训练出来的模型效果。