关于LLM 训练和推理的 padding

训练时候可以进行 左pad 或者 右pad,或者对 prompt 进行左pad,对 label 进行右pad。现在其实一般预训练或者微调的时候都不pad,否则会影响训练效率,大概的思路:假设 batch size = 2,max_seq_len = 16,sequence 1、2、3、4 分别有 7、9、6、10 个 token,那么就可以组成[[s1+s2], [s3+s4]] 进行训练,这个时候需要构造一个正确的 casual attention mask。flash_attn_varlen_qkvpacked_func 接口,就可以实现这样的计算而无需 padding。

batch 推理的时候一般只用 左pad。推理时也只有batch推理会有影响,另外左对齐方便所有行同时产生next token。在强化学习训练PPO/DPO/GRPO 的时候需要用到推理,所以也需要做左pad!!

padding_side 的影响

谈到 padding,我们自然要考虑 attention_mask,借助 attention_mask 可以在计算 attention weight 时将 padding 带来的影响屏蔽掉。下面是设置不同的 padding_side,tokenizer 的输出:

没有设置 padding_side 或者 padding_side=’right’:

>>> from transformers import LlamaForCausalLM, LlamaTokenizer
>>> tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> prompts = ["hello llama", "who are you?"]
>>> tokenizer(prompts, return_tensors="pt", padding=True)
{
    'input_ids': tensor([[    1, 22172, 11148,  3304,     2],                                                                                                                                                                                                                         │············
                         [    1,  1058,   526,   366, 29973]]),
    'attention_mask': tensor([[1, 1, 1, 1, 0],  [1, 1, 1, 1, 1]])
}

设置 padding_side=’left’:

>>> from transformers import LlamaForCausalLM, LlamaTokenizer
>>> tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", padding_side="left")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> prompts = ["hello llama", "who are you?"]
>>> tokenizer(prompts, return_tensors="pt", padding=True)
{
    'input_ids': tensor([[    2,     1, 22172, 11148,  3304],                                                                                                                                                                                                                         │············
                         [    1,  1058,   526,   366, 29973]]),
    'attention_mask': tensor([[0, 1, 1, 1, 1],  [1, 1, 1, 1, 1]])
}

要理解 padding_side=’right’ 为什么会导致结果不正确,关键的点是 next token 的预测是使用句子的最后一个 token 经过 transformer 层之后输出的 logit 来得到 next token 的。下面是 model.generate通过多次跳转后来到 next token 的处理逻辑:

# https://github.com/huggingface/transformers/blob/a7cab3c283312b8d4de5df3bbe719971e24f4281/src/transformers/generation/utils.py#L2411
        
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

# forward pass to get next token
outputs = self(
    **model_inputs,
    return_dict=True,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
)

next_token_logits = outputs.logits[:, -1, :]
# argmax
next_tokens = torch.argmax(next_tokens_scores, dim=-1)

从上面的代码可以看到,句子最后一个 token 所对应的 logit 会被用来计算 next token,因此,最后一个 token logit 的计算是否正确决定了推理的结果是否正确。
接下来,我们来看一下 padding_side=’left’ 和 padding_side=’right’,最后一个 token 所对应的 logit 是否是正确计算的。


我们先来看 padding_side=’left’ 的最后一个 logit 的计算过程,省略中间的具体细节,只给出关键的过程),这里只关注句子 “hello llama”:

从图 4 的计算过程可以看到,使用 padding_side=’left’ 的方式,attention score after masked 矩阵的最后一行和 V 的第一列进行内积后得到的值为正确且符合期望的值,即最后一个 token 所对应的 logit 的计算没有受 padding 的影响,该 logit 的计算过程正确。

因为最后一列计算得分时候,V第一行:pad token 的权重【 attention score】都是 0,且attention score 左下角权重为0那么计算结果最后一列的结果 只跟非pad的V有关

我们接下来看一下 padding_side=’right’ 的最后一个 logit 的计算过程:

从图 5 的计算过程可以看到,attention score after masked 矩阵的最后一行和 V 的第一列进行内积后得到的值是不符合期望的,即最后一个 token(pad token)所对应的 logit 的计算不正确,因为 pad token 也参与了计算,而正确预测 next token 的时候 pad token 是不应该参与计算的。

因为最后一列计算得分时候,V最后一行:pad token对应的的权重【 attention score 最后一行】不都是0,,那么计算结果最后一列的结果 跟非pad的V有关。

至此,我们弄清楚了为什么 padding_side=’right’ 会产生不正确的结果。

Prepacking-消除attention padding冗余计算

一个batch里,不同的request,其prompt长度不一样,这样计算attention时会做padding,确保所有的request长度相同。如下图所示,1个batch总共4个句子,后面3个句子做了padding,这样做的一个问题就是,会浪费计算。

,一个解决方法是去除padding,把这些句子放在一个句子里计算。如下图所示,去除了padding之后,所有的request放在一个句子里。但是带来的问题是,attention计算只有句子之内的不同token需要进行attention计算(红色计算attention、蓝色计算attention等等),句子之间是独立的。所以这种做法必须要进行适当的数据组织,让我们的attention算子能知道自己该把哪些token放在一个句子里计算。如下图所示,虽然输入到attention算子的是一个句子,包含10个token,但是需要一些额外数据,让attention算子去把红色部分当一个request计算attention、蓝色部分当一个request计算attention、绿色部分当一个request计算attention、黄色部分当一个request计算attention。

预打包在概念上很简单;我们不是将每个序列填充到相同的长度,而是使用现成的装箱算法将多个提示打包在一起,以代替填充标记

Qwen3 技术报告

  • MoE 模型:Qwen3-235B-A22B 和 Qwen3-30B-A3B;其中 235B 和 30B 分别是总参数量,22B 和 3B 分别是激活参数量。
  • 密集模型:Qwen3-32B、Qwen3-14B、Qwen3-8B、Qwen3-4B、Qwen3-1.7B 和 Qwen3-0.6B。

整体架构:

1) 包含num_experts个轻量级专家网络(Qwen3MoeMLP)的并行计算单元;

2) 基于注意力机制的路由网络(gate)。

在计算过程中,路由网络通过动态决策机制为每个输入Token生成路由决策,筛选出匹配度最高的top_k个专家节点。随后,系统将根据路由权重对选定专家的计算结果进行加权融合,最终生成该隐层的表征输出。

那么同样我们对比DeepSeekMOE,Qwen3MOE有两个点的改变:

1)没有shared expert。

2) 优化了MLP架构,变为Qwen3MoeSparseMoeBlock。

模型特性优化总结表

特性实现细节
注意力机制改进的Qwen3Attention(支持Flash Attention优化)
MoE路由策略Top-K专家选择(默认K=2),支持权重归一化
专家结构每个专家为标准MLP(hidden_size → moe_intermediate_size → hidden_size)
动态专家分配每间隔decoder_sparse_step层使用MoE(其他层使用标准MLP)
负载均衡机制通过router_logits计算辅助损失,防止专家极化
计算优化使用index_add操作实现零浪费的专家计算

对比传统MOE优化效果:

优化方向Qwen3-MoE实现方案对比传统MoE模型优势
路由机制Top-K + 动态权重归一化(norm_topk_prob)缓解专家利用不均衡问题,相比Mixtral的固定权重分配更灵活
稀疏模式分层动态稀疏(decoder_sparse_step控制MoE层间隔)混合密集与稀疏计算,相比全MoE结构降低计算开销
内存优化logits_to_keep参数支持部分logits计算长序列生成时内存占用减少,优于Mixtral的全序列计算
注意力机制改进的Flash Attention 3.0集成相比标准Attention实现,训练速度提升,显存占用减少
负载均衡改进的辅助损失函数(load_balancing_loss_func+自研调整系数)专家利用率从Mixtral的提升,防止专家极化
动态计算mlp_only_layers参数跳过MoE层支持按需切换密集/稀疏模式,相比固定结构推理灵活性提升

性能方面,在代码、数学、通用能力等基准测试中,旗舰模型 Qwen3-235B-A22B 与 DeepSeek-R1、o1、o3-mini、Grok-3 和 Gemini-2.5-Pro 等顶级模型表现相当

此外,小型 MoE 模型 Qwen3-30B-A3B 的激活参数数量是 QwQ-32B 的 10%,表现却更胜一筹。甚至像 Qwen3-4B 这样的小模型也能匹敌 Qwen2.5-72B-Instruct 的性能

性能大幅提升的同时,Qwen3 的部署成本还大幅下降,仅需 4 张 H20 即可部署满血版,显存占用仅为性能相近模型的三分之一

核心亮点

  • 多种思考模式

Qwen3 模型支持两种思考模式:

  1. 思考模式:在这种模式下,模型会逐步推理,经过深思熟虑后给出最终答案。这种方法非常适合需要深入思考的复杂问题。
  2. 非思考模式:在此模式中,模型提供快速、近乎即时的响应,适用于那些对速度要求高于深度的简单问题。

这种灵活性使用户能够根据具体任务控制模型进行“思考”的程度。例如,复杂的问题可以通过扩展推理步骤来解决,而简单的问题则可以直接快速作答,无需延迟。至关重要的是,这两种模式的结合大大增强了模型实现稳定且高效的“思考预算”控制能力。如上文所述,Qwen3 展现出可扩展且平滑的性能提升,这与分配的计算推理预算直接相关。这样的设计让用户能够更轻松地为不同任务配置特定的预算,在成本效益和推理质量之间实现更优的平衡。

下图为在 AIME24、AIME25、LiveCodeBech(v5)和 GPQA Diamond 等基准测试集中,非思考模式与思考模式的思考预算变化趋势。

  • 多语言

Qwen3 模型支持 119 种语言和方言。这一广泛的多语言能力为国际应用开辟了新的可能性,让全球用户都能受益于这些模型的强大功能。

语系语种&方言
印欧语系英语、法语、葡萄牙语、德语、罗马尼亚语、瑞典语、丹麦语、保加利亚语、俄语、捷克语、希腊语、乌克兰语、西班牙语、荷兰语、斯洛伐克语、克罗地亚语、波兰语、立陶宛语、挪威语(博克马尔语)、挪威尼诺斯克语、波斯语、斯洛文尼亚语、古吉拉特语、拉脱维亚语、意大利语、奥克语、尼泊尔语、马拉地语、白俄罗斯语、塞尔维亚语、卢森堡语、威尼斯语、阿萨姆语、威尔士语、西里西亚语、阿斯图里亚语、恰蒂斯加尔语、阿瓦德语、迈蒂利语、博杰普尔语、信德语、爱尔兰语、法罗语、印地语、旁遮普语、孟加拉语、奥里雅语、塔吉克语、东意第绪语、伦巴第语、利古里亚语、西西里语、弗留利语、撒丁岛语、加利西亚语、加泰罗尼亚语、冰岛语、托斯克语、阿尔巴尼亚语、林堡语、罗马尼亚语、达里语、南非荷兰语、马其顿语僧伽罗语、乌尔都语、马加希语、波斯尼亚语、亚美尼亚语
汉藏语系中文(简体中文、繁体中文、粤语)、缅甸语
亚非语系阿拉伯语(标准语、内志语、黎凡特语、埃及语、摩洛哥语、美索不达米亚语、塔伊兹-阿德尼语、突尼斯语)、希伯来语、马耳他语
南岛语系印度尼西亚语、马来语、他加禄语、宿务语、爪哇语、巽他语、米南加保语、巴厘岛语、班加语、邦阿西楠语、伊洛科语、瓦雷语(菲律宾)
德拉威语泰米尔语、泰卢固语、卡纳达语、马拉雅拉姆语
突厥语系土耳其语、北阿塞拜疆语、北乌兹别克语、哈萨克语、巴什基尔语、鞑靼语
壮侗语系泰语、老挝语
乌拉尔语系芬兰语、爱沙尼亚语、匈牙利语
南亚语系越南语、高棉语
其他日语、韩语、格鲁吉亚语、巴斯克语、海地语、帕皮阿门托语、卡布维尔迪亚努语、托克皮辛语、斯瓦希里语
  • 增强的 Agent 能力

我们优化了 Qwen3 模型的 Agent 和 代码能力,同时也加强了对 MCP 的支持。

预训练

在预训练方面,Qwen3 的数据集相比 Qwen2.5 有了显著扩展。Qwen2.5是在 18 万亿个 token 上进行预训练的,而 Qwen3 使用的数据量几乎是其两倍,达到了约 36 万亿个 token,涵盖了 119 种语言和方言。为了构建这个庞大的数据集,我们不仅从网络上收集数据,还从 PDF 文档中提取信息。我们使用 Qwen2.5-VL 从这些文档中提取文本,并用 Qwen2.5 改进提取内容的质量。为了增加数学和代码数据的数量,我们利用 Qwen2.5-Math 和 Qwen2.5-Coder 这两个数学和代码领域的专家模型合成数据,合成了包括教科书、问答对以及代码片段等多种形式的数据。

Qwen3模型采用三阶段预训练过程:

  1. 通用阶段 (S1): 在第一阶段,所有Qwen3模型使用4,096 token的序列长度,在超过30万亿token的数据上进行训练 。此阶段旨在建立模型的语言能力和通用世界知识基础,训练数据覆盖119种语言和方言 。
  2. 推理阶段 (S2): 为了进一步提升推理能力,此阶段的预训练语料库增加了STEM、编码、推理和合成数据的比例 。模型使用4,096 token的序列长度,在约5万亿高质量token上进行进一步预训练 。在此阶段还加速了学习率衰减 。
  3. 长上下文阶段: 在最后一个预训练阶段,收集高质量长上下文语料库,将Qwen3模型的上下文长度扩展到32,768 token 。长上下文语料库中,75%的文本长度在16,384到32,768 token之间,25%的文本长度在4,096到16,384 token之间 。报告提及沿用Qwen2.5的做法,使用ABF技术将RoPE的基础频率从10,000提高到1,000,000 。同时,引入YARN和Dual Chunk Attention (DCA)技术,在推理过程中实现序列长度容量的四倍增长 。

类似于Qwen2.5,Qwen3根据这三个预训练阶段开发了最优超参数(如学习率调度器和批次大小)预测的缩放律 。通过广泛实验,系统研究了模型架构、训练数据、训练阶段与最优训练超参数之间的关系 。最终为每个密集模型和MoE模型设定了预测的最优学习率和批次大小策略。

由于模型架构的改进、训练数据的增加以及更有效的训练方法,Qwen3 Dense 基础模型的整体性能与参数更多的Qwen2.5基础模型相当。例如,Qwen3-1.7B/4B/8B/14B/32B-Base 分别与 Qwen2.5-3B/7B/14B/32B/72B-Base 表现相当。特别是在 STEM、编码和推理等领域,Qwen3 Dense 基础模型的表现甚至超过了更大规模的 Qwen2.5 模型。对于 Qwen3 MoE 基础模型,它们在仅使用 10% 激活参数的情况下达到了与 Qwen2.5 Dense 基础模型相似的性能。这带来了训练和推理成本的显著节省。

后训练

为了开发能够同时具备思考推理和快速响应能力的混合模型,我们实施了一个四阶段的训练流程。该流程包括:(1)长思维链冷启动,(2)长思维链强化学习,(3)思维模式融合,以及(4)通用强化学习。

后训练部分详细介绍了Qwen3模型的后训练流程和评估结果 。Qwen3的后训练流程策略性地设计了两个核心目标:思维控制和强到弱蒸馏 。

思维控制 (Thinking Control):

思维控制涉及将“非思维”模式和“思维”模式集成到模型中,为用户提供灵活性,选择模型是否进行推理,并通过指定思维过程的token预算来控制思考的深度 。

强到弱蒸馏 (Strong-to-Weak Distillation):

强到弱蒸馏旨在优化轻量级模型的后训练过程 。通过利用大型模型的知识,显著降低了构建小型模型所需的计算成本和开发工作 。

如图1所示,Qwen3系列的旗舰模型遵循复杂的四阶段训练过程 。前两个阶段侧重于发展模型的“思维”能力 。后两个阶段旨在将强大的“非思维”功能整合到模型中 。

