Abstract

本项目旨在从零搭建一个基于 GPT-2 Medium 衍生架构的 LLaVA 多模态大模型，使其至少支持文本、图像两种模态输入，同时尽可能减少对 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

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() 这块）。

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

1.2 预训练

我选择的参数是：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• 12
• 13
• 14
• 15
• 16

vocab_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

架构则是：

1. Embedding ( 32000×102432000×1024 )
2. RoPE
3. 24 x Transformer Block，每个 Block 有：
   • 3.1 RMSNorm (pre‑norm)
   • 3.2 Multi‑Head Self‑Attention (16 头，d_model=1024，4 x 线性层)
   • 3.3 RMSNorm (pre‑norm)
   • 3.4 SwiGLU FFN (3 x 线性层)
4. RMSNorm
5. 线性层

架构之外，我使用了 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 is

• 1

Beijing. The central square is the center of every major city in China. China's metropolitan area is made up of hundreds of sub-provincial cities, each with its own municipal government and city councils....

部分偏差就很大：

The capital of America is

• 1

the 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 1：

Zaozhuang, Shandong province is

• 1

a number-one commercial and investment business center in the People's Republic of China.

但是不管逻辑、知识如何，目前的模型已经学会了基础的英语语法，有能力生成有信息的语句。

当然他和我们脑海中的那个 AI 还不一样。

you are a helpful assistant. answer me gently. Nice to meet you!

• 1

Fare 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 对论文的爬取。

1. HiRA: Parameter-Efficient Hadamard High-Rank Adaptation for Large Language Models
2. Entropic Distribution Matching for Supervised Fine-tuning of LLMs: Less Overfitting and Better Diversity
3. WizardLM: Empowering large pre-trained language models to follow complex instructions

其中 HiRA 是 LoRA 的一个衍生版本，旨在通过哈达玛积使更新的 Adapter 参数可以使用更高的秩；GEM 是使用最大熵原理改善传统交叉熵损失导致 SFT 过拟合的一种损失函数；Evol-Instruct 则是一种重写数据集的算法（这里我直接使用了 huggingface 上公开的产生的 训练集 而非跟着论文重写跑出来一份）。

所使用的数据集：

1. WizardLM_evol_instruct_70k
2. sharegpt-english
3. 自己通过 Deepseek 获取的短对话样本 - 30k

2.1 痛苦的训练过程

整个训练过程我投入了 2 个多星期，改了无数个版本。

由于 HiRA 的在高秩上的特色，我看了一俩篇 blog 就直接开了一个参数：r=512，并且替换了 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 之外的所有权重冻结，实现代码如下：

• 1
• 2
• 3
• 4
• 5
• 6

for 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 都是类似

• 1
• 2
• 3
• 4
• 5

<|user|> Hello! <|assistant|> Hello! Nice to meet you today! <endoftext>