探秘Transformer系列之（3）---数据处理

探秘Transformer系列之（3）—数据处理

接下来三篇偏重于工程，内容略少，大家可以当作甜点 ^_。

0x00 概要

有研究人员认为，大模型的认知框架看起来十分接近卡尔·弗里斯顿(Karl Friston)描绘的贝叶斯大脑。基于贝叶斯概率理论和生物物理学原理，大脑的主要目标是预测和控制外界的信息，以最大限度地降低不确定性和内部熵。

大脑通过不断收集和处理外部信息来构建内部模型，以预测和控制外界。而海量的文本或者多模态语料组成了大模型需要认知的外部世界的基本信息。预训练就是在不同尺度上提炼语料数据中的信息概率分布。增加训练数据量，模型参数量，训练时间都会丰富大模型在某一问题域的信息量，降低测试集上的信息熵，从而让它见多识广。微调则类似借助新语料和手段对模型内部的相关参数进行微扰，促进其进入更有序的空间，实现可控的可预测的涌现。因此，数据和数据处理决定了大模型的上限。

本章我们来分析哈佛代码的数据处理部分，藉此对Transformer的总体有进一步了解。另外，本篇会涉及到词表和分词器，所以我们提前把后续会讲到的一些概念提前做一下说明，以便读者可以更好的理解。

tokenize（分词）：把一句话按照一定规则来分成一个个词。比如按标点符号分词，或者按语法规则分词。
token（词元）：token是分词的结果，也就是最小语义单位。token既可以是一个单词、一个汉字，也可能是一个表示空白字符、未知字符、句首字符的特殊字符等。
词表（vocb）：词表是指LLM能够理解和识别的唯一单词或token的集合，用来定义token与整数之间的映射关系。在训练模型之前是需要构建好词表的。

0x01 总体流程

下图给出了LLM的常见数据处理流程，其包含质量过滤（Quality Filtering）、去重（De-deplication）、隐私擦除（Privacy Reduction）、Tokenization、数据混合等。其实这才是LLM工作中最繁杂的一部分。

哈佛的代码相对简单太多。我们首先给出训练代码的精简版。可以看到其主要分为两步：

建立分别用于训练数据和验证数据加载的数据加载器。
调用run_epoch()函数迭代运行训练步。每次运行都会利用数据加载器来加载数据。

具体代码如下。

def train_worker(
    gpu,
    ngpus_per_node,
    vocab_src,
    vocab_tgt,
    spacy_de,
    spacy_en,
    config,
    is_distributed=False,
):
    train_dataloader, valid_dataloader = create_dataloaders(
        gpu,
        vocab_src,
        vocab_tgt,
        spacy_de,
        spacy_en,
        batch_size=config["batch_size"] // ngpus_per_node,
        max_padding=config["max_padding"],
        is_distributed=is_distributed,
    )

    for epoch in range(config["num_epochs"]):
        _, train_state = run_epoch(
            (Batch(b[0], b[1], pad_idx) for b in train_dataloader),
            model,
            SimpleLossCompute(module.generator, criterion),
            optimizer,
            lr_scheduler,
            mode="train+log",
            accum_iter=config["accum_iter"],
            train_state=train_state,
        )

下图是对代码的进一步简化，后续在图上也只展示train数据集相关的代码，不展示验证集相关代码。

0x02 数据集

我们接下来看看数据集。

2.1 行业做法

常见数据集

在实践中，研究人员通常需要混合使用不同的数据源来进行LLM预训练，而不是使用单个语料库，通常会包括学术文献、书籍、网页内容和编程代码等。因此，现有的研究通常会混合几个现成的数据集（例如C4、OpenWebText和Pile），然后进行进一步处理以获得预训练语料库。此外，为了训练适应特定应用的LLM，从相关来源（如维基百科和BigQuery）提取数据以丰富预训练数据中的相应信息也很重要。只有提供足够语料，才可以降低概率空间的信息熵到一定阈值，从而对某类任务达成相变。下面是常见数据集。

下图给出来不同预训练模型的架构、训练语料、训练目标。

数据源比率

