从sft_clm_mlm三种训练方式来看data_collator——【transformers源码阅读】

背景

最近一直在做大模型的预训练（clm或者mlm），也在做sft（使用指令数据做有监督微调）。

我发现：从数据结构的角度来说，clm、mlm、sft其实本质上都是差不多的。

1. 工程上（或者叫代码上）98%都是相同的。
2. 2%的不同，体现在训练的数据结构上和data_collator部分。

之前也一直想好好写一写transformers包的data_collator部分，这个部分，给很多人的感觉：“不就是数据填充么”，其实没那么简单。他做了不少东西：

1. 比如mlm、clm的实现。
2. 如何在numpy、tensorflow、pytorch中丝滑切换的。

感觉这几个点，还是很有意思的。

因此，在本文中，将从数据结构的角度，分析sft、clm、mlm的异同点，把data_collator部分也顺带介绍了。如果有错误，也希望大佬不要吝啬时间，指导一下。

sft

sft-有监督微调，我现在做的类似于对齐的任务，基本上都是模仿这个仓库来了的https://github.com/tatsu-lab/stanford_alpaca

先不考虑像是lora、量化这样的训练技巧。如果你仔细阅读了上面这个仓库，然后从数据结构的角度来看，整体的思路是这样的：

1. 先用一个训练好的大模型。
2. 整理的数据有三列：instruction、input、output。
3. 然后使用使用一个prompt，将instruction和input搞在一起，变成source，将output直接转换成target。
4. 接下来，把source和target和token.eos_token_id直接拼接在一起，这个时候暂时叫sentence。
5. 然后把sentence通过tokenizer转换成input_ids。
6. 最后一步，要把input_ids复制一份，叫labels。然后把labels前面的，source对应的tokenid，全部变成-100。
7. 那么这个时候，一个面向sft任务的input_ids和labels就已经构造好了。
8. 剩下的就是常规操作，就不介绍了。

在这个任务里面，使用的就是transformers的DataCollatorForSeq2Seq。这个data_collator任务很简单：就是让每一个batch内的input_ids和labels都长度对齐。

https://github.com/tatsu-lab/stanford_alpaca这个仓库，有优缺点。

1. 优点： 提供的思想是非常好的，想法不错。
2. 缺点：但是代码写的是有点拉垮：当数量大的时候，完全没有进度条，完全不能多线程处理。

于是我就基于这个仓库的优点，解决了他的缺点，修改了一下数据处理部分，做了一个给bloom模型的sft代码。具体可以看这个https://github.com/yuanzhoulvpi2017/zero_nlp/tree/main/chinese_bloom

clm

clm-因果模型的训练方法，代码可以参考huggingface的transformers里面给到的代码，https://github.com/huggingface/transformers/examples/pytorch/language-modeling/run_clm.py

他的数据思路，大概是这样的，非常暴力。

1. 一大串的无监督数据，假设这些数据都叫content。
2. 然后把这个content放到tokenzier里面转换成input_ids。
3. 这个时候，有两种做法： 一种是直接将input_ids复制为labels，然后使用default_data_collator； 还有一种做法是使用DataCollatorForLanguageModeling，但是设置里面的参数为mlm=False。
4. 在上面那个步骤中，有个细节，要求要把labels里面所有pad_token_id都要替换成-100。

大家经常说：

1. clm的特征，就是在训练的时候，只能看到左边的词。
2. mlm的特征，就说在训练的时候，可以看到两边的词。

本来训练一个大模型，到这里基本上就可以了，但是我们的目的是把transformers的data_collator看懂，而mlm任务对应的data_collator可以说是这三个任务里面最难的。

mlm

mlm-遮蔽语言模型。代码可以参考huggingface的transformers里面给到的代码： https://github.com/huggingface/transformers/examples/pytorch/language-modeling/run_mlm.py

他的数据思路，大概是这样的。

1. 一大串的无监督数据，假设这些数据都叫content。
2. 然后把这个content放到tokenzier里面转换成input_ids。
3. 这个时候，使用DataCollatorForLanguageModeling。接下来，我将详细介绍这个类里面最重要的函数torch_mask_tokens。

    def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> Tuple[Any, Any]:
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
        """
        import torch
        
        # step 1
        labels = inputs.clone()
        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
        probability_matrix = torch.full(labels.shape, self.mlm_probability)
        if special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        # step 2
        labels[~masked_indices] = -100  # We only compute loss on masked tokens

        # step 3
        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
        # step 4
        # 10% of the time, we replace masked input tokens with random word
        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels

具体步骤已经在上面的代码中标记出来了。

1. step 1里面，就是把inputs复制，成为新的变量叫labels。
2. step 2里面，制作一个新的掩膜，这个掩膜和inputs大小一样。然后使用掩膜对labels部分进行遮盖。对没有盖住地方，设置为-100。
3. step 3里面，在掩膜的基础上，对inputs做操作：对被盖住的80%的token_id用mask_id替换掉。
4. step 4里面，在掩膜的基础上，对inputs继续做操作：对step 3中、没有被mask_id替换掉的token_id中，再用50%的概率，用随机的token_id替换原始的token_id。

因为我语文不太好，表达的有点难受。 但是本质上就是：创建掩膜，按照比例，随机的对labels的部分token_id做替换：一部分用mask_id替换；一部分用随机token_id替换。

所有的loss

最后，在回到模型的loss部分。都是叫自回归loss。 代码都是下面这个样子：

        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            batch_size, seq_length, vocab_size = shift_logits.shape
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)
            )

上面的代码中shift_logits = lm_logits[..., :-1, :]、shift_labels = labels[..., 1:]就可以看出来了。

最后

1. 本篇文章，就是想将sft、clm、mlm三种任务拿出来，从数据处理的角度来比较一下，他们的异同点。在抛开一些训练技巧（lora、量化等），其实可以发现，这三个任务，在代码层面，可以无缝切换。
2. transformers包的data_collator承担了大部分数据处理操作，并不只是承担pad操作。

自我介绍

1. 喜欢阅读transformers源码，对nlp和transformers包感兴趣。如果你对自然语言处理、文本转向量、transformers、大模型、gpt等内容感兴趣欢迎关注我～