文章目录
- DeepSeek 优化方式
- MoE架构的稀疏性
- 技术实现
- 基于计算强度的offload策略
- 引入CPU和GPU的高性能算子
- CUDA Graph的改进和优化
- 灵活高效的推理实验平台
这是系列博客,记录了我学习DeepSeek V3/R1时的学习笔记。其他博客:
- DeepSeek 简介
- DeepSeek R1原理
- DeepSeek V3原理
- DeepSeek 优化方式
- 在Deepseek-R1-ZERO出现前,为何无人尝试放弃微调对齐,通过强化学习生成思考链推理模型?
- MoE硬件部署
DeepSeek__8">DeepSeek 优化方式
参考资料:4090单卡跑满血版DeepSeek-R1,清华团队开源项目再破模型推理门槛
MoE架构的稀疏性
DeepSeek-R1/V3均采用了MoE(混合专家)架构,这种架构的核心是将模型中的任务分配给不同的专家模块,每个专家模块专注于处理特定类型的任务。MoE结构的模型具有很强的稀疏性,在执行推理任务的时候,每次只会激活其中一部分的模型参数。
因此,MoE架构需要大量的存储空间,但是并不需要很多的计算资源。
基于此,团队采用了GPU/CPU的异构计算划分策略:仅将非Shared部分的稀疏MoE矩阵放在CPU/DRAM上并使用llamafile提供的高速算子处理,剩余稠密部分放在GPU上使用Marlin算子处理。
- 矩阵部分划分
非 Shared 部分的稀疏 MoE 矩阵:稀疏矩阵是指矩阵中大部分元素为零的矩阵。在 MoE 模型的语境下,非 Shared 部分可能是指那些特定于某些 “专家”(expert)或者计算分支,不被其他部分共享的数据。这部分矩阵的特点是稀疏,即大部分元素为零,在存储和计算上有独特的优化需求。
剩余稠密部分:与稀疏矩阵相对,稠密矩阵指矩阵中大部分元素不为零。在整个稀疏 MoE 矩阵中,除了上述非 Shared 部分的稀疏矩阵外,剩下的就是稠密部分。
- 计算资源分配
CPU/DRAM 处理非 Shared 稀疏部分:CPU(中央处理器)和 DRAM(动态随机存取存储器)通常具有较高的通用性和相对较大的内存容量,适合处理稀疏矩阵这类不规则数据结构。因为处理稀疏矩阵时,需要频繁地进行随机访问和逻辑判断,CPU 在这方面具有优势。同时,llamafile 提供的高速算子专门用于处理这种稀疏矩阵,能够提高在 CPU/DRAM 上的处理效率。
GPU 处理稠密部分:GPU(图形处理器)擅长处理高度并行的计算任务,对于稠密矩阵的乘法、加法等运算,GPU 能够利用其大量的计算核心进行并行计算,从而快速完成计算。Marlin 算子是专门针对 GPU 上稠密矩阵处理进行优化的算子,能够更好地发挥 GPU 的性能。
在这样的情况下,同样使用4bit量化,GPU上的参数只需要24GB的显存环境,这样的消耗只需要一张4090就能满足。
此外通过这样的组合,还能够大幅度提升整个推理的性能,达到286 token/s的预填充和14 token/s的生成速度,比llama.cpp快28倍。
技术实现
具体到技术实现中,团队采用了基于计算强度的offload策略、高性能的CPU和GPU算子、CUDA Graph加速的多种方式来加速推理速度。
基于计算强度的offload策略
-
基于计算强度的offload策略
offload策略指的是将计算任务从一个计算资源转移到另一个计算资源的策略,通常是从计算能力较弱的资源转移到较强的资源,或者从当前资源转移到其他更适合处理该任务的资源。基于计算强度的offload策略就是根据任务的计算强度(即单位时间内执行的计算操作数量)来决定如何分配计算任务。计算强度高的任务优先分配到计算能力强的设备(如GPU)上执行,以充分利用设备的性能,提高整体计算效率。
-
Attention
Attention:在深度学习中,Attention机制(注意力机制)是一种让模型能够聚焦于输入数据不同部分的技术,它可以帮助模型在处理序列数据(如文本、语音)时,动态地分配对不同位置元素的关注度,从而更好地捕捉数据中的重要信息。例如在机器翻译中,Attention机制可以让模型在翻译一个单词时,关注原文中与之相关的部分,而不是同等对待所有单词。
-
MLA在官方开源v2版本中的问题
在官方开源的v2版本中,将MLA展开成MHA(Multi - Head Attention,多头注意力机制,是Attention机制的一种扩展,通过多个头并行计算注意力,然后将结果拼接起来,以捕捉更丰富的信息)进行计算。这个过程带来了两个负面结果:一是扩大了KV cache大小,KV cache是在计算Attention过程中用于存储键值对的数据缓存区,扩大它会占用更多的内存资源;二是降低了推理性能,推理性能指模型在对新数据进行预测时的速度和准确性,这里由于计算方式的改变导致速度变慢。
-
KTransformers推理框架中的优化
为了发挥MLA的真正性能,在KTransformers推理框架中,团队将矩阵直接吸收到q_proj(query投影,在Attention计算中,将输入数据投影到query空间的操作)和out_proj(output投影,将Attention计算后的结果投影到输出空间的操作)权重中。这样做的好处是压缩表示不需要解压缩来计算Attention,从而显著减少了KV缓存大小,并增加了该运算符的算术强度(单位时间内执行的算术运算数量)。算术强度的增加意味着在相同时间内可以执行更多的计算操作,进而优化了GPU计算能力的利用率。
-
计算强度与划分策略
由于MLA和Expert的计算强度相差数千倍,团队依据计算强度来决定划分策略。按照计算强度从高到低的顺序为MLA > Shared Expert > Routed Expert,优先将计算强度高的任务放入GPU,直到GPU的计算资源放不下为止。这样可以确保GPU始终在处理最能发挥其性能的高计算强度任务,提高整体计算效率。
引入CPU和GPU的高性能算子
在CPU算子中,团队使用llamafile作为CPU内核,使用expert并行和其他优化,组成高性能算子框架CPUInfer。此外增加多线程、任务调度、负载均衡、NUMA感知等优化。
在GPU算子的使用上,团队引入Marlin算子作为GPU计算的内核,它能够非常高效地进行量化后的矩阵计算,和torch这些计算量化后的矩阵乘法的库相比,使用Marlin算子完成在GPU上面的计算大概可以达到3.87倍的理想加速效果。
CUDA Graph的改进和优化
为了平衡推理性能和框架本身的易用性/可扩展性,基于Python构建KTransformers框架,同时使用CUDA Graph降低Python调用开销是一个必然的选择。
KTransformers中使用CUDA Graph过程中尽可能地减少了CPU/GPU通讯造成的断点,在CUDA Graph中掺杂和CPU异构算子通讯,最终实现一次decode仅有一个完整的CUDA Graph调用的结果。
灵活高效的推理实验平台
值得关注的是,KTransformers不止是一个固定的推理框架,也不只能推理DeepSeek的模型,它可以兼容各式各样的MoE模型和算子,能够集成各种各样的算子,做各种组合的测试。
此外还同时提供了Windows、Linux的平台的支持,方便运行。
GitHub 地址:https://github.com/kvcache-ai/ktransformers
具体技术细节指路:https://zhuanlan.zhihu.com/p/714877271