使用 TFLite 在边缘设备上简化大型语言模型的推理

八月 13, 2024
Quentin Khan Software Engineer
Linkun Chen Software Engineer

通过采用更智能的缓存策略,优化 XNNPack 中首次令牌的时间和峰值内存用量


XNNPack 是所有模型的默认 TensorFlow Lite CPU 推理引擎。该引擎为移动、桌面和 Web 平台带来了革命性的速度提升。XNNPack 所采用的一种优化技术会将卷积、深度卷积、转置卷积和全连接运算符的静态权重重新打包到一个针对推理计算优化的内部布局中。在推理过程中,这些重新打包后的权重将按照与处理器流水线结构相符的顺序模式得到访问。

降低推理延迟时间是以一定成本为代价的:重新打包实际上是在 XNNPack 内部创建了一个额外的权重副本。为了减少这一成本,我们已经在 XNNPack 中加入了内存缓存机制。此缓存机制使得独立运行相同模型的不同 TFLite 解释器之间可以共享打包后的权重。

我们已对 TFLite XNNPack 委托实现进行改进,以解决现有缓存的一些缺陷。


1. 缓存存放于匿名内存中,这在内存压力较大时会导致数据交换到磁盘,从而造成性能下降。

2. 每当进程启动时,都需要重新打包初始权重。

3. 由于重新打包会读取原始 TFLite 权重并将其写入新的缓冲区,这会导致在打包期间出现较高的峰值内存用量。

4. 通过 XNNPack 委托启用缓存需要繁琐的步骤和仔细的生命周期管理。

5. 无法跨进程共享权重。

TFLite XNNPack delegate architecture

新的 XNNPack 缓存提供程序接口

XNNPack 已进行更新,并提供一个可让您实现权重缓存提供程序的接口。权重缓存提供程序类似于一个字典,XNNPack 会填充这个字典,并从中查询以访问打包后的缓冲区。以下是提供程序的主要函数和功能。

  • look_up:查找打包后的缓冲区键并返回唯一标识符(或为 NotFound 保留的特殊标识符),该标识符稍后可用于检索缓冲区地址。

  • reserve_space:保留可用于存储给定大小信息的缓冲区。然后您需要使用 look_up_or_insert 提交该缓冲区。

  • look_up_or_insert:检查缓存提供程序中是否存在与给定键匹配的缓冲区。如果没有,给定的数据将被提交到缓存提供程序。此函数还会返回可用于检索缓冲区地址的标识符。

  • offset_to_addr:根据由 look_up 和 look_up_or_insert 函数返回的标识符返回缓冲区地址。

XNNPack 与权重缓存提供程序之间的交互如下图所示。

The interactions between XNNPack and the weight cache provider

在 TFLite 委托中使用 MMAP 从磁盘加载缓存

TFLite 委托现在使用这个新接口,并拥有自己的权重缓存提供程序。此提供程序能够直接将打包后的权重保存到磁盘,并从磁盘加载。TFLite 长期以来一直利用 Flatbuffer 以及与文件关联的 (file-backed) 内存映射。我们通过利用相同的技术解决了之前存在的问题,并获得了以下优势。


消除重新打包的开销。

将打包后的权重保留在磁盘上,可以避免每次加载模型时重新进行打包这一耗时的过程。这意味着启动延迟时间和峰值内存用量都会显著减少。即使在初次构建时,这种方法也能实现打包数据的重复信息删除,并通过避免再次重新打包相同的数据,进一步提高打包效率。


改善内存管理。

mmap 能够利用操作系统的虚拟内存管理功能,优化系统的整体内存用量和性能。对我们而言,这特别适用于随机访问大量只读文件的情况,比如神经网络操作中的常量权重。

将打包后的数据存储在磁盘上后,XNNPack 缓存不再依赖可能在内存压力下导致性能问题的匿名内存,而是利用操作系统的虚拟内存管理功能来实现更顺畅的操作。

通过消除在文件系统和内存之间复制数据的需要,mmap 显著减少了开销并加快了访问速度。

您可以从 mmap 的手册页面和其他相关阅读材料中找到更多关于文件映射和内存用量的信息。


支持跨进程协作。

基于 mmap 的文件加载为实现多个进程之间的无缝权重共享奠定了基础,因为每个进程的虚拟地址空间都映射到相同的物理内存页面。这不仅减少了整体内存占用(因为多个进程共享相同的内存),而且还加速了各个进程中的模型加载。

mmap-based file loading architecture

简化面向用户的 API。

用户无需在整个应用程序生命周期中设置和管理缓存对象,只需提供缓存文件的路径即可。

std::unique_ptr<tflite::Interpreter> interpreter;
// 设置 XNNPack 委托的选项。
TfLiteXNNPackDelegateOptions xnnpack_options = TfLiteXNNPackDelegateOptionsDefault();
xnnpack_options.weight_cache_file_path = "/tmp/cache_file.xnn_cache";
// 创建 XNNPack 委托并将其应用于 TFLite 解释器。
// 静态权重将在首次运行时被打包并写入 weights_cache。
// 在之后的所有运行中,打包后的权重将被自动加载。
TfLiteDelegate* delegate = TfLiteXNNPackDelegateCreate(&xnnpack_options);
interpreter->ModifyGraphWithDelegate(delegate);

维护缓存完整性

为了确保推理的准确性和高效性,关键是要在特定条件下使 XNNPack 缓存失效:

模型演变:如果模型的权重或结构发生变化,缓存的数据将过时,这时您必须使其失效。这意味着您需要删除提供的缓存路径下的文件。

XNNPack 升级:XNNPack 内部打包算法的更新可能会导致缓存的权重不兼容,这就需要重新计算缓存。幸运的是,XNNPack 能够检测这种情况,并自动替换现有的缓存。

总之,任何可能影响 XNNPack 打包或使用权重的方式的修改都应该导致缓存失效。


基准

会话初始化主要受权重打包过程的影响。对于 LLM,多个子图会重复使用相同的权重。构建缓存的速度更快,因为重复信息删除机制可避免多次打包相同的权重。更标准的模型(如稳定扩散模型)没有重复信息删除机制,因此初始化时间略高是由于需要将缓存保存到磁盘上。从第 2 次运行开始,重新加载缓存使得所有情况下的初始化时间都将显著减少,仅占之前时间的一小部分。

会话初始化的改进自然而然地影响到了 LLM 返回第一个令牌的时间,在基准测试中此时间大约缩短了一倍。

缓存实现带来的内存优化效果也是显而易见的。得益于重复信息删除机制,LLM 的峰值常驻内存大小(Resident Set Size,简称 RSS)得到了降低。对于无法从重复信息删除机制中受益的其他模型,内存用量没有变化。重新加载缓存进一步降低了峰值 RSS,因为 TFLite 原始模型不再被读取,因此不会被加载到内存中。


Pixel 8 Pro 上的 Gemma 2B

Benchmarks - Gemma 2B on a Pixel 8 Pro

Pixel 8 Pro 上的 Phi2

Benchmarks - Phi2 on a Pixel 8 Pro

Pixel 8 Pro 上的稳定扩散模型

Stable Diffusion on a Pixel 8 Pro

研究展望

当前的缓存机制与文件系统紧密相关。我们希望能够独立地利用数据重复信息删除机制,应对那些用户更愿意使用传统的分配内存,而不是与文件关联的映射的情况。mmap 允许创建匿名映射,这将使我们能够重用大部分现有的实现。