我們很高興地發(fā)布 AI Edge Torch 生成式 API，它能將開發(fā)者用 PyTorch 編寫的高性能大語言模型 (LLM) 部署至 TensorFlow Lite (TFLite) 運(yùn)行時(shí)，從而無縫地將新的設(shè)備端生成式 AI 模型部署到邊緣設(shè)備上。本文是 Google AI Edge 博客連載的第二篇。上一篇文章為大家介紹了 Google AI Edge Torch，該產(chǎn)品可以在使用 TFLite 運(yùn)行時(shí)的設(shè)備上實(shí)現(xiàn)高性能的 PyTorch 模型推理。

AI Edge Torch 生成式 API 使開發(fā)者能夠在設(shè)備上引入強(qiáng)大的新功能，例如摘要生成、內(nèi)容生成等。我們之前已經(jīng)通過 MediaPipe LLM Inference API 讓開發(fā)者們能夠?qū)⒁恍┳钍軞g迎的 LLM 部署到設(shè)備上?，F(xiàn)在，我們很高興能進(jìn)一步拓展對(duì)模型的支持范圍，并讓大家部署到設(shè)備，而且具備優(yōu)秀的性能表現(xiàn)。今天發(fā)布的 AI Edge Torch 生成式 API 是初始版本，提供以下功能:

簡(jiǎn)單易用的模型創(chuàng)作 API，支持自定義 Transformer。

在 CPU 上性能表現(xiàn)出色，并即將支持 GPU 和 NPU。

作為 AI Edge Torch 的擴(kuò)展，支持 PyTorch。

完全兼容現(xiàn)有的 TFLite 部署流程，包括量化和運(yùn)行時(shí)。

支持 TinyLlama、Phi-2 和 Gemma 2B 等模型。

兼容 TFLite 運(yùn)行時(shí)和 Mediapipe LLM 運(yùn)行時(shí)接口，支持 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

我們將在本文中為大家深入介紹該 API 的性能、可移植性、創(chuàng)作開發(fā)體驗(yàn)、端到端推理流水線和調(diào)試工具鏈。更具體的文檔和示例請(qǐng)查看: https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative/examples

性能表現(xiàn)

為了讓 MediaPipe LLM Inference API 順利支持最受歡迎的一些 LLM，我們的團(tuán)隊(duì)手工打造了幾款在設(shè)備上擁有最佳性能的 Transformer 模型。通過這項(xiàng)工作，我們確定了幾個(gè)主要課題: 如何有效地表示注意力機(jī)制、量化的使用以及良好的 KV 緩存。生成式 API 很好地完成了這些課題 (本文后面會(huì)具體提到)，而且依然能達(dá)到之前手寫版本性能的 90% 以上，并大大提高開發(fā)速度。

通過 MediaPipe 和 TensorFlow Lite 在設(shè)備上運(yùn)行大語言模型 https://developers.googleblog.com/en/large-language-models-on-device-with-mediapipe-and-tensorflow-lite/

下表顯示了三種模型樣本的關(guān)鍵基準(zhǔn)測(cè)試結(jié)果:

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

這些基準(zhǔn)測(cè)試是在大核上運(yùn)行，使用 4 個(gè) CPU 線程，并且使用了這些模型在所列設(shè)備上目前所知最快的 CPU 實(shí)現(xiàn)。

創(chuàng)作體驗(yàn)

核心創(chuàng)作庫提供了常見 Transformer 模型 (僅編碼器、僅解碼器或編碼-解碼器等樣式) 的基本構(gòu)建模塊。您可以用它從頭開始創(chuàng)作模型，或重新創(chuàng)作現(xiàn)有模型以提高性能。我們建議大多數(shù)用戶采用重新創(chuàng)作的方式，因?yàn)檫@樣就不需要訓(xùn)練 / 微調(diào)的步驟了。使用生成式 API 創(chuàng)作的核心優(yōu)勢(shì)如下:

一組針對(duì)可轉(zhuǎn)換性、性能和平臺(tái)可移植性進(jìn)行了優(yōu)化的核心 Transformer 構(gòu)建模塊，可以輕松與常規(guī) PyTorch 算子進(jìn)行混合和匹配。

一個(gè)簡(jiǎn)單的權(quán)重重映射機(jī)制。

直觀的量化 API。

支持多簽名導(dǎo)出，包括預(yù)填充、解碼或自定義簽名，并能無縫接入現(xiàn)成的 MP 任務(wù) / LLM Inference API。

