全网首篇从tensorRT-LLM MoE CUDA kernel角度理解Mixtral-8x7b的推理加速及展望

最新推荐文章于 2024-06-15 11:56:23 发布

不归牛顿管的熊猫

最新推荐文章于 2024-06-15 11:56:23 发布

阅读量1.9k

点赞数 45

文章标签：人工智能深度学习 pytorch nlp

本文链接：https://blog.csdn.net/weixin_43568400/article/details/137796873

版权

本文首发于我的公众号"AI不止算法", 原文链接位于此，欢迎关注公众号。

最近，LLM MoE非常火，Mixtral-8x7b是第一个开源的MoE，先说明一下，标题里面的“全网首篇”不能随便加，我一直在调研相关资料，发现确实没有从tensorRT-LLM角度来讲的，所以才加上了哈。

网上已经有了从模型结构和官方python源码和huggingface源码对Mixtral-8x7b作解析的优秀文章，解析得都非常好，然而，他们有两个不足点在于：

一是没有区分开训练和推理来解析，这导致很多朋友不知道哪部分是训练时候做的，哪部分是推理时候做的，看了后非常混乱，对MoE还是似懂非懂
二是多数都是停留在Pytorch python API的角度去解析，不是说不好，只是还比较偏上层，不利于读者深度理解大模型推理引擎层面该怎么样去实现并加速Mixtral-8x7b为代表的MoE模型

本文出于如下原因先讲讲MoE推理相关的内容：

训练比推理不仅内容方面更杂，而且理解起来门槛也更高，推理很适合对MoE建立起80%的认识
n本人做推理多一些

已经对MoE比较熟悉的朋友，可以直接划到最后看tensorRT-LLM的实现。

一句话讲明MoE的概念

如果大家对传统机器学习算法比较熟悉，看过李沐的统计学习方法或者吴恩达的机器学习，那么应该听过bagging、boosting等集成学习方法，MoE其实也是集成学习，相较于深度神经网络(Deep Neural Network), MoE更像是宽度神经网络，如下图所示，对于MoE的结果是多个expert的输出进行加权组合得到的，router又叫gating网络，包括一个linear和softmax，起到路由的作用，分发给不同expert权重。所以AAAI22也有一篇paper叫做“Go wider instead of deeper”
在这里插入图片描述

# moe的pytorch代码
import torch
import torch.nn as nn

class Expert(nn.Module):
    def __init__(self, input_dim):
        super(Expert, self).__init__()
        self.fc = nn.Linear(input_dim, 1)
    def forward(self, x):
        return self.fc(x)
        
class MoE(nn.Module):
    def __init__(self, input_dim, num_experts):
        super(MoE, self).__init__()
        self.experts = nn.ModuleList([Expert(input_dim) for _ in range(num_experts)])
        # gating的组成
        self.gating = nn.Sequential(
            nn.Linear(input_dim, num_experts),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        # 各个expert做forward前向推理
        expert_outputs = [expert(x) f