从零到一打造自己的大模型（四）SFT对话训练

前言

最近看了很多大模型，也使用了很多大模型。对于大模型理论似乎很了解，但是好像又缺点什么，思来想去决定自己动手实现一个 toy 级别的模型，在实践中加深对大语言模型的理解。

在这个系列的文章中，我将通过亲手实践，构建一个 1.2B 的模型，完成模型搭建、tokenizer 训练、模型预训练和指令微调这些流程。记录整个开发过程和其中遇到的各种挑战和对应解决方案。

最后这些内容并不以训练一个足够强大的模型为目标，更多的是走一遍流程，所以里面内容显得十分粗糙。所有的内容都是我对于大模型的理解形成的，如果您发现有任何过时或不准确的地方，请不吝指出。

完整代码请访问这个仓库。

在第一篇文章《从零到一打造自己的大模型（一）模型实现》中实现了一个 decoder only 的小模型。

在第二篇文章《从零到一打造自己的大模型（二）分词器》中成功训练了一个分词器，现在有了模型、数据集和分词器，我们可以开始训练自己的模型了。

在第三篇文章《从零到一打造自己的大模型（三）模型训练》中对模型进行预训练，向其中注入了知识，但是模型还不能灵活使用，现在我们教模型如何使用自己学会的知识。

对话格式

模型在训练的时候目标是预测下一个词，所以对话模型也是一个词一个词生成，为了能让模型像人一样聊天，或者说让模型能够正确处理人对话方式，我们需要构造一个对话模板。

对话模板将不同角色说的话按照先后顺序以一定格式拼成字符串，模型就可以根据字符串继续预测下一个词，这样就达到了生成对话的目的。

Qwen 模型的对话模板如下：

<|im_start|>system
系统提示词<|im_end|>
<|im_start|>user
用户问题<|im_end|>
<|im_start|>assistant
模型回答<|im_end|>

在训练分词器时，我们保留了几个特殊的 token，分别是 <|system|>，<|user|>，<|assistant|>，<|end|>。因此我们可以构造自己的对话模板，这里采用如下的对话模板：

<|system|>系统提示词<|end|>
<|user|>用户问题<|end|>
<|assistant|>模型回答<|end|>

数据集方面采用 stanford_alpaca 中微调数据集，这个数据集中每个数据有 instruction、input 和 output 三个字段，其中 instruction 和 input 是用户的输入，而 output 是模型输出，下面开始构造模板：

IGNORE_TOKEN_ID = -100

def _preprocess(
    source: Dict, tokenizer: PreTrainedTokenizer, max_len: int
) -> Dict[str, torch.Tensor]:
    system_message = "You are a helpful assistant."
    nl_tokens = tokenizer("\n").input_ids
    _system = tokenizer("<|system|>").input_ids
    _user = tokenizer("<|user|>").input_ids
    _assistant = tokenizer("<|assistant|>").input_ids
    _end = tokenizer("<|end|>").input_ids

    input_ids, labels = [], []

    # 系统指令
    system = _system + tokenizer(system_message).input_ids + _end + nl_tokens
    input_ids += system
    labels += _system + [IGNORE_TOKEN_ID] * (len(system) - 3) + _end + nl_tokens
    assert len(input_ids) == len(labels)

    # 输入指令
    if source["input"] != "":
        _input_ids = (
            _user
            + tokenizer(source["instruction"]).input_ids
            + nl_tokens
            + tokenizer(source["input"]).input_ids
            + _end
            + nl_tokens
        )
    else:
        _input_ids = (
            _user + tokenizer(source["instruction"]).input_ids + _end + nl_tokens
        )
    _labels = _user + [IGNORE_TOKEN_ID] * (len(_input_ids) - 3) + _end + nl_tokens
    input_ids += _input_ids
    labels += _labels
    assert len(input_ids) == len(labels)

    # 输出
    _output = _assistant + tokenizer(source["output"]).input_ids + _end
    input_ids += _output
    labels += _output
    assert len(input_ids) == len(labels)

    attention_mask = [1] * len(input_ids)

    if len(input_ids) < max_len:
        diff = max_len - len(input_ids)
        input_ids += [tokenizer.pad_token_id] * diff
        labels += [IGNORE_TOKEN_ID] * diff
        attention_mask += [0] * diff

    input_ids = input_ids[:max_len]
    labels = labels[:max_len]
    attention_mask = attention_mask[:max_len]

    return dict(
        input_ids=torch.tensor(input_ids, dtype=torch.int),
        attention_mask=torch.tensor(attention_mask, dtype=torch.int),
        labels=torch.tensor(labels, dtype=torch.int),
    )

