Abstract本项目旨在从零搭建一个基于 GPT-2 Medium 衍生架构的 LLaVA 多模态大模型使其至少支持文本、图像两种模态输入同时尽可能减少对Pytorch封装库的直接调用在此中熟练掌握基础Transformer的知识、模型的预训练微调等等处理技术和对多模态技术的了解。Apology由于成本 / 技术 / 与主题无过大关联 等原因我并没有手撕一些组件列举如下torchnumpyeinops最基础的函数如*这样的矩阵运算、广播机制还有arangezerosrearrange之类的基础数据处理函数torch自带的torch.nn.functional.scaled_dot_product_attention: 我独立实现了SDPA但是非常不幸由于不会手撕 FlashAttention又考虑到训练 / 推理时间成本重新用 Pytorch 优化过的注意力写了一个新的MultiHeadAttentionCLIP-ViT : 预算实在有限训练的成本还是太高了只能在了解其原理、感慨 OpenAI 财大气粗之后直接使用成品了。Info项目开源地址GithubPart 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()这块。然后就可以进入预训练了。1.2 预训练我选择的参数是12345678910111213141516vocab_size 32000 context_length 1024 batch_size 8 d_model 1024 num_layers 24 num_heads 16 d_ff 2752 rope_theta 10000 max_iters 3_000_000 max_learning_rate 2e-4 min_learning_rate 2e-5 weight_decay 0.1 betas (0.9, 0.95) eps 1e-8 max_grad_norm 1.0架构则是Embedding ( 32000×102432000×1024 )RoPE24 x Transformer Block每个 Block 有3.1 RMSNorm (pre‑norm)3.2 Multi‑Head Self‑Attention (16 头d_model10244 x 线性层)3.3 RMSNorm (pre‑norm)3.4 SwiGLU FFN (3 x 线性层)RMSNorm线性层架构之外我使用了AdamW优化器还有余弦学习率调节器并且使用 wandb 记录数据。这大概是一个 GPT2 - Medium 的参数参数量大概是 370M 。根据 Chinchilla 最优法则理想训练需要 20 倍参数量即 7B 的训练 Token 量。不过由于时间、预算有限我最终只跑了 3443​ 个 CS336 特供款 OpenWebText Sample 和大约 1331​ 个 Wikipedia English Chunks。然后在某炼丹平台租了一张 RTX PRO 6000 96GB开始了预训练。训练时发现训练的速度确实很慢我想大概是Scaled Dot-Product Attention的问题。我之前学过一点点Triton算子这种涉及大量张量融合计算的问题应该使用一个fused的版本。那么对于SDPA而言便是FlashAttention V2。但是这个的实现有些偏离主题了干脆直接拿‌torch.nn.functional.scaled_dot_product_attention代替了。请原谅我在这块的作弊 Orz我最初用的 OpenWebText 跑跑了一共大概 7 个小时分了两段。损失变化如下最初这里的学习率其实并没有收敛很好当时只训练到了橙色pretrain_owt_p3那里。。。我想这大概与那时只训练了 1.6B 左右的 token 有关。我最初并没有接着训练 Wikipedia 训练集而是直接开了接下来的 sft 微调训练。结果在疯狂碰壁一直改了两天 bug 耗了 50 多 RMB 后发现还是预训练不充足。于是继续训练了wiki保留 checkpoint 里面的模型权重但是重新初始化了AdamW优化器。这里补上了验证集评估的记录训练wiki阶段可见训练效果还不错。现在一共训练了大概 3B 的 token 量。部分测试还可以The capital of China is1Beijing. The central square is the center of every major city in China. Chinas metropolitan area is made up of hundreds of sub-provincial cities, each with its own municipal government and city councils....部分偏差就很大The capital of America is1the city of Dallas, Texas. Dallas is a small unincorporated community and census-designated place located in Dallas County in the central part of the state...还会把我的家乡描述成全国经济 TOP 1Zaozhuang, Shandong province is1a number-one commercial and investment business center in the Peoples Republic of China.但是不管逻辑、知识如何目前的模型已经学会了基础的英语语法有能力生成有信息的语句。当然他和我们脑海中的那个 AI 还不一样。you are a helpful assistant. answer me gently. Nice to meet you!1Fare your expectations for love. A cool and healthy girl. Did you know that this is the way she became? Did you know what it was like to be a beautiful woman? Did you understand that the idea of b...他不会回答你的问题因为他很少见到过这样问答式的样本。他遇到这样的问题第一反应是根据学习过的内容续写。可能输入的类似指令与训练集中的某篇散文类似所以他就会接着续写那个散文而非去和 “正常 AI” 一般回答。这就需要进行下一步——SFT 微调了。Part 2. Supervised Fine-Tuning这部分是对预训练模型进行监督微调目的是让 AI 学会对答模板。使用的技术栈是 HiRA GEM并且使用 LoRA CE 作为 Bench。2.0 技术选择使用了下面三篇论文使用的技术。论文的选择部分参考了 AI 对论文的爬取。HiRA: Parameter-Efficient Hadamard High-Rank Adaptation for Large Language ModelsEntropic Distribution Matching for Supervised Fine-tuning of LLMs: Less Overfitting and Better DiversityWizardLM: Empowering large pre-trained language models to follow complex instructions其中 HiRA 是 LoRA 的一个衍生版本旨在通过哈达玛积使更新的 Adapter 参数可以使用更高的秩GEM 是使用最大熵原理改善传统交叉熵损失导致 SFT 过拟合的一种损失函数Evol-Instruct 则是一种重写数据集的算法这里我直接使用了 huggingface 上公开的产生的 训练集 而非跟着论文重写跑出来一份。所使用的数据集WizardLM_evol_instruct_70ksharegpt-english自己通过 Deepseek 获取的短对话样本 - 30k2.1 痛苦的训练过程整个训练过程我投入了 2 个多星期改了无数个版本。由于 HiRA 的在高秩上的特色我看了一俩篇 blog 就直接开了一个参数r512并且替换了TransformerLM里面注意力模块的proj_q、proj_k、proj_v、proj_o四者的线性层。这在时候看来相当愚蠢首先r过大论文中和 LoRA 的对照组最大r才开到32这里开到512直接更新了整个模型接近 28%28% 的参数过拟合爆炸其次是 q k v o 四个层也不合适按照论文的做法应该只替换attn.proj_q、attn.proj_k、attn.proj_v、ffn.w2、ffn.w3这几个线性层。然而先抛开这些我遇到了一个问题预训练的语料并没有涉及特殊 token (|endoftext|、|assistant|、|user|、|pad|等 )而在预训练时我直接将除了 HiRALinear 之外的所有权重冻结实现代码如下123456for param in model.parameters(): param.requires_grad False for name, param in model.named_parameters(): if .A in name or .B in name: param.requires_grad True这导致特殊词汇的向量并未得到训练虽然构造的ids都是类似12345|user| Hello! |assistant| Hello! Nice to meet you today! endoftext