数据混合策略对于训练至关重要。为了平衡不同类型的数据，研究者往往使用大模型对数据进行分类，然后对不同类别的数据进行数据分布调整。比如会基于知识深度和帮助性等质量指标进行采样权重调整。或者采用平衡采样策略，确保高质量内容的优先性，同时保留多样化的类别。这样可以确保模型能够从各种类型的数据中学习，避免了某些领域数据过多而导致的偏差。下图是现有LLM预训练中数据源的比率。

数据治理

因为语料中的偏差与错误会让大模型学到扭曲的外部信息。因此语料的全面数据治理十分必要，既要丰富详实，又要不偏不倚。为了确保数据的高质量，很多LLM会在训练中采用了多种处理策略，比如：

数据质量增强。通过结合规则清洗和去重程序对文档质量进行严格评估。这往往通过前一代模型来对预训练数据进行智能过滤，评估文档的连贯性、简洁性、教育价值、帮助性、知识丰富性和类别相关性。这种方法不仅提高了数据质量，还增强了模型对多语言数据的处理能力。
数据格式优化。比如对于对话和问答数据，可以采用嵌套文档格式，使用灵活的模板来平衡自然理解与结构一致性。这种设计确保了模型在多种交互模式下的泛化能力。
数据合成。比如利用其它模型生成高质量的合成数据。而且会利用其它模型对这些合成数据进行进一步筛选，确保了合成数据的质量和相关性。这种方法不仅扩大了训练数据的规模，还保证了数据的高质量和多样性。

2.2 哈佛数据集

哈佛代码通过Multi30k数据集来训练模型，以此实现将德语句子翻译成英语的功能。Multi30K是Flickr30K数据集（Young等人，2014）的扩展，包含31014个德语翻译的英语描述和155070个独立收集的德语描述。因为此数据集极其简单，我们就不需要做很多处理。

数据集介绍参见 https://github.com/multi30k/dataset
PyTorch的相关说明参见 https://pytorch.org/text/stable/_modules/torchtext/datasets/multi30k.html

数据集由三个文件组成：mmt16_task1_test.tar.gz，training.tar.gz，validation.tar.gz。打开training.tar.gz，可以看到两个文件：train.de，train.en。里面分别是29000行德文和英文，摘录如下：

train.de

Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.
Mehrere Männer mit Schutzhelmen bedienen ein Antriebsradsystem.
Ein kleines Mädchen klettert in ein Spielhaus aus Holz.
Ein Mann in einem blauen Hemd steht auf einer Leiter und putzt ein Fenster.
Zwei Männer stehen am Herd und bereiten Essen zu.

train.en

Two young, White males are outside near many bushes.
Several men in hard hats are operating a giant pulley system.
A little girl climbing into a wooden playhouse.
A man in a blue shirt is standing on a ladder cleaning a window.
Two men are at the stove preparing food.

数据集用到了两次，分别是构建词表时和训练构建batch（批量/批次）时。

0x03 加载功能模块

在哈佛源码中构建了几个数据相关的全局变量，分别用来存储加载的分词器、字典和模型。

model = load_trained_model() # 加载模型
spacy_de, spacy_en = load_tokenizers() # 加载分词器
vocab_src, vocab_tgt = load_vocab(spacy_de, spacy_en) # 构建字典

3.1 加载模型

load_trained_model()函数负责加载模型，此处会设置batch_size等参数。如果此函数在执行过程中没有找到可以加载的模型，则会调用train_model()函数来训练一个模型。

def load_trained_model():
    config = {
        "batch_size": 2,
        "distributed": False, # 不进行分布式训练
        "num_epochs": 8,
        "accum_iter": 10, # 每训练10个批量后会更新一次模型参数
        "base_lr": 1.0, # 基础学习率
        "max_padding": 10, # 句子最大长度
        "warmup": 3000, # 依据基础学习率会预热3000次，此后学习率会下降
        "file_prefix": "multi30k_model_",
    }
    model_path = "multi30k_model_final.pt"
    if not exists(model_path):
        train_model(vocab_src, vocab_tgt, spacy_de, spacy_en, config)

    # 初始化模型
    model = make_model(len(vocab_src), len(vocab_tgt), N=6)
    # 从模型文件中加载模型参数
    model.load_state_dict(torch.load("multi30k_model_final.pt"))
    return model