在计算损失的时候，系统提示词和用户输入都是固定的，只需要对输出计算损失，因此在系统提示词和用户输入部分的 labels 设置为 -100 这样计算损失时可以忽略这个词。

有了上面的处理，可以构造自己的 dataset：

class SupervisedDataset(Dataset):
    def __init__(self, path, tokenizer: PreTrainedTokenizer, max_len: int):
        """
        Args:
            path (str): 路径可以是一个文件或者一个包含JSON文件的目录。
            transform (callable, optional): 一个用于进行数据转换的可选函数。
        """
        self.path = path
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.cache = {}
        self.data = []

        # 检查路径是文件还是目录
        if os.path.isfile(path):
            self._load_file(path)
        elif os.path.isdir(path):
            for filename in os.listdir(path):
                if filename.endswith(".json"):
                    file_path = os.path.join(path, filename)
                    self._load_file(file_path)

    def _load_file(self, file_path: str):
        """
        从指定的文件路径加载JSON数据。
        """
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                self.data.extend(data)
            elif isinstance(data, dict):
                self.data.append(data)

    def __len__(self):
        """
        返回数据集中的样本数量。
        """
        return len(self.data)

    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
        """
        返回索引idx对应的数据项。
        """
        if idx in self.cache:
            return self.cache[idx]

        res = _preprocess(self.data[idx], self.tokenizer, self.max_len)
        self.cache[idx] = res
        return res

现在我们测试一下构造的数据集：

tokenizer = CustomTokenizer.from_pretrained("tokenizer")
dataset = SupervisedDataset("alpaca_data.json", tokenizer, max_len=2048)
data_loader = DataLoader(dataset=dataset, batch_size=1)

for batch in data_loader:
    break

tokenizer.decode(batch["input_ids"][0])

<|system|>You are a helpful assistant.<|end|>
<|user|>Give three tips for staying healthy.<|end|>
<|assistant|>1.Eat a balanced diet and make sure to include plenty of fruits and vegetables. 2. Exercise regularly to keep your body active and strong. 3. Get enough sleep and maintain a consistent sleep schedule.<|end|>

至此完成了微调数据集的准备，接下来就可以进行训练了，因为训练目标任然是预测下一个词，因此代码与第三篇文章代码相似。

在训练好模型之后，我们希望模型能够按照文本顺利生成下一个词，首先我们先编写生成token的代码：

def generate(
    self,
    input_ids: torch.IntTensor,
    attention_mask: torch.IntTensor,
    max_new_tokens: Optional[int] = 50,
    temperature: float = 0.95,
):
    device = input_ids.device
    self.eval()
    _input_ids = input_ids.tolist()
    _attention_mask = attention_mask.tolist()
    generate_ids = []

    while len(generate_ids) < max_new_tokens:
        input_ids = torch.tensor(_input_ids, device=device)
        attention_mask = torch.tensor(_attention_mask, device=device)

        with torch.inference_mode():
            logtis = self(input_ids=input_ids, attention_mask=attention_mask)
            logits = logtis[0, -1, :] / temperature
            probs = nn.functional.softmax(logits, dim=-1)
            ids = torch.argmax(probs)

        _input_ids.append(ids.cpu().item)
        _attention_mask.append(1)
        generate_ids.append(ids.cpu().item)

    return generate_ids

这里我们使用 max_new_tokens 控制生成的数量，使用 temperature 控制生成的多样性。由于存在对话模板，我们不能直接送入原始文本而需要构造对应的模板，显然他和对话模板相似：

<|system|>系统提示词<|end|>
<|user|>用户问题<|end|>
<|assistant|>

这样模型按照预测下一个词就会顺利产生回复。

结语

至此我们成功完成了模型从构造到对话训练全过程，大致流程如此，后续可能会加入分布式训练过程。