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

67bb3e711571e6c5f0fb404ddf23cff9.png

作者 / 首席工程师 Cormac Brick,软件工程师 Haoliang Zhang

我们很高兴地发布 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

我们将在本文中为大家深入介绍该 API 的性能、可移植性、创作开发体验、端到端推理流水线和调试工具链。更具体的文档和示例请查看:

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

性能表现

为了让 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/

下表显示了三种模型样本的关键基准测试结果:

d692b53e324cba281ba0e0ca04eec04f.png

  • 三种模型样本

    https://github.com/google-ai-edge/ai-edge-torch/blob/main/ai_edge_torch/generative/examples/README.md

这些基准测试是在大核上运行,使用 4 个 CPU 线程,并且使用了这些模型在所列设备上目前所知最快的 CPU 实现。

创作体验

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

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

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

  • 直观的量化 API。

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

作为示例,下面展示如何使用新的生成式 API 以约 50 行 Python 代码重新创作 TinyLLama (1.1B) 的核心功能。

  • TinyLLama (1.1B)

    https://github.com/jzhang38/TinyLlama

步骤 1: 定义模型结构

import torch
import torch.nn as nnfrom 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 cfgclass 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

步骤 3: 映射旧模型权重

您可以使用库中的 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

量化

核心生成式 API 库还提供了一组量化 API,涵盖了常见的 LLM 量化模式。这些模式作为额外参数传递给 ai_edge_torch 转换器 API,由该 API 自动完成量化。我们会在未来的版本中提供更多的量化模式。

多签名导出

我们发现在实际推理场景中,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 的性能优化

我们在性能调查阶段发现了几个改善 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

使用 MediaPipe LLM Inference API

MediaPipe LLM Inference API 是一个高级 API,支持使用 prompt-in / prompt-out 接口进行 LLM 推理。它负责处理底层所有的 LLM 复杂流水线操作,让模型得以更轻松和顺畅地部署。要使用 MediaPipe LLM Inference 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)

在 TFLite 运行时使用纯 C++ 推理

我们还提供了一个简单易用的 C++ 示例 (无需 MediaPipe 依赖),来展示如何运行端到端的文本生成。如果您需要将导出的模型与自己独有的生产流水线和需求进行集成,这个示例是一个很好的起点,来帮助您实现更好的定制和灵活性。

跨平台支持

由于核心推理运行时都支持 TFLite,所以整个流水线都可以轻松集成到您的 Android (包括在 Google Play 中) 或 iOS 应用中,无需进行任何修改。这意味着用新的生成式 API 转换的模型只需添加几个自定义算子依赖即可立即部署。在未来的版本中,我们将为 Android 和 iOS 带来 GPU 支持,并支持 ML 加速器 (TPU、NPU)。

工具

最近发布的模型探索器 (Model Explorer) 是一款很好用的工具,可用于可视化诸如 Gemma 2B 之类的大型模型。分层查看和并排比较可以让您轻松查看和比较原始、重新创作和转换后的模型。我们也准备了专门的文章为您进一步介绍该工具,以及如何通过可视化基准信息来优化模型性能。

  • 模型探索器

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

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

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

以下是我们在编写 PyTorch TinyLlama 模型时使用该工具的示例。我们并排显示了 PyTorch export() 模型与 TFLite 模型。通过使用模型探索器,我们可以轻松比较每个层级 (如 RMSNorms、SelfAttention) 的表达情况。

d6a3aaba051b7068070714decc35c648.gif

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

总结以及下一步

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

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

  • GitHub repo

    https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative

在本次连载的第三篇文章中,我们将深入探讨模型探索器可视化工具,了解该工具如何帮助开发者们可视化、调试和探索模型。

  • 模型探索器

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


5e13c25da8e2c39d5e536bc00ad4d718.png

9a02db4d383a5f417307aa954cf086eb.png

a3f99636cb6d86070191d849b1b92855.png

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/pingmian/23741.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