初步实验表明,将教师模型的输出logit直接蒸馏到轻量级学生模型中,可以有效增强其性能,同时保持对其推理过程的细粒度控制 。这种方法避免了为每个小型模型单独执行详尽的四阶段训练过程 。它带来了更好的即时性能(通过更高的Pass@1分数体现),也提高了模型的探索能力(通过改进的Pass@64结果反映) 。此外,它以更高的训练效率实现了这些提升,所需的GPU小时仅为四阶段训练方法的1/10 。

在第一阶段,我们使用多样的的长思维链数据对模型进行了微调,涵盖了数学、代码、逻辑推理和 STEM 问题等多种任务和领域。这一过程旨在为模型配备基本的推理能力。后训练始于策划一个涵盖数学、代码、逻辑推理和通用STEM问题等广泛类别的综合数据集 。数据集中的每个问题都配有经过验证的参考答案或基于代码的测试用例 。该数据集作为长链式思维(long-CoT)训练“冷启动”阶段的基础 。数据集构建涉及严格的两阶段过滤过程:查询过滤和响应过滤 。报告详细描述了过滤过程,包括使用Qwen2.5-72B-Instruct识别和移除不易验证的查询,排除无需CoT推理即可正确回答的查询,以及对生成的候选响应进行多项标准的严格过滤 。此阶段的目标是在模型中注入基础的推理模式,而不过度强调即时推理性能 。

第二阶段的重点是大规模强化学习,利用基于规则的奖励来增强模型的探索和钻研能力。推理RL阶段使用的查询-验证对必须满足四个标准:未在冷启动阶段使用、对冷启动模型可学习、尽可能具有挑战性、涵盖广泛的子领域 。共收集了3,995对查询-验证对,并使用GRPO更新模型参数 。报告提及使用大批次大小和每次查询多次rollout,以及利用离线训练提高样本效率,对训练过程有益 。通过控制模型的熵,平衡探索和利用,实现了训练和验证性能的持续改进 。例如,Qwen3-235B-A22B模型的AIME’24分数在170个RL训练步骤中从70.1提高到85.1 

在第三阶段,我们在一份包括长思维链数据和常用的指令微调数据的组合数据上对模型进行微调,将非思考模式整合到思考模型中。确保了推理和快速响应能力的无缝结合。思维模式融合阶段的目标是将“非思维”能力整合到之前开发的“思维”模型中 。这允许开发者管理和控制推理行为,同时降低部署独立模型处理思维和非思维任务的成本和复杂性 。为此,在推理RL模型上进行持续监督微调(SFT),并设计聊天模板来融合两种模式 。

SFT数据构建:SFT数据集结合了“思维”和“非思维”数据 。为了不损害第二阶段模型的性能,“思维”数据是使用第二阶段模型本身通过对第一阶段查询进行拒绝采样生成的 。“非思维”数据则精心策划,涵盖编码、数学、指令遵循、多语言任务、创意写作、问答和角色扮演等广泛任务 。报告还提及使用自动生成的清单评估“非思维”数据的响应质量,并增加低资源语言翻译任务的比例以增强性能 。

聊天模板设计:为了更好地集成两种模式并允许用户动态切换模型的思维过程,Qwen3设计了聊天模板 。通过在用户查询或系统消息中引入/think/no think标志,模型可以根据用户的输入选择适当的思维模式 。即使在非思维模式样本中,也保留了空的思维块,以确保模型内部格式的一致性 。默认情况下,模型在思维模式下运行,因此也包含一些用户查询不含/think标志的思维模式训练样本 。对于更复杂的多轮对话,随机插入多个/think/no think标志,模型响应遵循最后遇到的标志 。

思维预算:思维模式融合的另一个优势是,一旦模型学会以非思维和思维模式响应,它自然会发展出处理中间情况的能力—基于不完整的思考生成响应 。这为实现模型思维过程的预算控制奠定了基础 。当模型的思考长度达到用户定义的阈值时,会手动停止思考过程并插入停止思考指令,然后模型根据其累积的推理生成最终响应 。报告指出,这种能力并非显式训练所得,而是思维模式融合应用自然产生的结果 。

最后,在第四阶段,我们在包括指令遵循、格式遵循和 Agent 能力等在内的 20 多个通用领域的任务上应用了强化学习,以进一步增强模型的通用能力并纠正不良行为。

通用RL阶段旨在广泛增强模型在各种场景下的能力和稳定性 。为此,建立了覆盖20多个不同任务的复杂奖励系统,每个任务都有定制的评分标准 。这些任务专门针对以下核心能力的提升:指令遵循、格式遵循、偏好对齐、Agent能力和专业场景下的能力(如RAG任务) 。

报告提及使用了三种不同类型的奖励来提供反馈:基于规则的奖励(用于推理RL阶段和通用任务,如指令遵循和格式遵循)、基于模型的奖励(带参考答案,允许更灵活地处理多样化任务)、基于模型的奖励(不带参考答案,利用人类偏好数据训练奖励模型,处理更广泛的查询并增强模型的互动性和帮助性)。

强到弱蒸馏 (Strong-to-Weak Distillation):

强到弱蒸馏流程专门为优化轻量级模型而设计,包括5个密集模型(Qwen3-0.6B、1.7B、4B、8B和14B)和1个MoE模型(Qwen3-30B-A3B)。这种方法在增强模型性能的同时,有效赋予了强大的模式切换能力 。蒸馏过程分为两个主要阶段:

  1. 离线蒸馏 (Off-policy Distillation): 在初始阶段,结合教师模型在/think/no think模式下生成的输出进行响应蒸馏 。这有助于轻量级学生模型发展基本的推理技能和在不同思维模式之间切换的能力 。
  2. 在线蒸馏 (On-policy Distillation): 在此阶段,学生模型生成在线序列进行微调 。具体来说,采样提示,学生模型以/think/no think模式生成响应 。然后通过将学生的logit与教师模型(Qwen3-32B或Qwen3-235B-A22B)的logit对齐,最小化KL散度来微调学生模型 。

通过评估Qwen3-32B模型在不同训练阶段的性能,报告得出结论:第三阶段将非思维模式整合到模型中,模型开始具备模式切换的初步能力 。第三阶段还增强了思维模式下的通用和指令遵循能力 。第四阶段进一步加强了模型在思维和非思维模式下的通用、指令遵循和Agent能力,确保了准确的模式切换 。

然而,对于知识、STEM、数学和编码等任务,思维模式融合和通用RL并未带来显著改进,甚至在一些挑战性任务上,思维模式下的性能有所下降 。报告推测这种性能下降是由于模型在更广泛的通用任务上进行训练,可能会损害其在处理复杂问题时的专业能力,并表示在Qwen3开发过程中接受了这种性能权衡以增强模型的整体多功能性 。

高级用法

我们提供了一种软切换机制,允许用户在 enable_thinking=True 时动态控制模型的行为。具体来说,您可以在用户提示或系统消息中添加 /think 和 /no_think 来逐轮切换模型的思考模式。在多轮对话中,模型会遵循最近的指令。

未来发展:

Qwen3 代表了我们在通往通用人工智能(AGI)和超级人工智能(ASI)旅程中的一个重要里程碑。通过扩大预训练和强化学习的规模,我们实现了更高层次的智能。我们无缝集成了思考模式与非思考模式,为用户提供了灵活控制思考预算的能力。此外,我们还扩展了对多种语言的支持,帮助全球更多用户。

展望未来,我们计划从多个维度提升我们的模型。这包括优化模型架构和训练方法,以实现几个关键目标:扩展数据规模、增加模型大小、延长上下文长度、拓宽模态范围,并利用环境反馈推进强化学习以进行长周期推理。我们认为,我们正从专注于训练模型的时代过渡到以训练 Agent 为中心的时代。

MCP-大模型上下文协议

MCP:大模型接入万物的通用协议,智能体时代的HTTP!

一、什么是MCP(Model Context Protocol)

定义

MCP(Model Context Protocol,模型上下文协议) ,2024年11月底,由 Anthropic 推出的一种开放标准,旨在统一大模型与外部数据源和工具之间的通信协议。MCP 的主要目的在于解决当前 AI 模型因数据孤岛限制而无法充分发挥潜力的难题,MCP 使得 AI 应用能够安全地访问和操作本地及远程数据,为 AI 应用提供了连接万物的接口。

Function Calling是AI模型调用函数的机制,MCP是一个标准协议,使大模型与API无缝交互,而AI Agent是一个自主运行的智能系统,利用Function Calling和MCP来分析和执行任务,实现特定目标。

MCP 的价值

举个栗子,在过去,为了让大模型等 AI 应用使用数据,要么复制粘贴,要么上传知识库,非常局限。

即使是最强大模型也会受到数据隔离的限制,形成信息孤岛,要做出更强的大模型,每个新数据源都需要自己重新定制实现,使真正互联的系统难以扩展,存在很多的局限性。

现在,MCP 可以直接在 AI 与数据(包括本地数据和互联网数据)之间架起一座桥梁,通过 MCP 服务器和 MCP 客户端,大家只要都遵循这套协议,就能实现“万物互联”。

有了MCP,可以和数据和文件系统、开发工具、Web 和浏览器自动化、生产力和通信、各种社区生态能力全部集成,实现强大的协作工作能力,它的价值远不可估量。

MCP 与 Function Calling 的区别

  • MCP(Model Context Protocol),模型上下文协议
  • Function Calling,函数调用

这两种技术都旨在增强 AI 模型与外部数据的交互能力,但 MCP 不止可以增强 AI 模型,还可以是其他的应用系统。

Function Calling

数据安全性

这样一个理想的“万物互联”生态系统看着很让人着迷。

但是大家是不是担心通过 MCP Server 暴露出来的数据会泄露或被非法访问,这个头疼的问题 MCP 也考虑到了。

MCP 通过标准化的数据访问接口,大大减少了直接接触敏感数据的环节,降低了数据泄露的风险。

还有,MCP 内置了安全机制,确保只有经过验证的请求才能访问特定资源,相当于在数据安全又加上了一道防线。同时,MCP协议还支持多种加密算法,以确保数据在传输过程中的安全性。

例如,MCP 服务器自己控制资源,不需要将 API 密钥等敏感信息提供给 LLM 提供商。这样一来,即使 LLM 提供商受到攻击,攻击者也无法获取到这些敏感信息。

不过,MCP 这套协议/标准,需要大家一起来共建,这个生态才会繁荣,现在,只是测试阶段,一切才刚刚开始,当然,还会涌现出更多的问题。

工作原理

MCP 协议采用了一种独特的架构设计,它将 LLM 与资源之间的通信划分为三个主要部分:客户端、服务器和资源。

客户端负责发送请求给 MCP 服务器,服务器则将这些请求转发给相应的资源。这种分层的设计使得 MCP 协议能够更好地控制访问权限,确保只有经过授权的用户才能访问特定的资源。

以下是 MCP 的基本工作流程:

  • 初始化连接:客户端向服务器发送连接请求,建立通信通道。
  • 发送请求:客户端根据需求构建请求消息,并发送给服务器。
  • 处理请求:服务器接收到请求后,解析请求内容,执行相应的操作(如查询数据库、读取文件等)。
  • 返回结果:服务器将处理结果封装成响应消息,发送回客户端。
  • 断开连接:任务完成后,客户端可以主动关闭连接或等待服务器超时关闭。

MCP 核心架构

MCP 遵循客户端-服务器架构(client-server),其中包含以下几个核心概念:

  • MCP 主机(MCP Hosts):发起请求的 LLM 应用程序(例如 Claude Desktop、IDE 或 AI 工具)。
  • MCP 客户端(MCP Clients):在主机程序内部,与 MCP server 保持 1:1 的连接。
  • MCP 服务器(MCP Servers):为 MCP client 提供上下文、工具和 prompt 信息。
  • 本地资源(Local Resources):本地计算机中可供 MCP server 安全访问的资源(例如文件、数据库)。
  • 远程资源(Remote Resources):MCP server 可以连接到的远程资源(例如通过 API)。

共分为了下面五个部分:

  • MCP Hosts: Hosts 是指 LLM 启动连接的应用程序,像 Cursor, Claude Desktop、Cline 这样的应用程序。
  • MCP Clients: 客户端是用来在 Hosts 应用程序内维护与 Server 之间 1:1 连接。
  • MCP Servers: 通过标准化的协议,为 Client 端提供上下文、工具和提示。
  • Local Data Sources: 本地的文件、数据库和 API。
  • Remote Services: 外部的文件、数据库和 API。

整个 MCP 协议核心的在于 Server,因为 Host 和 Client 相信熟悉计算机网络的都不会陌生,非常好理解,但是 Server 如何理解呢?

看看 Cursor 的 AI Agent 发展过程,我们会发现整个 AI 自动化的过程发展会是从 Chat 到 Composer 再进化到完整的 AI Agent。

AI Chat 只是提供建议,如何将 AI 的 response 转化为行为和最终的结果,全部依靠人类,例如手动复制粘贴,或者进行某些修改。

AI Composer 是可以自动修改代码,但是需要人类参与和确认,并且无法做到除了修改代码之外的其它操作。

AI Agent 是一个完全的自动化程序,未来完全可以做到自动读取 Figma 的图片,自动生产代码,自动读取日志,自动调试代码,自动 push 代码到 GitHub。

而 MCP Server 就是为了实现 AI Agent 的自动化而存在的,它是一个中间层,告诉 AI Agent 目前存在哪些服务,哪些 API,哪些数据源,AI Agent 可以根据 Server 提供的信息来决定是否调用某个服务,然后通过 Function Calling 来执行函数。

MCP Client

MCP client 充当 LLM 和 MCP server 之间的桥梁,MCP client 的工作流程如下:

  • MCP client 首先从 MCP server 获取可用的工具列表。
  • 将用户的查询连同工具描述通过 function calling 一起发送给 LLM。
  • LLM 决定是否需要使用工具以及使用哪些工具。
  • 如果需要使用工具,MCP client 会通过 MCP server 执行相应的工具调用。
  • 工具调用的结果会被发送回 LLM。
  • LLM 基于所有信息生成自然语言响应。
  • 最后将响应展示给用户。

Claude Desktop 和Cursor都支持了MCP Server接入能力,它们就是作为 MCP client来连接某个MCP Server感知和实现调用。

MCP Server

MCP server 是 MCP 架构中的关键组件,它可以提供 3 种主要类型的功能:

  • 资源(Resources):类似文件的数据,可以被客户端读取,如 API 响应或文件内容。
  • 工具(Tools):可以被 LLM 调用的函数(需要用户批准)。
  • 提示(Prompts):预先编写的模板,帮助用户完成特定任务。

这些功能使 MCP server 能够为 AI 应用提供丰富的上下文信息和操作能力,从而增强 LLM 的实用性和灵活性。

你可以在 MCP Servers Repository 和 Awesome MCP Servers 这两个 repo 中找到许多由社区实现的 MCP server。使用 TypeScript 编写的 MCP server 可以通过 npx 命令来运行,使用 Python 编写的 MCP server 可以通过 uvx 命令来运行。

通信机制

MCP 协议支持两种主要的通信机制:基于标准输入输出的本地通信和基于SSEServer-Sent Events)的远程通信。

这两种机制都使用 JSON-RPC 2.0 格式进行消息传输,确保了通信的标准化和可扩展性。

  • 本地通信通过 stdio 传输数据,适用于在同一台机器上运行的客户端和服务器之间的通信。
  • 远程通信利用 SSE 与 HTTP 结合,实现跨网络的实时数据传输,适用于需要访问远程资源或分布式部署的场景。

二、MCP的功能与应用:

如何使用 MCP

如果你还没有尝试过如何使用 MCP 的话,我们可以考虑用 Cursor(本人只尝试过 Cursor),Claude Desktop 或者 Cline 来体验一下。

当然,我们并不需要自己开发 MCP Servers,MCP 的好处就是通用、标准,所以开发者并不需要重复造轮子(但是学习可以重复造轮子)。

首先推荐的是官方组织的一些 Server:官方的 MCP Server 列表

目前社区的 MCP Server 还是比较混乱,有很多缺少教程和文档,很多的代码功能也有问题,我们可以自行尝试一下 Cursor Directory 的一些例子,具体的配置和实战笔者就不细讲了,大家可以参考官方文档。

