王全东-SUBLLM新架构文本下采样机制革新大语言模型效率

SUBLLM新架构：文本下采样机制革新大语言模型效率 王全东小米大模型团队 演讲嘉宾 王全东 小米大模型团队大模型高级算法工程师 中国科学院声学研究所博士、美国佐治亚理工访问学者、中科院认证高级工程师，长期从事大语言模型、多模态、语音识别等领域研究，曾获多项顶会竞赛冠亚军奖项，已发表顶会论文十余篇，拥有专列多项。深度参与了小米自研大模型从0到1的研发过程，荣获2024年度CCF计算机应用创新技术一等奖。近期和DanielPovey等提出SUBLLM新架构，被量子位等科技媒体报道。 1.长文本模型的技术挑战2.SUBLLM架构3.主要实验结果4.分析与讨论5.总结与展望 目录 长文本模型的技术挑战PART 01 长文本模型的技术挑战 长文本需求旺盛 长文本模型的技术挑战 长文本模型结构： Decoder only Transformer结构：Llama类似结构，attention的平方复杂度 长文本模型的技术挑战 训练成本高：attention的平方复杂度 长文本模型的技术挑战 模型窗长扩展方法：数据方向，训练成本不高1. Data Engineering for Scaling Language Models to 128K Context 长文本模型的技术挑战 模型窗长扩展方法：位置编码方向，训练成本不高2.YaRN: Efficient Context Window Extension of Large Language ModelsNTK-by-parts +温度控制 长文本模型的技术挑战 模型窗长扩展方法：位置编码方向，训练成本不高3.PoSE: Efficient Context Window Extension of LLMsvia Positional Skip-wise Training 长文本模型的技术挑战 训练成本主要在8k预训练阶段：attention的平方复杂度 长文本模型的技术挑战 Decoder only Transformer长文本模型训练infra开发1.DeepSpeed Ulysses： 优势：对Attention的实现不敏感，适合各种attention方法劣势：序列并行度不能超过头数 长文本模型的技术挑战 Decoder only Transformer长文本模型训练infra开发2. Ring-attention：“大号”的flash attention优势：并行度的扩展性较好劣势：对Attention变种不友好，eg.Sparse Attention 长文本模型的技术挑战 推理成本高：attention的平方复杂度推理速度角度，200k比20k贵28倍，比2k贵112倍 长文本模型的技术挑战 其他长文本模型结构： 1.Infini-Transformer：长期压缩记忆和局部因果注意力attention 长文本模型的技术挑战 其他长文本模型结构： 2. MEGALODON：继承MEGA(带有门控注意力的指数移动平均)并改进 长文本模型的技术挑战 加速方法很多备受关注SUBLLM应运而生 SUBLLM架构PART 02 SUBLLM架构 目标: 1.开发一种优化资源使用的架构，同时保持模型能力不变。 2.区分重要token和不重要token，重要token占主要算力 3.兼容现有attention based模型生态，模型广泛应用的关键 SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 语音识别中语音分帧后降采样4倍 SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 语音识别Zipformer，最高降采样16倍 SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 文本序列是否存在冗余？Yes! SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 文本序列中的token是否同等重要？No! CoLT5模型结构by Google, 2023 SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 SUBLLM：采样(Sub-sampling)-上采样(Up-sampling)-旁路(Bypass) LLM SUBLLM架构 受语音领域启发：语音信号下采样减少冗余保留必要信息 SUBLLM：采样(Sub-sampling)-上采样(Up-sampling)-旁路(Bypass) LLM SUBLLM架构 Learnable Subsampling Module1.通过去除不太重要的tokens来缩短序列长度 2.使用Score层来衡量token的重要性 序列采样，index选择预设采样保留比例d Score层给每个token打分再得到weight SUBLLM架构 Learnable Subsampling Module妙处1：解决训练和推理的不一致性 Inference Mode：推理时看不到未来token，设定阈值为0，正数分值的token保留 Balancer：通过对梯度进行惩罚来限制score的正数比例=预设下采样保留比例只保留正数分值的token SUBLLM架构 UpsamplingModule通过将采样后的tokens与原始序列合并来恢复序列长度 妙处2：用减法使得token的完全选择可导，使score层有token判别能力 实际基本为0，不影响weight大小，但是有梯度 SUBLLM架构 Subsampling andUpsamplingModule SUBLLM架构 Bypass Module：1.下采样前和上采样后的序列加权求和，per channel2.增强训练的收敛性和稳定性。 SUBLLM架构 对比Google: Mixture-of-Depths (MoD)：1.给重要token分配更多算力 2. MoD每间隔一个block做一次topk的token筛选3.加速50% 主要实验结果PART 03 主要实验结果 Pre-Training CorporaSlimPajama,100倍模型大小tokens Pre-Training DetailsSUBLLM 1.3B vs Llama 1.3B,SUBLLM仅仅多了8192个参数SUBLLM 2次下采样，最短的地方总保留比例仅40%BF16,Flash Attention2训练窗长2k4k 8k优化器为ScaledAdam 主要实验结果 加速比26%最大加速比31%推理加速比37%训练/推理显存减少10，1GB/GPU预训练valid loss持平Few-shot分数持平 分析与讨论PART 04 分析与讨论 预训练加速 分析与讨论 推理加速 分析与讨论 模型结构适用于不同的优化器 分析与讨论 下采样的有效性分析 总结与展望PART 05 总结与展望 SUBLLM贡献 1.提出SUBLLM新架构它结合了下采样、上采样和旁路模块动态地将资源分配给重要的token 2.提出了一种token序列子采样方法可以有效地测量token重要性分数并按预期控制分数值的分布在推理过程中实现所需的下采样保留率 3.与Llama模型相比SUBLLM在训练和推理方面分别实现了加速，训练加速34%，推理50%同时显着降低了内存成本，保持了模型能力 总结与展望 未来研究方向 1.不同的tokenizer有不同的压缩率如何确保SUBLLM结构的通用性并确定有效的保留比例 2.应用在长文本场景，比如200K窗长上，可以最低保留多少比例值得期待一下 3.多模态模型中的视觉标记具有冗余性在处理高分辨率图像时，大大降低了VLM的效率SUBLLM应用在这个领域也将具有落地前景 THANKS