我们很高兴能让开发者将全新的设备端生成式 AI 模型无缝应用到边缘设备上。为了满足这一需求，我们推出了 AI Edge Torch 生成式 API。该 API 可使开发者使用 PyTorch 创作高性能的 LLM，并利用 TensorFlow Lite (TFLite) 运行时进行部署。本文是 Google AI Edge 开发者版本系列博文中的第二篇文章。该系列的第一篇文章对 Google AI Edge Torch 进行了简单介绍，该产品可以在使用 TFLite 运行时的移动设备上实现高性能的 PyTorch 模型推理。
AI Edge Torch 生成式 API 使得开发者能够在设备端实现强大的新功能，比如摘要生成、内容创造等。我们已经通过 MediaPipe LLM 推理 API，让开发者能够将一些最热门的 LLM 部署到设备上。现在，我们非常兴奋地宣布，AI Edge Torch 生成式 API 将使开发者能够以极高的性能将任何受支持的模型部署至设备端。该 API 的初始版本提供以下功能：
在本篇博文中，我们将深入探讨性能、可移植性、创作开发体验、端到端的推理流水线以及调试工具链。如需更多文档和示例，请参见此处。
我们致力于通过 MediaPipe LLM 推理 API 无缝运行一些最热门的 LLM，在此过程中，我们的团队创作了几个全手写的 Transformer，它们具有最先进的设备端性能水平（参见 MediaPipe LLM 推理 API 博客）。从这项工作中我们总结出几个关键主题：有效表示注意力机制的方式、量化的使用，以及构建高效 KV 缓存表示的重要性。生成式 API 能轻松地实现上述主题（后文会详细介绍），同时其性能还超过手写版本 90% 及以上，极大提高了开发者的工作效率。
下表展示了 3 个模型实例的关键基准测试结果：
基准测试是在设备的大核心上进行的，使用了 4 个 CPU 线程，是这些模型在所列设备上目前所知最快的 CPU 实现。
核心创作库提供常见 Transformer 模型（如仅编码器、仅解码器或编码器 - 解码器样式等）的基本模块。您能够借助此库从零开始创建模型，或者重新创作现有模型以提升性能。我们建议大多数用户选择重新创作，因为这样不需要进行训练/微调步骤。生成式 API 创作的核心优势包括：
作为示例，我们在下面展示了如何使用新的生成式 API，以约 50 行 Python 代码重新创作 TinyLLama(1.1B) 的核心功能。
步骤 1：定义模型结构
import torch
import torch.nn as nn
from ai_edge_torch.generative.layers.attention import TransformerBlock
import ai_edge_torch.generative.layers.attention_utils as attn_utils
import ai_edge_torch.generative.layers.builder as builder
import ai_edge_torch.generative.layers.model_config as cfg
class TinyLLamma(nn.Module):
def __init__(self, config: cfg.ModelConfig):
super().__init__()
self.config = config
# Construct model layers.
self.lm_head = nn.Linear(
config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
)
self.tok_embedding = nn.Embedding(
config.vocab_size, config.embedding_dim, padding_idx=0
)
self.transformer_blocks = nn.ModuleList(
TransformerBlock(config) for _ in range(config.num_layers)
)
self.final_norm = builder.build_norm(
config.embedding_dim,
config.final_norm_config,
)
self.rope_cache = attn_utils.build_rope_cache(
size=config.kv_cache_max,
dim=int(config.attn_config.rotary_percentage * config.head_dim),
base=10_000,
condense_ratio=1,
dtype=torch.float32,
device=torch.device("cpu"),
)
self.mask_cache = attn_utils.build_causal_mask_cache(
size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
)
self.config = config
步骤 2：定义模型的前向函数
@torch.inference_mode
def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
B, T = idx.size()
cos, sin = self.rope_cache
cos = cos.index_select(0, input_pos)
sin = sin.index_select(0, input_pos)
mask = self.mask_cache.index_select(2, input_pos)
mask = mask[:, :, :, : self.config.kv_cache_max]
# forward the model itself
x = self.tok_embedding(idx) # token embeddings of shape (b, t, n_embd)
for i, block in enumerate(self.transformer_blocks):
x = block(x, (cos, sin), mask, input_pos)
x = self.final_norm(x)
res = self.lm_head(x) # (b, t, vocab_size)
return res
步骤 3：映射旧模型权重
您可以使用库中的 ModelLoader API 轻松映射权重，就像这样：
import ai_edge_torch.generative.utilities.loader as loading_utils
# 此映射将旧张量名称与新模型相关联。
TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
ff_up_proj="model.layers.{}.mlp.up_proj",
ff_down_proj="model.layers.{}.mlp.down_proj",
ff_gate_proj="model.layers.{}.mlp.gate_proj",
attn_query_proj="model.layers.{}.self_attn.q_proj",
attn_key_proj="model.layers.{}.self_attn.k_proj",
attn_value_proj="model.layers.{}.self_attn.v_proj",
attn_output_proj="model.layers.{}.self_attn.o_proj",
pre_attn_norm="model.layers.{}.input_layernorm",
pre_ff_norm="model.layers.{}.post_attention_layernorm",
embedding="model.embed_tokens",
final_norm="model.norm",
lm_head="lm_head",
)
完成这些步骤后，您可以运行一些示例输入来验证重新创作的模型的数值正确性（见链接）。如果数值检查达标，您可以继续进行转换与量化步骤。
利用 ai_edge_torch 提供的转换 API，您可以使用相同的 API 将（重新创作的）Transformer 模型转换为高度优化的 TensorFlow Lite 模型。转换过程包含以下关键步骤：
1) 导出到 StableHLO。通过 Torch Dynamo 编译器对 PyTorch 模型进行跟踪和编译，生成带有 Aten 算子的 FX 图，然后由 ai_edge_torch 将其降低为 StableHLO 图。
2) ai_edge_torch 在 StableHLO 上执行进一步的编译器传递，包括算子融合/折叠等，并生成高性能的 TFLite FlatBuffer（包含用于 SDPA、KVCache 的融合算子）。
核心生成式 API 库还提供一组量化 API，涵盖常见的 LLM 量化配方。这些配方作为额外参数传递给 ai_edge_torch 转换器 API，再由该 API 自动完成量化。我们将在未来版本中提供更多量化模式。
我们发现，在实际的推理场景中，LLM 为了达到最佳的服务性能，其模型需要具备明显区分的（解聚合）推理功能（预填充、解码）。这一定程度上源于此观察：预填充/解码可能需要采用不同形状的张量，其中预填充受计算限制，而解码则受内存限制。对于大型 LLM 而言，避免在预填充/解码之间重复模型权重至关重要。我们利用 TFLite 和 ai_edge_torch 中现有的多签名特性来实现这一点，使您能够轻松地为模型定义多个入口点，如下所示。
def convert_tiny_llama_to_tflite(
prefill_seq_len: int = 512,
kv_cache_max_len: int = 1024,
quantize: bool = True,
):
pytorch_model = tiny_llama.build_model(kv_cache_max_len=kv_cache_max_len)
# 用于在转换过程中跟踪模型图形的张量。
prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
prefill_input_pos = torch.arange(0, prefill_seq_len)
decode_token = torch.tensor([[0]], dtype=torch.long)
decode_input_pos = torch.tensor([0], dtype=torch.int64)
# 设置模型量化。
quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None
edge_model = (
ai_edge_torch.signature(
'prefill', pytorch_model, (prefill_tokens, prefill_input_pos)
)
.signature('decode', pytorch_model, (decode_token, decode_input_pos))
.convert(quant_config=quant_config)
)
edge_model.export(f'/tmp/tiny_llama_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite')
在性能调查阶段，我们发现了几个对于提升 LLM 性能至关重要的方面：
1) 高性能的 SDPA 和 KVCache：我们发现，如果不进行充分的编译器优化/融合，转换后的 TFLite 模型会因为这些函数中的算子粒度问题而无法展现出优越的性能。为解决这一问题，我们引入了高层函数边界和 StableHLO 复合算子。
2) 利用 TFLite 的 XNNPack 代理以进一步加速 SDPA：确保重负载的 MatMul/矩阵 - 向量计算得到良好优化至关重要。XNNPack 库能在广泛的移动 CPU 上以出色的性能完成这些基础运算。
3) 避免冗余计算：若静态形状模型在预填充阶段具有较长的固定输入消息大小，或在解码阶段具有较大的固定序列长度，则可能导致计算量大于最小必要计算量。
4) 运行时内存消耗：我们在 TFLite 的 XNNPack 代理中引入了权重缓存/预先打包机制，以显著降低内存的峰值用量。
LLM 推理通常涉及众多预处理/后处理步骤及复杂的编排工作，例如：令牌化、采样以及自回归解码逻辑。为此，我们提供了基于 MediaPipe 的解决方案，以及纯 C++ 推理示例。
MediaPipe LLM 推理 API 是高阶 API，支持通过提示输入/提示输出接口进行 LLM 推理。该 API 在后台处理实现 LLM 流水线的所有复杂操作，使部署变得更加简便和流畅。如要使用 MP LLM 推理 API 进行部署，您需要确保使用预期的预填充和解码签名来转换模型，并如下面代码所示创建一个软件包：
def bundle_tinyllama_q8():
output_file = "PATH/tinyllama_q8_seq1024_kv1280.task"
tflite_model = "PATH/tinyllama_prefill_decode_hlfb_quant.tflite"
tokenizer_model = "PATH/tokenizer.model"
config = llm_bundler.BundleConfig(
tflite_model=tflite_model,
tokenizer_model=tokenizer_model,
start_token="<s>",
stop_tokens=["</s>"],
output_filename=output_file,
enable_bytes_to_unicode_mapping=False,
)
llm_bundler.create_bundle(config)
同时，我们也为您提供了一个易于使用的 C++ 示例（没有 MediaPipe 依赖项），以展示如何运行端到端的文本生成实例。开发者可以将此示例作为起点，将导出的模型与他们独特的生产流水线和需求相结合，从而实现更好的定制化和灵活性。
由于核心推理运行时基于 TFLite，因此整个流水线可以轻松集成到您的 Android（已包含在 Google Play 中）或 iOS 应用中，无需进行任何修改。这意味着从新的生成式 API 转换的模型只需添加几个自定义算子依赖项，就能立即部署。在未来的版本中，我们将为 Android 和 iOS 引入 GPU 支持，同时还将支持 ML 加速器（如 TPU、NPU）。
最近发布的模型探索器是一个用于可视化大型模型（如 Gemma 2B）的实用工具。通过层次化查看和并排比较功能，开发者可以轻松地将原始模型/重新创作的模型/转换后的模型版本可视化。如要了解关于此工具的详细信息，以及如何可视化基准信息以进行性能调优，请参阅此博文。
下面是我们在创作 PyTorch TinyLlama 模型时使用该工具的一个示例。我们并排显示了 PyTorch export() 模型与 TFLite 模型。借助模型探索器，我们可以轻松对比每个层（例如 RMSNorms 和 SelfAttention）的表达情况。
AI Edge Torch 生成式 API 是为 MediaPipe LLM 推理 API 预构建优化模型的强大补充，适用于希望在设备上运行自己的生成式 AI 模型的开发者。我们会在接下来的几个月继续带来更新，包括对 Web 的支持、改进的量化技术，以及 CPU 之外的更广泛计算支持。同时，我们也正积极探索更好的框架集成方案。
目前发布的是该库的早期预览版，该版本仍处于实验阶段，其目的是与开发者社区进行互动。使用预览版 API 时，请预期其功能可能会有所不同，存在不完善之处，并且对量化和模型的支持有限。但我们已经在 GitHub repo 中为大家提供了很多用于上手的内容，欢迎大家测试和体验，并随时和我们分享 PR、问题和功能请求。
在本系列的第 3 篇文章中，我们将深入探讨模型探索器可视化工具，了解该工具如何帮助开发者们可视化、调试和探索模型。
这项工作是 Google 多个职能团队之间的合作成果。我们要感谢所有为这项工作做出贡献的团队成员：Aaron Karp、Advait Jain、Akshat Sharma、Alan Kelly、Andrei Kulik、Arian Afaian、Chun-nien Chan、Chuo-Ling Chang、Cormac Brick、Eric Yang、Frank Barchard、Gunhyun Park、Han Qi、Haoliang Zhang、Ho Ko、Jing Jin、Joe Zoe、Juhyun Lee、Kevin Gleason、Khanh LeViet、Kris Tonthat、Kristen Wright、Lin Chen、Linkun Chen、Lu Wang、Majid Dadashi、Manfei Bai、Mark Sherwood、Matthew Soulanille、Matthias Grundmann、Maxime Brénon、Michael Levesque-Dion、Mig Gerard、Milen Ferev、Mohammadreza Heydary、Na Li、Paul Ruiz、Pauline Sho、Pei Zhang、Ping Yu、Pulkit Bhuwalka、Quentin Khan、Ram Iyengar、Renjie Wu、Rocky Rhodes、Sachin Kotwani、Sandeep Dasgupta、Sebastian Schmidt、Siyuan Liu、Steven Toribio、Suleman Shahid、Tenghui Zhu、T.J. Alumbaugh、Tyler Mullen、Weiyi Wang、Wonjoo Lee、Yi-Chun Kuo、Yishuang Pang、Yu-hui Chen、Zoe Wang，以及 Zichuan Wei。