3.2 加载分词器

load_tokenizers()函数的功能是加载德语分词模型和英语分词模型。spacy是一个文本预处理的Python库，其提供了分词功能，具体信息可以参见 https://spacy.io/，https://github.com/explosion/spaCy。

import spacy

def load_tokenizers():
    try:
        spacy_de = spacy.load("de_core_news_sm")
    except IOError:
        os.system("python -m spacy download de_core_news_sm")
        spacy_de = spacy.load("de_core_news_sm")

    try:
        spacy_en = spacy.load("en_core_web_sm")
    except IOError:
        os.system("python -m spacy download en_core_web_sm")
        spacy_en = spacy.load("en_core_web_sm")

    return spacy_de, spacy_en

3.3 加载词表

load_vocab()函数会加载词表，然后构建词表。load_vocab()函数具体代码如下。

def load_vocab(spacy_de, spacy_en):
    # 如果文件不存在，则构建字典，否则直接加载词典
    if not exists("vocab.pt"):
        vocab_src, vocab_tgt = build_vocabulary(spacy_de, spacy_en)
        torch.save((vocab_src, vocab_tgt), "vocab.pt")
    else:
        vocab_src, vocab_tgt = torch.load("vocab.pt")
    return vocab_src, vocab_tgt

0x04 加载数据

create_dataloaders()函数定义了数据加载器，其中最主要的部分是collate_fn()函数，该函数利用collate_batch()函数来定义构建batch功能，即把若干数据聚集成一个batch。

def create_dataloaders(
    device,
    vocab_src, # 源词表（德语词表）
    vocab_tgt, # 目标词表（英语词表）
    spacy_de, # 德语分词器
    spacy_en, # 英语分词器    
    batch_size=12000, # batch size(批次大小）
    max_padding=128, # 句子最大填充长度
    is_distributed=True,
):
    
    # 德语分词函数，其会调用德语分词器对语句进行分词
    def tokenize_de(text):
        return tokenize(text, spacy_de)

    # 英语分词函数，其会调用英语分词器对语句进行分词
    def tokenize_en(text):
        return tokenize(text, spacy_en)

    # 定义构建batch功能，即把若干数据聚集成一个batch
    def collate_fn(batch):
        return collate_batch(
            batch,
            tokenize_de,
            tokenize_en,
            vocab_src, # 源词表（德语词表）
            vocab_tgt, # 目标词表（英语词表）
            device,
            max_padding=max_padding,
            pad_id=vocab_src.get_stoi()["<blank>"],
        )

    # 加载数据集
    train_iter, valid_iter, test_iter = datasets.Multi30k(
        language_pair=("de", "en")
    )

    # 将train_iter转换为map
    train_iter_map = to_map_style_dataset(
        train_iter
    )  # DistributedSampler needs a dataset len()
    train_sampler = (
        DistributedSampler(train_iter_map) if is_distributed else None
    )
    valid_iter_map = to_map_style_dataset(valid_iter)
    valid_sampler = (
        DistributedSampler(valid_iter_map) if is_distributed else None
    )

    # 构建训练数据加载器
    train_dataloader = DataLoader(
        train_iter_map,
        batch_size=batch_size,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        collate_fn=collate_fn,
    )
    # 构建验证数据加载器
    valid_dataloader = DataLoader(
        valid_iter_map,
        batch_size=batch_size,
        shuffle=(valid_sampler is None),
        sampler=valid_sampler,
        collate_fn=collate_fn,
    )
    return train_dataloader, valid_dataloader

前面加载词表时我们提到，会把词表作为参数传入collate_batch()函数。现在又提到数据加载器会利用collate_batch()函数加载数据。看来collate_batch()函数是核心所在，我们接下来就进行分析如何加载batch。

4.1 填充（Padding）

