Tiny-universe手戳大模型 Tiny-LLM--task2

该项目在于实现一个简单的大语言模型，从训练tokenizer开始，到训练模型，再到使用模型生成文本。仅使用Numpy和Pytorch即可实现一个简单的大语言模型训练，显存使用2G左右。以下为项目效果展示。

按照以下步骤进行训练。

1. 训练Tokenizer： python train_vocab.py --download True --vocab_size 4096
2. 数据预处理：python preprocess.py
3. 训练模型：python train.py
4. 使用模型生成文本：python sample.py --prompt "One day, Lily met a Shoggoth"

Step 1: 训练Tokenizer

首先，我们需要为文本处理训练一个Tokenizer。Tokenizer的作用是将文本转换为数字序列，以便模型能够理解和处理。我们使用的数据集是 TinyStory ，它是一个由GPT-3.5和GPT-4生成的小型故事数据集，包含简短的故事，且词汇量有限。在这个任务中，我们采用字符级Tokenizer，将文本中的每个字符映射为对应的数字。通过以下命令可以下载数据集并训练Tokenizer。

python train_vocab.py --download True --vocab_size 4096

LLaMA2 的词表大小为 32,000，但由于 TinyStory 数据集较小，词汇量有限，我们将词表大小设置为 4,096。训练完成后，我们得到的 Tokenizer 能够将文本转换为数字序列，也可以将数字序列还原为文本。

class Tokenizer:
    def __init__(self, tokenizer_model=None):
        """
        初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。

        参数:
        tokenizer_model: str, 可选，分词器模型的路径，如果不指定则使用默认路径 TOKENIZER_MODEL。
        """
        # 如果提供了分词器模型路径，使用该路径；否则使用默认模型路径
        model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
        # 确保模型文件存在
        assert os.path.isfile(model_path), model_path

        # 加载 SentencePiece 模型
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        self.model_path = model_path

        # 获取分词器的特殊token和词汇表大小
        self.n_words: int = self.sp_model.vocab_size()  # 词汇表大小
        self.bos_id: int = self.sp_model.bos_id()       # 句子开头 (BOS) 的ID
        self.eos_id: int = self.sp_model.eos_id()       # 句子结尾 (EOS) 的ID
        self.pad_id: int = self.sp_model.pad_id()       # 填充 (PAD) 的ID

        # 验证分词器词汇表大小是否正确
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        """
        将字符串编码为词元ID列表。可以选择是否添加句子开头 (BOS) 和句子结尾 (EOS) 标记。

        参数:
        s: str, 要编码的字符串。
        bos: bool, 是否在编码的词元列表前添加 BOS 标记。
        eos: bool, 是否在编码的词元列表末尾添加 EOS 标记。

        返回:
        List[int]: 编码后的词元ID列表。
        """
        # 确保输入是字符串类型
        assert type(s) is str
        # 使用SentencePiece将字符串编码为词元ID
        t = self.sp_model.encode(s)
        # 如果需要BOS标记，将其添加到词元列表开头
        if bos:
            t = [self.bos_id] + t
        # 如果需要EOS标记，将其添加到词元列表末尾
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t: List[int]) -> str:
        """
        将词元ID列表解码为字符串。

        参数:
        t: List[int], 词元ID列表。

        返回:
        str: 解码后的字符串。s
        """
        return self.sp_model.decode(t)

在这个 Tokenizer 类中，我们首先初始化了一些特殊的 token ID，这些特殊 tokens 在自然语言处理任务中有着重要作用，分别用于填充、处理未识别的词汇、表示句子的开头和结尾等。在模型训练和推理过程中，正确处理这些特殊 tokens 对于提升模型性能至关重要。

接着，我们定义了两个关键方法：

1. encode 方法：该方法负责将输入文本转换为 token ID 序列。通过加载预训练的 Tokenizer 模型，我们可以对文本进行分词，将其拆解为词或子词，并将其映射为相应的数字表示。这个数字序列可以被模型接受用于训练和推理。

2. decode 方法：与 encode 方法相反，decode 方法用于将 token ID 序列还原为可读的文本。它将数字序列转换回对应的 tokens，并拼接成完整的文本，从而可以对模型的输出进行解释和展示。

这些方法的定义使得我们在使用过程中，可以非常方便地在文本与数字序列之间进行转换，为模型的输入与输出提供接口。大家可以使用以下代码测试 Tokenizer 的功能，验证其是否能够正确地将文本转换为数字序列，或者将数字序列还原为文本。