MCP的功能

MCP通过引入多样化的MCP Server能力,显著增强了AI工具的功能,例如我们常用的Cursor和Claude。以下是一些官方参考服务器,展示了MCP的核心功能和SDK的应用:

数据与文件系统:

文件系统:提供安全文件操作,带可配置的访问控制。

PostgreSQL:提供只读数据库访问,具备架构检查功能。

SQLite:支持数据库交互和商业智能功能。

Google Drive:实现Google Drive的文件访问和搜索功能。

开发工具:

Git:工具用于读取、搜索和操作Git仓库。

GitHub:集成仓库管理、文件操作和GitHub API。

GitLab:支持项目管理的GitLab API集成。

Sentry:从http://Sentry.io获取并分析问题。

网络与浏览器自动化:

Brave Search:利用Brave的搜索API进行网络和本地搜索。

Fetch:为LLM优化的网络内容获取和转换。

Puppeteer:提供浏览器自动化和网页抓取功能。

生产力和通信:

Slack:支持频道管理和消息功能。

Google Maps:提供位置服务、路线和地点详情。

Memory:基于知识图谱的持久记忆系统。

AI与专业工具:

EverArt:使用多种模型进行AI图像生成。

Sequential Thinking:通过思维序列进行动态问题解决。

AWS KB Retrieval:使用Bedrock Agent Runtime从AWS知识库检索。

官方集成工具:

这些MCP服务器由公司维护,用于其平台:

Axiom:使用自然语言查询和分析日志、跟踪和事件数据。

Browserbase:云端自动化浏览器交互。

Cloudflare:在Cloudflare开发者平台上部署和管理资源。

E2B:在安全的云沙箱中执行代码。

Neon:与Neon无服务器Postgres平台交互。

Obsidian Markdown Notes:读取和搜索Obsidian知识库中的Markdown笔记。

Qdrant:使用Qdrant向量搜索引擎实现语义记忆。

Raygun:访问崩溃报告和监控数据。

Search1API:统一的API用于搜索、爬虫和网站地图。

Tinybird:与Tinybird无服务器ClickHouse平台交互。

集成工具:

Docker:管理容器、镜像、卷和网络。

Kubernetes:管理pod、部署和服务。

Linear:项目管理和问题跟踪。

Snowflake:与Snowflake数据库交互。

Spotify:控制Spotify播放和管理播放列表。

Todoist:任务管理集成。

三、怎么使用和开发MCP Server

使用

目前支持的部分工具列表(更多见这里):

LLM分词器-SentencePiece/tiktoken/subword-nmt

SentencePiece 简介

https://github.com/google/sentencepiece

SentencePiece 是一种无监督的文本 tokenizer 和 detokenizer,主要用于基于神经网络的文本生成系统,其中,词汇量在神经网络模型训练之前就已经预先确定了。 SentencePiece 实现了subword单元(例如,字节对编码 (BPE))和 unigram 语言模型),并可以直接从原始句子训练字词模型(subword model)。 这使得我们可以制作一个不依赖于特定语言的预处理和后处理的纯粹的端到端系统。

SentencePiece 在大模型领域主要用于文本的 分词 和 编码

分词 是将文本分割成一个个独立的词语或符号。传统的中文分词方法,例如 BMM 分词、HMM 分词,都是基于规则的,需要人工制定分词规则。而 SentencePiece 则是基于 无监督学习 的,它可以自动学习文本的语义和结构,并根据学习结果进行分词。

编码 是将分词后的词语或符号转换为数字形式,以便计算机能够处理。SentencePiece 使用了一种称为 字节对编码 的方法,它可以将每个词语或符号编码成一个或多个字节。字节对编码的优点是能够有效地利用空间,并且可以将词语或符号之间的关系编码到编码中。

SentencePiece 在大模型领域具有以下优势:

  • 分词效果好,能够准确地识别词语和符号的边界。
  • 编码效率高,能够节省空间。
  • 能够将词语或符号之间的关系编码到编码中,有利于模型学习。

因此,SentencePiece 已经被广泛应用于各大模型,例如 Google 的 BERTGPT-3,以及阿里巴巴的 M6 等。

简单来说,SentencePiece 就是大模型领域的一个 分词和编码工具。它可以帮助大模型更好地理解和处理文本。

自监督训练原理

SentencePiece 的自监督训练模型原理是通过 无监督学习 的方式,学习文本的语义和结构,并根据学习结果进行分词和编码。

具体来说,SentencePiece 的训练过程可以分为以下几个步骤:

  1. 数据准备:首先需要准备一个文本语料库,语料库中的文本可以是任何类型,例如新闻文章、书籍、代码等。
  2. 模型训练:使用无监督学习算法训练 SentencePiece 模型,模型的输入是文本语料库,输出是分词后的文本。
  3. 模型评估:使用评估指标评估模型的性能,例如分词准确率、召回率等。

SentencePiece 使用的无监督学习算法是一种称为 Masked Language Modeling (MLM) 的算法。MLM 的基本思想是:将文本中的部分词语或符号进行遮蔽,然后让模型预测被遮蔽的词语或符号。通过这种方式,模型可以学习文本的语义和结构。

在 SentencePiece 中,MLM 的具体实现如下:

  1. 随机选择文本中的部分词语或符号进行遮蔽。
  2. 将被遮蔽的词语或符号替换为一个特殊符号,例如 [MASK]
  3. 将处理后的文本输入模型,让模型预测被遮蔽的词语或符号。

通过这种方式,模型可以学习到被遮蔽词语或符号与周围词语或符号之间的关系,从而提高分词和编码的准确性。

SentencePiece 的自监督训练模型具有以下优势:

  • 不需要人工制定分词规则,能够自动学习文本的语义和结构。
  • 分词效果好,能够准确地识别词语和符号的边界。
  • 编码效率高,能够节省空间。

SentencePiece 的自监督训练模型已经被广泛应用于各大模型,例如 Google 的 BERT、GPT-3,以及阿里巴巴的 M6 等。

SentencePiece中包含有四个部分:

  • Normalizer: 规一化操作,类似Unicode把字符转为统一格式。
  • Trainer: 从规一化后的语料中学习subword的切分模型。
  • Encoder: 对应预处理时的tokenization操作,把句子转为对应的subword或id。
  • Decoder: 对应后处理时的detokenization操作,把subword或id还原为原有句子

和 SentencePiece 类似的工具

  • Jieba:Jieba 是一个中文分词工具,它使用了一种称为 最大似然法 的方法进行分词。Jieba 的分词效果较好,并且速度较快。
  • Hmmseg:Hmmseg 是另一个中文分词工具,它使用了一种称为 隐马尔可夫模型 的方法进行分词。Hmmseg 的分词效果较好,并且可以支持多种语言。
  • Stanford CoreNLP:Stanford CoreNLP 是一个自然语言处理工具包,它包含了分词、词性标注、句法分析等功能。Stanford CoreNLP 的分词效果较好,并且可以支持多种语言。

使用建议:

  • 如果需要对中文文本进行分词,并且对分词效果要求较高,可以选择 SentencePiece、Jieba 或 Hmmseg。
  • 如果需要对多种语言文本进行分词,可以选择 Stanford CoreNLP。
  • 如果需要对文本进行分词和编码,并且对速度要求较高,可以选择 Jieba。
  • 如果需要对文本进行分词和编码,并且对分词效果要求较高,可以选择 SentencePiece 或 Hmmseg。

TiToken 简介

Tiktoken 的功能与 SentencePiece 类似,都是用于文本的 分词 和 编码

Tiktoken 是一个基于 BPE 算法的 快速分词器,它专门针对 GPT-4 和 ChatGPT 等大模型进行了优化。Tiktoken 的主要特点如下:

  • 速度快:Tiktoken 的分词速度比 SentencePiece 快很多,可以满足大模型训练和推理的需求。
  • 效果好:Tiktoken 的分词效果与 SentencePiece 相当,能够准确地识别词语和符号的边界。
  • 易用性:Tiktoken 提供了简单的 API 接口,方便使用。

Tiktoken 与 SentencePiece 的主要区别 如下:

  • 分词算法:Tiktoken 使用 BPE 算法进行分词,而 SentencePiece 使用的是无监督学习算法
  • 速度:Tiktoken 的分词速度比 SentencePiece 快很多。
  • 模型:Tiktoken 专门针对 GPT-4 和 ChatGPT 等大模型进行了优化,而 SentencePiece 则是通用的。

具体使用建议:

  • 如果需要对文本进行快速分词,并且对分词效果要求较高,可以选择 Tiktoken。
  • 如果需要对文本进行分词和编码,并且需要支持多种语言,可以选择 SentencePiece。
  • 如果需要对 GPT-4 或 ChatGPT 等大模型进行训练或推理,可以选择 Tiktoken。

BPE分词算法-原理和代码实现

三种主流的Subword算法,它们分别是:Byte Pair Encoding (BPE)、WordPiece和Unigram Language Model。

执行分词的算法模型称为分词器(Tokenizer),划分好的一个个词称为Token(中文叫词元,为啥不直接叫Word?接着往后看),这个过程称为Tokenization

我们将一个个的token(可以理解为小片段)表示向量,我们分词的目的就是尽可能的让这些向量蕴含更多有用的信息,然后把这些向量输入到算法模型中。

由于一篇文本的词往往太多了,为了方便算法模型训练,我们会选取出频率(也可能是其它的权重)最高的若干个词组成一个词表(Vocabulary)

‼️古典分词方法的缺点

一个句子,使用不同的规则,将有许多种不同的分词结果。古典分词方法的缺点非常明显:

  • 对于未在词表中出现的词(Out Of Vocabulary, OOV),模型将无法处理(未知符号标记为 [UNK])。
  • 词表中的低频词/稀疏词在模型训无法得到训练(因为词表大小有限,太大的话会影响效率)。
  • ⭐️很多语言难以用空格进行分词,例如英语单词的多形态,”look”衍生出的”looks”, “looking”, “looked”,其实都是一个意思,但是在词表中却被当作不同的词处理,模型也无法通过old, older, oldest之间的关系学到smart, smarter, smartest之间的关系。这一方面增加了训练冗余,另一方面也造成了大词汇量问题。

Character embedding,是一种更为极端的分词方法,直接把一个词分成一个一个的字母和特殊符号。虽然能解决OOV问题,也避免了大词汇量问题,但缺点也太明显了,粒度太细,训练花费的成本太高,但这种思想或许我们后面会用到。

BERT算法的横空出世,NLP中的很多领域都被颠覆性的改变了,BERT也称为了一个非常主流的NLP算法。由于BERT的特性,要求分词方法也必须作出改变。这就对应提出了Subword算法(或成为WordPiece),该算法已经成为一种标配。

基于子词的分词方法(Subword Tokenization)

可见不论是传统分词算法的局限性,还是BERT的横空出世,都要求我们提出新的分词算法,下面就轮到本文的主角登场:基于子词的分词方法(Subword Tokenization),简称为Subword算法,意思就是把一个词切成更小的一块一块的子词。如果我们能使用将一个token分成多个subtokens,上面的问题就能很好的解决

这种方法的目的是通过一个有限的词表来解决所有单词的分词问题,同时尽可能将结果中token的数目降到最低。例如,可以用更小的词片段来组成更大的词,例如:

unfortunately” = “un” + “for” + “tun” + “ate” + “ly”。

可以看到,有点类似英语中的词根词缀拼词法,其中的这些小片段又可以用来构造其他词。可见这样做,既可以降低词表的大小同时对相近词也能更好地处理

Subword与传统分词方法的比较

  • 传统词表示方法无法很好的处理未知或罕见的词汇(OOV问题)。
  • 传统词tokenization方法不利于模型学习词缀之间的关系,例如模型学到的“old”, “older”, and “oldest”之间的关系无法泛化到“smart”, “smarter”, and “smartest”。
  • Character embedding作为OOV的解决方法粒度太细。
  • Subword粒度在词与字符之间,能够较好的平衡OOV问题。

目前有三种主流的Subword算法,它们分别是:Byte Pair Encoding (BPE)、WordPiece和Unigram Language Model。

字节对编码(BPE, Byte Pair Encoding)

字节对编码(BPE, Byte Pair Encoder),又称digram coding双字母组合编码,是一种数据压缩算法,用来在固定大小的词表中实现可变⻓度的子词。该算法简单有效,因而目前它是最流行的方法。

BPE首先将词分成单个字符,然后依次用另一个字符替换频率最高一对字符,直到循环次数结束。

接下来详细介绍BPE在分词中的算法过程:

算法过程

  1. 准备语料库,确定期望的subword词表大小等参数
  2. 通常在每个单词末尾添加后缀</w>,统计每个单词出现的频率,例如,low的频率为5,那么我们将其改写为"l o w </ w>”:5 注:停止符</w>的意义在于标明subword是词后缀。举例来说:st不加</w>可以出现在词首,如st ar;加了</w>表明该子词位于词尾,如we st</w>,二者意义截然不同
  3. 将语料库中所有单词拆分为单个字符,用所有单个字符建立最初的词典,并统计每个字符的频率,本阶段的subword的粒度是字符
  4. 挑出频次最高的符号对,比如说 th 组成的 th,将新字符加入词表,然后将语料中所有该字符对融合(merge),即所有 th 都变为 th。 注:新字符依然可以参与后续的merge,有点类似哈夫曼树,BPE实际上就是一种贪心算法
  5. 重复遍历 2和 3 操作,直到词表中单词数达到设定量下一个最高频数为1,如果已经打到设定量,其余的词汇直接丢弃

注:看似我们要维护两张表,一个词表,一个字符表,实际上只有一张,词表只是为了我们方便理解。</w> 是为了明确 subword 是否是词尾,它在训练、merge 过程中和普通字符一样对待,在最终输出时作为重建原词的标记使用。

我们举一个完整的例子,来直观地看一下这个过程:

1、获取语料库,这样一段话为例:FloydHub is the fastest way to build, train and deploy deep learning models. Build deep learning models in the cloud. Train deep learning models.

2、拆分,加后缀,统计词频:

3、建立词表,统计字符频率(顺便排个序):

4、以第一次迭代为例,将字符频率最高的de替换为de,后面依次迭代:

5、更新词表 继续迭代直到达到预设的subwords词表大小或下一个最高频的字节对出现频率为1。

如果将词表大小设置为10,最终的结果为:

d e
r n
rn i
rni n
rnin g</w>
o de
ode l
m odel
l o
l e

这样我们就得到了更加合适的词表,这个词表可能会出现一些不是单词的组合,但是其本身有意义的一种形式

BPE的优点

上面例子中的语料库很小,知识为了方便我们理解BPE的过程,但实际中语料库往往非常非常大,无法给每个词(token)都放在词表中。BPE的优点就在于,可以很有效地平衡词典大小和编码步骤数(将语料编码所需要的token数量)。

随着合并的次数增加,词表大小通常先增加后减小。迭代次数太小,大部分还是字母,没什么意义;迭代次数多,又重新变回了原来那几个词。所以词表大小要取一个中间值。

BPE的缺点

  • 对于同一个句子, 例如Hello world,如图所示,可能会有不同的Subword序列。不同的Subword序列会产生完全不同的id序列表示,这种歧义可能在解码阶段无法解决。在翻译任务中,不同的id序列可能翻译出不同的句子,这显然是错误的。
  • 在训练任务中,如果能对不同的Subword进行训练的话,将增加模型的健壮性,能够容忍更多的噪声,而BPE的贪心算法无法对随机分布进行学习。

个人理解:我感觉缺点直接可以忽略

BPE的适用范围

BPE一般适用在欧美语言拉丁语系中,因为欧美语言大多是字符形式,涉及前缀、后缀的单词比较多。而中文的汉字一般不用BPE进行编码,因为中文是字无法进行拆分。对中文的处理通常只有分词和分字两种。理论上分词效果更好,更好的区别语义。分字效率高、简洁,因为常用的字不过3000字,词表更加简短。

BPE的实现

实现代码如下:

import re, collections

def get_vocab(filename):
    vocab = collections.defaultdict(int)
    with open(filename, 'r', encoding='utf-8') as fhand:
        for line in fhand:
            words = line.strip().split()
            for word in words:
                vocab[' '.join(list(word)) + ' </w>'] += 1
    return vocab