JavaSE中的if语句、switch语句:如何控制程序流程?

哈喽&#xff0c;各位小伙伴们&#xff0c;你们好呀&#xff0c;我是喵手。运营社区&#xff1a;C站/掘金/腾讯云&#xff1b;欢迎大家常来逛逛 今天我要给大家分享一些自己日常学习到的一些知识点&#xff0c;并以文字的形式跟大家一起交流&#xff0c;互相学习&#xff0c;一…

MySQL与PostgreSQL关键对比一(整体篇)

目录 1 快速参考表&#xff1a;MySQL 与 PostgreSQL 功能表 2 快速参考表&#xff1a;MySQL 与 PostgreSQL 功能表 MySQL 和 PostgreSQL 提供许多相同的特性和功能 - 但是这两个关系数据库管理系统 (RDBMS) 之间存在不容忽视的关键差异。 如果您不熟悉这些差异&#xff0c;这…

RabbitMQ(五)集群配置、Management UI

文章目录 一、安装RabbitMQ1、前置要求2、安装docker版复制第一个节点的.erlang.cookie进入各节点命令行配置集群检查集群状态 3、三台组合集群安装版rabbitmq节点rabbitmq-node2节点rabbitmq-node3节点 二、负载均衡&#xff1a;Management UI1、说明2、安装HAProxy3、修改配置…

Windows安装运行elasticsearch服务

官方下载地址&#xff1a;Download Elasticsearch | Elastic 我在linux上执行的下载命令&#xff1a;wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-8.5.3-linux-x86_64.tar.gz Elasticsearch&#xff08;简称ES&#xff09;是一款基于Apache Lu…

hutool工具实践-缓存

简介 依赖引入 <dependency><groupId>cn.hutool</groupId><artifactId>hutool-cache</artifactId><version>5.8.17</version></dependency> hutool工具既可以像上一章hutool工具实践-验证码-CSDN博客所说直接全部引入&#x…

短剧小程序剧场短剧APP定制开发付费短剧之为什么自建?

在当今数字时代&#xff0c;拥有一个属于自己的小剧场短剧影视小程序不仅是追求创作梦想的新途径&#xff0c;也是与观众建立紧密联系的有效方式。这种新兴的平台为创作者提供了前所未有的自由和机会&#xff0c;使他们能够直接与广大观众交流和分享作品。 1、源码分享的重要性…

搭贝请假审批应用

在现代企业管理中&#xff0c;高效的请假审批系统至关重要。搭贝的请假审批应用通过简化员工的请假流程、提升管理层的工作效率&#xff0c;确保企业运作的连贯性和透明度。本文将介绍搭贝请假审批应用的主要功能模块&#xff1a;请假分析看板、请假申请审批流、请假类型维护和…

依赖注入方式和自动加载原理

依赖注入 Spring提供了依赖注入的功能&#xff0c;方便我们管理和使用各种Bean&#xff0c;常见的方式有&#xff1a; 字段注入&#xff08;Autowired 或 Resource&#xff09;构造函数注入set方法注入 在以往代码中&#xff0c;我们经常利用Spring提供的Autowired注解来实现…

elk:使用filebeat采集日志发送到kafka

# 安装 filebeat 下载 cd /chz/install/filebeat wget https://artifacts.elastic.co/downloads/beats/filebeat/filebeat-8.13.4-linux-x86_64.tar.gz解压 tar zxvf filebeat-8.13.4-linux-x86_64.tar.gz修改配置文件 cd /chz/install/filebeat/filebeat-8.13.4-linux-x86…

在Linux上的Java项目导出PDF乱码问题

在Linux上的Java项目导出PDF乱码问题 场景&#xff1a;一个Java项目导出PDF&#xff0c;在我本地导出是没有问题&#xff0c;但是部署上Linux上后&#xff0c;导出就出现了乱码了。 处理方案 我这里使用的处理方案是在Linux服务器上安装一些PDF需要使用的字体 1.把字体上传到…