# 测试 Tokenizer
enc = Tokenizer('./data/tok4096.model') # 加载分词器
tetx = 'Hello, world!' # 测试文本
print(enc.encode(text, bos=True, eos=True)) # 编码文本
print(enc.decode(enc.encode(text, bos=True, eos=True))) # 解码文本

OUTPUT:
[1, 346, 2233, 4010, 1475, 4021, 2]
Hello, world!

Step 2: 数据预处理

在训练模型之前，首先需要对数据进行预处理。这一步的核心任务是将文本数据转换为模型能够理解的数字序列。具体来说，文本中的每个字符、单词或子词都需要被映射为一个唯一的数字 ID，这样模型才能处理这些数据。

# 定义分片处理函数
def process_shard(args, vocab_size, tokenizer_model_path):
    """ 处理数据分片，将其中的文本进行分词并保存为二进制文件 """
    ···


# 定义预处理函数，用于对多个数据分片进行批量处理
def pretokenize(vocab_size):
    """ 预处理所有的数据分片，并将分词后的数据保存为二进制文件 """
    ···


class PretokDataset(torch.utils.data.IterableDataset):
    """从磁盘加载已预处理的分词数据，并将其以 PyTorch 张量的形式返回。"""

    def __init__(self, split, max_seq_len, vocab_size, vocab_source):
        """
        初始化数据集。

        参数:
        split: str, 数据集的分割方式（'train' 或 'test'）。
        max_seq_len: int, 最大序列长度，用于生成输入输出序列。
        vocab_size: int, 词汇表的大小。
        vocab_source: str, 词汇表的来源（'llama2' 或 'custom'）。
        """
        super().__init__()
        self.split = split  # 数据集划分（训练集或测试集）
        self.max_seq_len = max_seq_len  # 最大序列长度
        self.vocab_size = vocab_size  # 词汇表大小
        self.vocab_source = vocab_source  # 词汇表来源

    def __iter__(self):
        """
        返回迭代器，按批次加载数据并生成模型输入/输出。
        """
        # 获取DataLoader的worker信息（用于并行数据加载）
        worker_info = torch.utils.data.get_worker_info()
        worker_id = worker_info.id if worker_info else 0  # worker ID
        # 获取分布式训练的rank信息（用于多GPU训练）
        rank = dist.get_rank() if dist.is_initialized() else 0
        # 基于worker_id和rank生成唯一的随机数种子，确保数据在每个worker和rank之间是唯一的
        seed = 42 + worker_id + 1337 * rank
        rng = random.Random(seed)
        print(f"Created a PretokDataset with rng seed {seed}")

        # 根据词汇表来源决定数据路径
        if self.vocab_source == "llama2":
            # 如果使用 Llama 2 词汇表，.bin 文件和 .json 文件在同一目录下
            bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
        elif self.vocab_source == "custom":
            # 如果使用自定义词汇表，.bin 文件在 tok{N} 目录下
            bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}")
            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))

        # 根据数据集划分使用不同的分片文件
        # 训练集使用所有分片文件，测试集只使用第一个分片
        shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
        assert len(shard_filenames) > 0, f"在 {bin_dir} 中未找到任何 .bin 文件"

        while True:
            # 随机打乱分片文件
            rng.shuffle(shard_filenames)
            for shard in shard_filenames:
                # 使用 memmap 读取文件，使得数据留在磁盘上，减少内存占用
                m = np.memmap(shard, dtype=np.uint16, mode="r")
                # 计算该分片中的批次数量
                num_batches = len(m) // self.max_seq_len
                num_batches -= 1  # 去掉最后一个不完整的批次
                assert num_batches > 0, "这个分片文件太小了？请检查。"
                # 随机打乱批次索引
                ixs = list(range(num_batches))
                rng.shuffle(ixs)
                # 对每个批次生成输入 x 和目标输出 y
                for ix in ixs:
                    start = ix * self.max_seq_len  # 批次起始索引
                    end = start + self.max_seq_len + 1  # 批次结束索引
                    # 将数据转换为 NumPy 数组并拷贝到 RAM 中
                    chunk = torch.from_numpy((m[start:end]).astype(np.int64))
                    # 模型输入 x 是当前批次的前 max_seq_len 个词元
                    x = chunk[:-1]
                    # 模型输出 y 是下一个词元
                    y = chunk[1:]
                    # 生成 x, y 对
                    yield x, y