深度学习模型要求输入数据具有固定的尺寸，但是NLP（自然语言处理）领域中很难做到，因为其输入文本通常是变长的，很难找到多个长度一样的句子，因此难免会将长度不一的句子放到一个batch里。为了符合模型的输入方式，使这些模型能够处理不同长度的文本，在数据集的生成过程中，我们要对输入序列进行对齐，使同一个batch内所有序列的长度一致。具体来说就是：

但是如果输入的序列太长，则是截取左边的内容，把多余的直接舍弃。
需要将长度不足的句子用无意义的特殊字符补全到最大句子长度。

这样，所有的文本序列就会有相同的长度，从而可以被作为一个统一的batch输入到模型中进行处理。哈佛代码就是用了填充和截断。

改进

因为这些填充符号其实并不携带实际的语义信息，只是用来填充序列长度，所以还是对模型处理带来负面影响。因此目前也有相关优化工作，比如No Pad优化。研究人员修改了注意力算子的实现，参与运算的所有文本序列头尾相接，拼接成一个超长的输入序列。为了标记每一个文本序列的起止位置，我们还需要一个序列去记录拼接文本的长度。这项技术可以有效缩减模型推理时所需的计算量。

左填充

padding有左填充和右填充之分。bert采用的是右填充，目前大多数LLM使用左填充。这是因为大多数LLM总是选择最后一个 token 的 logits 来预测下一个 token，如果我们在右侧进行填充，则模型在某些情况下是使用的 logits 来预测下一个 token，这样会导致不正确的结构。假如要生成两个句子：”飞雪连天射白鹿，笑书神侠倚碧鸳“。目前已经生成了“飞雪连天射白鹿“，需要生成下一个单词。

左填充的结果是：”飞雪连天射白鹿“。会使用”鹿“来进行预测下一个单词。
右填充的结果是：”飞雪连天射白鹿“。会使用来进行预测下一个单词。

4.2 Batch类

源码中使用Batch类承载了batch概念。Batch类会把一个batch的源语言句子和目标语言句子整合在一起，并且依据数据来生成对应的掩码。在读取数据集的句子时，会在句首加入一个特殊 token（一般是或），在句尾也加入一个特殊 token（）。此处假定batch大小是8，句子最长的长度是32，特殊字符在词表中的序号是0，的序号是1，的序号是2。

成员变量

Batch类的关键成员函数如下：

src：源语言句子列表。src的形状为[batch size, max_seq_len]，其中每个句子内容是原始语句中token对应的词典序号。单个句子举例如下：[0, 3, 5, 6,…,7, 1,2,2]，其中0是，1是，2是。因此3,5,6,…,7是实际语句内容。max_seq_len代表句子最长的长度。
tgt：目标语言句子列表。逻辑和src类似，但可以为空，因为推理时候不需要传入目标语言句子。
tgt_y：目标语言句子真值列表。训练阶段，解码器需要将预测输出序列的最后一个字符和真实的结果作比较，因此需要把tgt复制为tgt_y作为真值。
src_mask：源语言句子的掩码，作用是把src中的盖住，这样就不会参与计算。
tgt_mask：目标语言句子的掩码，逻辑和src_mask类似。

Batch的代码如下。

