Abstract本项目旨在从零搭建一个基于 GPT-2 Medium 衍生架构的多模态大模型使其至少支持文本、图像两种模态输入同时尽可能减少对Pytorch封装库的直接调用在此中熟练掌握基础Transformer的知识、模型的预训练微调等等处理技术和对多模态技术的了解。Apology由于成本 / 技术 / 与主题无过大关联 等原因我并没有手撕一些组件列举如下torchnumpyeinops最基础的函数如*这样的矩阵运算、广播机制还有arangezerosrearrange之类的基础数据处理函数torch自带的torch.nn.functional.scaled_dot_product_attention: 我独立实现了SDPA但是非常不幸由于不会手撕 FlashAttention又考虑到训练 / 推理时间成本重新用 Pytorch 优化过的注意力写了一个新的MultiHeadAttentionCLIP-ViT : 预算实在有限训练的成本还是太高了只能在了解其原理、感慨 OpenAI 财大气粗之后直接使用成品了。Info项目开源地址Github模型权重开源地址ModelScope模型体验地址ModelScopePart 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_size32000。训练是很慢的跑了半个小时发现进度很差意识到需要优化了。BPE 训练无非两步一步是分词一步是合并。分词是很快的使用作业pdf给出的那个分块即可没啥优化的价值。问题就出在合并上。我最初实现的合并关键代码是1234for 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执行一次匹配了。但是还有一个很大的问题出在1best_pair max(pair_counts.items(), keylambda x: (x[1], x[0]))上面。当pair数量极多会查找极慢。干脆放弃 python 直接转向 C 使用 C 的std::unordered_map哈希表 搭配上PyBind11直接用 C 重写这部分。编译为动态链接库之后就能直接被我们的训练脚本当作module引入了。于是我们训练出了vocab_size32000的一个 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()这块。然后就可以进入预训练了。
