基于Core ML的本地Llama 3.1

许多应用开发者希望在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，利用用户硬件实现成本效益推理。