class Batch:
    """Object for holding a batch of data with mask during training."""

    def __init__(self, src, tgt=None, pad=2):  # 2 = <blank>
        self.src = src # 源语言句子列表
        # 创建源语言的掩码，这样可以忽略填充部分，unsqueeze()的作用是增加一个维度，因为后续要和注意力分数进行掩码计算，而注意力分数是三个维度，所以这里要保持一致。
        self.src_mask = (src != pad).unsqueeze(-2) 
        # 预测时候没有目标语言句子；训练时候有目标语言句子
        if tgt is not None: # 如果目标语言数据存在
            # 去掉tgt的最后一个单词<eos>。因为tgt存储的是解码器的输入，而解码器的输入不应该有<eos>。比如一个句子“<bos>新年好<eos>”，下面代码处理之后，self.tgt就应该是"<bos>新年好"。
            self.tgt = tgt[:, :-1] # 形状是torch.Size([batch size, 字数-1])
            # 去掉tgt的第一个词<bos>。因为tgt_y存储的是希望预测的结果，所以不需要<bos>。假设tgt是“<bos>新年好<eos>”，下面语句运行之后， self.tgt_y内容就是“新年好<eos>”，即我们希望模型预测出这几个token。
            self.tgt_y = tgt[:, 1:] # 形状是torch.Size([batch size, 字数-1])
            
            # 创建目标语言掩码，这样可以忽略填充部分和未来词汇
            self.tgt_mask = self.make_std_mask(self.tgt, pad)
            self.ntokens = (self.tgt_y != pad).data.sum() # 计算目标语言句子中非填充词的数量，<bos>，<eos>这些也算是句子的token，所以依然要计算

    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        # 生成填充词对应的掩码
        tgt_mask = (tgt != pad).unsqueeze(-2)
        # subsequent_mask()函数会生成未来词汇相关的掩码，然后填充词对应的掩码和未来词汇相关的掩码会做与操作，得到最终掩码
        # tgt.size(-1) 表示的是序列的长度
        tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
            tgt_mask.data
        )
        return tgt_mask

目标语句

源语言句子相关的成员变量只有src一个，而和目标句子相关的成员变量有两个：tgt和tgt_y，因此我们要特殊分析下。对于推理阶段，tgt可以为空，因为预测时候没有目标语言句子，只有源语言句子。对于训练阶段，模型的输入是src和tgt，模型处理之后输出out，然后需要把out和tgt_y进行比对，确定损失大小。需要注意两个细节：

解码器的输入要除掉最后一个 token（或）。这是因为我们最后一次的输入tgt是我爱你（没有），因此我们的输入tgt一定不会出现目标的最后一个token，所以一般tgt处理时，会将目标句子删掉最后一个token。
而解码器的预测目标要除掉第一个 token。因为我们不需要预测，即我们的label不包含。代码中把label命名为tgt_y。

上述操作分别通过如下语句完成：target_input=target[:-1, :] 和 target_out=target[1:, :] 。我们举例如下。假设原始目标语言句子是"新年好“，转换为tgt为”新年好"，假设计算之后得到out是"新年乐“，tgt_y则是"新年好”。

def run_epoch():
    """Train a single epoch"""
    for i, batch in enumerate(data_iter):
        out = model.forward(
            batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
        )
        loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)

生成掩码

Batch类最重要的作用就是生成句子的掩码。生成掩码有两个目的。

由于 padding token 是用来补全长度的，并没有实际意义，因此我们还希望能够尽可能的降低 padding 的影响。这样可以降低计算复杂度，也降低 padding 在模型对文本建模中的影响。
因为使用Teaching Forcing模式（后文会详细解读）进行训练，这也要添加掩码，使得 Self-Attention 不能访问未来的输入。

针对第一种目的的掩码，我们称之为Padding Mask（填充词对应的掩码）。针对第二种目的的掩码，我们称之为和Sequence mask（未来词汇相关的掩码）。而源语句需要Padding Mask，目标语句需要Padding Mask和Sequence mask的结合。我们接下来结合源语句和目标语句来一一介绍。

源语句的掩码对应的变量叫做src_mask。假设某个句子内容是[0, 3, 1, 2, 2]。生成src_mask的语句比较简单，只有self.src_mask = (src != pad).unsqueeze(-2) 这一行代码。主要起到两个作用：

把src中非pad的部分置为True，pad部分置为False。上述例句则对应的掩码是[True, True, True, False, False]。因为“”、“”和“”要算作句子成分，因此不做掩码处理。
使用unsqueeze()函数增加一维度，因为后续src_mask要和注意力分数进行掩码计算，而注意力分数是三个维度，所以这里要保持一致。所以最终src_mask的形状是[batch大小,1,句子最长长度]。

目标语句掩码对应的变量叫做tgt_mask。生成tgt_mask则比较复杂，具体逻辑在前面给出的Batch类的成员变量函数make_std_mask()中。tgt_mask与src_mask略有不同，除了需要盖住pad部分，还需要将对角线右上的也都盖住。就是要结合填充词对应的掩码和未来词汇相关的掩码。make_std_mask()函数的逻辑如下：