def get_stats(vocab):
    pairs = collections.defaultdict(int)
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols)-1):
            pairs[symbols[i],symbols[i+1]] += freq
    return pairs

def merge_vocab(pair, v_in):
    v_out = {}
    bigram = re.escape(' '.join(pair))
    p = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
    for word in v_in:
        w_out = p.sub(''.join(pair), word)
        v_out[w_out] = v_in[word]
    return v_out

def get_tokens(vocab):
    tokens = collections.defaultdict(int)
    for word, freq in vocab.items():
        word_tokens = word.split()
        for token in word_tokens:
            tokens[token] += freq
    return tokens

跑一个例子试一下,这里已经对原句子进行了预处理:

vocab = {'l o w </w>': 5, 'l o w e r </w>': 2, 'n e w e s t </w>': 6, 'w i d e s t </w>': 3}
print('==========')
print('Tokens Before BPE')
tokens = get_tokens(vocab)
print('Tokens: {}'.format(tokens))
print('Number of tokens: {}'.format(len(tokens)))
print('==========')

num_merges = 5
for i in range(num_merges):
    pairs = get_stats(vocab)
    if not pairs:
        break
    best = max(pairs, key=pairs.get)
    vocab = merge_vocab(best, vocab)
    print('Iter: {}'.format(i))
    print('Best pair: {}'.format(best))
    tokens = get_tokens(vocab)
    print('Tokens: {}'.format(tokens))
    print('Number of tokens: {}'.format(len(token

结果:

==========
Tokens Before BPE
Tokens: defaultdict(<class 'int'>, {'l': 7, 'o': 7, 'w': 16, '</w>': 16, 'e': 17, 'r': 2, 'n': 6, 's': 9, 't': 9, 'i': 3, 'd': 3})
Number of tokens: 11
==========
Iter: 0
Best pair: ('e', 's')
Tokens: defaultdict(<class 'int'>, {'l': 7, 'o': 7, 'w': 16, '</w>': 16, 'e': 8, 'r': 2, 'n': 6, 'es': 9, 't': 9, 'i': 3, 'd': 3})
Number of tokens: 11
==========
Iter: 1
Best pair: ('es', 't')
Tokens: defaultdict(<class 'int'>, {'l': 7, 'o': 7, 'w': 16, '</w>': 16, 'e': 8, 'r': 2, 'n': 6, 'est': 9, 'i': 3, 'd': 3})
Number of tokens: 10
==========
Iter: 2
Best pair: ('est', '</w>')
Tokens: defaultdict(<class 'int'>, {'l': 7, 'o': 7, 'w': 16, '</w>': 7, 'e': 8, 'r': 2, 'n': 6, 'est</w>': 9, 'i': 3, 'd': 3})
Number of tokens: 10
==========
Iter: 3
Best pair: ('l', 'o')
Tokens: defaultdict(<class 'int'>, {'lo': 7, 'w': 16, '</w>': 7, 'e': 8, 'r': 2, 'n': 6, 'est</w>': 9, 'i': 3, 'd': 3})
Number of tokens: 9
==========
Iter: 4
Best pair: ('lo', 'w')
Tokens: defaultdict(<class 'int'>, {'low': 7, '</w>': 7, 'e': 8, 'r': 2, 'n': 6, 'w': 9, 'est</w>': 9, 'i': 3, 'd': 3})
Number of tokens: 9
==========

编码与解码

上面的过程称为编码。解码过程比较简单,如果相邻子词间没有中止符,则将两子词直接拼接,否则两子词之间添加分隔符。 如果仍然有子字符串没被替换但所有token都已迭代完毕,则将剩余的子词替换为特殊token,如<unk>。例如:

# 编码序列
["the</w>", "high", "est</w>", "moun", "tain</w>"]

# 解码序列
"the</w> highest</w> mountain</w>"

如何调包使用BPE

BPE 可以直接用最经典的 subword-nmt 包,不需要自己实现

Dolphin -支持东方40语种+中国22方言的新SOTA语音大模型

在当今数字化时代,语音识别技术已成为人机交互的关键桥梁,广泛应用于智能客服、语音助手、会议转录等众多领域。然而,对于东方语言的识别如越南语、缅甸语等,现有模型往往表现不佳,难以满足用户的需求。为解决这一难题,海天瑞声携手清华大学电子工程系语音与音频技术实验室,共同推出了Dolphin —— 一款专为东方语言设计的语音大模型

Dolphin 采用的多任务格式,其主要沿用了 OpenAI Whisper的
格式。Dolphin 专注于自动语音识别 (ASR),不支持翻译任务。此外,Dolphin 引入了特定区域的标记,从而支持方言。

Dolphin 是由 Dataocean AI 与清华大学合作开发的多语言、多任务 ASR 模型。它支持东亚、南亚、东南亚和中东地区的 40 种东方语言,同时还支持 22 种中国方言。该模型基于超过 21 万小时的数据进行训练,其中包括 DataoceanAI 的专有数据集和开源数据集。该模型可以执行语音识别、语音活动检测 (VAD)、语音分割和语言识别 (LID)

二、创新技术架构 

  • 模型结构    

Dolphin网络结构基于CTC-Attention架构,E-Branchformer编码器和Transformer解码器,并引入了4倍下采样层,以实现高效的大规模多语言语音识别模型的训练。CTC-Attention架构结合了CTC的序列建模能力和注意力机制的上下文捕捉能力,能够有效提升模型的识别准确性和效率。E-Branchformer编码器采用并行分支结构,能够更有效地捕捉输入语音信号的局部和全局依赖关系,为模型提供了更丰富的特征表示。解码器部分则采用了在序列到序列任务中表现出色的Transformer,能够生成高质量的文本输出。为了进一步提高训练效率和性能,我们在模型中引入了4倍下采样层。这一层可以减少输入特征的序列长度,从而加速计算过程,同时保留关键的语音信息,确保模型的识别效果不受影响。

  • 多任务格式

Dolphin 借鉴了 Whisper 和 OWSM 的创新设计方法,但专注于ASR 进行了若干关键修改。Dolphin 不支持翻译任务,并且去掉了previous text及其相关标记的使用,这简化了输入格式并减少了潜在的复杂性Dolphin引入了两级语种标签系统,以便更好地处理语言和地区的多样性。第一个标签指定语种(例如: <zh> 、 <ja>),第二个标签指定地区(例如 <CN> 、 <JP>)。 比如:<ru><RU> 表示俄罗斯的俄语,而 <ru><BY> 表示白俄罗斯的俄语。这种分层方法使模型能够捕捉同一种语言内不同方言和口音之间的差异,以及同一地区内不同语言之间的相似性,从而提高了模型区分密切相关的方言的能力,并通过在语言和地区之间建立联系增强了其泛化能力。

三、强大的数据基础 

Dolphin的训练数据集整合了海天瑞声【Dataocean AI】的专有数据和多个开源数据集,总时长超过20万小时,涵盖40个东方语种。其中,海天瑞声数据集包含137,712小时的音频,覆盖38个东方语种。这些高质量、多样化的数据为模型的训练提供了坚实的基础,使其能够更好地适应不同语言和方言的语音特征。

清理后数据集中 40 种东方语言的数据时长分布(以对数刻度表示)。其中 36 种语言的数据时长超过 100 小时,16 种语言的数据时长超过 1000 小时。

数据处理:对于像 YODAS 这样包含人工注释和 ASR 生成的转录本的数据集,我们只使用人工注释的部分。因此,我们的大部分训练数据都是手动转录的,以确保更高的转录质量。这种数据质量,尤其是转录本的质量,是使模型即使在模型规模较小的情况下也能实现显著优于 Whisper 识别性能的关键因素。对于时间戳,采用与 Whisper 相同的句子级时间戳方法,其中时间戳标记标记每个句子的起始和结束。对于长音频录音(通常长达几分钟),会在数据预处理过程中将其分割成较小的片段,然后将它们合并为长音频序列。

训练优化:

在训练数据的初始版本中,我们直接使用了清理后的数据集。然而,一个主要问题是短音频样本的比例过高。大多数音频片段的时长约为 5 秒,导致跨多种语言的删除错误率过高。这个问题与大多数训练数据由短音频样本组成这一事实相符。

为了解决这个问题,尝试了一种替代方法,将清理后的音频数据连接成 25-30 秒的长片段。这显著降低了较高的删除错误率。虽然这种方法导致插入错误率略有增加,但整体识别性能有所提升,平均字词错误率 (WER) 降低了 9.01%。

四、卓越性能表现 

通过精心设计的架构和大规模的训练数据,Dolphin在多种语言上的词错误率(WER)显著低于现有开源模型。

例如,在海天瑞声数据集上,Dolphin 模型的平均WER为31.5%,small模型为24.5%,medium模型为22.2%;在CommonVoice数据集上,Dolphin 模型的平均WER为37.2%,small模型为27.4%,medium模型为25.0%。即使与Whisper large-v3模型相比,Dolphin在模型规模更小的情况下,性能也更为出色。以中文为例,Dolphin中模型的WER仅为9.2%,而Whisper large-v3模型为27.9%。 在KeSpeech (包含一个普通话子集和八个中国方言子集)测试集上,Dolphin模型表现出了卓越的效果.

五、技术挑战

内存占用问题

图 3: 数据加载策略优化。假设一个节点有 4 个 GPU,每个 GPU 分配一个对应的进程,称为 rank。优化前,每个 rank 加载数据集的完整副本,记为 {D0,D1,D2,D3}。优化后,每个 rank 仅分配其计算所需的数据集子集。

我们的训练集包含 1.6 亿条话语,在数据处理阶段遇到了内存不足 (OOM) 问题。我们对数据处理的 sampler、dataset、dataloader 模块进行了深入分析,发现大量的 utterances 导致了内存溢出。PyTorch 支持两种类型的数据集:map-style 和 iterable-style。ESPnet 使用的是 map-style。map-style 数据集将 utterance 的元数据(utterance id 与文本、音频的映射)加载到内存中,内存占用随着训练数据 utterances 的数量线性增长。为了提高数据加载速度,dataloader 内部会有多个 worker 进行数据预取,这进一步增加了物理机的内存占用,最终导致 OOM。

受 Zero-DP的启发,我们提出了图 3 中的数据分片策略。我们不再加载整个数据集副本,而是优化每个 Rank,使其仅加载数据集中必要的子集。这种方法显著减少了每个 Rank 的内存占用,从而降低了物理机上的整体内存消耗。此外,随着数据并行度的提高,单个节点的内存占用呈线性下降。

训练效率:

将短音频合并成长音频可以显著提高 GPU 的计算密度和利用率,从而显著提高训练效率。在我们的数据集中,音频时长呈现出明显的左偏分布,短音频(1-10 秒)占比较高,长音频(11-30 秒)占比较低。为了使音频时长分布更加均衡,我们将短音频合并,并将它们均匀地重新分配到 0-30 秒范围内以 5 秒为间隔的桶中。

在处理 21 万小时的大规模数据集时,使用 ffmpeg 将多个短音频物理合并成长音频会非常耗时。为此,我们采用了更高效的逻辑合并策略。具体来说,在数据准备阶段,我们使用字典来表示音频合并前后的映射关系,并在训练过程中动态地合并音频。

通过优化合并策略,小模型单次 epoch 训练时间从 64 小时大幅缩短至 28.6 小时,训练速度提升 123.78%,大大加速了模型迭代进程。

六、开源与社区贡献 

为促进语音识别技术的进一步发展,Dolphin的训练模型和推理源代码已公开发布。这一举措不仅为研究人员提供了宝贵的研究基础,也为开源社区注入了新的活力,鼓励更多创新与合作。通过共享技术成果,我们希望能够吸引更多的开发者和研究机构参与到东方语言语音识别的研究中来,共同推动技术的进步。 

 Dolphin,一个大规模多语言多任务自动语音识别 (ASR) 模型。Dolphin 构建于 Whisper 风格的架构之上,并基于 OWSM,集成了专有和公开可用的数据集。实验结果表明,Dolphin 在各种语言和模型规模上始终优于现有的 SOTA 模型,有效弥合了东西方语言之间的性能差距。值得一提的是,Dolphin 基础模型的性能甚至优于 Whisper large-v3 版本。通过开源 Dolphin 基础模型、小型模型以及推理代码,我们旨在为多语言语音处理的进一步发展做出贡献。

支持的语言列表:

Language code

Language CodeEnglish NameChinese Name
zhMandarin Chinese中文
jaJapanese日语
thThai泰语
ruRussian俄语
koKorean韩语
idIndonesian印度尼西亚语
viVietnamese越南语
ctYue Chinese粤语
hiHindi印地语
urUrdu乌尔都语
msMalay马来语
uzUzbek乌兹别克语
arArabic阿拉伯语
faPersian波斯语
bnBengali孟加拉语
taTamil泰米尔语
teTelugu泰卢固语
ugUighur维吾尔语
guGujarati古吉拉特语
myBurmese缅甸语
tlTagalog塔加洛语
kkKazakh哈萨克语
orOriya / Odia奥里亚语
neNepali尼泊尔语
mnMongolian蒙古语
kmKhmer高棉语
jvJavanese爪哇语
loLao老挝语
siSinhala僧伽罗语
filFilipino菲律宾语
psPushto普什图语
paPanjabi旁遮普语
kabKabyle卡拜尔语
baBashkir巴什基尔语
ksKashmiri克什米尔语
tgTajik塔吉克语
suSundanese巽他语
mrMarathi马拉地语
kyKirghiz吉尔吉斯语
azAzerbaijani阿塞拜疆语

Language Region Code

Language Region CodeEnglish NameChinese Name
zh-CNChinese (Mandarin)中文(普通话)
zh-TWChinese (Taiwan)中文(台湾)
zh-WUChinese (Wuyu)中文(吴语)
zh-SICHUANChinese (Sichuan)中文(四川话)
zh-SHANXIChinese (Shanxi)中文(山西话)
zh-ANHUIChinese (Anhui)中文(安徽话)
zh-TIANJINChinese (Tianjin)中文(天津话)
zh-NINGXIAChinese (Ningxia)中文(宁夏话)
zh-SHAANXIChinese (Shaanxi)中文(陕西话)
zh-HEBEIChinese (Hebei)中文(河北话)
zh-SHANDONGChinese (Shandong)中文(山东话)
zh-GUANGDONGChinese (Guangdong)中文(广东话)
zh-SHANGHAIChinese (Shanghai)中文(上海话)
zh-HUBEIChinese (Hubei)中文(湖北话)
zh-LIAONINGChinese (Liaoning)中文(辽宁话)
zh-GANSUChinese (Gansu)中文(甘肃话)
zh-FUJIANChinese (Fujian)中文(福建话)
zh-HUNANChinese (Hunan)中文(湖南话)
zh-HENANChinese (Henan)中文(河南话)
zh-YUNNANChinese (Yunnan)中文(云南话)
zh-MINNANChinese (Minnan)中文(闽南语)
zh-WENZHOUChinese (Wenzhou)中文(温州话)
ja-JPJapanese日语
th-THThai泰语
ru-RURussian俄语
ko-KRKorean韩语
id-IDIndonesian印度尼西亚语
vi-VNVietnamese越南语
ct-NULLYue (Unknown)粤语(未知)
ct-HKYue (Hongkong)粤语(香港)
ct-GZYue (Guangdong)粤语(广东)
hi-INHindi印地语
ur-INUrdu乌尔都语(印度)
ur-PKUrdu (Islamic Republic of Pakistan)乌尔都语
ms-MYMalay马来语
uz-UZUzbek乌兹别克语
ar-MAArabic (Morocco)阿拉伯语(摩洛哥)
ar-GLAArabic阿拉伯语
ar-SAArabic (Saudi Arabia)阿拉伯语(沙特)
ar-EGArabic (Egypt)阿拉伯语(埃及)
ar-KWArabic (Kuwait)阿拉伯语(科威特)
ar-LYArabic (Libya)阿拉伯语(利比亚)
ar-JOArabic (Jordan)阿拉伯语(约旦)
ar-AEArabic (U.A.E.)阿拉伯语(阿联酋)
ar-LVTArabic (Levant)阿拉伯语(黎凡特)
fa-IRPersian波斯语
bn-BDBengali孟加拉语
ta-SGTamil (Singaporean)泰米尔语(新加坡)
ta-LKTamil (Sri Lankan)泰米尔语(斯里兰卡)
ta-INTamil (India)泰米尔语(印度)
ta-MYTamil (Malaysia)泰米尔语(马来西亚)
te-INTelugu泰卢固语
ug-NULLUighur维吾尔语
ug-CNUighur维吾尔语
gu-INGujarati古吉拉特语
my-MMBurmese缅甸语
tl-PHTagalog塔加洛语
kk-KZKazakh哈萨克语
or-INOriya / Odia奥里亚语
ne-NPNepali尼泊尔语
mn-MNMongolian蒙古语
km-KHKhmer高棉语
jv-IDJavanese爪哇语
lo-LALao老挝语
si-LKSinhala僧伽罗语
fil-PHFilipino菲律宾语
ps-AFPushto普什图语
pa-INPanjabi旁遮普语
kab-NULLKabyle卡拜尔语
ba-NULLBashkir巴什基尔语
ks-INKashmiri克什米尔语
tg-TJTajik塔吉克语
su-IDSundanese巽他语
mr-INMarathi马拉地语
ky-KGKirghiz吉尔吉斯语
az-AZAzerbaijani阿塞拜疆语

DeepSpeed Ulysses: 训练极长序列Transformer模型的系统优化

从生成性AI到科研模型,长序列训练正在变得非常重要。 在生成性AI领域,会话式AI、长文档摘要和视频生成等任务都需要在空间和时间层面对长上下文进行推理。 例如,多模态基础模型,如同时处理语音、图像和波形的模型,需要对具有极长序列的高维输入进行长上下文推理。 同样,章节和书籍级别的摘要(数万甚至数十万字)在会话式AI和摘要任务中也非常重要。

对于科学AI来说,长序列同样至关重要,它为更好地理解结构生物学、医疗保健、气候和天气预测以及大分子模拟打开了大门。 例如,通过在基因序列上训练大型语言模型,我们可以创建可以使用极长序列(人类基因组有64亿个碱基对)学习基因组进化模式的语言模型。在医疗保健领域,以所有的患者护理记录为条件的诊断预测模型需要极长序列的上下文。

尽管对于生成性AI和科学AI来说,长序列长度的重要性逐渐增长,但现有的大型模型训练系统和底层的并行技术(数据、张量、流水线、序列并行)并不能支持高效的长序列训练。现有并行方法存在两个主要挑战。首先,现有的数据、张量和流水线等并行方法无法解决序列维度的扩展问题。其次,由于内存通信效率低下,现有的序列并行方法不够高效。此外,现有方法的易用性不足,需要进行侵入性和复杂易出错的代码重构。

为了解决这些问题,我们很高兴宣布推出DeepSpeed-Ulysses(或称为Ulysses,一个非常长的小说),这是一种简单、易用且高效的方法,用于支持具有极长序列长度的高效可扩展LLM训练。

DeepSpeed-Ulysses将各个样本在序列维度上分割给参与的GPU。然后,在attention计算之前,它对已分割的查询(Q)、键(K)和值(V)执行all-to-all通信操作,以使每个GPU接收完整的序列,但仅用于注意力头的非重叠子集。这使得参与的GPU可以并行计算不同的注意力头。最后,DeepSpeed-Ulysses还使用另一个all-to-all来在注意力头上收集结果,同时重新在序列维度上进行分区。

DeepSpeed-Ulysses及其与此博客一起发布的实现的关键特性如下:

  • 与现有系统相比,序列长度增加了4倍,支持训练超过百万个token的序列。
  • 与现有系统相比,通信减少了超过10倍,导致吞吐量提高了高达2.5倍,并且每个GPU的持续吞吐量超过175 TFlops(超过硬件峰值的54%)。
  • 完全通用的attention:DeepSpeed序列并行支持密集和稀疏的注意力,并可与高效的注意力实现(如FlashAttention v2)一起工作。
  • 支持大规模模型训练:DeepSpeed序列并行不仅支持大序列长度,还可以与ZeRO-3并用支持大模型尺寸。
  • 易于使用和迁移,最小化对现有训练框架的代码更改要求。

在接下来的章节中,我们详细讨论DeepSpeed-Ulysses的核心设计、通信复杂度分析、实验评估以及与现有工作的比较,并展示其可用性和使用指南。

DeepSpeed-Ulysses的核心设计

图1显示了DeepSpeed-Ulysses的核心设计。与已知的Transformer架构一样,设计由N个输入序列在P个可用设备上分区组成。每个本地N/P分区都被投影到查询(Q)、键(K)和值(V)嵌入中。接下来,(QKV) 嵌入通过参与计算设备之间的高度优化的全对全集合(all-to-all collectives)进行全局的 QKV 收集。在全对全集合后,每个头的注意力计算形式为:

注意力计算后,另一个全对全集合将注意力计算的输出上下文张量转换为序列(N/P)并行,用于Transformer模型层的剩余模块中的后续操作(MLP MatMul、层归一化等)。

显著的通信量减少

DeepSpeed-Ulysses与其他现有的长序列方法的区别在于其更小的累积通信量以及随着序列并行度增加而更好的可扩展性,如下所示:

在具有节点内NVSwitch互连和节点间胖树IB拓扑的现代集群上,针对一个聚合消息大小为M的全对全传输,传输到P个GPU上的每个链接的通信量为M/P。 对于隐藏层大小为h、序列长度为N且并行度为P的Transformer模型,DeepSpeed序列并行会在注意计算之前对QKV投影执行聚合消息大小为3Nh的全对全操作,并在注意计算之后对输出上下文投影执行大小为Nh的另一个全对全操作。因此,DeepSpeed序列并行每个链接的聚合通信量为4Nh/P(或O(N/P)复杂度)。值得注意的是,当N和P成比例增加时,这个通信量是恒定的。

相比之下,现有的方法,如Megatron-LM,在N线性增长的情况下会导致通信量线性增加,而与P无关,从而导致O(N)的通信复杂度。例如,Megatron-LM对每个Transformer模型层都执行两个大小为Nhall-gather操作,以及两个大小为Nhreduce-scatter操作。然而,当P >> 1时,大小为M的每个all-gather和reduce-scatter的成本仍然是M,而不是M/P。因此,Megatron-LM序列并行会导致每个链接的通信量为4Nh,这比DeepSpeed序列并行大P倍。这使得DeepSpeed序列并行可以在实现显著更高的训练效率的同时支持极长序列训练。我们的实验评估结果与此理论分析相符。

DeepSpeed-Ulysses的其他亮点

通用的注意力解决方案

DeepSpeed分布式注意力模块的实现足够通用,以支持任何类型的注意力,例如自注意、交叉注意和因果注意,无论是它们的密集还是稀疏版本,以及支持局部注意层级上的长序列的各种优化内核,例如不同版本的FlashAttention。

DeepSpeed-Ulysses的通用性来自其核心设计的模块化性质:一个以注意力为中心的序列并行设计。在注意力计算之前,序列并行性是对N/P分区的,而注意力计算是对每个头的并行性,每个头的注意力全都保留,但头的数量较少,因此注意力计算可以用任何类型的注意力机制替代,例如密集注意力和各种形式的稀疏注意力。

通过ZeRO-3集成实现更大的模型和更长的序列训练

尽管DeepSpeed序列并行在使用更长的序列进行训练时减少了激活内存的使用,但并不影响模型状态的内存占用。因此,为了支持具有大序列长度的大语言模型训练,我们实现了DeepSpeed序列并行与ZeRO-3的集成。

ZeRO Redundancy Optimizer Stage 3 (ZeRO-3) 是一种用于训练大模型的内存优化技术。与传统的神经网络数据并行训练中,模型状态在数据并行等级上进行复制不同,ZeRO-3通过将模型状态在数据并行等级之间进行分区来优化内存使用。然而,使用序列并行时,训练数据可以在批(样本)和序列维度上考虑,相关的并行群组可以组合成一个更大的群组以实现ZeRO并行。

因此,我们将ZeRO-3分区扩展到数据并行和序列并行等级的组合。换句话说,在DeepSpeed序列并行中,ZeRO将模型状态分区在序列和数据并行组之间,并在需要时收集每个等级分区(allgather)。类似地,梯度将在数据并行和序列并行等级之间进行减少,用于参数更新。ZeRO可以在序列和数据维度上实现巨大的内存节省,并且不仅可以扩展到大序列长度,还可以扩展到大模型。

评估

我们在GPT(用于许多NLP任务的基础模型)上使用最多64个A100 GPU(40GB显存)对DeepSpeed-Ulysses进行了评估。我们的评估分为四个方面:i) 序列长度可扩展性,ii) 密集注意力的吞吐量以及与现有系统的比较,iii) 稀疏注意力的吞吐量以及与现有系统的比较,iv) DeepSpeed序列并行的收敛性研究。接下来,我们将对每个类别讨论和展示评估结果。

