通过 AI Edge Torch 生成式 API 在设备上使用自定义大语言模型

我们很高兴地发布 AI Edge Torch 生成式 API，它能将开发者用 PyTorch 编写的高性能大语言模型 (LLM) 部署至 TensorFlow Lite (TFLite) 运行时，从而无缝地将新的设备端生成式 AI 模型部署到边缘设备上。本文是 Google AI Edge 博客连载的第二篇。上一篇文章为大家介绍了 Google AI Edge Torch，该产品可以在使用 TFLite 运行时的设备上实现高性能的 PyTorch 模型推理。

AI Edge Torch 生成式 API 使开发者能够在设备上引入强大的新功能，例如摘要生成、内容生成等。我们之前已经通过 MediaPipe LLM Inference API 让开发者们能够将一些最受欢迎的 LLM 部署到设备上。现在，我们很高兴能进一步拓展对模型的支持范围，并让大家部署到设备，而且具备优秀的性能表现。今天发布的 AI Edge Torch 生成式 API 是初始版本，提供以下功能:

• 简单易用的模型创作 API，支持自定义 Transformer。
• 在 CPU 上性能表现出色，并即将支持 GPU 和 NPU。
• 作为 AI Edge Torch 的扩展，支持 PyTorch。
• 完全兼容现有的 TFLite 部署流程，包括量化和运行时。
• 支持 TinyLlama、Phi-2 和 Gemma 2B 等模型。
• 兼容 TFLite 运行时和 Mediapipe LLM 运行时接口，支持 Android、iOS 和 Web。

• MediaPipe LLM Inference API

  https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference

• AI Edge Torch

  https://ai.google.dev/edge/lite/models/convert_pytorch

为了让 MediaPipe LLM Inference API 顺利支持最受欢迎的一些 LLM，我们的团队手工打造了几款在设备上拥有最佳性能的 Transformer 模型。通过这项工作，我们确定了几个主要课题: 如何有效地表示注意力机制、量化的使用以及良好的 KV 缓存。生成式 API 很好地完成了这些课题 (本文后面会具体提到)，而且依然能达到之前手写版本性能的 90% 以上，并大大提高开发速度。

• 通过 MediaPipe 和 TensorFlow Lite 在设备上运行大语言模型
  https://developers.googleblog.com/en/large-language-models-on-device-with-mediapipe-and-tensorflow-lite/

• 三种模型样本
  https://github.com/google-ai-edge/ai-edge-torch/blob/main/ai_edge_torch/generative/examples/README.md

核心创作库提供了常见 Transformer 模型 (仅编码器、仅解码器或编码-解码器等样式) 的基本构建模块。您可以用它从头开始创作模型，或重新创作现有模型以提高性能。我们建议大多数用户采用重新创作的方式，因为这样就不需要训练 / 微调的步骤了。使用生成式 API 创作的核心优势如下:

• 一组针对可转换性、性能和平台可移植性进行了优化的核心 Transformer 构建模块，可以轻松与常规 PyTorch 算子进行混合和匹配。

• 一个简单的权重重映射机制。

• 直观的量化 API。

• 支持多签名导出，包括预填充、解码或自定义签名，并能无缝接入现成的 MP 任务 / LLM Inference API。

• TinyLLama (1.1B)

  https://github.com/jzhang38/TinyLlama

步骤 1: 定义模型结构

import torchimport torch.nn as nn
from ai_edge_torch.generative.layers.attention import TransformerBlockimport ai_edge_torch.generative.layers.attention_utils as attn_utilsimport ai_edge_torch.generative.layers.builder as builderimport 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_modedef forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:B, T = idx.size()cos, sin = self.rope_cachecos = 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 itselfx = 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

您可以使用库中的 ModelLoader API 轻松映射权重，就像这样:

import ai_edge_torch.generative.utilities.loader as loading_utils

# This map will associate old tensor names with the new model.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",)

完成这些步骤后，您可以运行一些示例输入来验证重新创作过的模型的数值正确性。如果数值检查达标，您就可以继续进行后续的转换和量化操作。

• 验证重新创作的模型

  https://github.com/google-ai-edge/ai-edge-torch/blob/59946008def0ab867c2f4cd8931eaf607ac0d768/ai_edge_torch/generative/test/test_model_conversion.py#L132

通过 ai_edge_torch 提供的转换 API，您可以将 (重新创作的) Transformer 模型转换为高度优化的 TensorFlow Lite 模型。转换过程包含以下关键步骤:

1. 导出到 StableHLO。通过 torch dynamo 编译器对 PyTorch 模型进行追踪和编译，生成带有 Aten 算子的 FX 计算图，然后由 ai_edge_torch 将其降为 StableHLO 计算图。
2. ai_edge_torch 在 StableHLO 上执行进一步的编译器操作，包括算子融合 / 折叠等，生成高性能的 TFLite flatbuffer (包含用于 SDPA、KVCache 的融合算子)。

• StableHLO

  https://github.com/openxla/stablehlo

我们发现在实际推理场景中，LLM 模型需要有明确分离 (细分) 的推理函数 (预填充、解码)，才能实现最佳的服务性能。这部分基于这样的观察: 预填充 / 解码可能需要采用不同的 tensor 形状，预填充受到算力限制，而解码则受到内存限制。对于大型 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)# Tensors used to trace the model graph during conversion.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)
# Set up Quantization for model.quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else Noneedge_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 代理中引入了权重缓存 / 预打包机制，显著降低了内存的峰值使用量。

• SDPA

  https://github.com/google-ai-edge/ai-edge-torch/blob/7f52f70709bc12cf041b3b1fd4a49bc0d52c889a/ai_edge_torch/generative/layers/attention.py#L74

LLM 推理通常涉及许多预处理 / 后处理步骤和复杂的编排，例如令牌化、采样和自回归解码逻辑。为此，我们提供了基于 MediaPipe 的解决方案以及一个纯 C++ 推理示例。

• 基于 MediaPipe 的解决方案

  https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference

• 纯 C++ 推理示例

  https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative/examples/c%2B%2B

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)。

• 模型探索器

  https://ai.google.dev/edge/model-explorer

• 模型探索器: 大模型开发的计算图可视化工具

  https://research.google/blog/model-explorer/

△ 并排比较 TinyLlama PyTorch 和转换后的 TFLite

AI Edge Torch 生成式 API 是为 MediaPipe LLM Inference API 预构建优化模型的强大补充，适用于希望在设备上运行自己的生成式 AI 模型的开发者。我们会在接下来的几个月继续带来更新，包括 Web 支持、更好的量化和对 CPU 之外的硬件的支持。我们也会尝试探索更好的框架集成方案。

目前发布的是开发库的早期预览版本，该版本依然处于实验阶段，旨在与开发者社区进行开放互动。API 可能会发生变化，且存在不完善之处，并且对量化和模型的支持有限。但我们已经在 GitHub repo 中为大家提供了很多用于上手的内容，欢迎大家测试和体验，并随时和我们分享 PR、问题和功能需求。