首先生成填充词对应的掩码，即Padding Mas。上述例句则对应的掩码是[[[True, True, True, False, False]]]。

然后调用subsequent_mask()函数来生成未来词汇相关的掩码，即Sequence mask，这是一个对角线以及之下都是True的矩阵，具体掩码如下。

[[
  [ True, False, False, False, False ],
  [ True, True, False, False, False ],
  [ True, True, True, False, False ],
  [ True, True, True, True, False ],
  [ True, True, True, True, True ],
]]

最后填充词对应的掩码和未来词汇相关的掩码会做与操作，得到最终掩码如下

[[
  [ True, False, False, False, False ],
  [ True, True, False, False, False ],
  [ True, True, True, False, False ],
  [ True, True, True, False, False ],
  [ True, True, True, False, False ],
]]op

注意src_mask的shape是(batch,1,seq_len)，而trg_mask是(batch,seq_len,seq_len)。因为src_mask的每一个时刻都能attendto所有时刻(padding的除外)，一次只需要一个向量就行了，而trg_mask需要一个矩阵，这个矩阵代表若干个时刻。

subsequent_mask()函数对应的代码如下。

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )
    return subsequent_mask == 0

构建batch

函数data_gen()被用来构建batch，具体代码如下。

def data_gen(V, batch_size, nbatches):
    """
    生成一组随机数据。（该方法仅用于Demo）
    :param V: 词典的大小
    :param batch_size
    :param nbatches: 生成多少个batch
    :return: yield一个Batch对象
    """

    # 生成{nbatches}个batch
    for i in range(nbatches):
        # 生成一组输入数据
        data = torch.randint(1, V, size=(batch_size, 10))
        # 将每行的第一个词都改为1，即"<bos>"
        data[:, 0] = 1
        # 该数据不需要梯度下降
        src = data.requires_grad_(False).clone().detach()
        tgt = data.requires_grad_(False).clone().detach()
        # 返回一个Batch对象
        yield Batch(src, tgt, 0)

4.3 加载batch

collate_batch()函数是DataLoader类的collate_fn (Callable, optional)参数，其作用是将一个样本列表组合成一个张量的mini-batch。DataLoader内部会将”句子对的列表“传给collate_batch()函数来处理，然后把输入的batch发给模型。