序列长度可扩展性

第一组实验是在12亿参数的GPT模型上将序列长度扩展到100万token。这个评估的结果如图2所示。DeepSpeed序列并行允许随着GPU数量的增加线性增加序列长度,并且序列长度与GPU数量保持线性比例关系,适当的GPU数量下保持相似的计算吞吐量。

密集注意力评估

接下来,我们在300亿参数的密集注意力模型上对DeepSpeed序列并行进行了评估,并与Megatron序列并行在64个A100 GPU上进行了对比。这些评估的结果如图3所示。

我们将DeepSpeed序列并行与Megatron-LM在不同序列长度下的性能进行了比较。对于我们的评估,我们选择了能使DeepSpeed序列并行和Megatron-LM分别达到最佳性能(通过吞吐量或TFLOPs衡量)的序列长度-批大小组合,我们称之为最佳(批大小-序列长度)配置。对于DeepSpeed序列并行,我们始终使用64的ZeRO并行度。

图3显示,DeepSpeed序列并行在相同序列长度下始终优于Megatron-LM。此外,DeepSpeed序列并行可以运行比Megatron-LM更长的序列。DeepSpeed序列并行的性能优势在于两个方面:(1)DeepSpeed序列并行结合ZeRO-3的内存优化,可以容纳更多的样本,从而提高吞吐量;(2)相对于Megatron-LM序列并行中应用的all-gather通信,DeepSpeed序列并行使用更高效的全对全通信。

图3:DeepSpeed和Megatron LM序列并行在300亿参数模型上的密集注意力评估。

稀疏注意力评估

类似地,我们在300亿参数的稀疏注意力模型上对DeepSpeed序列并行进行了评估,并与Megatron序列并行进行了对比。我们的评估结果如图4所示。稀疏注意力的实验结果与密集注意力实验类似。我们观察到DeepSpeed序列并行的吞吐量性能相对于Megatron-LM提高了2倍以上。通过节省内存,DeepSpeed序列并行结合ZeRO-3可以扩展到比Megatron-LM更长4倍的序列长度。

DeepSpeed序列并行在相同序列长度下始终优于Megatron-LM。事实上,当前的DeepSpeed吞吐量受到本地稀疏注意力实现的瓶颈,因此DeepSpeed吞吐量随着序列长度的增加而降低。我们预计,随着未来局部稀疏注意力实现性能的改善,DeepSpeed与Megatron之间的性能差距将在更大的序列长度下进一步增加。

图4:DeepSpeed和Megatron LM序列并行在300亿参数模型上的稀疏注意力评估。

收敛性研究

最后,图5显示了1.3亿参数GPT模型在32K序列长度下,使用序列并行度设置为4的情况下,在8个A100 GPU上的收敛性。对于DeepSpeed序列并行,我们使用不同的ZeRO阶段进行了收敛性评估。DeepSpeed序列并行是一种纯系统优化技术,用于实现长序列Transformer模型的训练,因此在训练模型质量上没有(负面)影响,并通过实验得到了验证,如图5所示。

图5:使用不同ZeRO内存优化阶段的DeepSpeed序列并行的收敛性评估。

DeepSpeed-Ulysses软件可用性

DeepSpeed-Ulysses只需进行少量简单代码更改来集成到您的代码中。下面是一个启用它的示例:

from deepspeed.sequence.layer import DistributedAttention

# 将原始的自注意(attn)替换为DeepSpeed-Ulysses的自注意

dist_attn = DistributedAttention(attn, get_sequence_parallel_group())

与其他支持序列并行的库(如Megatron-LM)相比,DeepSpeed-Ulysses不需要进行模型重构。 DeepSpeed-Ulysses已经完全与Megatron-DeepSpeed代码库集成并经过测试。这意味着如果您已经 在使用这个代码库来训练大型语言模型,您可以无缝地使用DeepSpeed-Ulysses训练具有极长序列的模型。

DeepSeek-R1 技术报告

摘自:https://zhuanlan.zhihu.com/p/19744278380

Github: https://github.com/deepseek-ai/DeepSeek-R1

DeepSeek-R1:通过强化学习提升LLM的推理能力

R1训练流程:

•冷启动 •基于推理的强化学习 •Rejection Sampling •SFT •全场景强化学习

DeepSeek-R1-Zero 采用大规模强化学习(RL)进行训练,无需预先进行监督微调(SFT),表现出显著的推理能力。在强化学习过程中,DeepSeek-R1-Zero 展现出多种卓越且新颖的推理特性。但该模型仍面临可读性不足语言混杂等问题。

为解决这些问题并进一步增强推理性能,研究团队开发了 DeepSeek-R1,该模型在进行强化学习前引入了多阶段训练和冷启动数据。

DeepSeek-R1 在推理任务上实现了与 OpenAI-o1-1217 相当的性能水平

为促进学术研究发展,研究团队开源了 DeepSeek-R1-Zero、DeepSeek-R1,以及基于 Qwen 和 Llama 架构从 DeepSeek-R1 知识蒸馏获得的六个稠密模型(1.5B、7B、8B、14B、32B、70B)。

引言

近年来,LLM技术发展迅速,不断缩小与AGI的差距。后训练技术已成为完整训练流程中的关键环节,证实能够提升推理任务准确率,实现社会价值观对齐,适应用户偏好,同时相较于预训练所需计算资源较少。在推理能力方面,OpenAI的o1系列模型首次通过延长Chain-of-Thought(CoT)推理过程引入了推理时扩展机制,在数学、编程和科学推理等多个推理任务中取得显著进展。

然而,如何实现有效的测试时扩展仍是学术界面临的重要课题。前期研究探索了多种方法,包括过程型奖励模型、强化学习以及蒙特卡洛树搜索和束搜索等算法。但这些方法均未能达到与OpenAI的o1系列模型相当的通用推理水平。

本研究采用纯RL方法提升语言模型的推理能力。研究旨在探索LLM在无监督数据条件下通过纯RL过程实现自我进化的推理能力潜力。

具体而言,研究选用DeepSeek-V3-Base作为基础模型,采用群组相对策略优化(GRPO)作为RL框架提升模型推理性能。在训练过程中,DeepSeek-R1-Zero自然形成了多种高效且创新的推理特征。经过数千轮RL迭代,DeepSeek-R1-Zero在推理基准测试中展现出优异性能。例如,在AIME 2024测试中,pass@1得分从15.6%提升至71.0%,采用majority voting机制后,得分进一步提高到86.7%,达到OpenAI-o1-0912的性能水平。

然而,DeepSeek-R1-Zero仍面临可读性不足、语言混杂等挑战。

为解决这些问题并进一步提升推理性能,研究团队开发了DeepSeek-R1模型,该模型整合了初始训练数据和多阶段训练流程。具体实施步骤包括:首先收集数千条初始训练数据用于DeepSeek-V3-Base模型的微调;随后进行推理强化学习训练;在RL过程接近收敛时,通过拒绝采样(rejection sampling)方法从RL检查点生成新的SFT数据,并结合DeepSeek-V3在写作、事实QA和自我认知等领域的监督数据重新训练DeepSeek-V3-Base模型;最后,使用新数据完成微调后的检查点进行额外的RL训练,综合考虑各类场景的提示词。

经过上述步骤,最终获得的DeepSeek-R1模型达到了与OpenAI-o1-1217相当的性能水平。

