2025-11-10 01:26:00
月之暗面发布的 Kimi K2 Thinking，在 Humanity's Last Exam (HLE) 上达到了 44.9% 的成绩，在多个基准测试中表现优异，不过榜单简单看一眼即可；让我比较惊喜的是，K2 Thinking 可以执行 200-300 步连续的工具调用，有类似于 claude 一样的长程规划和自适应推理能力。
但是，K2 Thinking 的官方 blog 只展示了 benchmark 数据和 demo，并没有透露具体的技术细节。作为一个大模型从业者，看到 Twitter/知乎大家都在聊这个模型，所以我就比较好奇「模型的训练方法」以及「给我们工作学习中的启发」。
好在今年早些时候发布了 Kimi K2 的完整技术报告和 技术 blog。而 K2 Thinking 和 K2 师出同源，只是在 K2 的基础上增加了 thinking 能力，更强的工具调用能力，通过 test-time scaling 实现一个更强的 Thinking Agent。因此，通过深入研究 K2 的技术细节，我们就能理解 K2 Thinking 是如何炼成的。
我是朝发（CHAOFA）这篇文章会从 K2 的技术报告出发，结合 K2 Thinking 的特点，了解这个 SOTA 开源 thinking 模型是怎么训出来的。核心关注三个问题：
历史上此比较相关文章：
如果不喜欢看文字可以看视频解读，B 站-chaofa用代码打点酱油和 YouTube
算法视角深度解读 Kimi K2 和 K2 Thinking，从预训练优化到 Agentic 能力训练的完整流程（含MuonClip优化、Agentic 数据 --bilibili
Archiecture from: Sebastian Raschka
先来看一下上面的整体结构图，然后在深入技术细节之前，我们有必要先理解 K2 和 K2 Thinking 的关系。
Kimi K2 是一个 MoE (Mixture-of-Experts) 模型，拥有 32B 激活参数和 1T 总参数。它在非 thinking 模型中，在前沿知识、数学和编码任务上达到了 SOTA 性能。
K2 的核心特点是有比较强的 Agentic 能力。什么是 Agentic 任务？就是模型不仅要回答问题，还要主动使用工具、执行操作、完成复杂的多步骤任务。比如：
K2 发布了两个版本：
Kimi K2 Thinking 是在 K2 的基础上，通过额外的训练，让模型具备了 thinking 能力。它的核心特点是：
think → search → browse → think → code 的循环，动态生成和验证假设从架构上看，K2 Thinking = K2 + Thinking Ability + Test-Time Scaling。因此，理解 K2 的训练方法，就能理解 K2 Thinking 的 80%。
下面我们按照训练流程，依次讲解预训练、后训练和 test-time scaling 的关键技术。
预训练是 Agentic Intelligence 的关键基础，它建立了让强化学习探索变得可行、高效和可泛化的先验知识。但是，正如 Ilya Sutskever 所说，数据是有限的"化石燃料"，其增长速度远远落后于算力的增长。这使得预训练阶段的 token 利用效率成为 AI scaling laws 中的新关键系数。
给定一个大致有限的预训练数据集和固定的模型配置，更 token 高效的优化器能产生更多的智能。Moonshot 之前的工作 Moonlight 已经证明，Muon 优化器在 LLM 训练中显著优于广泛使用的 AdamW 优化器，即“相同配置训练下有更低的 loss”。
K2 的设计目标是进一步扩展 Moonlight，它采用了类似 DeepSeek-V3 的架构。基于 scaling-law 分析，他们做了两点改进（看图更清晰）：
原文这么写的：Based on scaling-law analysis, we reduce the number of heads for long-context efficiency, and increase MoE sparsity for greater token efficiency。
但在扩展过程中，他们遇到了一个持续的挑战：由 attention logits 爆炸引起的训练不稳定。这个问题在使用 Muon 时更频繁，而在 AdamW 中较少。现有的解决方案（如 Qwen3 用的 query-key normalization）都不够充分（防止数值溢出）。
为了解决这个问题，kimi 提出了 MuonClip 优化器，它通过 qk-clip 技术改进了 Muon。
核心思想：qk-clip 通过在 Muon 更新后直接重新缩放 query 和 key 投影的权重矩阵，从源头控制 attention logits 的规模，从而稳定训练。（注意：这里是更新完之后，所以不会改变这一次更新的 forward/backward 操作，影响的是下一步）。
具体来说，query 和 key 投影按如下方式缩放：
其中 是一个平衡超参数，因此 attention logit 变为：
自适应因子 （阈值为 ）在每一步之后根据该步的最大 attention logit 设置：
其中 是预设的阈值。这是一个通用技术，可能适用于其他稳定化场景。这里其实还有一些其他的细节，比如 每个 head 有不同的 。
实验表明，MuonClip 有效地防止了 logit 爆炸，同时保持了下游任务性能。在实践中，K2 使用 MuonClip 在 15.5T tokens 上进行预训练，实现了零训练尖峰（zero loss spike），证明了 MuonClip 是大规模 LLM 训练的稳健解决方案。
从 loss 曲线可以看出，MuonClip 的训练过程非常平滑，没有出现任何不稳定的情况。这为后续的 Agentic 能力训练打下了坚实的基础。
小结：MuonClip 优化器通过 qk-clip 技术，在保持 Muon 高 token 效率的同时，解决了训练不稳定问题，使得在同等条件下获得比 AdamW 更低的 loss，使得 K2 能够在有限的数据上训练出更强的基础模型。
K2 相比 K1.5 的一个关键进步是引入了合成数据生成策略来提高 token 利用率。核心思想是：通过精心设计的改写 pipeline，在不引入显著过拟合的情况下，扩大高质量 tokens 的数量。改写（Rephrasing） 就是数据合成的一种方式，主要是为了提高「高质量数据的占比」，尤其是「知识领域」和「数学领域」。：
在知识密集型文本上进行预训练面临一个权衡：单次 epoch 不足以全面吸收知识，而多次 epoch 重复会导致收益递减并增加过拟合风险。为了提高高质量知识 tokens 的利用率，K2 提出了一个合成改写框架，每个语料库最多改写两次，包含三个关键组件：
A. 风格和视角多样化的提示（Style- and perspective-diverse prompting）
通过精心设计的 prompts，引导大语言模型以不同的风格和视角生成原文的忠实改写。这样做的好处是：
B. 分块自回归生成（Chunk-wise autoregressive generation）
为了在长文档中保持全局连贯性并避免信息丢失，采用基于分块的自回归改写策略，一图胜千言：
C. 保真度验证（Fidelity verification）
为了确保原文和改写内容之间的一致性，进行保真度检查，比较每个改写段落与其源文本的语义对齐。这是训练前的初步质量控制步骤。
为了增强数学推理能力，K2 采用了两种策略：
A. "学习笔记"风格改写
将高质量的数学文档改写成"学习笔记"风格，遵循 SwallowMath 中引入的方法。这种风格更接近人类学习数学的方式，包含：
B. 多语言翻译
将其他语言的高质量数学材料翻译成英语，以增加数据多样性。这样可以：
小结：通过针对知识和数学领域的专门改写技术，K2 在不显著增加过拟合风险的情况下，大幅提高了高质量 tokens 的利用率。这种受控的数据增强策略是 K2 预训练成功的关键因素之一。
K2 的增强 Agentic 能力源于两个重要方面：
为了教会模型复杂的工具使用能力，kimi 是基于大规模模拟真实世界的工具使用场景，构建了数据 pipeline。
这个管道的核心思想是：系统地演化数百个领域，包含数千个工具（包括真实的 MCP 工具和合成工具），然后生成数百个具有不同工具集的 agents。
辅助看这个图：
具体流程如下：
这个可扩展的 pipeline 生成了多样化、高质量的数据，为大规模拒绝采样和强化学习铺平了道路。
传统的工具使用训练依赖于人工标注的数据，成本高、规模小、多样性有限。而 k2 的方法通过自动化合成，可以：
传统的强化学习主要应用于可验证奖励的任务，比如数学题（答案对错明确）和竞赛编程（能否通过测试用例）。但对于不可验证奖励的任务（如写研究报告、创意写作），传统 RL 就无能为力了。
是不是突然想起了 DeepSeek-GRM（通用奖励模型）。
核心思想是：模型作为自己的评判者，为不可验证的任务提供可扩展的、基于 rubric 的反馈。
具体做法：
但这里有个问题：模型的自我评估准确吗？这不还是 LLM as Judge 那一套吗？
kimi 的解决方案是：在可验证奖励的 on-policy rollouts 中，持续更新 critic，使 critic 在最新策略上不断提高评估准确性。
这可以看作是用可验证奖励来改进不可验证奖励的估计。通过这种方式，模型的自我评估能力会随着训练不断提升，从而支持更广泛的任务。
小结：通过大规模 Agentic 数据合成和通用强化学习，K2 学会了在各种场景下使用工具，并且能够处理可验证和不可验证的任务。这为 K2 Thinking 的长程推理能力打下了基础。
K2 Thinking 在 K2 的基础上，增加了 thinking 能力、更强的工具调用能力和 test-time scaling。这使得模型能够在推理时进行长程思考和工具调用，从而解决更复杂的问题。
Test-Time Scaling 是指在推理时增加计算量，以提升模型性能。对于 K2 Thinking，这体现在两个方面：
这两者结合，使得 K2 Thinking 能够解决需要深度推理和多步操作的复杂问题。
K2 Thinking 的核心能力是边思考边使用工具。它会进行动态的 think → search → browse → think → code 循环，这个循环可以重复数百次，直到找到答案：
在 BrowseComp benchmark 上，K2 Thinking 达到了 60.2% 的成绩，显著超越了 29.2% 的人类基线。
BrowseComp 是一个挺具有挑战性的 benchmark，旨在评估模型持续浏览、搜索和推理难以找到的真实世界网络信息的能力。
K2 Thinking 在编码任务上也表现出色：
从官网看，K2 Thinking 可以从单个 prompt 构建完整的应用，包括：
通过 test-time scaling，K2 Thinking 能够在推理时进行长程思考和工具调用，从而解决需要深度推理和多步操作的复杂问题。这使得它在 Agentic Reasoning、Agentic Search 和 Agentic Coding 任务上都达到了 SOTA 性能。（有点 claude 那味道了）
让我们总结一下 K2 和 K2 Thinking 的关键区别：
| 维度 | K2 (Instruct) | K2 Thinking | |
2025-11-03 04:52:00
我是朝发（chaofa），这是我 25 年第 9 次月度总结，希望还能坚持下去。
从信息传播角度，现在是一个好的社会，很多的大佬也愿意出来讲话，无论是播客还是文章，都很容易被传播到，从而得以瞥见他们的想法。这让我们看到很多。我有时候很感激，有时候又时长感慨：他们的成功真的可以复制吗？
罗永浩的播客是一个挺好的东西，我几乎是每一期都听的，10 月份和影视飓风 TIM 的采访让我很是羡慕，潘天鸿。96 年，还不到 30 岁，但现在已经彻底成功[1]。老罗和他的对话还挺有意思的，但我主要是聊四个东西对我的启发：
我觉得 TIM 给我的感觉就是另外一个维度的雷军，是那种特别努力的人。这种被称得上“努力”的人，不仅有着天生好的精力和体力，还配有强大的意志和成功欲望，这些可能都不是后天能够锻炼出来或者培养出来的东西[2]。我真的非常佩服 TIM 的精力，他几乎是每一期真人出镜，且有大量的口播，考虑到更新频率以及「影视飓风」那么大体量公司管理和沟通的事情，他必然是全情投入其中，不容其他人丝毫地质疑他的努力。
在对谈中，有一个词出现的频率特别高「野兽先生，Mr. Beast」，TIM 的对标账号是世界第一网红「野兽先生」，尽管这不是我所喜欢的风格[3]，但是我觉得这主要是他的野心太大了。「大众流量」、「声誉口碑」他全都要，并且要靠「流量变现」养活其他的探索与尝试，所以他选择了这种已经被证明的道路。野兽先生无疑是成功的，想要获得巨大的影响力这也是一个最好的办法，毕竟他要养着「冲击奥斯卡短片」的团队，靠着矩阵账号的方式，打造着目前中国（可能是）最成功的自媒体账号。
影视飓风很早之前就有一期飞书的广告视频，讲「影视飓风是如何通过数据驱动的方式增长」。作为任何一个「互联网行业从业者」，应该对此完全的不陌生，里面说的所有内容都是在职场中被宣贯无数次的准则，但是要真正在自己的生活或者「自己的事业」上落实这些道理却有巨大的挑战。很多时候，我们就是无法跨越这种「知道」和「做到」的巨大鸿沟。
就以我自己为例，作为一个曾经在中文最大的互联网公司的广告相关从业者，我自然知道「点击率」、「留存」以及「各种漏斗指标」对于一个视频的影响，但是我自己做「B 站/YouTube/小红书大模型教学视频（chaofa 用代码打点酱油）」的时候，我总是带有一种「类似于理想主义的天真」，希望别人能认真看下去[4]，从而收获真正的价值，所以我一直想要筛选掉那些「并不是想真正学习的人」，于是我通过「不剪辑、不配字幕、封面不设计、口播不写稿等」一系列愚蠢至极的操作把真正想学习的人筛选出来了，事实上我好像部分成功了，我成功筛选出了「高学历、高收入、低付费意愿」的能够自我学习和进化的人士，一方面为此感到欣慰，一方面又「因为自己的思想的转变」而感觉有点想笑。
自媒体似乎很光鲜亮丽，尤其是被「各种网红的造富神话」的冲击之后，现在越来越多的年轻人将「当网红」作为自己的最佳的职业选择[5]。但大多数人只要自己尝试之后，才会发现实际上是很难赚钱的（除了那些教别人做自媒体的），或者说很难赚到自己「付出同样精力打工上班」赚到的钱。
核心原因是自媒体赛道的变现问题，前段时间王自如限高做绿皮火车吃泡面的新闻就可见一斑（毕竟也是曾经的顶流）。这里主要有两个原因：
而 TIM 的选择是通过卖货，即现在大多数账号的首要变现方式一样（包括像 B 站的大物是也等）也都是通过卖货来变现的，只是可能 TIM 这种自建电商品牌的方式更正规军一点，并且在很早开始尝试，而很多账号可能还在犹豫、彷徨、迷茫，而影视飓风已经带着 TIM 的同款内裤把钱赚了个够。
回到我的账号——chaofa用代码打点酱油，变现效率可能说是奇差，刚刚提到的「高学历、高收入、低付费意愿」的人其实就是我自己（我的同事、我的同学们），我也是属于那种付费意愿比较低的人🤣，思想正在转变中。想要卖货是不可能了，但是我发现「网络上很多盗版我视频的人」，在我视频不到40 的情况下盗版视频可能有近百个了，因此我的视频绝对是有价值的，所以我一定要快点写书，等写完书我就可以尝试卖课了（严肃脸😠），也许是 26 年底，也许是永远不会（懒惰可能会占领智商的高地）。
很多人都说 TIM 的成功当然也离不开家庭的支持，这里指的是那份底气和关键性的指导，他可以有足够的尝试，可以自由的探索。这是显然的，毕竟做成这样缺失路上任何一环都不会有现在的「影视飓风」，不过这可能是最不重要的，因为他真的很强，没有这样的家庭他也能成功，只是可能不一定有现在成功[6]。但是无论如何，他的成功都是自己争取得到的，其他的东西只是景上添花，就算不是 29 岁，39 岁的 TIM 也一定会走上人生巅峰。RESPECT！
最后欢迎来探讨对世界的认知，基本全网同名 chaofa用代码打点酱油 (推荐)
大约在 21 年才关注到影视飓风，真的是一路看着他成功，以后他会更成功，真的是太羡慕了，无论是才华、认知、努力还是天生的精力。 ↩︎
尽管 TIM 称自己并不是一个高能量的人，但是他所展现出来的精力，都是大多数人不可比拟的。在媒体前的表现自然会有一定的表演性质，但是只要能一直持续下去（10 年啊，人生能有几个十年），那就是常人难以企及的能量。 ↩︎
我还是喜欢有温度一点的内容，不喜欢被当做数字，想要被认为是一个真正的观众，所以我永远喜欢那种真诚深度对话。备注：这个月和一个网友聊了个天还挺好的。 ↩︎
我的视频都特别的长，视频质量一定是远超 B 站同类视频，但是数据比较一般。我经常能在评论区有一些比较搞笑的对话，问我资料/问我代码（GitHub 链接/个人 Blog/公众号差点就怼脸上了），我有时候真的很想笑， Google 搜一下「chaofa 用代码打点酱油」就这么难吗？ ↩︎
备注：我觉得以 TIM的认知，他是不可能不能推测出自己家庭情况是比较有钱的，因为这东西很容易感受到，爸妈能陪你飞到国外去看看学校（校长？）一定不是 TIM 口中普通家庭可以做到的，至少不差，没 TIM 口中感受的那么差吧。 ↩︎
2025-01-28 03:30:00
本次课一共讲解三个不同版本的 MOE，分别是基础版MOE，大模型训练用的 SparseMoE，还有 DeepSeek 用的比较多的 shared_expert 的 SparseMoE。
输入是一个 Token, 输出是一个 Token Embedding。暂时先不考虑 MOE 得到的 Embedding 怎么使用。
因为 MOE 网络对应着 Expert，这个 Expert 一般是一个 FeadFoward Network，FFN。而为了简化，后续我们都用一层的 Linear 代替，更高级版本的 Expert 留给大家当做课后作业。下面是一个专家的定义。
class BasicExpert(nn.Module):
# 一个 Expert 可以是一个最简单的， linear 层即可
# 也可以是 MLP 层
# 也可以是 更复杂的 MLP 层（active function 设置为 swiglu）
def __init__(self, feature_in, feature_out):
super().__init__()
self.linear = nn.Linear(feature_in, feature_out)
def forward(self, x):
return self.linear(x)基础版本的 MOE 可以看这个图，非常的简单。
class BasicMOE(nn.Module):
def __init__(self, feature_in, feature_out, expert_number):
super().__init__()
self.experts = nn.ModuleList(
[
BasicExpert(feature_in, feature_out) for _ in range(expert_number)
]
)
# gate 就是选一个 expert
self.gate = nn.Linear(feature_in, expert_number)
def forward(self, x):
# x 的 shape 是 （batch, feature_in)
expert_weight = self.gate(x) # shape 是 (batch, expert_number)
expert_out_list = [
expert(x).unsqueeze(1) for expert in self.experts
] # 里面每一个元素的 shape 是： (batch, ) ??
# concat 起来 (batch, expert_number, feature_out)
expert_output = torch.cat(expert_out_list, dim=1)
# print(expert_output.size())
expert_weight = expert_weight.unsqueeze(1) # (batch, 1, expert_nuber)
# expert_weight * expert_out_list
output = expert_weight @ expert_output # (batch, 1, feature_out)
return output.squeeze()
def test_basic_moe():
x = torch.rand(2, 4)
basic_moe = BasicMOE(4, 3, 2)
out = basic_moe(x)
print(out)
test_basic_moe()这个一般我们用 switch transformers 这篇文章的图作为演示，详情看：
和 Basic 区别是，MOE 选择 topK 个专家，然后对这 topK 个专家的输出进行加权求和，并且把输入样本变成了大模型中真实的输入 Shape，(batch, seq_len, hidden_dim)
# 主要参考自 mistral MOE 的实现
class MOERouter(nn.Module):
def __init__(self, hidden_dim, expert_number, top_k):
super().__init__()
self.gate = nn.Linear(hidden_dim, expert_number)
self.expert_number = expert_number
self.top_k = top_k
def forward(self, hidden_states):
# 计算路由logits
router_logits = self.gate(hidden_states) # shape is (b * s, expert_number)
# 计算专家经过softmax之后的概率
routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float)
# 计算topk的专家的输出
router_weights, selected_experts = torch.topk(
routing_probs, self.top_k, dim=-1
) # shape都是 (b * s, top_k)
# 专家权重归一化
router_weights = router_weights / router_weights.sum(dim=-1, keepdim=True)
router_weights = router_weights.to(hidden_states.dtype)
# 生成专家掩码
expert_mask = F.one_hot(
selected_experts,
num_classes=self.expert_number
) # shape是 (b * s, top_k, expert_number)
expert_mask = expert_mask.permute(2, 1, 0) # (expert_number, top_k, b * s)
return router_logits, router_weights, selected_experts, expert_mask
class MOEConfig:
def __init__(
self,
hidden_dim,
expert_number,
top_k,
shared_experts_number=2,
):
self.hidden_dim = hidden_dim
self.expert_number = expert_number
self.top_k = top_k
self.shared_experts_number = shared_experts_number
class SparseMOE(nn.Module):
# 稀疏 MOE 模型，这里每一个 token 都会过 topk 个专家，得到对应token 的 hidden_embeddings
def __init__(self, config):
super().__init__()
self.hidden_dim = config.hidden_dim
self.expert_number = config.expert_number
self.top_k = config.top_k
self.experts = nn.ModuleList(
[
BasicExpert(self.hidden_dim, self.hidden_dim) for _ in range(self.expert_number)
]
)
self.router = MOERouter(self.hidden_dim, self.expert_number, self.top_k)
def forward(self, x):
# x shape is (b, s, hidden_dim)
batch_size, seq_len, hidden_dim = x.size()
# 合并前两个维度，因为不是 Sample 维度了，而是 token 维度
hidden_states = x.view(-1, hidden_dim) # shape is(b * s, hidden_dim)
router_logits, router_weights, selected_experts_indices, expert_mask = self.router(hidden_states)
# 其中 selected_experts_indices shape 是 (b * s, top_k)
# 其中 expert_mask shape 是 (expert_number, top_k, b * s)
final_hidden_states = torch.zeros(
(batch_size * seq_len, hidden_dim),
dtype=hidden_states.dtype,
device=hidden_states.device
)
for expert_idx in range(self.expert_number):
expert_layer = self.experts[expert_idx]
# expert_mask[expert_idx] shape 是 (top_k, b * s)
idx, top_x = torch.where(expert_mask[expert_idx])
# idx 和 top_x 都是一维 tensor
# idx 的值是 0 或 1, 表示这个 token 是作为当前专家的 top1 还是 top2
# top_x 的值是 token 在 batch*seq_len 中的位置索引
# 例如对于 batch_size=2, seq_len=4 的输入:
# top_x 的值范围是 0-7, 表示在展平后的 8 个 token 中的位置
# idx 的值是 0/1, 表示这个 token 把当前专家作为其 top1/top2 专家
# hidden_states 的 shape 是 (b * s, hidden_dim)
# 需要取到 top_x 对应的 hidden_states
current_state = hidden_states.unsqueeze(
0
)[:, top_x, :].reshape(-1, hidden_dim) # （selected_token_number, hidden_dim）
# router_weight 的 shape 是 (b * s, top_k)
current_hidden_states = expert_layer(
current_state
) * router_weights[top_x, idx].unsqueeze(-1) # （selected_token_number, 1） 这里有广播
# 把当前专家的输出加到 final_hidden_states 中
# 方式1 的写法性能更好，并且方式1容易出现
final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
# 方式2
# final_hidden_states[top_x] += current_hidden_states.to(hidden_states.dtype)
# 方式2 的写法性能更差，并且方式2容易出现错误，+= 操作在处理重复索引时需要多次读写内存，可能会导致竞争条件
# 把 final_hidden_states 还原到原来的 shape
final_hidden_states = final_hidden_states.reshape(batch_size, seq_len, hidden_dim)
return final_hidden_states, router_logits # shape 是 (b * s, expert_number)
def test_token_level_moe():
x = torch.rand(2, 4, 16)
config = MOEConfig(16, 2, 2)
token_level_moe = SparseMOE(config)
out = token_level_moe(x)
print(out[0].shape, out[1].shape)
test_token_level_moe()备注：这里是参考 deepseek moe 思想，写的一个共享 expert 的 MOE 网络，有一定的简化，但是可以方便理解训练过程。
和 版本2 的 SparseMOE 区别是，这里多了一个 shared experts 的模型，这个模型是所有 token 共享的，也就是说，所有 token 都过这个 shared experts 模型，然后每个 token 会用计算的 Router 权重，来选择 topK 个专家，然后和共享的专家的输出一起加权求和。
具体结构图为：
class ShareExpertMOE(nn.Module):
def __init__(self, config):
super().__init__()
self.moe_model = SparseMOE(config)
self.shared_experts = nn.ModuleList(
[
BasicExpert(
config.hidden_dim, config.hidden_dim
) for _ in range(config.shared_experts_number)
]
)
def forward(self, x):
# x shape 是 (b, s, hidden_dim)
# 首先过 moe 模型
sparse_moe_out, router_logits = self.moe_model(x)
# 针对的还是 x 的每一个
# 然后过 shared experts
shared_experts_out = [
expert(x) for expert in self.shared_experts
] # 每一个 expert 的输出 shape 是 (b, s, hidden_dim)
shared_experts_out = torch.stack(
shared_experts_out, dim=0
).sum(dim=0)
# 把 sparse_moe_out 和 shared_experts_out 加起来
return sparse_moe_out + shared_experts_out, router_logits
def test_share_expert_moe():
x = torch.rand(2, 4, 16)
config = MOEConfig(16, 2, 2)
share_expert_moe = ShareExpertMOE(config)
out = share_expert_moe(x)
print(out[0].shape, out[1].shape)
test_share_expert_moe()用于测试上面的代码是否可以跑通？
def switch_load_balancing_loss(router_logits: torch.Tensor, num_experts: int) -> torch.Tensor:
"""
计算 Switch Transformers 的负载均衡损失
Args:
router_logits: shape [batch_size * sequence_length, num_experts]
num_experts: 专家数量
Returns:
total_loss: 总损失 = auxiliary_loss + z_loss
"""
# 计算路由概率
router_probs = torch.softmax(router_logits, dim=-1) # [b*s, num_experts]
# 获取每个token的最优专家
_, selected_experts = torch.topk(router_probs, k=2, dim=-1)
# 创建one-hot矩阵表示选中的专家
mask = torch.nn.functional.one_hot(selected_experts, num_experts).float()
# 计算每个专家的期望负载 (理想情况下应该是 1/num_experts)
expected_load = torch.ones_like(router_probs) / num_experts
# 计算实际负载 (每个专家处理的token数量除以总token数量)
# 在batch维度上计算平均值
actual_load = mask.mean(dim=0)
# 计算auxiliary loss
# 这会惩罚负载分布与期望负载的差异
aux_loss = torch.sum(actual_load * router_probs.mean(dim=0)) * num_experts
# 计算z_loss (可选)
# 这会惩罚过大的路由logits
z_loss = torch.mean(torch.square(router_logits))
z_loss_weight = 0.001 # 可调整的超参数
# 总损失
total_loss = aux_loss + z_loss * z_loss_weight
return total_loss
def test_moe_training():
# Create a simple dataset
batch_size = 32
seq_len = 16
hidden_dim = 32
num_batches = 100
# Initialize model and optimizer
config = MOEConfig(hidden_dim=hidden_dim,
expert_number=4,
top_k=2,
shared_experts_number=2)
model = ShareExpertMOE(config)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# Training loop
model.train()
for batch in range(num_batches):
# Generate random input data
x = torch.randn(batch_size, seq_len, hidden_dim)
target = torch.randn(batch_size, seq_len, hidden_dim)
# Forward pass
output, router_logits = model(x)
# Compute losses
# MSE loss for prediction
mse_loss = F.mse_loss(output, target)
aux_loss = switch_load_balancing_loss(router_logits, config.expert_number)
# Combined loss
total_loss = mse_loss + 0.01 * aux_loss
# Backward pass and optimize
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if batch % 10 == 0:
print(f"Batch {batch}, Loss: {total_loss.item():.4f} "
f"(MSE: {mse_loss.item():.4f}, Aux: {aux_loss.item():.4f})")
# Run the training test
test_moe_training()最后欢迎关注我，基本全网同名 chaofa用代码打点酱油
2025-01-28 02:58:00
自 chatGPT 22年底问世以来，大模型（Large Language Model, LLM）一般使用 Causal Language Model 的形式，属于 Transformers 中的 Decoder 部分，其中在 Decoder 的 Block 中有一个 FFN(FeadForward) 层，一般认为这部分参数用于存储知识。而标准的 FFN 一般有一个升维度和降维度的过程，一共有两个权重矩阵，用公式表示为
其中 x shape 是 ，w1 shape 是 ，w2 shape 是 , w1 是升维（up），w2 是降维(down)
激活函数主要是为了实现神经网络学习输入和输出之间的复杂非线性关系而使用的一个函数。在公式 (1) 中，ReLU 是一个激活函数（Transfromers原版），可以替换成其他的激活函数，比如 BERT 开始用 Gaussian Error Linear Unit，GELU 比较多，随后就成了激活函数的主流选择，但是随着大模型的爆火以及 PaLM 模型的发布，大家开始慢慢使用 swishGLU 作为激活函数，并且作为一个主要的优化点。
具体可以看下面一段代码即可清楚的理解 FFN 模型是什么实现的。
class FeedForward(nn.Module):
# 实际上就是 MLP
def __init__(self, config):
super().__init__()
self.net = nn.Sequential(
nn.Linear(config.n_embd, 4 * config.n_embd),
# 激活函数
nn.ReLU(),
# 可以替换成 nn.GELU(),
# 但是 如果是 SwishGLU 则实现方式有所不同，接下来就会介绍 swishGLU 是怎么实现的
nn.Linear(4 * config.n_embd, config.n_embd),
nn.Dropout(config.dropout)
)
def forward(self, x):
return self.net(x)ReLU 深度学习以来最常用的激活函数，其公式非常的简单。
从 GPT、BERT 以来，GELU 似乎成了新时代取代 ReLU 的激活函数，具体形式如下：
其中 是标准正态分布的累计分布函数，定义为
这里的 erf 是误差函数
但是这个函数由于计算成本较高，因此有两个初等函数作为近似计算（但目前【2025年1月27日】其实很多框架已经可以精确计算 erf 函数）。
近似计算分析详细可以参见苏神的文章，GELU的两个初等函数近似是怎么来的
SwiGLU（或者swishGLU，以下可能混用） 是 swish 激活函数和 GLU 门控单元的结合体，因此需要分别介绍两者的不同。
其中需要注意的是：在 T5 开始，很多模型（比如 PaLM ）在FFN层都不用 bias 了，也就是说 FFN的公式变成了
注意公式 6 和公式 1 的区别，一共没有 bias 一个有 bias，但具体得看不同模型的实现，并不能一概而论。
swish 是一个非线性函数（激活函数都是如此，笑🤣），具体公式为：
其中 是一个超参数，当 时，Swish 就变成了 SiLU (Sigmoid Linear Unit)，大多数框架的默认实现（如 PyTorch、TensorFlow 的 nn.SiLU()）使用的是 的固定版本。
因此如果采用 swish 激活函数，FFN 的公式变成了
共有两个可学习的矩阵，其中 是升维矩阵， 是降低维度的矩阵。
GLU，Gated Linear Units，是一种门控结构（有参数，因此相对于普通的激活函数多了一个 gate 矩阵），通过 sigmoid 控制不同维度的激活。公式如下[1]：
这里是不是熟悉 LSTM, GRU 的同学一下就理解，其中需要注意的是，b, c 对应的 bias 不是必须的。
对比公式 7 和公式 9，公式 9 中的 对应 公式 7 中的 ，而 对应公式 7 中的 矩阵。
而 SwiGLU 就是把门控函数替换成了 swish，并且去除掉了 bias 部分，以及把 FFN 层的一个 Linear 层替换成了 GLU 层，因此一共有三个可训练的参数矩阵, w1, w2, w3。
因此最终的公式表达为，
而我们都知道 FFN 是一个升高维度，然后降低维度的过程，因此可以写成，W2 是一个降低维度的参数，W1 是升高维度的过程，而 W3 是一个 Gate 需要用到的参数矩阵。
通过这个公式整体就非常的清晰理解使用 swiGLU 的 FFN。
而我们都知道在 basic 版本的 FFN，见公式（1）， 只有 和 分别是 (h, 4h) 和（4h, h），因此整体参数是 。
而公式9 中，一共有三个矩阵，如果想要实现总参数 ，那么每一个参数矩阵的大小应该是 ，因此 的shape应该是 ， 的 shape 是 。
假设输入的 hidden_dim 大小是 hidden_dim，那么中间层（up 后的维度）大小是 mid_dim， 具体计算逻辑如下：
mid_dim = int(8 * hidden_dim / 3)
# multiple_of：make SwiGLU hidden layer size multiple of large power of 2
mid_dim = multiple_of * ((mid_dim + multiple_of - 1) // multiple_of)
# multiple_of 一般设置为 256， LLaMA 和 GPT等模型注意，在 LLM (大语言模型) 架构中，multiple_of 是一个用于优化计算效率的参数，通常设置为 256 或其他 2 的幂次方数（如 128、512 等），最终让 mid_dim 调整为 multiple_of 的整数倍。这样做有几个原因：
class FFNExpert(nn.Module):
def __init__(self, hidden_dim, dropout): # LLM 进化之路， FFN 激活函数从 GELU -> SwiGLU
super().__init__()
# 有一个 magic number 叫做 8/3
hidden_dim = hidden_dim
# 这里可以自己去优化成 multiple_of 的倍数
mid_dim = hidden_dim * 8 // 3
self.up = nn.Linear(hidden_dim, mid_dim, bias=False)
self.down = nn.Linear(mid_dim, hidden_dim, bias=False)
self.gate = nn.Linear(hidden_dim, mid_dim, bias=False)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
out = self.dropout(
self.down(
# up 之后的 Shape 是(b, s, mid_dim)
# gate 和 up 之后的Shape都是 (b, s, mid_dim)
# 两者是 element-wise 相乘
F.silu(
self.gate(x)
) * self.up(x)
)
)
return out最后欢迎关注我，基本全网同名 chaofa用代码打点酱油
https://zhuanlan.zhihu.com/p/693332639 ↩︎
2024-12-29 05:00:00
如果今年要挑一个最重大的事情，那只能是点点(我妻子)怀孕了，我明年就要当爹了。这事情的影响是非常巨大的，不仅涉及到家庭，而且是一个需要长期付出不可逆的过程。有无孩子一定会是两个截然不同的世界，所以我是有些恐慌的（教育孩子的难度不言而喻）。
由于点点想生两个孩子，因此在去年结婚之后，就开始盘算着什么时候生娃，不然到时候得做高龄产妇。左思右想后决定从 4 月份开始备孕，并且一定要让孩子在次年 9 月之前出生。之所以有这么荒诞的想法是因为我国的入学政策是：“凡年满6周岁的儿童，其父母或者其他法定监护人应当送其入学接受并完成义务教育。”
这里有一个很有意思的内容，关于点点的高度不自洽，一方面天天讲以后孩子打螺丝能养活自己就好了，一方面又很焦虑孩子的入学时间，不然可能导致她/他以后读博可能有年龄压力（就半年也不至于吧，而且孩子一定想读博吗🤣）。所以未来会怎么样呢，这种极大的不确定真是让人又担心又期待。
经过半年的时间，周末各种跑医院，所幸预产期会是在 25年的 9 月前，缓解了点点下半年最大的烦恼，但随之而来的孕早期一系列的孕反，嗳气、呕吐、尿频等，十分不易。我也在反思自己，我的关心不够，好像真的只会抖机灵逗她开心。。。
如果要说今年最有意义的是一件事情，那就是【从十月份开始时不时录制一些技术视频】，并把它分享在 B站上。正是因为这样的公开表达，最终收获了一些朋友的关注，尤其部分比较热心的人甚至会私信我表示感谢，这里面充满了正反馈，也让我感受到了一点点意义。这里简单讲一下相关的数据（以后全网基本都叫【chaofa用代码打点酱油】了）
我写博客其实还挺早的，但是根本没有人在看，没有什么反馈，基本属于自嗨。第一次自建博客是 17年，那个时候写了一个关于自己学习 React 的一些记录，但是后面读研之后不做前端了就把对应的内容删除了。后面受到【极客兔兔】在 V站发帖自建博客过程的影响，又开始第二次写自己的博客，这时候是 19年的 6 月份，也就是当前的博客：chaofa用代码打点酱油，改过很多次名字，但唯一值得高兴的时候，这个博客持续存在了 5年，里面记录了自己很多的碎碎念。
那么为什么今年我却把它称为【公开表达元年】，因为这一次不一样了。以前我想过写公众号，想过做视频，想过回答知乎问题，但是大多数都没有坚持下去。为此我今年反思了一下为什么以前没有持续下去？
大多数人都知道“公开表达是难却正确的事情”，但是真正的领悟者却不多，李笑来在《把时间当做朋友》一书中提到，互联网用户行为遵循 "90-9-1原则"。
不是说因为创作者特别稀疏我们采取成为创作者的，而是这其中有巨大的好处。
那么后面我应该怎么做呢？
去年换工作之后，高强度的工作了一段时间，加上和岗位、老板的风格不是很适应，没干多久就感觉天天精疲力尽的，很快就想要辞职，但是迫于职业生涯的延续性，我自然是不敢真的就裸辞，因此在苦苦坚持，想要寻求一些方法延续自己的职业生涯，比如：自我鼓励——《工作，再坚持坚持》，理性分析——《如何在大厂工作六个月以上且保持一定的心理健康?》，只能说收效甚微。毕竟饿了就想吃饭，累了就会想休息，天经地义。
差不多待一年之后，就开始考虑活水转岗 or 换工作，不过深圳的就业机会还是较少，思来想去还是觉得活水合适一些。然后开始内部看一些机会，这个时候又涉及到去干什么业务的问题，所以很纠结到底干什么？继续和在腾讯一样做广告相关的业务，还是去做搜索，还是去做推荐，还是去做纯AIGC的业务，还是去做NLP相关的业务，最后兜兜转转又到了与我最有缘分的客服。
活水可能也算是一次跳槽吧，毕竟也要经过3轮技术面试，因此我基本把它当作全新的工作，工作方式也适当地做出一些改变。工作上一直在向表现好的大佬学习，希望明年工作上有一些突破。
2024-12-09 06:00:00
nums_key_value_head 参数就可nums_key_value_head 设置等于 1 就是 MQAnums_key_value_head 设置等于 nums_head 就是 MHA如果不喜欢看文字的同学可以查看 B站 或者 YouTube 视频。
备注：也可以直接由 GQA 中修改参数得到。但是本代码更完整一些
import math
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, hidden_dim, nums_head) -> None:
super().__init__()
self.nums_head = nums_head
# 一般来说，
self.head_dim = hidden_dim // nums_head
self.hidden_dim = hidden_dim
# 一般默认有 bias，需要时刻主意，hidden_dim = head_dim * nums_head，所以最终是可以算成是 n 个矩阵
self.q_proj = nn.Linear(hidden_dim, hidden_dim)
self.k_proj = nn.Linear(hidden_dim, hidden_dim)
self.v_proj = nn.Linear(hidden_dim, hidden_dim)
# gpt2 和 bert 类都有，但是 llama 其实没有
self.att_dropout = nn.Dropout(0.1)
# 输出时候的 proj
self.o_proj = nn.Linear(hidden_dim, hidden_dim)
def forward(self, X, attention_mask=None):
# 需要在 mask 之前 masked_fill
# X shape is (batch, seq, hidden_dim)
# attention_mask shape is (batch, seq)
batch_size, seq_len, _ = X.size()
Q = self.q_proj(X)
K = self.k_proj(X)
V = self.v_proj(X)
# shape 变成 （batch_size, num_head, seq_len, head_dim）
q_state = Q.view(batch_size, seq_len, self.nums_head, self.head_dim).permute(
0, 2, 1, 3
)
k_state = K.view(batch_size, seq_len, self.nums_head, self.head_dim).transpose(
1, 2
)
v_state = V.view(batch_size, seq_len, self.nums_head, self.head_dim).transpose(
1, 2
)
# 主意这里需要用 head_dim，而不是 hidden_dim
attention_weight = (
q_state @ k_state.transpose(-1, -2) / math.sqrt(self.head_dim)
)
print(type(attention_mask))
if attention_mask is not None:
attention_weight = attention_weight.masked_fill(
attention_mask == 0, float("-1e20")
)
# 第四个维度 softmax
attention_weight = torch.softmax(attention_weight, dim=3)
print(attention_weight)
attention_weight = self.att_dropout(attention_weight)
output_mid = attention_weight @ v_state
# 重新变成 (batch, seq_len, num_head, head_dim)
# 这里的 contiguous() 是相当于返回一个连续内存的 tensor，一般用了 permute/tranpose 都要这么操作
# 如果后面用 Reshape 就可以不用这个 contiguous()，因为 view 只能在连续内存中操作
output_mid = output_mid.transpose(1, 2).contiguous()
# 变成 (batch, seq, hidden_dim),
output = output_mid.view(batch_size, seq_len, -1)
output = self.o_proj(output)
return output
attention_mask = (
torch.tensor(
[
[0, 1],
[0, 0],
[1, 0],
]
)
.unsqueeze(1)
.unsqueeze(2)
.expand(3, 8, 2, 2)
)
x = torch.rand(3, 2, 128)
net = MultiHeadAttention(128, 8)
net(x, attention_mask).shape备注：以下代码省略了 attention_dropout attention_mask等情况的处理，真实实现过程中需要考虑。
import torch
import torch.nn as nn
import math
# 忽略了 attention_mask, attention_dropout;
class GroupQueryAttention(nn.Module):
def __init__(self, hidden_dim, nums_head, nums_key_value_head):
super().__init__()
assert hidden_dim % nums_head == 0 # 可以整除
assert nums_head % nums_key_value_head == 0 # N 个 query head 为一组
self.hidden_dim = hidden_dim
self.nums_head = nums_head
self.nums_key_value_head = nums_key_value_head
self.head_dim = hidden_dim // nums_head
# 初始化 qkv o
self.q_proj = nn.Linear(hidden_dim, nums_head * self.head_dim) # out feature_size (nums_head * head_dim)
# k v out shape (nums_key_value_head * head_dim)
self.k_proj = nn.Linear(hidden_dim, nums_key_value_head * self.head_dim)
self.v_proj = nn.Linear(hidden_dim, nums_key_value_head * self.head_dim)
self.o_proj = nn.Linear(hidden_dim, hidden_dim) # input_size nums_head * head_dim
def forward(self, X, attention_mask=None):
# X shape (batch, seq, hidden_dim)
batch_size, seq, _ = X.size()
# qkv projection
q = self.q_proj(X) # （batch, seq, hidden_dim)
k = self.k_proj(X)
v = self.v_proj(X)
# attention_weight 目标shape 是 (batch, nums_head, seq, seq)
q = q.view(batch_size, seq, self.nums_head, self.head_dim)
k = k.view(batch_size, seq, self.nums_key_value_head, self.head_dim)
v = v.view(batch_size, seq, self.nums_key_value_head, self.head_dim)
# 关注: nums_head 和 nums_key_value_head 的关系
q = q.transpose(1, 2) # (b, nums_head, seq, head_dim)
k = k.transpose(1, 2) # (b, nums_key_value_head, seq, head_dim)
v = v.transpose(1, 2) # (b, nums_key_value_head, seq, head_dim)
# k v repeat； （广播操作）
k = k.repeat_interleave(self.nums_head // self.nums_key_value_head, dim=1)
v = v.repeat_interleave(self.nums_head // self.nums_key_value_head, dim=1)
attention_score = (q @ k.transpose(2, 3)) / math.sqrt(self.head_dim)
attention_weight = torch.softmax(attention_score, dim=-1)
# （attention_mask 忽略） # 可以看前面的视频
output = attention_weight @ v # (b, nums_head, seq, head_dim)
# output projection 变成 (b, seq, hidden_dim)
output = output.transpose(1, 2).contiguous()
final_output = self.o_proj(output.view(batch_size, seq, -1))
return final_output
# 测试
x = torch.rand(3, 2, 128)
net = GroupQueryAttention(128, 8, 4)
net(x).shape由于 MQA 是 GQA 的一种特殊形式，因此只要在参数设置的时候将 nums_key_value_head = 1 就是 Multi Query Self-Attention。
最后欢迎关注我，基本全网同名 chaofa用代码打点酱油