def collate_batch(
    batch, # 句子对的列表。比如[(源句子1, 目标句子1),(源句子2, 目标句子2),.....]，列表大小为batch size
    src_pipeline, # 德语分词功能，即spacy_de的封装器
    tgt_pipeline, # 英语分词功能，即spacy_en的封装器
    src_vocab, # 德语词典，Vocab对象
    tgt_vocab, # 英语词典，Vocab对象
    device,
    max_padding=128, # 句子最大长度
    pad_id=2,
):
    # <bos>和<eos>在词典中的index
    bs_id = torch.tensor([0], device=device)  # <s> token id
    eos_id = torch.tensor([1], device=device)  # </s> token id
    src_list, tgt_list = [], []
    for (_src, _tgt) in batch: # 遍历句子对列表
        # 首先调用src_vocab(src_pipeline(_src))对源句子处理，具体是利用分词器src_pipeline和词表src_vocab把句子转换为词表index的序列；其次调用torch.cat在句子前面加上<bos>，句子后面加上<eos>。
        processed_src = torch.cat(
            [
                bs_id,
                torch.tensor(
                    src_vocab(src_pipeline(_src)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        # 首先调用tgt_vocab(tgt_pipeline(_tgt))对源句子处理，具体是利用分词器tgt_pipeline和词表tgt_vocab把句子转换为词表index的序列；其次调用torch.cat在句子前面加上<bos>，句子后面加上<eos>。
        processed_tgt = torch.cat(
            [
                bs_id,
                torch.tensor(
                    tgt_vocab(tgt_pipeline(_tgt)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        # 如果processed_src大于max_padding，则截断；如果小于max_padding，则填充
        src_list.append(
            # warning - overwrites values for negative values of padding - len
            pad(
                processed_src,
                (
                    0,
                    max_padding - len(processed_src),
                ),
                value=pad_id,
            )
        )
        # 如果processed_tgt大于max_padding，则截断；如果小于max_padding，则填充
        tgt_list.append(
            pad(
                processed_tgt,
                (0, max_padding - len(processed_tgt)),
                value=pad_id,
            )
        )

    src = torch.stack(src_list) # 把列表堆叠在一起
    tgt = torch.stack(tgt_list) # 把列表堆叠在一起
    return (src, tgt)

4.3 训练使用

train_worker()函数在训练时，在每个epoch中，会调用把从数据集之中获取的数据构建成一个Batch，然后调用run_epoch()函数进行具体训练。

_, train_state = run_epoch(
    # 拿到Batch类的实例
    (Batch(b[0], b[1], pad_idx) for b in train_dataloader),
    model,
    SimpleLossCompute(module.generator, criterion),
    optimizer,
    lr_scheduler,
    mode="train+log",
    accum_iter=config["accum_iter"],
    train_state=train_state,
)

run_epoch()函数代码如下。

def run_epoch(
    data_iter, # 可迭代对象，一次返回一个Batch对
    model, # Transformer模型，EncoderDecoder类对象
    loss_compute, # SimpleLossCompute对象，用于计算损失
    optimizer, # Adam优化器。验证时，optimizer是DummyOptimizer
    scheduler, # LambdaLR对象，用于调整Adam的学习率，实现WarmUp
    mode="train", 
    accum_iter=1, # 多少个batch更新一次参数，默认为1，也就是每个batch都对参数进行更新
    train_state=TrainState(), # TrainState对象，用于保存一些训练状态
):
    """Train a single epoch"""
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    n_accum = 0
    
    # 遍历数据集中的每个batch
    for i, batch in enumerate(data_iter):
        # 对每个batch进行前向传播，等价于model(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)。这里的out是Decoder的输出，并不是Generator的输出，因为在EncoderDecoder的forward中并没有使用generator。generator的调用放在了loss_compute中        
        out = model.forward(
            batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
        )
        
        """
        调用loss_compute()函数来计算每个批次的损失，传入的三个参数分别为：
        1. out: EncoderDecoder的输出
        2. tgt_y: 要被预测的所有token，例如src为`<bos> I love you <eos>`，则`tgt_y`则为`我 爱 你 <eos>`
        3. ntokens：这批batch中有效token的数量，用于对loss进行正则化。
        """   
        loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
        # loss_node = loss_node / accum_iter
        if mode == "train" or mode == "train+log":
            loss_node.backward() # 计算梯度
            train_state.step += 1 # 记录step次数
            train_state.samples += batch.src.shape[0] # 记录样本数量。batch.src.shape[0]获取的是Batch size
            train_state.tokens += batch.ntokens # 记录处理过的token数
            # 如果达到了accum_iter次，就进行一次参数更新
            if i % accum_iter == 0:
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                n_accum += 1
                train_state.accum_step += 1
            # 更新学习率    
            scheduler.step()

        # 累计loss
        total_loss += loss
        # 累计处理过的tokens
        total_tokens += batch.ntokens
        # 累计从上次打印日志开始处理过得tokens
        tokens += batch.ntokens
        if i % 40 == 1 and (mode == "train" or mode == "train+log"):
            lr = optimizer.param_groups[0]["lr"]
            elapsed = time.time() - start
            start = time.time()
            tokens = 0
        del loss
        del loss_node
    # 返回平均损失和训练状态    
    return total_loss / total_tokens, train_state # 返回平均损失

小结

我们用一个完整的图展示训练的总体数据流程如下。以这个数据处理流程为基础和起点，LLM被注入足够的信息量，从而构建了海量自然语言和代码的概率分布空间，成各种复杂关联的模式，涵盖自然语言和代码中各种知识与结构。这些知识和结构会体现为概率分布的距离与关系，从而为对比、类比、归纳、演绎等推理步骤提供支撑，也就是“涌现出”这些推理能力。