LLM中损失函数解析

在GPT系列大语言模型中损失函数采用的是自回归语言建模任务，即根据前K-1个token预测第K个token，本质上都是交叉熵分类损失，在实现上预训练和监督微调稍有不同，本文分别进行介绍

预训练Pretrain

数据集

由于预训练数据集一般很大，因此一般全部处理完成之后，在代码中直接加载最终的token进行训练。一般经过以下流程：

1. 依次读入所有的预训练语料，对每一个预训练语料的每一个样本进行分词处理tokenizer，并添加结束符，例如<eos>。
2. （optional）将经过分词处理之后的所有预训练语料，拼接成一个整的大语料文件
3. 对预训练语料进行维度变换，最终预训练样本的shape=[语料token总数//max_length, max_length]，max_length是指模型输入token的最大长度
4. 返回最终的训练语料，每条预训练样本的长度=max_length - 1或者max_length，X是输入，Y在自监督标签。代码如下：

def __getitem__(self, index: int):
        sample = self.data[index]  # 一个样本
        # X,Y前后错一位，seq_len = max_length - 1
        X = np.array(sample[:-1]).astype(np.int64)
        Y = np.array(sample[1:]).astype(np.int64)
        return torch.from_numpy(X), torch.from_numpy(Y)

例如：

原始预训练语料为：我爱吃饭   (分词结果为[100, 101, 102, 103])

则， X=[我，爱，吃] = [100, 101, 102]

        Y=[爱，吃，饭] = [101, 102, 103]

注：以上就是采用Teacher Forcing并且进行并行化训练，模型直接输入整个单词序列和对应的真值标签。

计算损失

最终计算损失函数的次数 = batch size * seq_len。即相当于以下计算过程：

X = [我]， Y=[爱]

X = [我，爱]， Y=[吃]

X = [我，爱，吃]， Y=[饭]

注：但每条预训练样本只完整前向计算1次，为了保证看不到后面的token，需要在网络结果中采用masked self-attention

采用分类中常用的交叉熵损失函数，类别数=词典大小

logits = self.output(h)  # logits.shape=[bs, seq_len, vocab_size]
# logits_view.shapre=[bs * seq_len, vocab_size],  targets.shape=[bs * seq_len]
self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

监督微调SFT

数据集

由于SFT样本一般较小，因此也可以只读取原始文本，在dataloader中进行分词处理，一般经过以下步骤：

1. 构建完整SFT原始数据集，例如可以以csv/json等格式存储，最终只包括prompt和answer两个字段
2. 在SFTDataset类中分别对prompt和answer进行分词处理，（还可以在answer前后增加特殊字符<bos>,<eos等>）
3. 设置每一个样本的loss_mask：prompt部分位置的值为0，其余部分位置的值为1（下面示例代码只是一种策略）
4. 返回X, Y, loss_mask。 其中X和Y也是同一条样本前后错一位得到（注意：并不是X是prompt，Y是answer）

def __getitem__(self, index: int):
        #
        sample = self.df.iloc[index] # 一行样本
        prompt = self.tokenizer.encode(sample['prompt'],add_special_tokens=False)
        answer = self.tokenizer.encode(sample['answer'],add_special_tokens=False)
        if len(prompt) > self.prompt_max_len:
            prompt = prompt[:self.prompt_max_len-2]
        if len(answer) > self.answer_max_len:
            answer = answer[:self.answer_max_len-2]
        #
        input_id = prompt + [self.bos] + answer + [self.eos]
        context_length = input_id.index(self.bos)  # 答案的开始位置
        mask_position = context_length - 1         # mask的结束位置
        pad_len = self.max_length - len(input_id)  # pad长度
        input_id = input_id + [self.pad] * pad_len

        if pad_len == 0:
            loss_mask = [0]*context_length + [1]*(len(input_id[mask_position+1:]))
        else:
            loss_mask = [0]*context_length + [1]*(len(input_id[mask_position+1:-pad_len])) + [0]*pad_len
        #
        input_id = np.array(input_id)
        X = np.array(input_id[:-1]).astype(np.int64)
        Y = np.array(input_id[1:]).astype(np.int64)
        loss_mask = np.array(loss_mask[:-1])
        #
        return torch.from_numpy(X),torch.from_numpy(Y),torch.from_numpy(loss_mask)

计算损失

仍采样交叉熵损失，但参数reduce=False，计算返回结果与输入的维度相同，然后采用loss_mask不计算Prompt部分的损失，最后求在answer个数上的均值。

logits = model(X, Y)
# logits.shape = [bs, seq_len, vocab_size],  (seq_len=max_len - 1),    Y.shape = [bs, seq_len]
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), Y.view(-1), ignore_index=0, reduce=False)  # loss.shape = [bs * seq_len]
loss_mask = loss_mask.view(-1) # [bs * seqlen]
loss = torch.sum(loss*loss_mask) / loss_mask.sum()