class Task:
    @staticmethod
    def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs):
        ds = PretokDataset(**dataset_kwargs)
        dl = torch.utils.data.DataLoader(
            ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers
        )
        for x, y in dl:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            yield x, y

在这部分中，首先定义了 process_shard 函数，用于处理数据分片。该函数的主要功能是将文本数据分词后，转换为更高效的二进制文件格式，以便后续更快速地加载和处理数据。接下来，我们定义了 pretokenize 函数，用于批量处理多个数据分片。通过这一函数，所有数据可以并行处理，进一步加快预处理的速度。然后，我们设计了一个 PretokDataset 类，用于加载已预处理好的数据集。我们继承了 torch.utils.data.IterableDataset 来定义该数据集，这使得我们可以更灵活、高效地处理数据。在这个类中，核心是 __iter__ 方法，它负责生成用于训练的数据批次。最后，我们还定义了一个 Task 类，专门用于迭代数据集，并生成模型所需的输入和目标输出。这一部分的设计确保了数据流的顺畅对接，为模型训练提供了标准化的数据输入。可以通过以下代码来测试预处理后的数据集。

Step 3: 训练模型

在数据预处理完成后，我们就可以开始训练模型了。我们使用的模型是一个和LLama2结构一样的 Decoder only Transformer模型，使用Pytorch实现。相关代码在model.py文件中。此处不再赘述，源码中有详细的中文注释，且我们在之前的文章中也有详细的介绍。

在模型这一部分可以重点看一下生成式模型是如何实现生成token的，可以查看model.py文件中的Transforerm类中的generate方法。

在完成数据预处理后，我们就可以开始训练模型了。我们使用的模型是一个与 LLaMA2 结构相同的 Decoder-only Transformer 模型，采用 PyTorch 实现。具体的实现细节已经包含在 model.py 文件中，在此不再赘述。该源码中包含详细的中文注释，此外我们在之前的文章中也对模型架构进行了深入介绍。

在模型部分，建议重点关注生成式模型如何生成 token 的过程。可以参考 model.py 文件中的 Transformer 类，尤其是 generate 方法的实现，它展示了模型如何基于已有的上下文生成后续 token 的机制。

