使用PyTorch从零训练自己的大模型

LLM是如今大多数AI聊天机器人的核心基础，例如ChatGPT、Gemini、MetaAI、Mistral AI等。这些LLM背后的核心是Transformer架构。

本文介绍如何一步步使用PyTorch从零开始构建和训练一个大型语言模型（LLM）。该模型以Transformer架构为基础，实现英文到马来语的翻译功能，同时也适用于其他语言翻译任务。

（本文以论文 "Attention is all you need " 来构建 transformer 架构。）

1.加载数据集

为了让LLM模型能够执行从英文到马来语的翻译任务，需要使用含有英马双语对照的数据集。

为此，这里选择了Huggingface提供的“Helsinki-NLP/opus-100”数据集（https://huggingface.co/datasets/Helsinki-NLP/opus-100）。包含百万级的英文-马来语对照句对，足以确保模型训练的准确性。此外，该数据集还包含了2000条验证和测试数据，且已经预先完成了分割工作，省去了手动分割的繁琐步骤。

# 导入必需的库
# 如果还没安装datasets和tokenizers库，请先安装(!pip install datasets, tokenizers)。
import os
import math
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from pathlib import Path
from datasets import load_dataset
from tqdm import tqdm

# 将device值设置为“cuda”以在GPU上训练，如果GPU不可用则默认回退为“cpu”。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  

# 从以下huggingface路径加载训练、验证和测试数据集。
raw_train_dataset = load_dataset("Helsinki-NLP/opus-100", "en-ms", split='train')
raw_validation_dataset = load_dataset("Helsinki-NLP/opus-100", "en-ms", split='validation')
raw_test_dataset = load_dataset("Helsinki-NLP/opus-100", "en-ms", split='test')

# 用于存储数据集文件的目录。
os.mkdir("./dataset-en")
os.mkdir("./dataset-my")

# 用于保存模型的目录，在训练模型期间每个EPOCHS后保存。
os.mkdir("./malaygpt")

# 用于存储源代码和目标tokenizer的目录。
os.mkdir("./tokenizer_en")
os.mkdir("./tokenizer_my")

dataset_en = []     
dataset_my = []
file_count = 1      

# 为了训练tokenizer（在步骤2中），将训练数据集分成英文和马来语。 
# 创建多个大小为50k数据的小文件，并将其存储到dataset-en和dataset-my目录中。
for data in tqdm(raw_train_dataset["translation"]):
    dataset_en.append(data["en"].replace('\n', " "))
    dataset_my.append(data["ms"].replace('\n', " "))
    if len(dataset_en) == 50000:
        with open(f'./dataset-en/file{file_count}.txt', 'w', encoding='utf-8') as fp:
            fp.write('\n'.join(dataset_en))
            dataset_en = []

        with open(f'./dataset-my/file{file_count}.txt', 'w', encoding='utf-8') as fp:
            fp.write('\n'.join(dataset_my))
            dataset_my = []
        file_count += 1

2.创建分词器

Transformer模型不处理原始文本，只处理数字，因此需要将原始文本转换为数字格式。

这里使用名为BPE（Byte Pair Encoding）的流行分词器来完成这一转换过程。这是一种子词级别的分词技术，已在GPT-3等先进模型中得到应用。

分词器流程

通过训练数据集来训练这个BPE分词器，生成英马双语的词汇表，这些词汇表是从语料中提取的独特标记的集合。

分词器的作用是将原始文本中的每个单词或子词映射到词汇表中的相应标记，并为这些标记分配唯一的索引或位置ID。

这种子词分词方法的优势在于，它能有效解决OOV问题，即词汇表外单词的处理难题。

通过这种方式，我们能够确保模型在处理翻译任务时，无论是常见词汇还是生僻词汇，都能准确无误地进行编码，为后续的嵌入表示打下坚实基础。

# 导入tokenizer库的类和模块。
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace

# 用于训练tokenizer的训练数据集文件路径。
path_en = [str(file) for file in Path('./dataset-en').glob("**/*.txt")]
path_my = [str(file) for file in Path('./dataset-my').glob("**/*.txt")]

