Abstract

本项目旨在从零搭建一个基于 GPT-2 Medium 衍生架构的多模态大模型，使其至少支持文本、图像两种模态输入，同时尽可能减少对 Pytorch 封装库的直接调用，在此中熟练掌握基础 Transformer 的知识、模型的预训练微调等等处理技术和对多模态技术的了解。

Apology

由于成本 / 技术 / 与主题无过大关联 等原因，我并没有手撕一些组件，列举如下：

• torch numpy einops 最基础的函数，如 @ * 这样的矩阵运算、广播机制，还有 arange zeros rearrange 之类的基础数据处理函数
• torch 自带的 torch.nn.functional.scaled_dot_product_attention : 我独立实现了SDPA，但是非常不幸由于不会手撕 FlashAttention，又考虑到训练 / 推理时间成本，重新用 Pytorch 优化过的注意力写了一个新的 MultiHeadAttention
• CLIP-ViT : 预算实在有限，训练的成本还是太高了，只能在了解其原理、感慨 OpenAI 财大气粗之后直接使用成品了。

Info

项目开源地址：Github

模型权重开源地址：ModelScope

模型体验地址：ModelScope

Part 1. Pre-training

这部分是预训练，目的是让 AI 简单学会基础语法，能够续写语料。

1.0 开始的开始

这个项目的起源在于斯坦福的 AI 神课 CS336 。个人认为这门课的作业部分教育指导意义远大于讲课视频，Transformer、优化器等等等等都是要亲手实现才能理解，看视频、看blog、单纯调用 torch.nn.functional 库中的函数与类永远无法弥补此点。

基础 Transformer 模块中的大部分代码（即 modules.py 中的大多数模块）都是直接从我亲自跟着 Assignment 1 敲下的代码摘除或改编的。

Assignment 1 在带着我们亲手实现手搓完整的 BPE分词器、线性层、Embeddings、RMSNorm、SwiGLU、RoPE、SDPA及MHA、CE损失函数、SGD、AdamW、学习率调节器、梯度裁剪、DataLoader、ckpt设计和基础文本生成后，便让我们自己训练一个模型。然而从各个模块的分别实现到组装为一个完整的训练循环是另一码事。训练期间遇到了很多问题。

1.1 预训练的前戏

首先是训练集获取问题。显然为了模型的能力上限、应用场景考虑，不能使用简单的 TinyStories 训练集。但是由于算力、成本有限，我选择了 CS336 特供 - OpenWebText 训练集。

我先训练了一个 BPE 分词器，设定 vocab_size=32000 。

训练是很慢的，跑了半个小时发现进度很差意识到需要优化了。

BPE 训练无非两步，一步是分词，一步是合并。

分词是很快的，使用作业pdf给出的那个分块即可，没啥优化的价值。
问题就出在合并上。我最初实现的合并关键代码是：

• 1
• 2
• 3
• 4

for segment in segments: it = re.finditer(PAT, segment) for i in it: list_origin_words[i.group()] += 1

这是对作业pdf原理直接的转写，但是显然对于真实实际应用场景很糟糕。

我在实现作业之初用 flameprof 记录分析了一下对 TinyStories bpe 训练过程，发现整个火焰图几乎都是 regex 相关的内容，而其他函数的占比极小：

显然，训练脚本花了大部分的时间放在 finditer() 正则匹配上面了。那么我们直接空间换时间，只需先进行一次预匹配，创一个词频表，然后就无需在每一个 merge 执行一次匹配了。

但是还有一个很大的问题出在

• 1

best_pair = max(pair_counts.items(), key=lambda x: (x[1], x[0]))

上面。当 pair 数量极多，会查找极慢。干脆放弃 python 直接转向 C++ ，使用 C++ 的 std::unordered_map 哈希表 ，搭配上 PyBind11 ，直接用 C++ 重写这部分。编译为动态链接库之后就能直接被我们的训练脚本当作 module 引入了。

于是，我们训练出了 vocab_size=32000 的一个 BPE 分词器，并且将词典表和合并序列直接用 pickle 保存下来。

那么下一步的当务之急是将这 12G 的纯文本 Tokenize 一下。显然，我在 CS336 中 Tokenizer 的 encode() 函数对长文本的效果很不好，于是我用 AI 快速跑了一个快速 Tokenizer，依旧使用的了 C++ 来写核心的循环合并部分（毕竟 C++ 的哈希表还是太强了），又加入了并行计算，然后就以 8 核全部 98.7% 以上榨干极致性能实现了快速的 Token 化。经过不精确估算，高性能 Tokenizer 的 encode 使 Token 化的速度提升了大约 46000%46000%。将训练语料 Token 化后，我用 np.uint16 格式 + torch.save 储存到了本地的 .npy 文件（def make_npy() 这块）。

然后就可以进入预训练了。