研究进一步探索了从DeepSeek-R1到较小dense模型的知识蒸馏。以Qwen2.5 32B为基础模型,直接从DeepSeek-R1进行知识蒸馏的效果优于直接应用RL训练,表明大型基础模型所发现的推理模式对提升推理能力具有关键作用。研究团队已开源蒸馏后的Qwen和Llama系列模型。

值得注意的是,14B蒸馏模型的性能显著超越了当前最先进的开源模型QwQ-32B-Preview,而32B和70B蒸馏模型则在稠密模型推理基准测试中创造了新的记录

主要贡献

后训练:基础模型的大规模强化学习应用

  • 本研究直接将RL应用于基础模型,无需将SFT作为前置步骤。这种方法使模型能够通过CoT探索复杂问题的解决方案,最终开发出DeepSeek-R1-Zero模型。DeepSeek-R1-Zero具备自我验证、反思和生成长CoT等能力,为学术界提供了重要研究成果。这是首个验证LLM推理能力可纯粹通过RL提升而无需SFT的开放研究,为该领域未来发展奠定基础
  • 研究提出了DeepSeek-R1的开发流程,包含两个RL阶段用于优化推理模式和人类偏好对齐,以及两个SFT阶段用于构建模型的推理和非推理基础能力。该流程将有助于行业开发更高性能的模型。

知识蒸馏:小型模型的性能提升

  • 研究表明大型模型的推理模式可通过知识蒸馏迁移至小型模型,其效果优于直接对小型模型进行RL训练。开源的DeepSeek-R1及其API将支持学术界开发更优秀的小型模型
  • 利用DeepSeek-R1生成的推理数据,研究团队对学术界广泛使用的多个稠密模型进行了微调。评估结果显示,经过知识蒸馏的小型dense模型在基准测试中表现优异。DeepSeek-R1-Distill-Qwen-7B在AIME 2024上达到55.5%的性能,超越QwQ-32B-Preview。DeepSeek-R1-Distill-Qwen-32B在AIME 2024、MATH-500和LiveCodeBench上分别达到72.6%、94.3%和57.2%的成绩,显著优于现有开源模型,达到与o1-mini相当的水平。研究团队已向学术界开源基于Qwen2.5和Llama3系列的1.5B、7B、8B、14B、32B和70B蒸馏检查点

研究方法

概述

传统研究主要依赖大规模监督数据提升模型性能。本研究证实,即使在无需监督微调(SFT)作为初始训练的情况下,通过大规模强化学习(RL)也能显著提升推理能力。此外,引入适量初始训练数据可进一步优化性能。后续章节将介绍:(1)DeepSeek-R1-Zero:直接对基础模型应用RL,无需任何SFT数据;(2)DeepSeek-R1基于经数千个长CoT样例微调的检查点进行RL训练;(3)将DeepSeek-R1的推理能力通过知识蒸馏迁移至小型稠密模型

DeepSeek-R1-Zero:基础模型的强化学习应用

前期相关研究表明强化学习在推理任务中具有显著效果。然而,这些研究高度依赖耗时的监督数据采集。本节探索LLM在无监督数据条件下通过纯强化学习实现推理能力自我进化的潜力。研究首先概述强化学习算法,随后展示实验结果,以期为学术界提供研究参考。

强化学习算法

群组相对策略优化(GRPO): 为优化RL训练成本,研究采用GRPO算法,摒弃了通常与策略模型规模相当的评论家模型,转而通过群组评分估计基线。具体而言,对每个问题 q ,GRPO从旧策略 πθold 采样输出组{ o1,o2,…,oG },通过最大化以下目标优化策略模型 πθ :

其中 ε 和 β 是超参数, Ai 是优势函数,使用组内每个输出对应的奖励组{ r1,r2,…,rG }计算得到:

奖励建模

奖励机制作为训练信号来源,决定RL的优化方向。DeepSeek-R1-Zero采用基于规则的双重奖励系统

  • 准确性奖励:评估响应正确性。如对确定性数学问题,要求模型以特定格式(如方框内)提供最终答案,实现基于规则的可靠验证。对LeetCode问题,则通过编译器基于预设测试用例生成反馈。
  • 格式奖励:要求模型将推理过程置于指定标签对内。研究未采用结果或过程神经奖励模型,原因在于神经奖励模型可能在大规模RL过程中产生奖励欺骗问题,且重训奖励模型需额外资源,增加训练流程复杂度。

训练模板

DeepSeek-R1-Zero的训练始于简洁指令模板的设计。

如表1所示,模板要求模型首先生成推理过程,随后给出最终答案。研究刻意将约束限定于结构格式,避免引入内容偏见(如强制反思推理或特定问题解决策略),以准确观测模型在RL过程中的自然演化。

DeepSeek-R1-Zero的性能分析、演化过程及关键突破

性能分析 图2记录了DeepSeek-R1-Zero在RL训练过程中AIME 2024基准测试的性能变化轨迹。

图2 | DeepSeek-R1-Zero训练过程中的AIME准确率变化。为确保评估稳定性,对每个问题采样16个响应并计算总体平均准确率。

数据显示,随着RL训练的深入,模型性能呈现稳定上升趋势。在AIME 2024测试中,平均pass@1得分从初始的15.6%显著提升至71.0%,达到OpenAI-o1-0912的性能水平,充分证实了RL算法在模型性能优化方面的有效性。

表2 | DeepSeek-R1-Zero与OpenAI o1模型在推理相关基准测试上的性能对比。

表2详细对比了DeepSeek-R1-Zero与OpenAI o1-0912模型在各类推理基准测试上的表现。结果表明,纯RL训练使DeepSeek-R1-Zero获得了出色的推理能力,无需借助监督微调数据,这证实了模型通过单一RL机制实现有效学习和泛化的能力。通过引入majority voting机制,模型性能得到进一步提升。例如,在AIME基准测试中,采用majority voting后性能从71.0%提升至86.7%,超越OpenAI-o1-0912。这种优异表现凸显了模型的基础能力和推理潜力。

演化过程分析 DeepSeek-R1-Zero的演化过程展示了RL在推理能力自主优化方面的显著效果。通过直接对基础模型实施RL训练,研究得以在无监督微调影响下观测模型进展。

图3 | 展示DeepSeek-R1-Zero在RL训练过程中训练集的平均响应长度变化,反映模型自主习得延长推理时间的能力。

如图3所示,模型的推理时长在训练过程中持续优化,这种进展源于模型的内生发展而非外部干预。DeepSeek-R1-Zero通过扩展测试计算时间,自然形成了解决复杂推理任务的能力。其计算规模从数百到数千个推理token不等,实现了深度的思维探索和优化。随着测试计算时间的延长,模型展现出复杂的行为特征,包括反思机制(重新评估先前推理步骤)和多元问题解决策略的探索。这些行为模式并非预设,而是源于模型与RL环境的交互作用,显著增强了其处理高难度任务的效率和准确性。

关键突破与局限性 研究过程中观察到模型出现重要突破,如表3所示,体现在中期版本中。

表3:记录DeepSeek-R1-Zero中期版本的重要突破,展示模型获得自主思考复核能力的过程,体现RL在模型能力提升方面的有效性。

此阶段,DeepSeek-R1-Zero习得了重新评估初始方法并延长思考时间的能力。这一进展不仅体现了模型推理能力的提升,也展示了RL在实现复杂学习成果方面的潜力。这种现象验证了RL的核心优势:通过适当的激励机制,促使模型自主发展高级问题解决策略。

然而,DeepSeek-R1-Zero仍存在若干局限性。尽管具备强大的推理能力和创新的推理行为,但在可读性和语言一致性方面仍面临挑战。为提高推理过程的可读性并促进开放社区交流,研究团队开发了DeepSeek-R1模型,该模型结合了RL和用户友好的初始训练数据。

DeepSeek-R1:基于冷启动的强化学习方法

基于DeepSeek-R1-Zero的成功实践,研究聚焦两个核心问题:

  1. 通过引入少量高质量数据作为冷启动,是否能够进一步提升推理性能或加速收敛?
  2. 如何开发既能生成清晰连贯的CoT,又具备强大通用能力的用户友好型模型?

为解决上述问题,研究团队设计了四阶段训练流程

冷启动机制

区别于DeepSeek-R1-Zero,DeepSeek-R1采用少量长CoT数据对模型进行预微调作为初始RL策略网络,以避免基础模型RL训练早期的不稳定性。数据收集采用多种方法:

  • 基于长CoT示例的少样本提示
  • 直接提示生成包含反思验证的详细答案
  • 整理DeepSeek-R1-Zero的规范化输出
  • 人工标注后处理优化

研究收集数千条冷启动数据用于DeepSeek-V3-Base的预训练。相较于DeepSeek-R1-Zero,冷启动数据具有以下优势:

  • 可读性增强:克服了DeepSeek-R1-Zero输出内容可读性差的局限。通过设计标准化输出模式,包括响应末尾的总结性内容,并筛除不符合阅读友好性要求的输出。输出采用|special_token|<reasoning_process>|special_token|<summary>格式,包含查询的推理过程和结果摘要。
  • 性能提升:基于人类认知模式优化的冷启动数据设计,展现出优于DeepSeek-R1-Zero的性能表现,验证了迭代训练对推理模型的优越性。

推理强化学习优化

完成冷启动数据预训练后,采用与DeepSeek-R1-Zero类似的大规模RL训练流程,重点提升模型在编码、数学、科学和逻辑等明确定义问题域的推理能力。在训练过程中发现Chain-of-Thought存在语言混杂现象,尤其是多语言提示场景下。为此引入语言一致性奖励机制,基于目标语言词占比计算。尽管消融实验显示该机制略微影响模型性能,但提升了人类使用体验。最终将任务准确率和语言一致性奖励合并计算总体奖励,持续RL训练直至模型在推理任务上收敛。

拒绝采样与监督微调

推理RL收敛后,利用检查点生成后续SFT数据。不同于专注推理的冷启动阶段,此阶段整合多领域数据以增强模型的写作、角色扮演等通用能力。具体实施如下:

推理数据构建 通过对RL训练检查点执行拒绝采样生成推理轨迹。扩展了评估机制,除规则型奖励外,引入基于DeepSeek-V3判断的生成式奖励模型。优化输出质量,过滤混杂语言、冗长段落和代码块。对每个提示词进行多样本采样,保留正确结果。最终获得约60万条推理训练样本。

非推理数据整合 在写作、事实QA、自我认知和翻译等领域,采用DeepSeek-V3流程和部分SFT数据。对复杂非推理任务,通过提示DeepSeek-V3生成前置CoT;对简单查询则直接响应。累计获取约20万条非推理训练样本。使用总计约80万样本数据对DeepSeek-V3-Base执行两轮微调。

全场景强化学习

优化人类偏好对齐,实施第二阶段RL训练,着重提升模型实用性、安全性和推理能力。采用多元奖励信号和多样化提示分布:

  • 推理数据:延续DeepSeek-R1-Zero方法,在数理逻辑领域应用规则型奖励
  • 通用数据:采用奖励模型捕捉复杂场景下的人类偏好
  • 实用性评估:专注于响应摘要,确保输出的实用性和相关性
  • 安全性保障:全面评估推理过程和摘要,识别并降低潜在风险

通过奖励信号和数据分布的系统整合,实现了推理能力和用户体验的均衡发展。

知识蒸馏:增强小型模型的推理能力

本研究采用DeepSeek-R1生成的80万训练样本,对Qwen和Llama等开源模型进行直接SFT微调,旨在将DeekSeek-R1的推理能力迁移至计算效率更高的小型模型。

实验结果表明,这种直接知识蒸馏方法显著提升小型模型的推理性能

研究选用的基础模型包括:Qwen2.5-Math-1.5B、Qwen2.5-Math-7B、Qwen2.5-14B、Qwen2.5-32B、Llama-3.1-8BLlama-3.3-70B-Instruct

选择Llama-3.3的原因在于其推理能力较Llama-3.1略有优势

蒸馏过程中仅采用SFT,未纳入RL阶段,尽管引入RL可能带来显著的性能提升。研究重点在于验证知识蒸馏技术的有效性,为后续学术界对RL优化的深入研究奠定基础。

实验设计与评估

研究采用多维度基准测试体系评估模型性能:

标准评估基准 8类16个评估标准如下所示:

  • 知识理解类:MMLU、MMLU-Redux、MMLU-Pro
  • 跨语言评估:C-Eval、CMMLU
  • 格式理解:IFEval
  • 长文本处理:FRAMES
  • 专业知识:GPQA Diamond
  • 事实问答:SimpleQA、C-SimpleQA
  • 编程能力评估: SWE-Bench Verified、Aider、LiveCodeBench、Codeforces
  • 数学能力测试: CNMO 2024、AIME 2024

除标准基准测试外,研究还使用LLM作为评估器评估模型在开放式生成任务上的表现。具体而言,遵循AlpacaEval 2.0Arena-Hard的原始配置,使用GPT-4-Turbo-1106作为成对比较的评估器。评估时仅输入最终摘要以避免长度偏差。对于蒸馏模型,报告其在AIME 2024、MATH-500、GPQA Diamond、Codeforces和LiveCodeBench上的代表性结果。

评估用prompt 不同的评估标准采用不同的prompt,具体如下所示:

  • 基础评估:采用simple evals框架标准prompt评估MMLU、DROP、GPQA Diamond和SimpleQA
  • 特殊处理: MMLU-Redux采用Zero-Eval prompt格式实现零样本评估,MMLU-Pro、C-Eval、CLUE-WSC将原少样本prompt改造为零样本形式
  • 编程评估: HumanEval-Mul覆盖8种主流编程语言,LiveCodeBench采用CoT格式,Codeforces基于10个Div.2竞赛题目与专家测试用例,SWE-Bench通过无代理框架验证

值得注意的是,DeepSeek-R1的输出在每个基准测试上限制为最多32,768个token。

基准模型 研究与多个强基准模型进行全面对比,包括DeepSeek-V3、Claude-Sonnet-3.5-1022、GPT-4o-0513、OpenAI-o1-miniOpenAI-o1-1217。鉴于在中国大陆访问OpenAI-o1-1217 API的限制,其性能数据来源于官方报告。对于蒸馏模型,额外与开源模型QwQ-32B-Preview进行比较。

生成配置 所有模型的最大生成长度设置为32K token。对需要采样的基准测试,采用0.6的temperature参数、0.95的top-p值,并为每个查询生成64个响应以估算pass@1。

DeepSeek-R1评估结果

表4 | DeepSeek-R1与其他代表性模型的比较。

在面向教育的知识基准测试(如MMLU、MMLU-Pro和GPQA Diamond)中,DeepSeek-R1相较于DeepSeek-V3展现出优越性能。这一进步主要归因于STEM相关问题准确率的提升,这得益于大规模RL带来的显著进步。

此外,DeepSeek-R1在依赖长文本理解的问答任务FRAMES上表现卓越,展示了其强大的文档分析能力。这凸显了推理模型在AI驱动的搜索和数据分析任务中的潜力

在事实性基准测试SimpleQA上,DeepSeek-R1的表现优于DeepSeek-V3,证明了其处理基于事实查询的能力。类似地,在该基准测试中也观察到OpenAI-o1超越GPT-4o的趋势。

然而,DeepSeek-R1在中文SimpleQA基准测试中的表现不如DeepSeek-V3,主要是由于安全性RL后倾向于拒绝回答某些查询。若不考虑安全性RL,DeepSeek-R1可以达到超过70%的准确率。

DeepSeek-R1在IF-Eval(一个用于评估模型遵循格式指令能力的基准测试)上也取得了令人瞩目的成果。这些改进可归因于在最终阶段的SFT和RL训练中引入了指令遵循数据。

此外,在AlpacaEval 2.0和ArenaHard上的出色表现表明DeepSeek-R1在写作任务和开放域问答方面具有优势。其显著优于DeepSeek-V3的表现凸显了大规模RL的泛化效益,不仅提升了推理能力,还改善了各个领域的性能。

而且DeepSeek-R1生成的摘要长度简洁,在ArenaHard上平均为689个token,在AlpacaEval 2.0上平均为2,218个字符。这表明DeepSeek-R1在基于GPT的评估中避免了引入长度偏差,进一步证实了其在多任务场景下的稳健性。