# [ 创建源语言tokenizer - 英语 ].
# 创建额外的特殊标记，如 [UNK] - 表示未知词，[PAD] - 用于维持模型序列长度相同的填充标记。
# [CLS] - 表示句子开始的标记，[SEP] - 表示句子结束的标记。
tokenizer_en = Tokenizer(BPE(unk_token="[UNK]"))
trainer_en = BpeTrainer(min_frequency=2, special_tokens=["[PAD]","[UNK]","[CLS]", "[SEP]", "[MASK]"])

# 基于空格分割标记。
tokenizer_en.pre_tokenizer = Whitespace()

# Tokenizer训练在步骤1中创建的数据集文件。
tokenizer_en.train(files=path_en, trainer=trainer_en)

# 为将来使用保存tokenizer。
tokenizer_en.save("./tokenizer_en/tokenizer_en.json")

# [ 创建目标语言tokenizer - 马来语 ].
tokenizer_my = Tokenizer(BPE(unk_token="[UNK]"))
trainer_my = BpeTrainer(min_frequency=2, special_tokens=["[PAD]","[UNK]","[CLS]", "[SEP]", "[MASK]"])
tokenizer_my.pre_tokenizer = Whitespace()
tokenizer_my.train(files=path_my, trainer=trainer_my)
tokenizer_my.save("./tokenizer_my/tokenizer_my.json")

tokenizer_en = Tokenizer.from_file("./tokenizer_en/tokenizer_en.json")
tokenizer_my = Tokenizer.from_file("./tokenizer_my/tokenizer_my.json")

# 获取两个tokenizer的大小。
source_vocab_size = tokenizer_en.get_vocab_size()
target_vocab_size = tokenizer_my.get_vocab_size()

# 定义token-ids变量，我们需要这些变量来训练模型。
CLS_ID = torch.tensor([tokenizer_my.token_to_id("[CLS]")], dtype=torch.int64).to(device)
SEP_ID = torch.tensor([tokenizer_my.token_to_id("[SEP]")], dtype=torch.int64).to(device)
PAD_ID = torch.tensor([tokenizer_my.token_to_id("[PAD]")], dtype=torch.int64).to(device)

3.准备数据集和数据加载器

在构建模型的第三步，着手准备数据集及其加载器。这一阶段的目标是为源语言（英语）和目标语言（马来语）的数据集做好训练与验证的准备。

为此，需要编写一个类，能够接收原始数据集，并利用英语和马来语的分词器（分别为tokenizer_en和tokenizer_my）对文本进行编码处理。编码后的数据会通过数据加载器进行管理，该加载器将按照设定的批次大小（本例中为10）来迭代处理数据集。

如有需要，还可以根据数据量和计算资源的实际情况，对批次大小进行调整。

# 此类以原始数据集和max_seq_len（整个数据集中序列的最大长度）为参数。
class EncodeDataset(Dataset):
    def __init__(self, raw_dataset, max_seq_len):
        super().__init__()
        self.raw_dataset = raw_dataset
        self.max_seq_len = max_seq_len
    
    def __len__(self):
        return len(self.raw_dataset)

    def __getitem__(self, index):
        
        # 获取给定索引处的原始文本，其包含源文本和目标文本对。
        raw_text = self.raw_dataset[index]
        
        # 分离文本为源文本和目标文本，稍后将用于编码。
        source_text = raw_text["en"]
        target_text = raw_text["ms"]

        # 使用源 tokenizer（tokenizer_en）对源文本进行编码，使用目标 tokenizer（tokenizer_my）对目标文本进行编码。
        source_text_encoded = torch.tensor(tokenizer_en.encode(source_text).ids, dtype = torch.int64).to(device)    
        target_text_encoded = torch.tensor(tokenizer_my.encode(target_text).ids, dtype = torch.int64).to(device)

        # 为了训练模型，每个输入序列的序列长度都应等于 max seq length。 
        # 因此，如果长度少于 max_seq_len，则会向输入序列添加额外的填充数量。
        num_source_padding = self.max_seq_len - len(source_text_encoded) - 2 
        num_target_padding = self.max_seq_len - len(target_text_encoded) - 1 

        encoder_padding = torch.tensor([PAD_ID] * num_source_padding, dtype = torch.int64).to(device)
        decoder_padding = torch.tensor([PAD_ID] * num_target_padding, dtype = torch