Apple Machine Learning Research
·
基于Core ML的本地Llama 3.1
💡
原文英文，约5800词，阅读约需22分钟。
📝
内容提要
许多应用开发者希望在Apple硅上本地运行大型语言模型（LLMs），以提高推理效率和保护用户隐私。本文介绍了如何优化和部署Llama-3.1-8B-Instruct模型，利用Apple的Core ML框架实现约33个tokens/s的解码速度。通过引入状态化的键值缓存和4位量化，显著提升了模型性能，适用于其他基于变换器的LLMs。
🎯
关键要点
- 许多应用开发者希望在Apple硅上本地运行大型语言模型（LLMs），以提高推理效率和保护用户隐私。
- 本文介绍了如何优化和部署Llama-3.1-8B-Instruct模型，利用Apple的Core ML框架实现约33个tokens/s的解码速度。
- 模型必须经过仔细优化，以有效利用可用的系统资源，因为LLMs通常对内存和处理能力有很高的需求。
- 使用Core ML框架和优化措施，Llama-3.1-8B-Instruct模型可以在M1 Max上本地运行。
- 基线模型导出和性能评估是优化过程的第一步，了解模型的基本性能和改进空间。
- 导出PyTorch模型到Core ML需要定义一个薄包装器，并使用Core ML工具进行转换。
- 模型输入包括token化的文本序列和二进制注意力掩码，输出为每个token的原始未归一化概率分数。
- 模型执行分为“提示”和“扩展”两个阶段，使用贪婪采样策略生成下一个token。
- 评估模型性能的指标包括提示延迟和扩展吞吐量，后者反映了模型生成输出token的速度。
- 基线模型的扩展吞吐量较低，主要由于静态输入形状和缺乏键值缓存。
- 通过引入灵活的输入形状和状态化的键值缓存，可以显著提高模型性能。
- 使用融合的缩放点积注意力（SDPA）和块状int4权重量化进一步优化模型性能。
- 状态化的键值缓存允许模型在每次预测调用后更新缓存，减少了数据复制的开销。
- 块状int4量化将模型大小从16GB减少到4.2GB，解码速度提高到约33 tokens/s。
- 使用Core ML框架和所述优化，开发者可以在Apple硅上本地部署LLMs，利用用户硬件实现成本效益推理。
➡️