作為示例，下面展示如何使用新的生成式 API 以約 50 行 Python 代碼重新創(chuàng)作 TinyLLama (1.1B) 的核心功能。

TinyLLama (1.1B)

https://github.com/jzhang38/TinyLlama

步驟 1: 定義模型結(jié)構(gòu)

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: 定義模型的前向函數(shù)

@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)
returnres

您可以使用庫中的 ModelLoaderAPI 輕松映射權(quán)重，就像這樣:

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",
)

完成這些步驟后，您可以運(yùn)行一些示例輸入來驗(yàn)證重新創(chuàng)作過的模型的數(shù)值正確性。如果數(shù)值檢查達(dá)標(biāo)，您就可以繼續(xù)進(jìn)行后續(xù)的轉(zhuǎn)換和量化操作。

驗(yàn)證重新創(chuàng)作的模型

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

轉(zhuǎn)換和量化

通過 ai_edge_torch 提供的轉(zhuǎn)換 API，您可以將 (重新創(chuàng)作的) Transformer 模型轉(zhuǎn)換為高度優(yōu)化的 TensorFlow Lite 模型。轉(zhuǎn)換過程包含以下關(guān)鍵步驟:

導(dǎo)出到 StableHLO。通過 torch dynamo 編譯器對(duì) PyTorch 模型進(jìn)行追蹤和編譯，生成帶有 Aten 算子的 FX 計(jì)算圖，然后由 ai_edge_torch 將其降為 StableHLO 計(jì)算圖。

ai_edge_torch 在 StableHLO 上執(zhí)行進(jìn)一步的編譯器操作，包括算子融合 / 折疊等，生成高性能的 TFLite flatbuffer (包含用于 SDPA、KVCache 的融合算子)。

StableHLO

https://github.com/openxla/stablehlo

量化

核心生成式 API 庫還提供了一組量化 API，涵蓋了常見的 LLM 量化模式。這些模式作為額外參數(shù)傳遞給 ai_edge_torch 轉(zhuǎn)換器 API，由該 API 自動(dòng)完成量化。我們會(huì)在未來的版本中提供更多的量化模式。

多簽名導(dǎo)出

我們發(fā)現(xiàn)在實(shí)際推理場(chǎng)景中，LLM 模型需要有明確分離 (細(xì)分) 的推理函數(shù) (預(yù)填充、解碼)，才能實(shí)現(xiàn)最佳的服務(wù)性能。這部分基于這樣的觀察: 預(yù)填充 / 解碼可能需要采用不同的 tensor 形狀，預(yù)填充受到算力限制，而解碼則受到內(nèi)存限制。對(duì)于大型 LLM，避免在預(yù)填充 / 解碼之間重復(fù)模型權(quán)重至關(guān)重要。我們使用 TFLite 和 ai_edge_torch 中現(xiàn)有的多簽名特性來實(shí)現(xiàn)這一點(diǎn)，使得開發(fā)者能輕松地為模型定義多個(gè)入口，如下所示:

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

針對(duì) LLM 的性能優(yōu)化

我們?cè)谛阅苷{(diào)查階段發(fā)現(xiàn)了幾個(gè)改善 LLM 性能的關(guān)鍵要素:

高性能的 SDPA 和 KVCache: 我們發(fā)現(xiàn)，如果沒有足夠的編譯器優(yōu)化 / 融合，轉(zhuǎn)換后的 TFLite 模型會(huì)因?yàn)檫@些函數(shù)中算子的粒度問題，性能不會(huì)很好。為了解決這個(gè)問題，我們引入了高級(jí)函數(shù)邊界和 StableHLO 復(fù)合算子。

利用 TFLite 的 XNNPack 代理進(jìn)一步加速 SDPA: 確保大量 MatMul / 矩陣-向量計(jì)算得到很好的優(yōu)化至關(guān)重要。XNNPack 庫能在廣泛的移動(dòng) CPU 上以出色的性能完成這些基礎(chǔ)計(jì)算。

避免不必要的計(jì)算: 靜態(tài)形狀模型如果在預(yù)填充階段有長(zhǎng)且固定的輸入消息大小，或者在解碼階段有大的固定序列長(zhǎng)度，則帶來的計(jì)算量會(huì)大于該模型需要的最小計(jì)算量。