数学任务上,DeepSeek-R1展现出与OpenAI-o1-1217相当的性能,大幅超越其他模型。在LiveCodeBench和Codeforces等编码算法任务上也观察到类似趋势,其中注重推理的模型在这些基准测试中占据主导地位。

在面向工程的编码任务上,OpenAI-o1-1217在Aider上优于DeepSeek-R1,但在SWE Verified上表现相当。考虑到目前相关RL训练数据量仍然非常有限,研究团队认为DeepSeek-R1的工程性能将在下一版本中得到改善。

蒸馏模型评估

表5 | DeepSeek-R1蒸馏模型与其他可比模型在推理相关基准测试上的比较。

如表5所示,仅通过蒸馏DeepSeek-R1的输出,高效的DeepSeek-R1-7B(即DeepSeek-R1-Distill-Qwen-7B,以下类似缩写)就能在各方面超越GPT-4o-0513等非推理模型。

DeepSeek-R1-14B在所有评估指标上超越QwQ-32B-Preview,而DeepSeek-R1-32B和DeepSeek-R1-70B在大多数基准测试中显著超越o1-mini。这些结果展示了知识蒸馏的巨大潜力

此外,研究发现对这些蒸馏模型应用RL能带来显著的进一步提升。考虑到这值得进一步探索,此处仅呈现简单SFT蒸馏模型的结果。

讨论

蒸馏与强化学习对比

通过蒸馏DeepSeek-R1,小型模型能够取得出色的结果。然而,仍有一个问题待解答:

模型是否可以通过本文讨论的大规模RL训练而不依赖蒸馏来达到相当的性能?

为回答这个问题,研究团队对Qwen-32B-Base使用数学、代码和STEM数据进行了超过10K步的大规模RL训练,得到DeepSeek-R1-Zero-Qwen-32B。

如表6所示的实验结果表明,32B基础模型经过大规模RL训练后,达到了与QwQ-32B-Preview相当的性能。然而,从DeepSeek-R1蒸馏得到的DeepSeek-R1-Distill-Qwen-32B在所有基准测试中的表现都显著优于DeepSeek-R1-Zero-Qwen-32B

因此,可以得出两个结论:

首先,将更强大的模型蒸馏到较小的模型中可以产生优异的结果,而较小的模型依靠本文提到的大规模RL需要巨大的计算力,甚至可能无法达到蒸馏的性能水平

其次,虽然蒸馏策略既经济又有效,但要突破智能的边界可能仍需要更强大的基础模型和更大规模的强化学习

未成功的尝试

在开发DeepSeek-R1的早期阶段,研究也遇到了失败和挫折。在此分享这些失败经验以提供见解,但这并不意味着这些方法无法开发出有效的推理模型。

过程奖励模型(PRM)

PRM是一种合理的方法,可以引导模型采用更好的方法解决推理任务。然而,在实践中,PRM有三个主要限制可能阻碍其最终成功。

首先,在一般推理中明确定义细粒度步骤具有挑战性。其次,确定当前中间步骤是否正确是一项具有挑战性的任务。使用模型的自动标注可能无法产生令人满意的结果,而手动标注不利于规模化。第三,一旦引入基于模型的PRM,必然导致奖励欺骗,重新训练奖励模型需要额外的训练资源,并使整个训练流程变得复杂。

总之,虽然PRM在对模型生成的前N个响应重新排序或辅助引导搜索方面表现良好,但在实验中,相比其在大规模强化学习过程中引入的额外计算开销,其优势有限

蒙特卡洛树搜索(MCTS)

AlphaGoAlphaZero的启发,研究探索使用MCTS来增强测试时计算的可扩展性。这种方法包括将答案分解为更小的部分,使模型能够系统地探索解决方案空间。为此,提示模型生成多个标签,对应搜索所需的具体推理步骤。在训练方面,首先使用收集的提示通过预训练值模型引导的MCTS寻找答案。随后,使用产生的问答对来训练actor模型和值模型,不断改进过程。

然而,这种方法在扩大训练规模时遇到几个挑战。首先,与搜索空间相对明确的象棋不同,token生成呈现指数级更大的搜索空间。为解决这个问题,为每个节点设置最大扩展限制,但这可能导致模型陷入局部最优。其次,值模型直接影响生成质量,因为它指导搜索过程的每个步骤。训练细粒度值模型本质上是困难的,这使得模型难以迭代改进。虽然AlphaGo的核心成功依赖于训练值模型来逐步提升性能,但由于token生成的复杂性,这一原则在团队的设置中难以复制。

总之,虽然MCTS在与预训练值模型配对时可以改善推理性能,但通过自搜索迭代提升模型性能仍然是一个重大挑战

结论、局限性和未来工作

本文分享了通过RL增强模型推理能力的探索历程。DeepSeek-R1-Zero代表了一种不依赖冷启动数据的纯RL方法,在各种任务中取得了出色的表现。DeepSeek-R1通过结合冷启动数据和迭代RL微调展现出更强的性能,最终在多个任务上达到与OpenAI-o1-1217相当的水平。

研究进一步探索了将推理能力蒸馏到小型稠密模型的可能性。以DeepSeek-R1作为教师模型生成80万条数据,并对多个小型稠密模型进行微调。

结果令人鼓舞:DeepSeek-R1-Distill-Qwen-1.5B在数学基准测试中超越GPT-4o和Claude-3.5-Sonnet,在AIME上达到28.9%,在MATH上达到83.9%的成绩。其他稠密模型也取得了显著成果,大幅超越基于相同基础检查点的其他指令微调模型。

未来,计划在以下方向继续推进DeepSeek-R1的研究:

  • 通用能力:目前DeepSeek-R1在函数调用、多轮对话、复杂角色扮演和json输出等任务上的能力仍不及DeepSeek-V3。后续研究将探索如何利用长CoT增强这些领域的任务表现。
  • 语言混杂:DeepSeek-R1当前针对中文和英文进行了优化,在处理其他语言的查询时可能出现语言混杂问题。例如,即使查询使用非英文或中文的语言,DeepSeek-R1可能使用英语进行推理和响应。未来更新将着力解决这一限制。
  • 提示词工程:在评估DeepSeek-R1时发现,模型对prompt较为敏感。少样本提示会持续降低其性能。因此,建议用户直接描述问题并使用零样本设置指定输出格式以获得最佳结果
  • 软件工程任务:由于评估时间较长影响RL过程效率,大规模RL尚未在软件工程任务中广泛应用。因此,DeepSeek-R1在软件工程基准测试上相比DeepSeek-V3未显示出显著改进。未来版本将通过对软件工程数据实施拒绝采样或在RL过程中引入异步评估来提高效率。

Baichuan-Omni-1.5 多模态语音大模型

Baichuan-Omni-1.5 🤗 | Baichuan-Omni-1.5-Base 🤗 | 技术报告 📖

OpenMM-Medical 🤗 | OpenAudioBench 🤗

Github: https://github.com/baichuan-inc/Baichuan-Omni-1.5/

Baichuan-Omni-1.5 是 Baichuan-omni 系列的最新、性能一流模型。该模型通过端到端方式训练和推理。与其他开源模型相比,Baichuan-Omni-1.5 在文本、图像、音频和视频输入的理解能力有显著提升,并支持了可控的实时语音对话和多模态实时交互的新功能。此外,Baichuan-Omni-1.5 也是目前最出色的开源医疗多模态模型。Baichuan-Omni-1.5 的主要特性包括:

  • 多模态理解和交互能力: Baichuan-Omni-1.5 接受图像、视频、文本、音频作为输入,并生成高质量文本和语音输出,能够在不损害任何模态能力的情况下实现无缝的高质量跨模态交互,并和用户进行实时语音对话。在针对全模态理解的综合评测基准 OminiBench 中,Baichuan-Omni-1.5 取得开源社区一流水平,并超过了 GPT-4o-mini
  • 优秀的视觉能力: Baichuan-Omni-1.5 在 OpenCompass 常用的十个视觉评测集上平均得分 73.3,在7B量级的大小下,在图像理解方面超越了 GPT-4o-mini,比GPT-4o-mini平均高出6分,并且和GPT-4o的结果相近。此外,视频理解表现也优于GPT-4V
  • 出色的语音能力: Baichuan-Omni-1.5 通过一个 8 层 RVQ 音频Tokenizer(Baichuan-Audio-Tokenizer),在 12.5 Hz 帧率下实现了语义和声学信息捕获的最佳平衡,支持高质量可控制声音的中英双语实时对话。Baichuan-Omni-1.5 在语音理解任务优于 GLM-4-Voice,并在语音对话的语义和声学评估中展现了开源模型中最高的语音生成性能。同时,我们还开源了音频理解和生成基准(OpenAudio-Bench),以评估音频的端到端能力。
  • 领先的医疗图像理解能力: 我们从开源数据集中收集了一个比较全面的医学理解基准(OpenMM-Medical)用于评估模型的医学能力。Baichuan-Omni-1.5 在 GMAI-MMBench 以及 OpenMM-Medical 上取得了最佳的表现。在 OpenMM-Medical 上,Baichuan-Omni-1.5 仅使用 7B 的 LLM 取得了 83.8% 的高分,超过 Qwen2-VL-72B 的 80.7%。
在图像、视频和音频模态上的评估。(左)Baichuan-Omni-1.5覆盖了比Qwen2 VL更多的模态,并且超越了当前领先的全模态模型VITA-1.5和MiniCPM-o 2.6。(右)各模态在所有基准测试中的平均得分。

简介

Baichuan-Omni-1.5是一款全模态模型,具备全模态理解能力,并提供端到端的音频生成能力。为了实现跨模态的流畅高质量互动,同时不妥协任何模态的能力,优先优化了三个关键方面。首先,建立了一个全面的数据清洗和合成管道,用于多模态数据处理,获得了约5000亿条高质量数据(包括文本、音频和视觉数据)。其次,设计了一种音频分词器(Baichuan-Audio-Tokenizer),能够同时捕捉音频中的语义和声学信息,从而实现无缝集成,并提升与多模态大语言模型(MLLM)的兼容性。最后设计了一个多阶段的训练策略,逐步整合多模态对齐和多任务微调,确保各模态之间的有效协同。Baichuan-Omni-1.5在全模态能力方面领先于当前的模型(包括GPT4o-mini和MiniCPM-o 2.6)。值得注意的是,在多个多模态医学基准测试中,它取得了与领先模型(如Qwen2-VL-72B)相媲美的成绩。

与开源对手相比,Baichuan-Omni-1.5在文本、图像、音频和视频输入的理解能力上表现出显著的提升。特别地,该模型在可控实时语音互动和跨模态实时协同理解方面展示了令人印象深刻的能力。除了其通用能力外,Baichuan-Omni-1.5在医疗领域作为最出色的多模态大语言模型(MLLM)脱颖而出。这为人工智能在促进人类社会福祉方面开辟了令人兴奋的新可能性。Baichuan-Omni-1.5的架构如下图所示。根据评估结果,我们总结了Baichuan-Omni-1.5的关键优势和贡献:

全模态互动:Baichuan-Omni-1.5旨在处理文本、图像、音频和视频输入,输出高质量的文本和语音。它能够实现无缝、高质量的跨模态互动,且不会妥协任何模态的能力。

卓越的视觉-语言能力:Baichuan-Omni-1.5在十个图像理解基准测试中平均得分为73.3,超越GPT-4o-mini平均6分。

统一且出色的语音能力:我们设计了一种8层的RVQ音频分词器(Baichuan-Audio-Tokenizer),在捕捉语义和声学信息之间达到了最佳平衡,帧率为12.5Hz,支持高质量的可控双语(中文和英文)实时对话。同时,我们还开源了音频理解和生成基准(OpenAudioBench),用于评估音频的端到端能力。

领先的医学图像理解能力:我们收集了一个全面的医学理解基准:OpenMM-Medical,这是多个现有数据集的整合。我们的模型在GMAI-MMBench和OpenMM-Medical上取得了最先进的表现。具体来说,在OpenMM-Medical上,Baichuan-Omni-1.5使用7B大语言模型取得了83.8%的得分,超过了Qwen2-VL-72B的80.7%。

模型架构:

Baichuan-Omni-1.5 的架构。我们的模型旨在处理纯文本/音频输入以及视频/图像与文本/音频的组合。在生成音频时,Baichuan-Omni-1.5 LLM Decoder 交替预测文本标记和音频标记。然后,音频解码器对音频令牌进行解码,以生成最终音频。

Baichuan-Omni-1.5

高质量预训练数据

Baichuan-Omni-1.5的预训练数据集展示:我们构建了一个广泛的全模态数据集,包括文本、图像-文本、视频-文本、音频-文本及它们之间的交互。我们的数据集还包含了交织的图像-音频-文本和视频-音频-文本数据。

构建了全面且高质量的跨模态数据集,包含文本、图像-文本、视频-文本、音频-文本及它们之间的交互

图像数据:我们将图像训练数据分为三种类型:交织的图像-文本数据、图像描述数据和问答数据。为了提高数据的多样性并提升模型性能,我们采用了以下两种图像数据合成策略:

1)我们利用公司内收集的书籍和论文,并解析它们生成交织的图像-文本、OCR数据和图表数据。这些数据具有高度的完整性、专业性和知识密集性。

2)我们根据[19]的方案,训练了一个专用的描述模型,能够生成所需的图像描述,如OCR提示。这些描述提供了图像内容的深入说明。

3)目前,大量的开源数据集主要是英文数据。为了避免模型的中文能力下降,我们合成了大量的中文描述和交织数据。

视频数据:视频数据集包含了广泛的公开资源,涵盖了视频分类 、动作识别 和时间定位等多种任务。视频-文本数据源可以分为视频描述数据和视频问答(QA)数据。

音频数据:音频数据可以大致分为两种主要类型:音频理解数据和音频生成数据。音频理解数据包括自动语音识别(ASR)、音频问答(AQA)、语音转文本翻译和音频-文本交织数据。音频生成数据包括文本转语音(TTS)、交织的文本转语音数据和纯音频数据。交织数据由文本和音频模态交替组成,并以标点符号分隔,以促进跨模态知识的迁移。完全对齐的生成数据由完全对齐的文本和音频内容组成,旨在增强模型在文本监督下生成音频标记的能力。音频-文本配对数据(例如ASR和TTS数据)提高了基础语音任务的表现。纯音频数据则增强了独立处理音频模态的能力。

文本数据:为了构建一个高质量的文本语料库,我们从各种来源聚合了数据,包括网页、书籍、学术论文、代码等。我们遵循之前研究中的数据处理指南,采用严格的选择方法,旨在提高文本语料库的多样性和质量。这种多样性确保训练语料库涵盖了广泛的主题和语言风格,适用于不同的应用。同时,我们的高质量处理技术旨在消除冗余和过滤噪声,从而丰富数据集的信息密度和整体效用。最终,我们获得了1.507亿条纯文本数据。

跨模态交互数据:为了增强模型的跨模态交互能力,我们合成了一系列跨模态交互数据集,涵盖了图像-音频-文本视频-音频-文本格式。图像-文本数据来源包括两种类型:图像-文本描述数据和图像-文本交织数据。具体而言,文本数据首先在句子级别进行分割。然后,将四分之一的文本通过我们的内部文本转语音(TTS)接口转换为音频元素。随后,我们利用生成的音频元素替换原始图像-文本数据中的相应文本句子。这种方法通过将多样化的音频元素融入现有的文本内容,促进了丰富的跨模态交互框架。我们的音频数据包含44种不同的语音类型,确保了语调的多样性。此设置配有任务提示,如“请听以下音频,描述图像的内容。您的任务是在听完后结合音频和图像补充额外信息”,旨在预测剩余的三分之二文本描述。对于视频-文本数据集,音频成分直接从原始视频中提取,作为跨模态的音频元素。总计,我们生成了1000亿个用于跨模态交互的标记数据。

模型结构:

Baichuan-Omni-1.5是一个统一的全模态模型,由视觉分支、音频分支和一个预训练的大型语言模型(LLM)主干组成,支持文本、音频、视觉输入以及端到端的文本和音频输出。