@torch.inference_mode()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
        在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
        """
        for _ in range(max_new_tokens):
            # 如果序列上下文过长，截断它到最大长度
            idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
            
            # 前向传播获取序列中最后一个位置的 logits
            logits = self(idx_cond)
            logits = logits[:, -1, :] # 只保留最后一个时间步的输出
            
            if temperature == 0.0:
                # 选择最有可能的索引
                _, idx_next = torch.topk(logits, k=1, dim=-1)
            else:
                # 缩放 logits 并应用 softmax
                logits = logits / temperature
                if top_k is not None:
                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits[logits < v[:, [-1]]] = -float('Inf')
                probs = F.softmax(logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            
            # 将采样的索引添加到序列中并继续
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

在 generate 方法中，我们首先获取序列中最后一个位置的 logits，然后基于这些 logits 生成新的 token。接着，生成的新 token 会被添加到序列中，模型随后会继续生成下一个 token。通过这种迭代过程，我们能够生成完整的文本。接下来，您可以使用以下命令开始训练模型。

python train.py

迭代了2000次以后，loss还有1.9685，val loss还有2左右。

# -----------------------------------------------------------------------------
# I/O 配置，用于定义输出目录和训练时的日志记录与评估设置
out_dir = "output"  # 模型输出保存路径
eval_interval = 2000  # 评估间隔步数
log_interval = 1  # 日志记录间隔步数
eval_iters = 100  # 每次评估时迭代的步数
eval_only = False  # 如果为True，脚本在第一次评估后立即退出
always_save_checkpoint = False  # 如果为True，在每次评估后总是保存检查点
init_from = "scratch"  # 可以选择从头开始训练（'scratch'）或从已有的检查点恢复（'resume'）

# 数据配置
batch_size = 8  # 每个微批次的样本数量，如果使用梯度累积，实际批次大小将更大
max_seq_len = 256  # 最大序列长度
vocab_size = 4096  # 自定义词汇表大小

# 模型配置
dim = 288  # 模型的隐藏层维度
n_layers = 8  # Transformer的层数
n_heads = 8  # 注意力头的数量
n_kv_heads = 4  # 模型分组
multiple_of = 32  # 在某些层的维度必须是该数的倍数
dropout = 0.0  # Dropout概率

# AdamW优化器配置
gradient_accumulation_steps = 4  # 梯度累积步数，用于模拟更大的批次
learning_rate = 5e-4  # 最大学习率
max_iters = 100000  # 总的训练迭代次数
weight_decay = 1e-1  # 权重衰减系数
beta1 = 0.9  # AdamW优化器的β1参数
beta2 = 0.95  # AdamW优化器的β2参数
grad_clip = 1.0  # 梯度裁剪阈值，0表示不裁剪

# 学习率衰减配置
decay_lr = True  # 是否启用学习率衰减
warmup_iters = 1000  # 学习率预热的步数

# 系统设置
device = "cuda:0"  # 设备选择：'cpu'，'cuda'，'cuda:0'等
dtype = "bfloat16"  # 数据类型：'float32'，'bfloat16'，'float16'

Step 4: 使用模型生成文本

在模型训练完成后，会在output目录下生成一个ckpt.pt文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。

python sample.py --prompt "One day, Lily met a Shoggoth"

遇到的问题

在训练过程中torch出现了各种报错，一开始在数据预处理运行python preprocess.py的代码时候报了AttributeError: module ‘torch.utils’ has no attribute ‘data’，然后我在 preprocess.py脚本文件里面加了import torch.utils.data 这一行代码，就运行正常了，紧接着又报了attributeerror module ‘torch.utils.data‘ has no attribute ‘iterabledataset‘这个错误,pytorch1.2之前的版本都是没有torch.utils.data.IterableDatset这个类的.
查阅了pytorch 1.2有关torch.utils.data.DataLoade 的技术文档。
At the heart of PyTorch data loading utility is the torch.utils.data.DataLoader class. It represents a Python iterable over a dataset, with support for map-style and iterable-style datasets,customizing data loading order,automatic batching,single- and multi-process data loading,automatic memory pinning.
class torch.utils.data.IterableDataset
An iterable Dataset.

All datasets that represent an iterable of data samples should subclass it. Such form of datasets is particularly useful when data come from a stream.

All subclasses should overrite iter(), which would return an iterator of samples in this dataset.

When a subclass is used with DataLoader, each item in the dataset will be yielded from the DataLoader iterator. When num_workers > 0, each worker process will have a different copy of the dataset object, so it is often desired to configure each copy independently to avoid having duplicate data returned from the workers. get_worker_info(), when called in a worker process, returns information about the worker. It can be used in either the dataset’s iter() method or the DataLoader ‘s worker_init_fn option to modify each copy’s behavior.
Example 1: splitting workload across all workers in iter():

Example 2: splitting workload across all workers using worker_init_fn:

因为cuda版本太旧了，就把原来的cuda10.0版本卸载了，去https://developer.nvidia.com/cuda-toolkit-archive下载cuda 11.8，然后在https://developer.nvidia.com/cudnn-downloads下载cudnn 9.4.0

安装好以后就能

在安装torch时候又遇到时候遇到问题了

输入pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118，安装以后，运行python train.py又报错了，显示AssertionError Torch not compiled with CUDA enabled这个问题，直接用不了gpu训练，然后回忆之前安装Pytorch（GPU版）的时候，我首先创建了一个虚拟环境，而这个虚拟环境的作用就是隔绝外界操作，相当于构建了一个独立空间。然后，我是在这个虚拟环境内安装的Pytorch（GPU版）。这就意味着，我在Pycharm上引入conda environment，是没有办法使用Pytorch（GPU版）的，因为Pytorch（GPU版）已经被虚拟环境隔离了。

然后使用官网的conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

安装完以后又遇到新的错误，究其原因其实是，anaconda的环境下存在两个libiomp5md.dll文件。所以直接去虚拟环境的路径下搜索这个文件，可以看到在环境里有两个dll文件：解决方式就是把虚拟环境本身路径下的libiomp5md.dll删除了，然后就可以正常运行了


完结撒花

参考时候

• llama2.c
• https://github.com/KMnO4-zx/tiny-llm
• https://pytorch.org/docs/1.2.0/data.html#torch.utils.data.IterableDataset