運(yùn)行時(shí)內(nèi)存消耗: 我們?cè)?TFLite 的 XNNPack 代理中引入了權(quán)重緩存 / 預(yù)打包機(jī)制，顯著降低了內(nèi)存的峰值使用量。

SDPA

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

部署

LLM 推理通常涉及許多預(yù)處理 / 后處理步驟和復(fù)雜的編排，例如令牌化、采樣和自回歸解碼邏輯。為此，我們提供了基于 MediaPipe 的解決方案以及一個(gè)純 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

使用 MediaPipe LLM Inference API

MediaPipe LLM Inference API 是一個(gè)高級(jí) API，支持使用 prompt-in / prompt-out 接口進(jìn)行 LLM 推理。它負(fù)責(zé)處理底層所有的 LLM 復(fù)雜流水線操作，讓模型得以更輕松和順暢地部署。要使用 MediaPipe LLM Inference API 進(jìn)行部署，您需要使用給定的預(yù)填充和解碼簽名來轉(zhuǎn)換模型，并創(chuàng)建一個(gè)任務(wù)包，如下方代碼所示:

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="",
      stop_tokens=[""],
      output_filename=output_file,
      enable_bytes_to_unicode_mapping=False,
  )
llm_bundler.create_bundle(config)

在 TFLite 運(yùn)行時(shí)使用純 C++ 推理

我們還提供了一個(gè)簡(jiǎn)單易用的 C++ 示例 (無需 MediaPipe 依賴)，來展示如何運(yùn)行端到端的文本生成。如果您需要將導(dǎo)出的模型與自己獨(dú)有的生產(chǎn)流水線和需求進(jìn)行集成，這個(gè)示例是一個(gè)很好的起點(diǎn)，來幫助您實(shí)現(xiàn)更好的定制和靈活性。

跨平臺(tái)支持

由于核心推理運(yùn)行時(shí)都支持 TFLite，所以整個(gè)流水線都可以輕松集成到您的 Android (包括在 Google Play 中) 或 iOS 應(yīng)用中，無需進(jìn)行任何修改。這意味著用新的生成式 API 轉(zhuǎn)換的模型只需添加幾個(gè)自定義算子依賴即可立即部署。在未來的版本中，我們將為 Android 和 iOS 帶來 GPU 支持，并支持 ML 加速器 (TPU、NPU)。

工具

最近發(fā)布的模型探索器 (Model Explorer) 是一款很好用的工具，可用于可視化諸如 Gemma 2B 之類的大型模型。分層查看和并排比較可以讓您輕松查看和比較原始、重新創(chuàng)作和轉(zhuǎn)換后的模型。我們也準(zhǔn)備了專門的文章為您進(jìn)一步介紹該工具，以及如何通過可視化基準(zhǔn)信息來優(yōu)化模型性能。

模型探索器

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

模型探索器: 大模型開發(fā)的計(jì)算圖可視化工具

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

以下是我們?cè)诰帉?PyTorch TinyLlama 模型時(shí)使用該工具的示例。我們并排顯示了 PyTorch export() 模型與 TFLite 模型。通過使用模型探索器，我們可以輕松比較每個(gè)層級(jí) (如 RMSNorms、SelfAttention) 的表達(dá)情況。

△ 并排比較 TinyLlama PyTorch 和轉(zhuǎn)換后的 TFLite

總結(jié)以及下一步

AI Edge Torch 生成式 API 是為 MediaPipe LLM Inference API 預(yù)構(gòu)建優(yōu)化模型的強(qiáng)大補(bǔ)充，適用于希望在設(shè)備上運(yùn)行自己的生成式 AI 模型的開發(fā)者。我們會(huì)在接下來的幾個(gè)月繼續(xù)帶來更新，包括 Web 支持、更好的量化和對(duì) CPU 之外的硬件的支持。我們也會(huì)嘗試探索更好的框架集成方案。

目前發(fā)布的是開發(fā)庫的早期預(yù)覽版本，該版本依然處于實(shí)驗(yàn)階段，旨在與開發(fā)者社區(qū)進(jìn)行開放互動(dòng)。API 可能會(huì)發(fā)生變化，且存在不完善之處，并且對(duì)量化和模型的支持有限。但我們已經(jīng)在 GitHub repo 中為大家提供了很多用于上手的內(nèi)容，歡迎大家測(cè)試和體驗(yàn)，并隨時(shí)和我們分享 PR、問題和功能需求。