“ 做了很多的优化工作,比较硬核,记录下这么个新框架用来推理moe模型的把。
https://arxiv.org/pdf/2401.14361.pdf
https://github.com/TorchMoE/MoE-Infinity
这篇文章介绍了一个名为MOE-INFINITY的系统,它是一个针对Mixture-of-Experts(MoE)模型的高效服务系统。MOE模型在AI服务中越来越受欢迎,如聊天机器人和机器翻译。这些模型通常包含大量的参数,需要大量的GPU资源来部署。MOE-INFINITY旨在通过优化专家参数的卸载来减少部署成本,同时保持低延迟。
MOE-INFINITY的核心算法原理包括以下几个关键部分:
- 序列级专家激活追踪(Sequence-level Expert Activation Tracing) :
- MOE-INFINITY通过追踪每个输入序列在推理过程中专家的激活模式,来识别稀疏激活和MoE推理的时间局部性。这种追踪方法能够捕捉到在序列级别上专家激活的稀疏性和时间局部性,这对于传统的基于模型的追踪方法来说是不可见的。
- 激活感知专家预取(Activation-Aware Expert Prefetching) :
- 系统根据当前的激活模式(EAM)和激活模式集合(EAMC)来评估专家的预取优先级。这种方法能够根据预测的激活比率和当前执行层的接近程度来优先预取专家,从而显著减少预取流量并及时开始预取。
- 激活感知专家缓存(Activation-Aware Expert Caching) :
- MOE-INFINITY采用一种新的缓存策略,它优先考虑在当前序列处理中显示出更高激活比率的专家,并倾向于缓存位于MoE初始层的专家,因为这些专家不太可能从预取中受益。这种MoE特定的缓存设计可以显著提高缓存命中率,有效提升专家缓存的性能。
- 系统实现和优化(System Implementation and Optimizations) :
- MOE-INFINITY在多GPU服务器性能优化方面做了很多工作,包括NUMA(Non-Uniform Memory Access)优化、内存池优化以及为每个NUMA节点分配预取线程等,以提高并行PCI链接的利用率。
- 集群部署和专家并行(Cluster Deployment and Expert Parallelism) :
- MOE-INFINITY支持通过专家并行来部署MoE模型,这允许系统在多服务器上分布式执行MoE模型,并通过战略性地分配专家到不同服务器来优化通信成本。
- 序列级专家激活追踪的实践考虑(Practical Concerns of Sequence-Level Expert Activation Tracing) :
- 文章还讨论了在实际部署中进行序列级专家激活追踪时的一些实际问题,如EAMC(Expert Activation Matrix Collection)的容量限制、处理分布偏移(distribution shift)的机制等。
MOE-INFINITY通过这些算法原理,能够在GPU集群中高效地服务MoE模型,显著降低了延迟并减少了部署成本。