Llama2推理引擎压缩至1277字节汇编
「开发者将Llama2推理引擎压缩至1277字节x86汇编代码,可直接从引导扇区运行,展示了极致的代码精简技术。」
在AI模型推理领域,效率和资源占用一直是开发者关注的焦点。近日,一个名为SectorLlama的开源项目在Hacker News上引起了广泛关注。该项目实现了一个完整的Llama2推理引擎,但令人惊讶的是,其全部代码仅占用1277字节的x86实模式汇编代码。这一成果不仅展示了极致的编程技巧,也为AI在资源受限环境下的应用提供了新的思路。
根据项目介绍,SectorLlama能够直接从磁盘引导扇区启动,加载量化后的模型,并在任何操作系统加载之前生成文本。这意味着,它可以在没有操作系统支持的环境下独立运行,例如在嵌入式系统或老旧硬件上。目前,该项目支持运行stories260K模型,这是一个经过专门训练的小型模型,拥有26万个参数,包含5层、8个注意力头,词汇表大小为512个token。
为了实现如此极致的代码压缩,SectorLlama采用了多种优化策略。首先,项目包含一个Python脚本(quantize.py),负责将模型打包成一种自定义的二进制格式。这种格式专为最小化解码开销而设计:权重被量化为int8,并采用全局absmax缩放;指数函数(exp)和SiLU激活函数的查找表被预先计算并直接嵌入代码中;此外,Q/K/V和gate/up权重矩阵被融合,使得汇编代码只需调用一次矩阵乘法,而不是三次。这些措施大大减少了代码体积和运行时计算量。
在内存管理方面,项目通过将CPU切换到“非真实模式”(unreal mode),使得KV缓存能够存储在更大的内存空间中。此前,KV缓存曾被量化为int8以最小化内存使用,但在追求极致代码精简的过程中,这一优化被移除。目前,采样仅采用贪婪argmax方式,项目作者表示理论上还有空间实现更复杂的采样技术,但当前目标是最小化代码体积。
值得注意的是,SectorLlama的代码编写以“代码高尔夫”(code golfing)为原则,即追求最短的代码长度,因此性能和精度并非最优。模型架构和提示词(prompt)也被硬编码在代码中。项目作者指出,理论上可以通过修改代码来加载更大的模型(如stories15M),但这需要切换到保护模式或非真实模式。
该项目在GitHub上开源,并欢迎汇编语言高手贡献代码,进一步缩小二进制体积。项目明确禁止使用自解压技巧或将推理逻辑移至Python脚本等“作弊”手段。SectorLlama的诞生,不仅是对编程极限的挑战,也为AI推理在极端环境下的应用提供了宝贵的参考,展示了在极小资源下实现AI推理的可能性。
来源:Heooo AI工具导航