-
「DeepSeek-V3 技术解析」：DeepSeek-V3-Base 预训练阶段解析 编者按： 这篇技术解析详细阐述了 DeepSeek-V3-Base 的预训练阶段所采用的关键技术。 文章重点介绍了三项核心技术：Document Packing 技术有效解决了输入序列长度差异导致的资源浪费问题；Fill-in-the-Middle（FIM）采用 PSM 框架和特殊 tokens，使模型具备上下文感知的中间内容生成能力；基于 YaRN 的长上下文窗口扩展技术则通过频率插值策略解决了位置编码的扩展挑战。 随后，文章详细描述了 DeepSeek-V3-Base 的预训练过程，包括数据构建、训练策略和评估结果。 评估显示，这些技术组合使 DeepSeek-V3 每训练 1T token 仅需 180K NVIDIA H800 GPU 小时数，并在“大海捞针”测试中展现卓越的长文本理解能力，为后续 RL 阶段奠定了优质基座。 作者 | Shirley Li 编译 | 岳扬 这是 DeepSeek 系列文章的第五篇，也是首篇聚焦 DeepSeek-V3 [1, 2] 训练流程的文章。 如下图所示，DeepSeek-V3 的训练分为多个阶段： 产出 DeepSeek-V3-Base 基础模型的预训练阶段 基于 DeepSeek-V3-Base，通过大规模强化学习（RL）分别训练出 DeepSeek-R1-Zero（无需监督式微调冷启动）和 DeepSeek-R1（含有监督式微调） 利用 DeepSeek-R1 生成推理数据，用于 DeepSeek-V3 的监督式微调（SFT），接着是未在图中展示的 RL 阶段。 图 1. DeepSeek-V3 训练流程示意图（由原文作者绘制） 本文将重点关注产出 DeepSeek-V3-Base 的预训练阶段，阐述该阶段实现高效预训练的关键技术。后续文章将涵盖： 群组相对策略优化（GRPO）[7] DeepSeek-R1-Zero 和 DeepSeek-R1 的训练细节 DeepSeek-V3 的后训练阶段（监督式微调与 RL 阶段） 目录 技术背景：解析 DeepSeek-V3 预训练阶段的相关技术，包括 Document Packing，Fill-in-Middle 和 long context extension。 预训练阶段：详解如何构建预训练数据、强调一些关键的训练策略，并回顾评估结果。 总结 参考文献 01 技术背景 本节将介绍预训练 DeepSeek-V3 过程中使用的几种技术，包括 document packing、Fill-in-the-Middle（FIM）和基于 YaRN 的长上下文窗口扩展技术。 1.1 Document Packing 要理解为什么需要 document packing，我们首先需要回顾一下 Transformer 模型是如何构建输入序列 tokens 的。 Transformer 模型默认情况下需要固定长度的 token 序列作为输入，然而同一 batch 的文本输入往往长度不同。为了适应这种情况，文本输入通常需要经过以下预处理步骤： 将所有原始文本输入分词为 token 序列 将 token 序列截断或填充到预定义的固定长度（max_seq_len）：若原始序列过长则截断，否则用特殊 [PAD] token 进行填充 生成掩码 IDs 使模型在训练时能忽略填充的 token 为了更清晰地展示这个过程，以下这个示例我们将使用 GPT-2 [10]的分词器处理两个句子： 运行上述脚本后，会得到如下输出，其中： 第一句话被填充了 4 个额外的 padding token，体现在 input_ids 和 mask_ids 中； 第二句被截断，因此无需添加 padding token。 图 2. 填充操作示例（此图由作者绘制） 上述截断和填充方法虽然能让模型处理不同长度的输入，但当输入序列长度差异过大时（这在 LLM 训练中非常常见）会引发一系列问题： 对超长序列，截断可能导致有用信息丢失 对较短的序列，填充过多 token 会造成计算资源浪费 因此，LLM 训练通常采用 document packing 技术来处理输入序列。 更具体地说，如果给定若干长度不同的文档，我们首先将其分割为较小的块（chunk），如下图所示（用不同颜色代表不同文档）： 图 3. 文档分割（图片改编自文献[3]） 随后，我们将不同文档的块（chunk）进行拼接，以避免对长文档进行截断和对短文档进行填充： 图 4. 传统拼接方式（图片改编自文献[3]） 在上例中： 第一个输入（译者注：图 4 第一行）仅包含文档 1 的 tokens 第二个输入（译者注：图 4 第二行）拼接自文档 1 和文档 2 的 tokens 第三个输入（译者注：图 4 第三行）拼接自文档 2 和文档 3 的 tokens 第四个输入（译者注：图 4 第四行）拼接自文档3、4、5 的 tokens 这种方法虽能在一定程度上避免进行填充和截断，但由于仅按数据中的相对顺序拼接来自不同文档的块（chunks），无法控制最终输入序列的构建方式。 例如：文档 3（紫色）被不必要地分割为两部分，尽管其实际长度小于 max_seq_len，可以完整放入。 为了解决这个问题，文献 [3] 提出了 Best-fit Packing 技术，通过两个步骤完全消除不必要的分割： Step 1：将每个文档分割为更小的块。 Step 2：以一种智能的方式将这些块（chunks）分组为训练序列，确保在不进一步分割任何块（chunks）的前提下生成最少量的序列。 图 5. Best-fit packing技术（此图改编自文献[3]） 1.2 Fill-in-the-Middle（FIM） 在传统的自回归生成中，只能以从左到右的方式训练模型，即模型只能根据前面的 tokens 预测下一个 token。然而在实际应用中，模型常需根据上下文生成中间缺失的内容。 尤其在代码生成场景中 —— 我们常会给定输入/输出和部分代码片段，要求模型填充中间逻辑，如下例所示： 为了适配此类需求，文献 [4] 提出了一种简单有效的方法，称为 “fill-in-the-middle”：即将文档随机切分为 prefix、middle 和 suffix 三部分，然后将 middle 部分移至末尾： 由于数据组织形式为 "Prefix-Suffix-Middle"，该方法常被称为 PSM 框架。实际实现时通过添加特殊 token 来标记各部分的边界： 其中： <|fim_begin|>和<|fim_hole|>标记 prefix 部分 <|fim_hole|>和<|fim_end|>标记 suffix 部分 <|fim_end|>和<|eos_token|>标记 middle 部分 以如下输入为例： 若需模型预测第二行代码，可将该行作为 middle 部分，并构造 FIM 输入如下： 图 6. PSM 框架示意图（此图由作者绘制） 此时模型的预期输出应为： 1.3 基于 YaRN 的长上下文窗口扩展技术 现代 LLM 常需处理极长的提示词（如整个代码仓库），但直接使用 128K 等长上下文窗口进行预训练并不现实。多数 LLM 采用分阶段渐进式扩展策略：先在较小的上下文窗口进行预训练，再分多个阶段逐步扩展到更长的上下文窗口，从而大大降低训练成本。 例如，在 DeepSeek-V3 中，模型首先使用 4K 的上下文窗口完成预训练，然后再分两阶段扩展到 128K： 第一阶段：从 4K 到 32K（1000 steps） 第二阶段：从 32K 到 128K（再 1000 steps） 需特别指出的是，这种扩展不能通过简单调大上下文窗口实现，而需借助基于旋转位置编码（RoPE）改进的 YaRN（Yet another RoPE extensioN）技术对位置编码进行修改。 关于 RoPE 的详细介绍，请参阅我们之前的文章《「DeepSeek-V3 技术解析」：多头潜在注意力机制（MLA）》。 RoPE 是一种相对位置编码方法，其核心思想是通过使用复杂的旋转嵌入修改 Query 和 Key，使得二者的内积依赖于它们的相对位置： 然而，由于余弦函数和正弦函数是周期性的，(pos_i, pos_j) 之间的内积可能看起来与 (pos_i, pos_k) 之间的内积相似，因此在固定 θ 的情况下，仅使用 1K tokens（即位置索引 1~1000） 进行预训练的模型在测试时可能会混淆，因为测试时遇到的位置索引（如 5K 或 10K）可能远远超出了预训练时的上下文窗口。 下图展示了这种现象：当 32K 上下文窗口的预训练模型在超出该窗口的位置测试时，困惑度（Perplexity）急剧上升： 图 7. 困惑度与上下文窗口的关系（此图由作者绘制） 那么，YaRN 是如何应对这一挑战的呢？ 既然外推法（extrapolate）效果欠佳，YaRN 转而采用插值频率（interpolate the frequency）的策略。 假设我们有一个在 4 个 token 长度的输入上训练的模型，希望将其扩展到 8 个 token，且基础频率 θ=0.5。 对于原始 RoPE，直接使用 cos(θ×pos) 和 sin(θ×pos) 对 Query 和 Key 进行旋转即可。 而对于 YaRN： 首先，计算扩展后的上下文长度与原始长度的比值作为缩放因子，本例中为 2。 然后，生成新频率 θ' = θ / 2 = 0.25。 再使用新频率对 Query 和 Key 进行旋转，即 cos(θ'×pos) 和 sin(θ'×pos)。 下图对比了 RoPE 与 YaRN 的 cos 和 sin 值： 图 8. YaRN 工作原理示意图（此图由作者绘制） 通过该图可观察到： 在 RoPE 中，cos 和 sin 值会随位置索引的增加而快速振荡，导致扩展到更长的上下文时出现问题。 而在 YaRN 中，原始的余弦和正弦函数通过频率缩放被插值到扩展后的上下文长度（如蓝色高亮区域所示），实现了更平滑的过渡，使得模型能够更有效地处理长序列。 下图展示了 DeepSeek-V3 在"大海捞针"（Needle In A Haystack，NIAH）测试中的表现，表明其在 128K 以下的上下文窗口长度中均表现出色： 图 9. DeepSeek-V3 的"大海捞针"测试结果（引自文献[2]） 02 预训练阶段 本节将介绍 DeepSeek-V3-Base 的训练方法，重点解析数据构建流程，并强调预训练阶段中的一些关键策略。 2.1 数据构建 数据规模与质量对 LLM 训练至关重要。DeepSeek-V3 的预训练语料库通过持续优化策略构建，具体优化路径如下： 在 DeepSeek 67B [8] 中，训练语料采用去重-过滤-再混合策略构建。 首先对 Common Crawl 语料进行去重，随后通过严格的文档质量评估标准进行过滤，最后通过数据再混合阶段解决数据不平衡问题。 在 DeepSeek-V2 [9] 中，通过以下方式扩展训练语料：1) 增加更多中文数据及来自不同来源的高质量数据；2) 通过优化数据清洗流程，恢复大量此前在文献 [8] 的策略中被删除的数据。同时，通过改进基于质量的过滤算法提升数据质量。 在 DeepSeek-V3 [2] 中，预训练语料进一步扩充，加入更多数学与编程样本，以及除中英文之外的多语言样本。 收集的预训练语料会通过前文提出的 Prefix-Suffix-Middle（PSM）框架结合 FIM（Fill-in-Middle）策略进行预处理，并应用 document-packing 技术。 2.2 训练策略 原论文[2]对预训练参数进行了详细描述，此处我们仅强调几个关键点： 长上下文窗口扩展：首先在 14.8T token 上以 4K 上下文窗口进行预训练，随后通过 1000 steps 扩展到 32K 上下文，最终再通过 1000 steps 扩展到 128K 上下文。 多词元预测：如我们本系列前一篇文章《「DeepSeek-V3 技术解析」：多词元预测技术（Multi-Token Prediction, MTP）》所述，DeepSeek-V3 采用了优化版的多词元预测机制，允许模型同时解码多个词元（tokens），以加速训练中的解码过程。 以 FP8 精度进行训练：DeepSeek-V3 采用混合精度计算提升效率，对部分计算使用低精度格式（如 8-bit 浮点数），在不过度影响精度的前提下减少内存占用并加速计算。 学习率的调度：在前 2K steps 中，学习率（learning rate）从 0 线性增长至 2.2e–4，并在 10T token 的训练过程中保持恒定；随后在 4.3T token 的训练过程中按照余弦曲线下降至 2.2e-5；在最后 500B token 的训练过程中，前 333B token 保持恒定的学习率，剩余 167B token 进一步降至 7.3e-6。 Batch size 的调度：在前 469B token 的训练过程中，Batch size 从 3072 逐步提升至 15360，后续训练中保持恒定。 2.3 评估结果 下表对比了 DeepSeek-V3 与其他开源基座模型在不同任务上的表现。其中 DeepSeek-V3 在多数数据集上都取得了最佳性能，尤其是在数学与代码相关的任务中表现突出。 需特别说明，得益于本系列文章中介绍的各项创新技术，DeepSeek-V3 的优异性能是在极高的训练效率下实现的。具体而言，DeepSeek-V3 每训练 1T token 仅需 180K H800 GPU hours，远低于训练 72B 或 405B 稠密模型的成本。 文献[2]中的表 3 文献 [2] 还通过全面的消融实验验证了无辅助损失函数的负载均衡、多词元预测等关键技术。由于我们已在前文中讨论过相关内容，此处不再赘述。 03 总结 本文探讨了 DeepSeek-V3 预训练策略中的关键创新，旨在提升效率、可扩展性与性能。由此产生的 DeepSeek-V3-Base 模型成为更高级推理模型（如 DeepSeek-R1-Zero 和 DeepSeek-R1）的基础，而这些模型又通过知识蒸馏反哺优化 DeepSeek-V3。 除此前讨论的架构创新 —— 多头潜在注意力（Multi-head Latent Attention）、DeepSeekMoE、无辅助损失函数的负载均衡及多词元预测（Multi-token Prediction）外，本文还引入了包括 document packing、Fill-in-the-Middle（FIM）和基于 YaRN 的长上下文窗口扩展在内的多项技术。 这些技术共同推动了大语言模型效率与可扩展性边界的突破，为高性能 AI 模型设立了新标杆。 参考文献 [1] DeepSeek（https://www.deepseek.com/） [2] DeepSeek-V3 Technical Report（https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf） [3] Fewer Truncations Improve Language Modeling（https://arxiv.org/abs/2404.10830） [4] Efficient Training of Language Models to Fill in the Middle（https://arxiv.org/abs/2207.14255） [4] DeepSeek-Coder: When the Large Language Model Meets Programming — The Rise of Code Intelligence（https://arxiv.org/abs/2401.14196） [5] DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence（https://arxiv.org/abs/2406.11931） [6] YaRN: Efficient Context Window Extension of Large Language Models（https://arxiv.org/abs/2309.00071） [7] DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models（https://arxiv.org/abs/2402.03300） [8] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism（https://arxiv.org/pdf/2401.02954） [9] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（https://arxiv.org/abs/2405.04434） [10] Language Models are Unsupervised Multitask Learners（https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf） Thanks for reading! Hope you have enjoyed and learned new things from this blog! About the author Shirley Li I am a Machine Learning Engineer working on building multi-modality models to solve real-world problems. END 本期互动内容 当前位置编码方案（RoPE/YaRN）已支持 128K 上下文，但人类书籍平均长度约 200K tokens。要实现真正无损的长文档理解，您认为下一代位置编码需要突破哪些理论瓶颈？ 原文链接： https://medium.com/data-science-collective/deepseek-explained-5-deepseek-v3-base-86c078ed5504 -
DeepSeek模型MOE结构代码详解 其实在DeepSeek-R1爆火之前，DeepSeek V2在我们行业就已经妇孺皆知了，它独特的MOE结构值得研究一下。这篇文章是基于 ZOMI酱 的2个视频写的，这2个视频讲的很好，建议大家都学习一下：《MOE终于迎来可视化解读！傻瓜都能看懂MoE核心原理！》和《使用昇腾NPU手撕MoE单机版代码！没想到如此简单！》。 这篇文章是把我自己的理解梳理一下，加强自己的理解和记忆。 MOE结构概述 我们可以从zomi酱视频里面的这张图开始： MOE是mixture of experts 的缩写，简单来说，就是把传统transformer结构中decoder层里面的单个线性层替换层多个并列的线性层。在这些线性层前面还有一个Router，Router会选择并列线性层里面的一部分进行计算。这样的话，既能让模型学习更多的知识（多个“专家”），又能减少推理计算量（选择部分“专家”进行计算）。 MOE计算代码 接下来我们参考zomi酱提供的代码来详细看一下MOE的计算过程是怎样的： import torch import torch.nn as nn import torch.nn.functional as F import torch_npu from torch_npu.contrib import transfer_to_npu class Expert(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super().__init__() self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, output_dim)) def forward(self, x): return self.net(x) class MoE(nn.Module): def __init__(self, input_dim, num_experts, top_k, expert_capacity, hidden_dim, output_dim): super().__init__() self.num_experts = num_experts self.top_k = top_k self.expert_capacity = expert_capacity # 路由网络 self.gate = nn.Linear(input_dim, num_experts) # 专家集合 self.experts = nn.ModuleList( [Expert(input_dim, hidden_dim, output_dim) for _ in range(num_experts)]) def forward(self, x): batch_size, input_dim = x.shape device = x.device # 路由计算 logits = self.gate(x) probs = torch.softmax(logits, dim=-1) print("probs: ", probs) topk_probs, topk_indices = torch.topk(probs, self.top_k, dim=-1) print("topk_probs: ", topk_probs) print("topk_indices: ", topk_indices) # 辅助损失计算 if self.training: # 重要性损失（专家利用率均衡） importance = probs.sum(0) importance_loss = torch.var(importance) / (self.num_experts ** 2) # 负载均衡损失（样本分配均衡） mask = torch.zeros_like(probs, dtype=torch.bool) mask.scatter_(1, topk_indices, True) routing_probs = probs * mask expert_usage = mask.float().mean(0) routing_weights = routing_probs.mean(0) load_balance_loss = self.num_experts * (expert_usage * routing_weights).sum() aux_loss = importance_loss + load_balance_loss else: aux_loss = 0.0 # 专家分配逻辑 flat_indices = topk_indices.view(-1) flat_probs = topk_probs.view(-1) sample_indices = torch.arange(batch_size, device=device)[:, None]\ .expand(-1, self.top_k).flatten() print("sample_indices: ", sample_indices) # 初始化输出 outputs = torch.zeros(batch_size, self.experts[0].net[-1].out_features, device=device) # 处理每个专家 for expert_idx in range(self.num_experts): print("expert_idx: ", expert_idx) # 获取分配给当前专家的样本 expert_mask = flat_indices == expert_idx print("expert_mask: ", expert_mask) expert_samples = sample_indices[expert_mask] print("expert_samples: ", expert_samples) expert_weights = flat_probs[expert_mask] print("expert_weights: ", expert_weights) # 容量控制 if len(expert_samples) > self.expert_capacity: expert_samples = expert_samples[:self.expert_capacity] expert_weights = expert_weights[:self.expert_capacity] if len(expert_samples) == 0: continue # 处理专家计算 expert_input = x[expert_samples] print("expert_input: ", expert_input) expert_output = self.experts[expert_idx](expert_input) weighted_output = expert_output * expert_weights.unsqueeze(-1) # 累加输出 outputs.index_add_(0, expert_samples, weighted_output) return outputs, aux_loss # 测试示例 if __name__ == "__main__": input_dim = 5 output_dim = 10 num_experts = 8 top_k = 3 expert_capacity = 32 hidden_dim = 512 batch_size = 10 # add device = torch.device("npu:4" if torch.npu.is_available() else "cpu") moe = MoE(input_dim, num_experts, top_k, expert_capacity, hidden_dim, output_dim).to(device) x = torch.randn(batch_size, input_dim).to(device) moe.eval（) output, _ = moe(x) print(f"Eval output shape: {output.shape}") # torch.Size([64, 256]) 初始化函数定义 首先，定义了Expert类，也就是“专家”，可以看到，专家是由线性层和激活函数构成的简单模型。 然后开始定义MOE类。在初始化函数中，定义了这样几个变量： self.num_experts：专家的数量，也就是上面提到的“并列线性层”的个数，训练后的每个专家的权重都是不同的，代表它们所掌握的“知识”是不同的。 self.top_k：每个输入token激活的专家数量。 self.expert_capacity：代表计算每组token时，每个专家能被选择的最多次数。 self.gate：路由网络，一般是一个线性层，用来计算每个专家被选择的概率。 self.experts：实例化Expert类，生成多个专家。 前向计算逻辑 接下来看一下forward函数。为了方便大家理解，我们把上面代码的执行打印结果也一起附上。 首先是输入x，shape是（batch_size, input_dim），batch_size我们可以看作是token的数量，也就是序列长度。然后通过self.gate和softmax计算每个token在每个专家上的激活概率： logits = self.gate(x) probs = torch.softmax(logits, dim=-1) probs的打印结果如下：我们设置的batch_size是10，num_experts是8，所以probs是个10行8列的矩阵。 probs: tensor([[0.1710, 0.1348, 0.0746, 0.1714, 0.0594, 0.2695, 0.0251, 0.0940], [0.1556, 0.0776, 0.1658, 0.1489, 0.1152, 0.1679, 0.0565, 0.1124], [0.1077, 0.1154, 0.1564, 0.1317, 0.0630, 0.2026, 0.0518, 0.1715], [0.0681, 0.0680, 0.1236, 0.1030, 0.1707, 0.2827, 0.0627, 0.1211], [0.0453, 0.0648, 0.2313, 0.0781, 0.1026, 0.1304, 0.1326, 0.2149], [0.1394, 0.2278, 0.0625, 0.1832, 0.0395, 0.1512, 0.0691, 0.1274], [0.1096, 0.1462, 0.1302, 0.1397, 0.0607, 0.1898, 0.0639, 0.1598], [0.1200, 0.1952, 0.0970, 0.1648, 0.0360, 0.1072, 0.1018, 0.1779], [0.0650, 0.0501, 0.1463, 0.1025, 0.2219, 0.1446, 0.1439, 0.1257], [0.0641, 0.0813, 0.0579, 0.1348, 0.1170, 0.0631, 0.3554, 0.1264]], device='npu:4', grad_fn=<SoftmaxBackward0>) 接着，再用topk算子把每个token的激活专家选出来： topk_probs, topk_indices = torch.topk(probs, self.top_k, dim=-1) topk_probs和topk_indices 的打印结果如下，因为我们设置的top_k=3，所以每个token都把排名前三的概率选出来了，同时topk_indices把这些概率对应的专家编号也选出来了，比如第0个token，激活了5号专家、3号专家、0号专家。 topk_probs: tensor([[0.2695, 0.1714, 0.1710], [0.1679, 0.1658, 0.1556], [0.2026, 0.1715, 0.1564], [0.2827, 0.1707, 0.1236], [0.2313, 0.2149, 0.1326], [0.2278, 0.1832, 0.1512], [0.1898, 0.1598, 0.1462], [0.1952, 0.1779, 0.1648], [0.2219, 0.1463, 0.1446], [0.3554, 0.1348, 0.1264]], device='npu:4', grad_fn=<TopkBackward0>) topk_indices: tensor([[5, 3, 0], [5, 2, 0], [5, 7, 2], [5, 4, 2], [2, 7, 6], [1, 3, 5], [5, 7, 1], [1, 7, 3], [4, 2, 5], [6, 3, 7]], device='npu:4') self.training分支对应的是训练过程中计算损失函数的部分，我们后面再讲。 选择好专家后，就要开始计算了。计算规则是，对于每个token，假如它选择的专家是e1、e2、e3，概率分别是p1、p2、p3，那么这个token的计算结果就是p1e1_out+p2e2_out+p3*e3_out。 由于计算个体是每个专家，所以代码中用for循环遍历每个专家。我们以第0个专家为例，看看它的计算过程是怎样的。 首先需要确定0号专家的输入。由于不是每个token都选择了0号专家，所以不能把x直接作为输入，而是要确定一个下标向量idxes，把x[idxes]作为0号专家的输入，idxes的值就是激活了0号专家的所有token编号，那么怎么得到idxes呢？代码里面是这样做的： 首先计算一个mask： expert_mask = flat_indices == expert_idx 打印结果如下： expert_mask: tensor([False, False, True, False, False, True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False], device='npu:4') flat_indices是topk_indices平铺之后的向量。通过对比，可以看到expert_mask中True的位置和topk_indices中0的位置铺平之后是一致的，代表第0个专家被第0个token和第1个token激活了。 而且expert_mask代表的含义是：只要它的第0-2的位置是True的话，就代表被第0个token激活了，只要它的第3-5的位置是True的话，就代表被第1个token激活了，以此类推，我们可以声明一个sample_indices向量： sample_indices: tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9], device='npu:4 再通过下面的代码就可以把idxes取出来了： expert_samples = sample_indices[expert_mask] 也顺便把概率权重取出来： expert_weights = flat_probs[expert_mask] 接着把输入取出来： expert_input = x[expert_samples] 打印结果如下： expert_samples: tensor([0, 1], device='npu:4') expert_weights: tensor([0.1710, 0.1556], device='npu:4', grad_fn=<IndexBackward0>) expert_input: tensor([[-1.4382, -1.5939, 0.0802, -0.5614, 0.2586], [-1.2631, 1.0266, 0.1806, -0.7280, -0.6913]], device='npu:4') 再进行专家计算： expert_output = self.experts[expert_idx](expert_input) weighted_output = expert_output * expert_weights.unsqueeze(-1) 最后还需要把计算结果叠加到对应的token上面去： outputs.index_add_(0, expert_samples, weighted_output) 完成上面的for循环之后，就把所有专家的计算任务完成了，通过index_add_的操作，把每个token的计算结果也汇总了。 损失函数 损失函数包含2部分：专家利用率均衡和样本分配均衡。 首先是专家利用率均衡，如果每个专家被选择的概率相近，那么说明分配越均衡，损失函数越小： importance = probs.sum(0) importance_loss = torch.var(importance) / (self.num_experts ** 2) 然后是样本分配均衡，首先得到每个token、每个专家的分配概率矩阵： mask = torch.zeros_like(probs, dtype=torch.bool) mask.scatter_(1, topk_indices, True) routing_probs = probs * mask 然后按照token维度（样本维度）求平均，得到每个专家被分配的token平均数量和平均概率： expert_usage = mask.float().mean(0) routing_weights = routing_probs.mean 两者相乘求和得到负载均衡损失： load_balance_loss = self.num_experts * (expert_usage * routing_weights).sum() 样本分配越均衡，这个损失函数越小。举个例子，10个专家，10个样本，如果所有样本都分到1个专家，那么损失函数值为10x1+0+0...+0=10，如果平均分给10个专家，那么损失函数值为1x0.1+1x0.1+...+1x0.1=1。 本文由博客一文多发平台 OpenWrite 发布！