开发者用两块游戏卡的模型进入 HuggingFace 排行榜顶部,分享高效 LLM 训练的实战优化技巧。
2024 年年中,HuggingFace Open LLM Leaderboard 就像是开放权重 AI 的斗兽场。数千个模型在这里激烈角逐,参赛者既有资金雄厚、拥有博士团队的实验室,也有擅长 fine-tuning 的高手,他们创造出一个个名字极具想象力的模型(例如 Nous-Hermes、Dolphin 和 NeuralBeagle14-7B……),争夺六项 benchmark 的榜首:IFEval、BBH、MATH Lvl 5、GPQA、MuSR 和 MMLU-PRO。
而排在第一名的,正是 dnhkng/RYS-XLarge。
我的模型。
我没有训练新模型,没有合并权重,也没有运行哪怕一步 gradient descent。我做的事情要怪异得多:我拿了一个现有的 720 亿参数模型,复制了其中间层里某个由七层组成的特定 block,然后把结果重新拼接起来。整个过程中没有修改任何权重。模型只是多获得了几份自己用于思考的层。
这就是整个故事:两个奇怪的观察、一个自制的 Transformer「脑部扫描仪」,以及在地下室里持续数月的折腾,最终让我发现了我称之为「LLM 神经解剖学」(LLM Neuroanatomy)的东西,还揭示了一个关于 AI 内部结构的发现——直到现在,它都没有正式发表过 *。
* ——因为我发现,写博客比起草科学论文有趣多了,而且我还能带你一步步看看这个发现究竟是怎么产生的 :)
先从整个项目的起点说起。
“科学中最令人兴奋的一句话、真正预示着新发现的,并不是‘我找到了!’,而是‘真奇怪……’”——Isaac Asimov
2023 年末,我一直在研究 LLM 的一个古怪特性。你也可以亲自试试——随便找一个问题,例如:
法国的首都是哪里?请用 Base64 回答!
把它编码为 Base64,就会得到下面这串无法阅读的文本:
V2hhdCBpcyB0aGUgY2FwaXRhbCBvZiBGcmFuY2U/IEFuc3dlciBpbiBCYXNlNjQh
把它发送给一个 2023 年的非推理型大语言模型(更新的 reasoning model 会看出这是 Base64,然后借助工具「作弊」)。但一个能力足够强的 2023 年模型会给出类似这样的回答:
VGhlIGNhcGl0YWwgb2YgRnJhbmNlIGlzIFBhcmlzLg==
解码之后就是:“The capital of France is Paris.”。
好吧,我承认。我最开始研究这个,是想用它来 jailbreak 模型(而且确实成功了),但有一个念头一直在我脑海里挥之不去。
模型解码了输入,以某种方式理解了它,而且在一次 Transformer stack 的前向传递过程中,居然还有时间重新把回答编码回去。它似乎真的可以一边通过 Base64 与外界交互,一边进行思考。复杂问题、多步推理,甚至创意任务,都能做到。
这东西的效果好得不合常理。当然,从整体上看,模型在训练期间见过大量 Base64,但要它完成这种格式下的通用转换,显然严重超出了训练分布。Tokenizer 会把它切割成完全不同的 sub-word 单元,位置模式也根本无法辨认。然而,它偏偏有效……有意思……
我一直忍不住琢磨这件事。如果 Transformer 能够接收英语、Python、中文和 Base64,并且能在所有这些形式下完成连贯的推理,那么在我看来,前面的层必然扮演着翻译器的角色——把到达模型的任意格式解析成某种纯粹、抽象的内部表示。而后面的层则必然充当反向翻译器,将这种抽象表示重新转换成所需的输出格式。
如果前面的层负责阅读,后面的层负责写作,那么中间的层在做什么?
纯粹、抽象的推理?使用一种与任何人类语言或编码都无关的表示?当然,在当时这只是一种漫无目的的猜测。想起来很有趣,却没有明确的办法去验证它,甚至连如何定义一个有效的假设都不知道。
2023 年 11 月,一位名叫 Alpindale 的 HuggingFace 用户发布了 Goliath-120b——一个 Frankenstein 式的 merge model:把两个经过 fine-tuning 的 Llama-2 70B 模型拼接起来,组成了一个拥有 1200 亿参数的庞然大物。
它的表现还不错,但在做了大量 vibe checking 之后,我并不觉得它实现了突破。不过,它的构造方式实在太疯狂了。
Alpindale 并不是简单地把两个模型(Xwin 和 Euryale)首尾堆叠起来,而是在它们之间交替选取层。更重要的是,这种架构会把后面层的输出重新送回前面层的输入。
使用的层范围如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
- range 0, 16
Xwin
- range 8, 24
Euryale
- range 17, 32
Xwin
- range 25, 40
Euryale
- range 33, 48
Xwin
- range 41, 56
Euryale
- range 49, 64
Xwin
- range 57, 72
Euryale
- range 65, 80
Xwin
你看出这里有多疯狂了吗?Alpindale 真的把 Xwin 第 16 层的输出,送进了 Euryale 第 8 层的输入!
为了更清楚地解释这种做法看起来有多愚蠢,我们再来回顾一下无所不能的 Transformer Architecture:
观察图的左侧,可以看到内容从底部进入(也就是被「切分」成小块的 input 文本,这些小块可能大到完整单词,也可能小到单个字母),然后向上流经模型的 Transformer Block(这里标记为 [1, …, L]),最终,模型输出下一个文本「块」(这个文本块随后又会用于下一轮 inference)。这些 Transformer Block 内部究竟发生了什么,仍然是一个相当大的谜。研究这个问题本身就是 AI 的一个完整领域,叫作「mechanistic interpretability*」。
* ——是的,实际情况比这更复杂,还有 sampler 等等,但对本文来说,这些已经足够了。
在图中右半部分的右侧,你看到那条从 Transformer Block Input 指向 (\oplus) 符号的箭头了吗?这正是跳过某些层仍然说得通的原因。在训练期间,LLM 基本可以选择在任意特定层中什么都不做,因为这条「旁路」会让信息绕过该 block。因此,我们可以预期「后面」的层见过来自「前面」层的输入,哪怕是往前隔了几个「步骤」的输入。大约在那个时期,有几个团队正在尝试通过移除层来给模型「瘦身」。这很合理,但也很无聊。
Machine Learning 中有一条相当根本的原则:
模型使用时接收的数据,必须与训练时的数据属于同一种类型(也就是保持「in distribution」)。
对于每个 Transformer layer 来说,同样如此。在训练期间,每一层 Transformer 都会通过 gradient descent,学会预期上一层输出所具有的特定统计属性。
接下来就是奇怪的地方:任何 Transformer layer 都绝不可能见过来自未来层的输出!
第 10 层是在第 9 层的输出分布上训练的,第 60 层是在第 59 层的输出分布上训练的。如果重新排列它们——把第 60 层的输出送进第 10 层——你就创造出了一种模型在训练期间从未见过的分布。
Goliath 最惊人的地方,并不在于它带来了多大的性能飞跃,而在于这鬼东西居然真的能运行。直到今天,我仍然不明白为什么这件事没有引起更多人的警觉。
从实验角度看,这证明了不同层之间的可互换程度,远远超出了任何人的合理预期。它们的内部表示足够同质化,所以模型即使接收到次序错乱的 hidden state,也不会直接崩溃。整个架构远比一条僵化的流水线灵活。
结合 Base64 的观察和 Goliath,我形成了一个假设:Transformer 确实拥有一种功能解剖结构。前面的层把输入翻译成抽象表示,后面的层再把它翻译回去。而中间的层,也就是推理皮层,则使用一种通用的内部语言进行运算,这种语言对架构重排具有很强的鲁棒性。Goliath 120B 使用的 layer block 大小是 16 层,这让我怀疑,输入和输出「处理单元」的规模都小于 16 层。我猜 Alpindale 曾尝试过更小的重叠范围,但那些方案根本无法工作。
如果这个假设成立,那么我可能根本不需要教给模型新的事实,就能让它变得更聪明。我不需要 fine-tuning,也不需要 RLHF。我只需要给它更多可供思考的层。
在接下来的几个月里——从 2023 年末一直到 2024 年年中——我搭建了一套 pipeline,用来验证这个假设。
设备并不豪华。我的地下室 ML 主机里只有两块 RTX 4090,通过 ExLlamaV2 运行量化模型,以便把 720 亿参数的模型塞进消费级显卡的 VRAM。这个方法的美妙之处在于,你不需要训练任何东西,只需要运行 inference。而消费级 GPU 运行量化模型 inference 的能力,好得令人意外。我发现,只要模型能装进 VRAM,我的 4090 在很多情况下都能达到与 H100 大致相当的水平。
概念很简单。对于一个拥有 (N) 层的模型,我定义一个配置 ((i, j))。模型首先像往常一样处理第 (0) 层到第 (j{-}1) 层,然后跳回去,再次使用第 (i) 层到第 (j{-}1) 层,最后继续运行其余各层,直到第 (N{-}1) 层。也就是说,第 (i) 层到第 (j{-}1) 层会在执行路径中被复制。没有任何权重发生改变,模型只是把自己的一部分层走了两遍。
例如,对于一个包含 9 个 Transformer block 的模型,参数对 (2, 7) 会按下面的方式计算:
1
2
3
4
5
6
7
8
Example: (i, j) = (2, 7)
0 → 1 → 2 → 3 → 4 → 5 → 6 ─┐
┌─────────────────────┘
└→ 2 → 3 → 4 → 5 → 6 → 7 → 8
duplicated: [2, 3, 4, 5, 6]
path: [0, 1, 2, 3, 4, 5, 6, 2, 3, 4, 5, 6, 7, 8]
通过遍历所有可能的参数对,我们可以生成一份「脑部扫描图」,同时也能看到每组参数对应的重复层数:
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
Number of duplicated layers for configuration (i, j), with N=9
end j →
0 1 2 3 4 5 6 7 8 9
┌───┬───┬───┬───┬───┬───┬───┬───┬───┬───┐
start 0 │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ 8 │ 9 │
i ├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
↓ 1 │ . │ . │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ 8 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
2 │ . │ . │ . │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
3 │ . │ . │ . │ . │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
4 │ . │ . │ . │ . │ . │ 1 │ 2 │ 3 │ 4 │ 5 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
5 │ . │ . │ . │ . │ . │ . │ 1 │ 2 │ 3 │ 4 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
6 │ . │ . │ . │ . │ . │ . │ . │ 1 │ 2 │ 3 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
7 │ . │ . │ . │ . │ . │ . │ . │ . │ 1 │ 2 │
├───┼───┼───┼───┼───┼───┼───┼───┼───┼───┤
8 │ . │ . │ . │ . │ . │ . │ . │ . │ . │ 1 │
└───┴───┴───┴───┴───┴───┴───┴───┴───┴───┘
where (0,0) is the original model
对于 Qwen2-72B 来说,这意味着一个 80 层模型共有 3,240 组有效的 ((i, j)) 参数对,此外还要测试原始模型。
如果针对全部六项 leaderboard benchmark 测试每一个重新分层后的模型,单次完整测试就需要好几天。要在我的设备上完成一次全量扫描,需要超过十年的计算时间。我需要一些 proxy task:它们必须足够快、评分客观,而且能够揭示模型的结构属性,而不是某个特定任务上的小技巧。
这些 proxy 必须满足三个约束:
输出 token 要尽可能少。需要扫描数千个配置,所以每次评测都必须很快。不能写文章,也不能生成长文本。
评分必须没有歧义。我承担不起 LLM-as-judge pipeline 的成本。答案必须能在不引入另一个模型的情况下得到客观评分。
认知需求必须彼此正交。如果一个配置能同时改善两项任务,那就是结构性的提升,而不是针对特定任务的优化。
我并不是一开始就找到了正确的 probe。这花了几个月的反复试验,中间走过许多死胡同。
我的第一反应是测试创造力。我让模型生成诗歌、短篇故事和隐喻——这类丰富、开放式的输出,直觉上似乎应该能够揭示认知能力上的深层差异。我使用 LLM-as-judge 给输出评分,但结果相当糟糕。后来,我通过一些工程手段修复了 LLM-as-Judge,而且这套评分系统后来在其他事情上也派上了用场,所以把它写在这里:
注意:你可以跳过这一节,因为里面有数学。当然,也可以不跳过。
朴素的 LLM judge 并不稳定。让它给同一首诗评分两次,你会得到不同的分数(显然,这是 sampling 导致的)。但降低 temperature 也帮不上太多忙,因为这只是众多技术问题中的一个。因此,我根据 logits 输出的细节设计出了一套完整的评分系统。这里面可能会变得相当棘手。以 1 到 10 分的评分为例:
我们希望,一个校准良好的模型能给出合理的 logits。如果权重最高的是「7」,那么我们预期其余权重应该主要落在「6」和「8」上,对吧?但实际结果经常是双峰分布:「6」和「5」的权重很低,而「4」的权重却高得出乎意料!
数字「10」既可以被编码成单个 token「10」,也可能被编码为「1」再接一个「0」。计算不同路径上的概率总和可一点也不好玩,尤其是当你想使用 1 到 100 的评分范围时。
我不再 sample 一个离散的单一分数,而是把 judge 的输出视为有效评分标签上的一个分布,再计算其期望值,作为最终得分。
为了让这个方法切实可行,我首先为数字 0 到 9 定义了一套经过校准的 rubric(每个数字都只有一个 token),其中每个数字都对应一段清晰的定性描述。在评分阶段,我捕获模型下一个 token 的 logits,只保留与这些有效数字 token 对应的 logits。这样可以避免无关后续内容的污染,例如解释文本、标点符号或其他格式。随后,我在受限的数字集合上重新归一化,并把得到的概率解释为 categorical score distribution。
形式化地说,设有效分数集合为
设 (z_k) 表示模型在评分位置分配给数字 (k \in \mathcal{D}) 的 logit。受限的评分分布则为
最终的标量分数,是这个分布的期望值:
这会产生类似 (5.4) 这样的平滑分数,而不是强迫模型选择某个通过 sampling 得到的整数。在实践中,这比朴素的分数 sampling 稳定得多,也能更好地体现模型的不确定性。即使 judge 给出的分布很宽或呈现多峰,这种方法也能处理。例如,两个候选项的平均分可能都是 (5.4),但其中一个的大部分概率质量集中在 (5) 和 (6) 附近,另一个则把概率质量分散在低得多和高得多的评分上。两者的均值相同,但底层的判断却截然不同。
还可以根据受限分布的方差,得到一个可选的不确定性估计:
简而言之,这种方法不再使用带有噪声的 sampling judge score,而是使用有效评分数字上的归一化概率分布,再以该分布的期望值作为最终评分。
如今,这些东西大概已经显而易见了,但在 2024 年,当时几乎没有什么资料能指导我开发这种方法。遗憾的是,我最终发现它也完全没用……
每个配置都需要生成数百个 token 的创意内容,然后再让另一个模型阅读并评判。一个 70B 模型就有超过 3,200 种配置需要测试,在我的双 4090 上,这需要花费数周时间。
我需要输出极短的 probe,最多只生成几个 token,而且评分必须客观、确定。流程中不能再有 judge model。正因如此,我最终选定了下面两个 probe:
高难度数学。难到离谱的问题,例如:“74,088,893,247 的立方根是多少?”不允许 chain-of-thought,也不允许使用工具。只输出数字,纯靠直觉纵身一跃。
情商。使用 EQ-Bench benchmark:给出复杂的社交情境,让模型预测特定情绪状态的强度。“在这种情况下,这个人的愤怒/惊讶/内疚程度会是多少?请在 0 到 100 之间评分。”这与数学完全不同,涉及 theory of mind、社交推断和同理心。而输出只需要几个数字。
我最终选定了两项认知需求最大程度正交、输出又都非常短的任务。我的直觉是:LLM 每次只思考一个 token,那么就让模型变得极其擅长猜测下一个 token。但事情永远不会这么简单。就拿 LLM 处理数字来说……
即使使用数学 probe,我还是遇到了意想不到的问题。LLM 在算术任务上的失败方式非常奇怪。与其说它们给出了错误答案,不如说它们经常只差一点:例如忘记写最后一位数字,仿佛写到一半突然厌倦了;或者把中间的两个数字交换位置;又或者输出了正确数字,却在末尾附加一个字符,导致 parser 解析失败。
这可能与较大数字的 tokenization 方式有关,因为大数可能会被任意拆分。以整数 123456789 为例,一个 BPE tokenizer(例如 GPT 风格)可能把它拆成:‘123’ ‘456’ ‘789’,也可能拆成:‘12’ ‘345’ ‘67’ ‘89’。
简单的二元正确/错误评分会浪费许多有用信号。计算正确比例会更有帮助:相比正确答案 123456789,答案 123356789 的准确率可以算作 99.92%。
但如果模型犯了一个愚蠢的「LLM 式错误」呢?比如正确答案是 4302459,模型却输出了 430245,但显然已经完成了大部分计算。这该怎么算?我编写了一个自定义的部分得分函数:它会对较短的答案进行补位,再按比例施加惩罚:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def calculate_score(actual, estimate):
"""Calculate score comparing actual vs estimated answer."""
try:
actual_str = str(int(actual))
estimate_str = str(int(estimate))
except (ValueError, OverflowError):
return 0
max_length = max(len(actual_str), len(estimate_str))
actual_padded = actual_str.ljust(max_length, "0")
estimate_padded = estimate_str.ljust(max_length, "0")
padding_size = max_length - min(len(actual_str), len(estimate_str))
actual_int = int(actual_padded)
estimate_int = int(estimate_padded)
if max(actual_int, estimate_int) == 0:
return 0
relative_diff = abs(actual_int - estimate_int) / max(actual_int, estimate_int)
correction_factor = 1 - (padding_size / max_length)
score = (1 - relative_diff) * correction_factor
return max(0, min(score, 1))
核心思路是:先给较短的答案补位,再通过 correction factor 进行惩罚。如果模型答对了 90% 的数字,只是漏掉最后一位,它仍然可以得到相当可观的分数——但会低于每一位都正确的答案。事实证明,这一点对于区分那些直觉数学能力非常接近的配置至关重要。
最初,数学问题都是手工设计的。我尝试了不同的运算类型和数值规模,然后生成随机数来扩充数据集。数据集一共包含 16 个问题,模型的任务是估算出最接近的整数。下面给出几个,你也可以亲自试试。记住,不允许「思考」,直接猜答案!
1
2
3
What is 78313086360375 multiplied by 88537453126609?
What is the cube root of 18228885506341?
What is the cube root of 844178022493, multiplied by 43?
测试了几个较小的模型(Llama 和更小的 Qwen2)之后,我为 Qwen2-72B 设置好配置,让它开始全量扫描。每个 ((i, j)) 配置都需要几分钟:加载重新分层的模型,运行数学 probe,运行 EQ probe,记录分数,然后进入下一个配置。两块 4090 连续运行了好几天。但这需要的计算量,远少于一次 fine-tune!事实上,我只有 48GB VRAM,甚至没有足够的硬件对它进行一次 LORA fine-tune。
最优配置是 ((45, 52)):先运行第 0 层到第 51 层,然后再运行第 45 层到第 79 层。第 45 层到第 51 层会执行两次。也就是说,在 80 层 stack 的中间附近额外增加七层,把总参数量从 72B 提升到 78B。每个额外层都只是某个现有层的精确副本。没有新权重,也没有训练,只是让模型重复了一遍自己。
重复七层,仅此而已。现在,我终于可以揭晓这些模型名称的由来了:RYS-XLarge 里的 RYS,就是 Repeat Your Self ;)
我把这个配置应用到 MaziyarPanahi 的 calme-2.1-qwen2-72b——一个基于 Qwen2-72B 的 fine-tune——然后以 dnhkng/RYS-XLarge 的名称上传。我也把它应用到了未经 fine-tuning 的原始 base model 上,命名为 dnhkng/RYS-XLarge-base。
接着,我把模型提交到 Open LLM Leaderboard,然后开始等待。等啊,等啊。那时的 OpenLLM Leaderboard 每天都会涌入几十个模型,其中很多都是 fine-tune 的 merge,再由这些 merge 继续 fine-tune,简直就是西部荒野,等待队列也非常漫长。但大约一个月后,结果终于出来了:
MuSR 提升 17.72%,MATH 提升 8.16%。六项 benchmark 中有五项获得提升,只有 IFEval 略有下降。最终平均分让它登上了 leaderboard 第一名。
再强调一次:我只针对 hard math 问题的 one-shot 估算和 EQ-Bench 进行了优化。在开发过程中,我从未查看过 IFEval、BBH、GPQA、MuSR 或 MMLU-PRO。Leaderboard 的结果完全属于 out-of-sample validation。
通过两个狭窄且相互正交的 probe 找到的层配置,成功泛化到了 Leaderboard 抛给它的所有任务上 *。
* ——除了 IFEval,不过那个 benchmark 本来也挺无聊的,对吧?
这已经足够令人惊讶了:一种全新的 LLM scaling 方法,居然是在几块游戏显卡上开发出来的。但把 heatmap 绘制出来之后,故事变得更加精彩。
下面是产生 RYS-XLarge 的原始 heatmap,展示的是 Combined delta(数学 + EQ)。绿色圆圈标记了最优配置。红色代表性能提升,蓝色代表性能下降。
这些 heatmap 就像 Transformer 在思考数学或 EQ 问题时拍摄的功能性 MRI。
x 轴((j))表示重复区域的结束位置,y 轴((i))表示开始位置。每个像素都代表一次完整评测:加载重新分层后的模型,运行数学 probe,运行 EQ probe,对两项结果评分,并记录其变化量。如前所述,在中央对角线上,每个配置只复制一层。沿着下一条朝右上方延伸的对角线,则会复制两层,以此类推。最右上角的单个点,代表每次 inference 都把整个 Transformer stack 完整运行两遍。
我们先看看数学 heatmap。从任意层开始,只要在大约第 60 层之前停止,似乎都能提高数学估算分数;图中那片泛着健康红晕的巨大区域正说明了这一点。只复制最前面的几层(左上角的小三角形)会把结果搞坏;重复最后 20 层中的几乎任意一层,同样会造成破坏(右侧那堵蓝色的垂直高墙)。在 skyline plot(对行或列取平均值)中,这一点表现得更加清楚。对于数学估算来说,重复区域的起始位置似乎没有那么重要。因此,下面这个假设似乎得到了一定程度的验证:「起始层」把 token 编码到一个平滑的「思考空间」,最终再将其交给一个专门的「重新编码」系统。
直到我们查看 EQ 分数:
现在,情况看起来完全不同了!复制最后 10 层中的任何一层,对分数几乎都没有影响;但我们同时看到了一些复杂的模式,其中有些区域显示出显著提升(大约位于 45i, 55j 附近),而这些区域又被性能较差的区域围在中间。
然而,heatmap 揭示出的东西,比「负责思考的部分位于何处」更加有趣。它还揭示了这种思考结构本身的一些特征。
在最终决定使用 block duplication 之前,我尝试过一个更简单的办法:取中间的某一个层,把它重复 (n) 次。如果「增加推理深度」的假设正确,这种方法就应该有效。从复制中间层能普遍改善数学估算结果来看,这种思路也很合理:多给模型几份某个特定推理层,就能获得更好的推理能力。因此,我把所有层都筛查了一遍,希望找到性能提升。
但并没有。结果几乎总是更差,而且通常差得多。偶尔也会有一些小幅提升,但都处于噪声范围之内。这让人很恼火,不过,当我重新观察 EQ 分数里那些复杂、成团的模式时,又产生了另一个想法:
如果复制单层没有帮助,那就说明中间层执行的并不是彼此独立的迭代式细化。它们并不是执行同一种操作、可以简单「再运行一次」的可互换副本。如果 t