鲲鹏腾破解国外高端硬件的依赖

说起来大语言模型现在的规模都搞到了几千亿甚至几万亿参数,MoE这种稀疏架构挺吃香,毕竟既能把容量弄大又不费劲儿。南京大学李猛博士的团队这回靠着鲲鹏昇腾科教创新孵化中心的算力支持,愣是在这方面拿下了突破。团队给MoE提了个软硬件一起干的方案,把内存占用砍掉了50%,还让推理速度快了2倍。 主要是大模型部署的时候太麻烦了,显存吃紧只能塞几个热专家进去,剩下的大量参数来回折腾内存和显存,延迟直接就上去了。而且国产化适配的呼声也越来越高。 针对这种情况,团队先是琢磨出了专家等价性的点子,然后设计了四级流水线并行架构,让CPU和GPU一起干活,参数传输压力也就小了。他们还搞了个新的路由机制,把专家分成高、中、低三类。高分的直接缓存着用,低分的要是没加载上来就用可替代的中等分专家顶上。这样一来冗余传输少了,显存占用自然也就降下来了。 为了避免加载参数慢的问题,他们还首创了一种共享专家引导的在线预取机制,不用离线训练就能把后面要用的专家给预先准备好。最后他们还弄了个双指针动态负载均衡算法,根据任务大小来分配CPU和GPU的算力。 依托华为的昇腾硬件还有CANN软件栈,他们做了个专门的加速引擎叫MoE-Ascend。这个引擎不光能管动态路由和预测预取,连算子适配都搞定了。在昇腾平台上跑起来以后不仅把显存消耗给砍了一半还多,推理速度也翻倍了。关键是把缓存命中率提升到了70%以上。 这就把原来对国外高端硬件的依赖给破解了。这套优化工具链打算开源到昇腾社区、DeepModeling社区和GitHub上让大伙儿用。未来这东西就能用在AI推理平台、智能客服还有语音生成这些场景里了。尤其是适配像NPU这种自主硬件的环境时效果特别好。