Apache POI(使用Java读写Excel表格数据)

1.Apache POI简介 Apache POI是一个开源的Java库&#xff0c;用于操作Microsoft Office格式的文件。它支持各种Office文档的读写功能&#xff0c;包括Word文档、Excel电子表格、PowerPoint演示文稿、Outlook电子邮件等。Apache POI提供了一组API&#xff0c;使得Java开发者能够…

course-nlp——8-translation-transformer

本文参考自https://github.com/fastai/course-nlp。 注意力机制和 Transformer Nvidia AI 研究员 Chip Huyen 写了一篇很棒的文章《Top 8 trends from ICLR 2019》&#xff0c;其中的趋势之一是 RNN 正在失去研究人员的青睐。 这是有原因的&#xff0c;RNN 可能很麻烦&#…

【Qt】Qt QTreeWidget隐藏列名称(横向表头)

1. 效果 未隐藏 隐藏 2. 方法 方法1 ui->treeWidget->header()->hide();方法2 ui->treeWidget->header()->setVisible(false);

Go语言升级1.22.0版本VSCode启动调试报 Version of Go is too old for this version of Delve

因为项目中调用了其它同事的服务&#xff0c;该服务同事的Go环境是1.22.0&#xff0c;我本地go环境是Go1.20.6&#xff0c;不得已也升级了go的版本到1.22.0&#xff0c;其实升级很简单&#xff0c;只需要去go官网下载windows环境的zip包&#xff0c;把原来配置的环境变量GOROOT…

计算机科学(学习笔记三)

内容来源&#xff1a;计算机科学 指令和程序 指令&#xff1a;指示计算机要做什么的代码&#xff0c;多条指令共同组成程序。 计算机指令长度 由于早期计算机每个字只有8位&#xff0c;指令只占4位&#xff0c;意味着只能有16个指令&#xff0c;这远远不够。 现代计算机有两…

[个人总结]-java常用方法

1.获取项目根路径 user.dir是一个系统属性&#xff0c;表示用户当前的工作目录&#xff0c;大多数情况下&#xff0c;用户的当前工作目录就是java项目的根目录&#xff08;src文件的同级路径&#xff09; System.getProperty("user.dir") 结果&#xff1a;D:\code…

C++实现,简单的命令行交互框架

目录 背景 背景 在实际开发中&#xff0c;经常需要有对端测试程序&#xff0c;配合自己的程序&#xff0c;验证功能、逻辑等。面对繁杂、多变的需求&#xff0c;如果对端程序设计得不够灵活&#xff0c;则无法提升工作效率&#xff0c;如果能够与对端程序交互&#xff0c;通过…

Spring Cloud 微服务集成Sentinel实现服务熔断降级

文章目录 一、前言二、技术思路及方案2.1 实现思路2.2 实现方案2.2.1 nacos动态数据源实现类关系图 三、功能实现3.1 快速集成方案3.1.1 引入依赖3.1.2 服务端熔断降级3.1.3 feign调用降级 四、扩展4.1 SPI机制4.2 自定义Slot实现4.3 基于 Sentinel 实现 Feign 全局异常兜底4.3…

Html5如何播放hls格式的视频

目录 一、什么情况下需要播放 二、hls的原理 1、 切片 2、编码 3、分发 4、M3U8文件 5、客户端请求 6、动态码率调整 7、缓存 8、实时性 三、方法一&#xff1a;使用hls.js播放 四、方法二&#xff1a;使用video.js播放 一、什么情况下需要播放 当前端需要嵌入设备视…

[数组] 子数组最大平均数

给你一个由 n 个元素组成的整数数组 nums 和一个整数 k 。 请你找出平均数最大且 长度为 k 的连续子数组&#xff0c;并输出该最大平均数。 任何误差小于 10-5 的答案都将被视为正确答案。 示例 1&#xff1a; 输入&#xff1a;nums [1,12,-5,-6,50,3], k 4 输出&#xff…