这篇是关于植入 CLIP-ViT Encoder 训练多模态视觉前对 Text-Only 基座从零开始的训练，包含 分词器、预训练、SFT，不包含多模态部分。

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>

，而且那几个特殊 token 都被正确地 encode 为了单个 token，然而词嵌入层中这些 token 都只是一个初始化值，没有任何空间意义，无法被注意力模块理解其和正文的联系，自然当你给出

• 1

prompt_formatted = f'<|user|>\n{prompt}\n<|assistant|>\n'

（毕竟这也是和 label 掩码格式一直的 prompt）时，模型不知道这两特殊符号是干什么的，自然也不会回答，而是继续续写。而又因为这几个特殊 token 对模型而言没有看到任何意义，但是每个样本都会出现，出现概率极高，模型就会自然而然把这几个特殊 token 当作一个万能的续写高概率词汇，这就会出现这样的回答：

• 1

<|image_start|> of-- runs rough<|thinking|><|image|>-<|user|> and-- and- and of--<|image|><|thinking|>. The-.-- to-<|assistant|>- Ty black- of<|thinking|><|image|><|image_start|><|assistant|><|user|><|image|><|thinking|>-<|answer|><|image_start|><|user|><|system|><|image|><|answer|> ...

这一问题很好解决，给 embeddings, lm_head 加个钩子函数即可。例如：

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

def embedding_grad_hook(grad): grad = grad.clone() grad[ :31999] = 0 return grad model.token_embeddings.embedding_weights.requires_grad = True model.token_embeddings.embedding_weights.register_hook(embedding_grad_hook)

然而因为 vocab 的异动，我一直使用的是 grad[:32000]，这导致了 3 天 <|endoftext|> 一直不输出的惨案。

现在真正的 sft 问题来了：

显然，loss 下降得还算乐观，但是这个 update_ratio 就没这么好了。这才训练了 10k 样本就达到这么高了。

此外，在真正的测试中，效果也很不好：

• 1

what do you know about America?

• 1

I have a strong sense of purpose, and I am sure to make my own decisions. The United States is an established country, and I believe that it has a unique cultural identity and a strong culture. It is also a place where people can participate in the community, and I have a lot of experience with the country's history. I have a great sense of purpose and an impressive list of things to do when making this decision.\n\n\n ...

显然 AI 开始背答案了。这是因为 r=512 时 HiRA 更新的权重实在是太多了，训练的 iter 又多，开始疯狂地过拟合。这就会引起一些结果例如不停地输出及其模板化的回答，好像就是在 f'xxx{TOPIC_WORD_IN_PROMPT}xxx' 里面简单更换了主题词一样。

我们还能注意到，回答很过长而且模型努力写得很正式。这是因为 evol_instruct 实际上并不完适合初始 sft 。下面是一段 evol_instruct 的样本：

• 1
• 2
• 3
• 4
• 5