视觉分支
像当前主流的MLLM一样,视觉分支旨在将图像和视频输入处理为视觉标记,并将其与文本标记一起输入到LLM中。我们使用Qwen2-VL的NaViT作为视觉编码器,该编码器能够动态处理任意分辨率和纵横比的图像和视频。然后,我们应用一个由两层MLP组成的视觉投影器,将视觉特征压缩为2×2的因子,从而在性能和效率之间取得平衡。

音频分支
音频分支扩展了LLM,使其能够支持端到端的语音输入和输出。这通过引入Baichuan-Audio-Tokenizer和基于流匹配的解码器来实现,前者负责将音频信号转换为离散标记,后者负责将音频标记解码为语音波形。我们在下图中展示了详细信息。

Baichuan-Audio-Tokenizer基于残差向量量化(RVQ)和多目标训练,帧率为12.5 Hz。在使用Whisper Large Encoder 从Mel谱图特征中提取高级特征后,残差卷积网络执行下采样以获得低帧率序列特征。然后使用8层残差向量量化器对这些特征进行量化,生成音频标记。这些标记随后被输入到音频解码器和预训练的LLM中,分别执行Mel谱图重建和转录预测。音频解码器采用与Whisper编码器对称的结构,并使用多尺度Mel损失来增强声音重建的质量。在训练过程中,预训练LLM的参数保持不变,以确保音频标记器和文本空间之间的语义对齐。

除了传统的任务如ASR、AQA和S2TT外,我们还将一定比例的交织文本-音频数据融入其中,以提高VQ模块建模复杂上下文场景的能力。

为了进一步增强合成音频的质量和感知逼真度,音频解码器模块通过流匹配模型进行优化。借鉴Matcha-TTS 和CosyVoice 的设计,U-Net包括一个单独的下采样块、一个上采样块和12个中间块。具体而言,流匹配解码器在24 kHz音频数据上进行训练,以生成目标Mel谱图,然后使用HiFi-GAN 声码器将其转换为语音波形。

多阶段模型训练:

图像-文本预训练

图像-文本预训练阶段扩展了LLM,使其能够处理和理解视觉输入,使用3000亿图像-文本样本,该阶段可以分为两个部分。

第一阶段:在第一阶段,我们训练视觉投影器,利用开源图像描述数据(例如LAION-5B数据集),建立图像表示与文本之间的初步对齐。在此阶段,我们冻结LLM和视觉编码器,仅训练视觉投影器,学习率为1e−3。

第二阶段:在第二阶段,我们解冻视觉编码器和LLM,以促进图像和文本表示之间更好的对齐。具体来说,我们以学习率1e−5训练LLM和视觉投影器,并以更低的学习率1e−6训练视觉编码器。我们使用公共和内部图像文本数据,包含交织数据和图像描述数据,以增强视觉-语言的表现力。具体来说,我们收集并标注高质量的OCR数据和图表数据,以增强文本/图表识别和理解能力。此外,我们还使用高质量的纯文本数据,这些数据占总数据的40%,以更好地保持语言模型的原始能力。

图像-音频-文本预训练

图像-音频-文本预训练阶段扩展了一个预训练在视觉数据上的LLM,使其能够以端到端的方式理解音频数据,使用887k小时的语音-文本数据,并结合我们的Baichuan-Audio-Tokenizer、新引入的音频嵌入层和独立音频头

具体来说,Baichuan-Audio-Tokenizer生成的音频令牌首先通过音频嵌入层转化为音频嵌入。音频LLM交替生成对齐的文本令牌和音频令牌,使用一个特殊令牌实现文本与音频之间的模态切换。生成的音频令牌由独立的音频头处理,该音频头基于先前的工作设计,包含3层深度变换器和8个分类头。

为了缓解语音和文本特征之间显著差异带来的冲突,我们参考了之前的工作,采用音频和文本数据交织的方法进行预训练。此外,采用了两阶段训练策略,以保持原始LLM的文本知识,同时有效地整合音频模态。

第一阶段:在第一阶段,我们冻结LLM、视觉模块和音频标记器的参数,只有音频嵌入层和音频头的参数更新,学习率为1e−4。我们在这一阶段使用包括ASR、TTS、INTLV和ITTS数据的音频数据。

第二阶段:在第二阶段,训练扩展到除视觉编码器和音频标记器之外的所有参数,学习率为1e−5。具体来说,我们使用音频数据、图像数据和纯文本数据,分别占比0.2、0.4和0.4,这可以更好地提升音频能力,同时保持视觉和语言能力。

全模态预训练

基于之前预训练阶段获得的视觉和音频能力,我们继续使用高质量的跨模态交互数据集进行训练,数据集涵盖图像-音频-文本和视频-音频-文本格式,并将最大序列长度扩展到64k,以支持长时间的语音和视频流。具体来说,输入的视频帧以每秒1帧的速率进行采样,每个视频最多包含32帧。每个输入帧被调整为最大分辨率560×1120像素,以保持最佳质量和细节。这一精心配置在性能和效率之间达到了平衡,促进了有效的模型训练,同时管理了计算负载。此训练过程使用4e−6的低学习率,以进一步优化与语言模态和跨模态交互的对齐。

多模态监督微调

在本节中,我们描述了全模态监督微调(SFT)阶段,旨在增强模型在各种任务中执行复杂全模态指令的能力。我们收集了包括开源、合成和内部注释数据在内的综合数据集。这些数据集涵盖了多个任务,包含大约1700万个数据对,跨越文本、音频、图像-文本、视频-文本和图像-音频组合等多种模态。关于这些数据类型和数量的详细信息见表4。

实验

如表6所示,Baichuan-Omni-1.5 在纯文本基准测试中表现出色,特别是在与仅专注于语言模态的开源LLM模型相比时。例如,在通用MMLU基准测试中,Llama3-Instruct的得分为67.1%,而Baichuan-Omni-1.5则达到了72.2%。Baichuan-Omni-1.5在语言模态上的成功,主要归功于我们在训练策略上的调整以及多模态训练数据的平衡比例,其中保持了适当比例的纯文本数据。这些结果表明,我们的数据合成与平衡方法,以及多阶段的训练策略,能够有效解决在多模态训练过程中纯语言任务性能下降的问题。此外,相较于最新的开源多模态模型MiniCPM-o 2.6,Baichuan-Omni-1.5在中文基准测试中展示了明显的优势,例如CMMLU(63.3%对75.5%)和C-Eval(61.5%对73.1%),并且在通用基准测试中也大大超过了MiniCPM-o 2.6,MMLU(65.3%对72.2%)和AGIEval(50.9%对54.4%)。这些结果表明,相较于当前的全模态模型,这些模型在训练非文本模态数据后可能会导致文本理解能力下降,而我们的模型在理解纯文本方面依然保持强大能力。

如表7和表8所示,显然,我们的模型在大多数基准测试中优于最新的开源模型VITA-1.5和MiniCPM-o 2.6。例如,与最近的MiniCPM-o 2.6相比,我们的模型在包括MMBench、SEED-IMG、MME和MMMU在内的十个基准测试中的六个上表现更好,这些基准测试要求具备专家级的感知和推理能力。这表明,我们的全模态模型已经处于开源模型的前沿。此外,与其他非全模态模型相比,Baichuan-Omni-1.5也取得了相当或更优的表现。例如,与MiniCPM-Llama3-V 2.5相比,我们的模型在大多数视觉问答(VQA)任务中表现更好。总体而言,与Qwen2-VL-7B相比,我们的模型在各类图像理解基准测试中的表现相当。我们的模型在MMBench-CN(81.9%对83.6%)、MMMU(52.7%对53.9%)、MathVista-mini(58.2%对63.6%)和ChartQA(83.0%对84.9%)等方面取得了更好的表现。此外,值得注意的是,在MMBench-EN/CN和OCRBench上,我们的模型已经超越了像GPT4o这样的闭源模型。

总结

在这项工作中,我们介绍了Baichuan-Omni-1.5,一个全模态模型,代表了朝着开发一个涵盖所有人类感官的综合框架迈出的重要一步。通过使用高质量的多模态数据以及多阶段的全模态预训练和微调策略,Baichuan-Omni-1.5在处理视频、图像、文本和音频理解方面取得了优异的表现。Baichuan-Omni-1.5的关键特点包括:(1) 强大的纯文本和多模态理解能力;(2) 全模态输入(文本、图像、视频、文本)和双模态输出(文本和音频)的端到端并行处理;(3) 在医疗场景中的卓越表现;以及(4) 高质量的可控音频生成。

尽管这些结果很有前景,但每种模态的基础能力仍有相当大的改进空间。即:(1) 增强文本理解能力;(2) 支持更长的视频帧理解;以及(3) 改进音频理解和生成,不仅能识别人类声音,还能识别自然环境中的声音,如流水声、鸟鸣声和碰撞声等。

我们的未来研究将专注于完善这些领域,以确保开发出更复杂、更通用的模型,能够理解和与复杂环境互动。我们预计,在这些领域的持续进展将对实现人工通用智能(AGI)的更广泛目标做出重要贡献。

Qwen2 Lora LLM微调训练教程

Github:https://github.com/datawhalechina/self-llm/tree/master

 本项目是一个围绕开源大模型、针对国内初学者、基于 Linux 平台的中国宝宝专属大模型教程,针对各类开源大模型提供包括环境配置、本地部署、高效微调等技能在内的全流程指导,简化开源大模型的部署、使用和应用流程,让更多的普通学生、研究者更好地使用开源大模型,帮助开源、自由的大模型更快融入到普通学习者的生活中。

  本项目的主要内容包括:

  1. 基于 Linux 平台的开源 LLM 环境配置指南,针对不同模型要求提供不同的详细环境配置步骤;
  2. 针对国内外主流开源 LLM 的部署使用教程,包括 LLaMA、ChatGLM、InternLM 等;
  3. 开源 LLM 的部署应用指导,包括命令行调用、在线 Demo 部署、LangChain 框架集成等;
  4. 开源 LLM 的全量微调、高效微调方法,包括分布式全量微调、LoRA、ptuning 等。

环境配置

在完成基本环境配置和本地模型部署的情况下,你还需要安装一些第三方库,可以使用以下命令:

python -m pip install --upgrade pip
# 更换 pypi 源加速库的安装
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

pip install modelscope==1.18.0
pip install transformers==4.44.2
pip install streamlit==1.24.0
pip install sentencepiece==0.2.0
pip install accelerate==0.34.2
pip install datasets==2.20.0
pip install peft==0.11.1

模型下载

使用 modelscope 中的 snapshot_download 函数下载模型,第一个参数为模型名称,参数 cache_dir 为模型的下载路径。

在 /root/autodl-tmp 路径下新建 model_download.py 文件并在其中输入以下内容,粘贴代码后请及时保存文件,如下图所示。并运行 python /root/autodl-tmp/model_download.py 执行下载,模型大小为 15GB,下载模型大概需要 5 分钟。

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import os
model_dir = snapshot_download('qwen/Qwen2.5-7B-Instruct', cache_dir='/root/autodl-tmp', revision='master')

指令集构建

LLM 的微调一般指指令微调过程。所谓指令微调,是说我们使用的微调数据形如:

{
  "instruction": "回答以下用户问题,仅输出答案。",
  "input": "1+1等于几?",
  "output": "2"
}

其中,instruction 是用户指令,告知模型其需要完成的任务;input 是用户输入,是完成用户指令所必须的输入内容;output 是模型应该给出的输出。

即我们的核心训练目标是让模型具有理解并遵循用户指令的能力。因此,在指令集构建时,我们应针对我们的目标任务,针对性构建任务指令集。例如,在本节我们使用由笔者合作开源的 Chat-甄嬛 项目作为示例,我们的目标是构建一个能够模拟甄嬛对话风格的个性化 LLM,因此我们构造的指令形如:

{
  "instruction": "你是谁?",
  "input": "",
  "output": "家父是大理寺少卿甄远道。"
}

我们所构造的全部指令数据集在根目录下。

数据格式化

Lora 训练的数据是需要经过格式化、编码之后再输入给模型进行训练的,如果是熟悉 Pytorch 模型训练流程的同学会知道,我们一般需要将输入文本编码为 input_ids,将输出文本编码为 labels,编码之后的结果都是多维的向量。我们首先定义一个预处理函数,这个函数用于对每一个样本,编码其输入、输出文本并返回一个编码后的字典:

def process_func(example):
    MAX_LENGTH = 384    # Llama分词器会将一个中文字切分为多个token,因此需要放开一些最大长度,保证数据的完整性
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(f"<|im_start|>system\n现在你要扮演皇帝身边的女人--甄嬛<|im_end|>\n<|im_start|>user\n{example['instruction'] + example['input']}<|im_end|>\n<|im_start|>assistant\n", add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens
    response = tokenizer(f"{example['output']}", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]  # 因为eos token咱们也是要关注的所以 补充为1
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
    if len(input_ids) > MAX_LENGTH:  # 做一个截断
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

Qwen2 采用的 Prompt Template格式如下:

<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
你是谁?<|im_end|>
<|im_start|>assistant
我是一个有用的助手。<|im_end|>

加载 tokenizer 和半精度模型

模型以半精度形式加载,如果你的显卡比较新的话,可以用 torch.bfolat形式加载。对于自定义的模型一定要指定 trust_remote_code参数为 True

tokenizer = AutoTokenizer.from_pretrained('/root/autodl-tmp/qwen/Qwen2.5-7B-Instruct/', use_fast=False, trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained('/root/autodl-tmp/qwen/Qwen2.5-7B-Instruct/', device_map="auto",torch_dtype=torch.bfloat16)

定义 LoraConfig

LoraConfig这个类中可以设置很多参数,但主要的参数没多少,简单讲一讲,感兴趣的同学可以直接看源码。

  • task_type:模型类型
  • target_modules:需要训练的模型层的名字,主要就是 attention部分的层,不同的模型对应的层的名字不同,可以传入数组,也可以字符串,也可以正则表达式。
  • rlora的秩,具体可以看 Lora原理
  • lora_alphaLora alaph,具体作用参见 Lora 原理

Lora的缩放是啥嘞?当然不是 r(秩),这个缩放就是 lora_alpha/r, 在这个 LoraConfig中缩放就是 4 倍。

config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    inference_mode=False, # 训练模式
    r=8, # Lora 秩
    lora_alpha=32, # Lora alaph,具体作用参见 Lora 原理
    lora_dropout=0.1# Dropout 比例
)

自定义 TrainingArguments 参数

TrainingArguments这个类的源码也介绍了每个参数的具体作用,当然大家可以来自行探索,这里就简单说几个常用的。

  • output_dir:模型的输出路径
  • per_device_train_batch_size:顾名思义 batch_size
  • gradient_accumulation_steps: 梯度累加,如果你的显存比较小,那可以把 batch_size 设置小一点,梯度累加增大一些。
  • logging_steps:多少步,输出一次 log
  • num_train_epochs:顾名思义 epoch
  • gradient_checkpointing:梯度检查,这个一旦开启,模型就必须执行 model.enable_input_require_grads(),这个原理大家可以自行探索,这里就不细说了。
args = TrainingArguments(
    output_dir="./output/Qwen2.5_instruct_lora",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    logging_steps=10,
    num_train_epochs=3,
    save_steps=100,
    learning_rate=1e-4,
    save_on_each_node=True,
    gradient_checkpointing=True
)

使用 Trainer 训练

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_id,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
)
trainer.train()

加载 lora 权重推理

训练好了之后可以使用如下方式加载 lora权重进行推理:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from peft import PeftModel

model_path = '/root/autodl-tmp/qwen/Qwen2.5-7B-Instruct/'
lora_path = 'lora_path'

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 加载模型
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto",torch_dtype=torch.bfloat16)

# 加载lora权重
model = PeftModel.from_pretrained(model, model_id=lora_path, config=config)

prompt = "你是谁?"
messages = [
    {"role": "system", "content": "现在你要扮演皇帝身边的女人--甄嬛"},
    {"role": "user", "content": prompt}
]

text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

model_inputs = tokenizer([text], return_tensors="pt").to('cuda')

generated_ids = model.generate(
    model_inputs.input_ids,
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)