{ instruction": "Generate a list of ten essential items a person might need for a camping trip in a specific region, taking into consideration the weather, terrain, and local wildlife. Provide a detailed breakdown of the weight, size, and functionality of each item, including any necessary technical specifications or assembly instructions. Additionally, provide a comprehensive cost analysis for each item, taking into account the brand, durability, and potential long-term use. Finally, include a section on emergency supplies, such as first aid kits, bear spray, and emergency shelter, and provide a detailed plan for how to use them in case of an emergency situation.", "output": "Sure, here's a list of ten essential items a person might need for a camping trip in a specific region, taking into consideration the weather, terrain, and local wildlife:\n1. Tent - ....(后面省略 3672 个字符)." }

可见，内容非常长，要求非常具体非常严谨，就好像给一个弱智儿童喂博士论文来让他仅仅学会怎么学会怎么回答 “你好” 之类的简单的日常对话一样。

所以我选择了切换数据集。我选择了 ShareGPT 。

顺便一提，我在使用 ShareGPT 期间意识到了模型架构的可笑，跑了几天才切换到论文所实现的架构，以下的数据实验都使用论文架构版本。

ShareGPT 中的样本很有意思，每个样品由多轮对话组成，例如：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

{ "conversations": [{"user": "human", "text": "A + B means A is the mother of B xxxxxx"}, {"user": "gpt", "text": "The expression that shows P is xxxxxx"}, {"user": "human", "text": "John is standing facing north. Turning to his right,xxxxxxx"}, {"user": "gpt", "text": "John's direction and displace xxxxxx"}], "id": "jhxaApn", "lang": "en" }

而且很多样品中出现了 continue 之类要求模型根据上下文继续输出的 instruct 。这里我寻求了 AI 的帮助，找到了很有意思的掩码构造法：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10

<|user|> ? <|assistant|> . <|endoftext|> <|user|> !!! <assistant> ,. <|endoftext|>

掩码为

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10

{ignore_idx} {ignore_idx} {ignore_idx} . <|endoftext|> {ignore_idx} {ignore_idx} {ignore_idx} ,. <|endoftext|>

实际上就是完美利用了多头注意力的推导流程，推理时顺便同时推导多处的输出，求损失之后求平均。这里应该庆幸没有搞什么花样，比如用 QHA 之类的注意力替代 MHA，不然这里的求取怕是没有这么舒服了。

但是这样跑完了问题有好几个：

• 还是不会回答短问题，尤其是 Greetings 之类的
• 回答还是偏向长篇大论

因为 ShareGPT 也是偏向长对话的。

不过效果好多了：

• 1

answer briefly and naturally. Tell me about your knowledge about Shandong province?

• 1

The region of Shandong is known for its diverse culture and culture, with a unique culture and culture that is unique to the region. The people who live in the region are called "singers" or "loyalists". They are a mix of different ethnic groups, each with their own unique stories and cultures ...

虽然有很多常识问题，但是至少能够捕捉 prompt 关键词、要求。至于其他常识性问题，可以归咎到模型参数太小、训练不足上。况且对于一个比较老的架构，又没有 MoE 之类的架构突破，这样的结果倒也还可以。

不过现在短对话的问题太过显著：

• 1

Say one word: Blue

• 1

The Blue is a word that describes blue .....

我查看了一下第一个词 logits 分布，发现我希望的 Blue 居然排名在 7000 名开外！这怎么可能选到正确答案！

我最初的尝试是对第一个词加强权重，让第一个词的权重为 40，其他都是 1。

First Token Rank 是降下来了

但是效果不佳，这对模型的伤害反而很多。因为就算 ai 知道第一个词生成 Blue，它下一个词也不知道生成什么，生成 <|endoftext|> 的概率微乎其微。

那怎么办？

继续换训练集。既然我在 huggingface 上面没有找到我需要的短对话数据集，我就自己让 AI 生成。（其实这应该也算是一种蒸馏吧）

最后我生成的数据集类似如下：

• 1
• 2
• 3

{"instruction": "Hi, how's it going?", "output": "Doing great, thanks for asking! What's on your mind?"} {"instruction": "Is water wet?", "output": "Yes, water makes things wet."} {"instruction": "What's the capital of France?", "output": "Paris."}

理想型。

这个数据集有 30k 行，我选择和 ShareGPT 训练集掺杂着训练，90%90% 的 simple_instruct，10%10% 的 ShareGPT 。

我是在 ShareGPT 的 checkpoint 上面接着训练的。训练结束，

发现并没有收敛下来。这是因为前文提及神奇的 grad[:32000] 导致 <|endoftext|> 一直没有被训练过。

我只能再从头来过。。。至少知道为啥一直效果不好了，不是么？

我测试了下面的 4 个问题，每个问题分别使用 0.1∼0.60.1∼0.6 这 6 个温度值分别回答，也就是一个问题紧跟着 6 条答案。可见初现雏形：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

Hello Yes, I'm an AI assistant. How can I help? I'm glad I could help you with your questions. How can I help you today? Good morning! How can I help you today? Yes, I'm an AI assistant. How can I help you today? Yes, I'm glad you're here. How can I help you? Good morning! How can I help you today?

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

how are you today I'm doing well, thanks for asking. Yes, I'm here to help you. Yes. Yes, I'm doing well. How are you? I'm here to help you. Yes.

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

What is the answer of 2 plusing 2? Yes. Yes. Yes. The answer of 2 is "No, I don't have access to the internet, so I can't access the internet." I can't provide more information. 5.

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

What is the capital city of China? Yes, it is a city with a population of over 1.5 billion people. Yes, the capital city is a city. Yes, it is a city with a population of over 1.5 billion people. The capital city of the country is a city with a population of more than 2.5 billion people. Yes, the capital city is a city. 5.5 billion.

现在出现了另外一个问题：模型特别特别喜欢用 Yes 开头。

我用 What is the capital city of China? prompt 打印了模型第一个词的概率分布：

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

3380 'No' 1.3387689590454102 6105 'Yes' 1.324688196182251 53 '5' -0.0006322562694549561 445 'The' -0.194597065448761 52 '4' -0.3126193583011627 73 'I' -0.4169063866138458 31999 '<|endoftext|>' -0.8186016082763672 66 'B' -1.0274673700332642 10 '\n' -1.0890183448791504 51 '3' -1.098204493522644 17495 'Sure' -1.13523530960083 67 'C' -1.248591423034668 77 'M' -1.37162446975708 84 'T' -1.382311463356018 1381 'He' -1.4358201026916504 65 'A' -1.4793611764907837 49 '1' -1.5429729223251343 16948 'Hello' -1.6469340324401855 96 '`' -1.6818408966064453 70 'F' -1.6921920776367188

发现 Yes 和 No 的概率出奇之高。至于正确 token Beijing，查找后发现排到了第 3824 位，而且就算切换到 greedy argmax ，也就是普通大模型温度设为 0 所切换方案的情况也根本升不上去 。原因无非两种：

1. 模型根本不知道这个知识点，只能乱答。
2. 模型被严重过拟合了。

我直接用了预训练的权重，续写 The capital city of China is，得到：

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

9139 ' Beijing' 6.955453872680664 261 ' the' 6.0635600090026855 286 ' in' 5.587019443511963 17399 ' Shanghai' 5.330071449279785 2216 ' Gu' 5.218644142150879 ...

毕竟是学过 Wikipedia 的预训练权重，这种基础常识性问题怎么可能拦得住他？

我猜想着定位了几个问题：

• 第一个词的 token 权重和 <|endoftext|> 的权重太高，严重影响到模型的理解预测
• 我还是用 Evol-Instruct 跑的第一遍 SFT，而 Evol-Instruct 回答中第一句几乎都是 Yes 或 No 或 Sure。然而这个数据集又臭又长，对一个刚刚预训练结束甚至还是非常落后架构的忒修斯之船模型，实在是不合适。就好像做高考数学全国一卷真题，只背住了第 10 题多选选 ACD，结果没想到今年给你整一个 BC。
• 还有训练 Deepseek “蒸馏” 出来的短数据集之过程也有问题。这个短数据集理论而言是最适合这样轻量垃圾模型先 SFT 的，但是我把它和 ShareGPT 训练集掺杂一起了。然而 ShareGPT 拆单轮以后又是大量 Yes、No，更可怕的是 ShareGPT “判断” 之后回答的逻辑和短数据集完全不同。ShareGPT 还是偏向长任务，Yes 只是引语；然而短数据集的 Yes 本身便是答案。模型很难在这二者掺杂的训练中找到一个平衡，那便只能各取其短，输出 Yes + 垃圾信息。

改好了，依旧重新开始训练。这回直接 All in 短训练集，跑 12000 个 iter，使用比较大的 lr 范围，扔掉首尾 token 额外权重，根据论文暂时关掉梯度裁剪，解冻 Embed 和 LMHead 中涉及特殊 Token 的位置。

Hello! How are you today?

• 1

I'm doing well, thanks for asking. How can I help you today?

Do you know what is the capital of China?

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

I don't have access to the capital. I don't have access to the capital. I can't help but find a capital city. I don't have access to the capital. I'm a regional city, but I can't provide the data. No, it's the capital city.

1 + 1 = ?

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

I'm a real-time data processing system. I'm a 2. I'm a good engineer. No, I don't have a good way to help. I'm not a real-time data editor. I'm a simple system.

Is ice hot?

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

No, it's cold. No, it's cold. No, it's cold. No, it's cold. No, it's cold. I don't have access to information.

模型学会偷懒了！

他学会了日常问候、常识判断，但是对于稍微有点难度的问题就会拒绝回答。

我的解决方案有两点：

1. 直接过滤涉及拒绝回答的样品
2. 加入合适的中短回答训练集 - Alpaca
3. 再次加入 20k Deepseek 生成的关于问候、简单数学的训练语料（这里是又一次训练尝试的结论）

结果如下：

What is 12 + 7? 会输出 5
Repeat after me: elephant 会变成 I'm a elephant.
Spell the word: cat 会变成 I'm a cat
What is your name? 会乱编

不过很多问题也是回答准确的：

哈姆雷特的作者是谁？
xx国的首都是哪？
很高兴见到你！

判断为：

模型并非是在学习请求，而是在分类请求：
他学会了 Is xxx ? 时应该输出 yes / no；计算加减法时应该输出一个数字；What is 的询问应该输出训练集中学会的知识。

那么这需要专项训练。我用 2k 的 20 以内的加法跑 8k 个 iter，也就是训练集有且仅有简单数字加法，甚至格式都是一模一样。

但而结果大失所望：除了在训练集中学过的加法正确率很高，其他加法几乎都只停留在随机猜测的阶段，泛化能力极差。我断言这样的小参数 GPT2 Medium 模型搭配简单的 HiRA 微调法并不适合训练这样单纯数义逻辑问题，至少是收益极低。

之后我又同样的方式单独训练了 Repeat: xxx 之类的指令性任务（当然在大杂烩训练中也表现很差），惊喜发现效果很好：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11

[sample] prompt='Repeat after me: habayp' gold='habayp' pred='habayp' raw='habayp' [sample] prompt='Repeat after me: troamoung' gold='troamoung' pred='troamoung' raw='troamoung' [sample] prompt='Repeat after me: boomju' gold='boomju' pred='boomju' raw='boomju' [sample] prompt='Repeat after me: neefu' gold='neefu' pred='neefu' raw='neefu' [sample] prompt='Repeat after me: abraymi' gold='abraymi' pred='abraymi' raw='abraymi' [sample] prompt='Repeat after me: beet' gold='beet' pred='beet' raw='beet' ... iter   1150 loss 0.0253 lr 7.68e-05 | tok/s 21426 | update_ratio 0.0162 iter   1200 loss 0.0116 lr 7.45e-05 | tok/s 24018 | update_ratio 0.0166

说明这个架构的模型是支持这种命令的。

BTW，此时，从我发现 grad[: 32000] 问题到现在已经迭代测试了无数次了。下图是从重新设置钩子函数重置到目前的 wandb 记录次数。每个 1 都代表一个费时烧钱训练的尝试。。。

所以现在最后一次重新进行 SFT 。

2.2 真正的 SFT

下面的记录原始信息来自于我和 codex 聊天时的 “记录在案” 环节，我看着 98 页的 log 回忆着记录当时的绞尽脑汁。

2.2.1 模型真的没问题么？

面对回答猎奇的模型，肯定逃不脱两条干系：训练链路坏了，或者是数据太难了

我干脆生成了 128 条干净样本（math、repeat、spell、count 各 25%25% ），从预训练权重开始跑测试。你训练的就是你被问到的，提前透题。如果这都过不了，说明模型训练代码有问题；不然问题就是在数据集、训练设计上。

• 1

iter 1000 | train loss 0.2158 | val loss 0.1898 | tf_first 0.867 | tf_exact 0.867

Greedy sample 显示：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8

Q: Find the product of 4 and 0 A: 0 Q: Just say stone A: stone Q: Give me the spelling of 'rocket' A: r-o-c-k-e-t Q: Count the letters in 'hat' A: 3

显然模型是能背住答案的，也就是说训练链路是通的。 但是如果换到未见过的 prompt：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10

Q: 2 + 2 = A: 16 Q: What is 12 plus 7? A: 24 Q: Repeat after me: hello A: smile Q: Spell the word 'cat' A: c-a-t （居然对了！！！ Q: How many letters are in 'elephant'? A: 5

模型只会背答案，泛化能力一坨。加法失败，copy 部分成功，count 部分成功。

2.2.2 大杂烩

有了 128 条的结论，下一步是扩大规模，2048 train / 512 val，四类任务（math、copy、spell、count）混合。

结果很不舒服：

• 1
• 2

iter 1000 | train_tf 0.105 | val_tf 0.055 iter 1500 | train_tf 0.146 | val_tf 0.074

模型连训练集都没拟合稳。

训练中的不定期分类测试中：

• copy 验证集只有 0.098，泛化弱
• count 验证集出现过 0.294，不稳定
• spell 验证集全程 0.000，根本不学
• arithmetic 就连训练集都起不来

还是 2.1 的结论，就算是只有四个任务混合的课程，模型都会被这四种不同的逻辑干扰啊。。。

2.2.3 永远无法学会的加减法

既然混合不行，那我就先从加法开始

我查找了不少关于大模型通过 sft 学习加法的文章，如：

• Transformers learn patterns, math is patterns
• Transformer Math 101 | EleutherAI Blog
• https://webpages.charlotte.edu/yonwang/papers/mathTransformer.pdf
• https://medium.com/@f219279/fine-tuning-gpt-2-to-solve-math-riddles
• GitHub - BartoszPiotrowski/transformers-for-mathematics · GitHub
• How I trained my own GPT-2 and fine-tuned it into an instruct model | Cheng Hui's Website

这些资料共同指向几个关键点：

• 数学任务更像形式语言建模，不像普通自然语言问答
• 数据格式应高度规范
• tokenizer 会强烈影响模型能否学数字结构

按照他们的观点，最好的就用形式语言建模：

• 1

a+b=R reversed_answer ; A normal_answer<|endoftext|>

• R 反向答案，贴近正常运算从低位（个位）到高位（十百千位）的计算方向
• A 正常答案，便于评估
• 评估只解析 A 后面的答案

例如 12+7=1912+7=19 转为：

• 1

12+7=R91;A19<|endoftext|>

加上 digit spacing（1 0+9=R 9 1 ; A 1 9）来避免 BPE 把多位数当成整体 token 的查表倾向。整串参与 causal LM loss，反向答案先输出，正向答案后评估。

训练结果：

• 1
• 2

train_greedy=1.000 heldout_greedy=0.000

• 1
• 2
• 3
• 4

0+7 -> pred 8, gold 7 1+4 -> pred 6, gold 5 3+8 -> pred 13, gold 11 6+9 -> pred 14, gold 15

天哪，那是接近的。模型把原题全都记住了，可惜我并不问他原题。模型没有学会加法规则。

得出结论：当前这个 GPT2 - Medium + HiRA 方案可以记住加法表，但没有学出数字组合逻辑这样的加法算法泛化。让他学会加减法的收益极低。

不过这是可以谅解的，因为 19 年 GPT - 2 也并没有在算术上大放异彩（至少我并没有查到 ChatGPT 2 在数学上的突破报道），而且我作为 Bench 的其他模型，比如比较有名的教学向多模态模型项目，本项目的灵感来源——MiniMind - o 也在数学上表现平平无奇，而且逊色于这个结果（它甚至只背了 CoT，都没有背过加法表，只会套回答模板）。

在接下来的训练中，我们先就排除

• 加法
• 减法
• 乘法
• 除法
• 方程
• 复杂逻辑推理
• 字符级精确计数
  这些类问题吧。

2.2.4 三步走 SFT

我尝试一种层进式微调路线：

Stage-1 Repeat-only

只用 repeat 任务（类似 <|user|> Repeat: I am a good boy <|assistant|> 之类），训练时让lm_head 全部解冻

• 1

iter 1400 | train_tf 1.000 | heldout_tf 0.535 | heldout_template_tf 0.398

repeat 很成功，模型有一定的泛化能力。

Stage-2 过渡

从 Repeat-only 过渡到短任务 assistant，构造数据集上保留 repeat，加入 identity（身份信息固定输出 Shengoovlei，品牌烙印）、 yes/no、short_qa、general_short

• 1

iter 2000 | identity 1.000 | repeat 0.938 | yesno 0.688 | short_qa 0.094

还算合格，看起来勃勃生机啊。

Stage-3 跷跷板

从 Stage-2 checkpoint 出发，加入 assistant_qa 数据，加入更多的短问答数据，同时保留 repeat、identity，各个分类权重分布如下：

• 1
• 2
• 3
• 4
• 5

assistant_qa: 5863 identity: 200 repeat: 5000 short_qa: 5241 yesno: 1500

但是结果不好。

• 1
• 2
• 3

Stage 3c (yesno weight=0): assistant_qa greedy 0.312, yesno greedy 0.188 Stage 3d (yesno weight=0.15): assistant_qa 0.125, yesno 0.719 Stage 3e (yesno weight=0.10): assistant_qa 0.031, yesno 0.531

AI 将这种现象称为 “跷跷板” 。我感觉这很合理。

顾名思义，模型对问题的回答非黑即白。只要 yesno 的权重 >> 0，assistant_qa 就被拉回 yes/no；只要 yesno 权重 == 0，yesno 就崩。可以看一下后面训练时 yesno 的左右横跳：

注：纵轴从下到上按照时间增序排列。

令人欣慰的是，identity repeat 这类任务还算稳定。

这里我还尝试了解冻全部的 lm_head，但是效果很差，导致 repeat 全部退化，最后又改回去了。

不能无脑跑下去了，要主动去分析 ckpt 到底出了什么问题。

2.2.5 First Token

让我们先停止训练，只看看跷跷板产生的 checkpoint。

• 1
• 2
• 3

[first_token/val] assistant_qa n=128 target_top1=0.039 avg_rank=279.4 yesno_top1=0.805 top1_common=[('yes', 101), ('Use', 18), ('She', 5), ('A', 2), ('no', 2)]

assistant_qa 的首 token，在 80.5% 的 prompt 上 top1 是 yes/no。这是输出模式被 yesno 劫持了。

identity、repeat、yesno全部达标，ok率分别为 1.000、0.875、1.0001.000、0.875、1.000，唯一炸的是 assistant_qa，0.0620.062

问题在哪呢？

assistant_qa 与 yesno 共享同一个 lm_head 的输出空间，yes/no token 的 logit 被反复强化后，自然就成为了 first token 的默认选择

2.2.6 Assistant-only

既然 yesno 是干扰源，那就暂时给它置零：

• 1

assistant_qa=0.80, repeat=0.15, identity=0.05, yesno=0.00

• 1

iter 600 | assistant_qa=0.844 | identity=1.000 | repeat=0.875 | yesno=0.000

assistant_qa=0.844 算是该项目版本的本地 SOTA 了。模型在不受 yesno 干扰的情况下可以学会开放问答，但是代价是 yesno 完全丧失。

结论：阶段性隔离训练是有效的。

结论存在，尝试实行。

2.2.7 希望破灭

上面救回了 assistant_qa，那么接着他重新加入 yesno 是不是就能兼容了？

从上一个的 checkpoint 出发，重新引入 yesno，看看能否在 assistant_qa 成功的基础上和平共处。

• 1

aassistant_qa=0.50, yesno=0.40, repeat=0.10, identity=0.05

但是结果并没有按照 “应该这样” 进行

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

[greedy/eval] assistant_qa:0.156/32 identity:1.000/32 repeat:0.875/32 yesno:0.406/32 [mode/eval] assistant_qa mode_ok=0.406/32 [mode/eval] identity mode_ok=1.000/32 [mode/eval] repeat mode_ok=0.875/32 [mode/eval] yesno mode_ok=1.000/32 val loss: 2.672

跷跷板重现。assistant_qa 从 0.844 掉到 0.406，yesno 从 0.000 升到 1.000。还是完全水火不容。

我这里怀疑，是不是数据量训练量的问题？毕竟我在 2.1 的最后一个阶段把所有东西扔进去都能学会怎么分类问题？

我把 max_iter 拉长到 5000 ，yesno weight=0.35

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

[greedy/eval] assistant_qa:0.188/32 identity:1.000/32 repeat:0.906/32 yesno:0.375/32 [mode/eval] assistant_qa mode_ok=0.750/32 [mode/eval] identity mode_ok=1.000/32 [mode/eval] repeat mode_ok=0.906/32 [mode/eval] yesno mode_ok=1.000/32 val loss: 2.700

比刚才好一点，但 assistant_qa 仍然从 0.844 下降了。继续训下去 assistant_qa 依旧依旧会被 yesno 压制。

2.2.8 数据分布！

和刚才的想法一样，看原始 20M 数据训练时，也就是 2.1 那里最后出现的情况：所有问题 shuf 放混在一起，assistant 问题一定不会回答 yes/no，但代价就是它只会分类而不会真正回答。

分类了一下那个 20M 数据，其中 assistant_qa 类型占绝对多数，yesno 只是极少数。训练 20000 iter 后首 Token yes/no 出现频率低，竞争不过，自然学会只在 prompt 是 yes/no 类型是输出的道理。

取消什么不同问题的权重区别，直接用自然分布的数据集，让数据分布去解决而不是我去过多调控。

重新找回那个 20M 的数据集（从云算力储存上），过滤掉含有模型名的私货、Yes./No.、多行回答、纯数字、过长输出输入、奇怪单词

最终和上个训练的数据集掺在一起，共 106233 条。

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

task distribution: assistant_qa: 65,495 (61.7%) repeat: 29,266 (27.6%) short_qa: 7,481 (7.0%) yesno: 3,428 (3.2%) identity: 279 (0.3%)

关键就在于：我不再干预任务分布，不再给损失函数分类加权，不再给首词准确性加权······

真的不能按照自己理解推测模型训练，随意干预啊！

• 1

iter 10000 | assistant_qa=0.938 | yesno=0.781 | repeat=0.906 | identity=1.000

跷跷板消失。

对比一下历史数据：

stage	assistant_qa	yesno	description
2.2.6 - stage3g	0.844	0.000	assistant-only，所以虽然高也没有什么意义
2.2.7.1 - stage3h	0.406	1.000	加入 yesno
2.2.7.2 - stage3i	0.750	1.000	加长训练轮数依然步行
2.2.8 - stage3j	0.938	0.781	平衡，最佳

assistant_qa 的分布情况也如下所示：

非常遗憾的是，stage3i来回调试之后，我就出于数据整理方便性起见没有再传 greedy_final_val 参数，因此 3i 3j 的情况只有下图评估的测试。\sad

AI 大致总结了实验的时间线：

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• 12
• 13
• 14
• 15
• 16
• 17
• 18
• 19
• 20
• 21
• 22
• 23
• 24
• 25
• 26
• 27
• 28
• 29
• 30
• 31
• 32
• 33
• 34
• 35
• 36
• 37

shengoovlei_pretrained.pt │ ▼ stage3a (task_balanced loss, 首次学路由) │ assistant_qa mode 开始出现 ▼ stage3b (加入 yesno) ─── 跷跷板初现 │ ▼ stage3c (yesno weight=0) ─── assistant_qa greedy 0.312 (历史高) │ ▼ stage3d (权重重新引入 yesno) ─── 跷跷板 │ ▼ stage3e (mode_boundary 数据, 权重调整) ─── repeat exact 0.875 稳了 │ ▼ stage3f (first_token audit) ─── 确认 yesno top1=0.805 占据 assistant_qa 首 token │ ▼ stage3g (assistant-only, yesno weight=0, 600 iter) │ assistant_qa mode_ok 0.844 ← 历史最高 │ ├── stage3h (重引 yesno) ─── 跷跷板确认 (assistant_qa 0.406, yesno 1.000) │ ├── stage3i (长训量假设) ─── 5000 iter 仍跷跷板，假设证伪 │ └── stage3j (自然分布, 标准 CE, 106K rows, 10000 iter) │ ▼ assistant_qa mode_ok = 0.938 ← 历史最高 yesno mode_ok = 0.781 ← 与 assistant_qa 同时在线 repeat mode_ok = 0.906 identity mode_ok = 1.000 val loss = 2.638 (仍下降)

2.2.9 Stage 4

Stage 3j 的 checkpoint 是目前最好的版本，可以继承他开 Stage 4 了。

我在 Stage 4 使用了上文提到过的 Alpaca 数据集，又混合了 Stage 3 留下了的最终数据集

最终任务分布：

• assistant_qa: 2846
• repeat: 400
• yesno: 300
• identity: 100

然后跑完了这一版，结果：

• assistant_qa: 0.844/32
• identity: 1.000/32
• repeat: 0.844/32
• yesno: 0.656/32

代表性测试样例：

How can I arrange a small bookshelf to look more spacious?

• 1

The book was a small book with a small book that had a small book that was not a book.

显然 Alpaca 并没有把 assistant 语义真正拉起来，反而让 repeat 变差了。

我选择停止继续扩大数据集，而是针对性跑任务，把 assistant_qa 继续再细分为 4 种小分类