大型语言模型 (LLM) 应用的快速增长与能源需求的快速增长息息相关。根据国际能源署 (IEA) 的数据,主要受 AI 驱动,数据中心电力消耗预计到 2026 年将大致翻一番。这是由于大规模 LLM 训练的能源密集型需求造成的——然而,AI 推理工作负载的增加也起到了作用。例如,与传统搜索查询相比,单次 AI 推理可能消耗大约 多 10 倍的能源。
作为开发者,我们直接影响 AI 解决方案的能源密集度。我们可以采取技术决策,帮助使我们的 AI 解决方案更具环境可持续性。最大限度地减少计算资源以提供 LLM 解决方案并不是创建可持续 AI 用途的唯一要求。例如,可能需要政策干预等系统性变革,但利用能源效率高的解决方案是一个重要因素,也是我们可以立即采取的有效干预措施。
话虽如此,最大限度地减少 LLM 推理的云计算需求也导致降低您的云账单,并使您的应用程序更节能,创造一个双赢局面。在本博客中,我们将引导您完成通过优化和在 PyTorch 上部署 Llama 3.1 模型来创建 LLM 聊天机器人的步骤,量化特定架构决策带来的计算效率优势。
我们将评估什么?
在本博客中,我们的目标是创建一个沉浸式的幻想故事应用程序,用户通过与生成式 AI 聊天进入幻想世界。第一个地点是邪恶之地(Land of Wicked),允许人们角色扮演漫步翡翠城(Emerald City),实时观察景象和场景。我们将通过聊天机器人和自定义系统提示来实现这一点。
我们将评估 LLM 在 CPU 上的性能。您可以在此处 查看 CPU 与 GPU 推理的优势。通常,在云端利用 CPU 进行 LLM 推理对于参数大约 10B 或更少(例如 Llama 系列)的模型是一个不错的选择。
我们还将使用基于 Arm 的 CPU,特别是 AWS Graviton 系列。根据研究,基于 Arm 的 Graviton3 服务器内置可提供降低 67.6% 的工作负载碳强度。虽然这项研究基于模拟,但它是展示最大限度减少应用程序能源需求的可能性的一个很好的开端。
首先,您将看到如何在 PyTorch 上运行一个简单的 LLM 聊天机器人,然后探索三种提高计算效率的应用程序优化技术
- 模型优化:利用 4 位量化和新增的 KleidiAI 内核。
- 捷径优化:实现向量数据库以处理常见查询。
- 架构优化:采用无服务器架构。
让我们开始吧。
在 AWS Graviton4 上通过 PyTorch 运行 Llama-3.1
为了最大限度地提高能源效率,我们将只使用支持此 LLM 聊天机器人所需的最低服务器资源。对于这个 Llama-3.1 80 亿参数模型,需要 16 核、64GB RAM 和 50GB 磁盘空间。我们将使用运行 Ubuntu 24.04 的 r8g.4xlarge Graviton4 实例,因为它符合这些规格。
启动此 EC2 实例,连接到它,并开始安装所需的依赖项
sudo apt-get update
sudo apt install gcc g++ build-essential python3-pip python3-venv google-perftools -y
然后安装 Torchchat,这是 PyTorch 团队开发的库,可在各种设备上运行 LLM
git clone https://github.com/pytorch/torchchat.git
cd torchchat
python3 -m venv .venv
source .venv/bin/activate
./install/install_requirements.sh
接下来,通过 CLI 从 Hugging Face 安装 Llama-3.1-8b 模型。您需要先在您的 HF 帐户上创建一个 Hugging Face 访问令牌。这将把 16GB 模型下载到您的实例,可能需要几分钟。
pip install -U "huggingface_hub[cli]"
huggingface-cli login
<enter your access token when prompted>
python torchchat.py export llama3.1 --output-dso-path exportedModels/llama3.1.so --device cpu --max-seq-length 1024
现在您已准备好运行 LLM 模型,添加一个系统提示,使其成为邪恶之地(Land of Wicked)的引导式故事讲述者
LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libtcmalloc.so.4 TORCHINDUCTOR_CPP_WRAPPER=1 TORCHINDUCTOR_FREEZING=1 OMP_NUM_THREADS=16 python torchchat.py generate llama3.1 --device cpu --chat
输入 ‘y’ 进入系统提示模式,然后输入以下提示
您是一个幻想冒险应用的引导式故事讲述者。让用户沉浸在邪恶之地的迷人世界中,引导他们在翡翠城进行互动式、实时体验。描述生动的景象、动态的场景,并让用户参与感觉鲜活且响应迅速的故事讲述。允许用户做出选择来塑造他们的旅程,同时保持邪恶世界的神奇基调。
然后输入您的用户查询
我穿过翡翠城的大门,抬头望去
输出将显示在屏幕上,生成第一个 token 大约需要 7 秒,速度低于每秒 1 token。
这个例子花了 245 秒,也就是 4 分钟,才生成完整的回复——速度不是很理想。我们将要查看的第一个优化将加速 LLM 的生成,减少其计算占用空间。
优化 1:KleidiAI 和量化
上述基本实现可以进行几种优化。最简单、最快速的一种是将模型从 FP16 量化到 INT4。这种方法牺牲了一些精度,但将模型大小从 16GB 减小到约 4GB,从而提高了推理速度。
另一种常见的优化是利用 TorchAO (Torch Architecture Optimization),它是 PyTorch 库,与 TorchChat 无缝协作,通过各种量化和稀疏化方法增强模型性能。
最后,我们将使用 Arm KleidiAI 优化。这些是用汇编语言编写的微内核,可显著提高 Arm CPU 上 LLM 推理的性能。如果您有兴趣,可以 在此处阅读有关 KleidiAI 内核如何工作的更多信息。
要实现这些优化,请启动一个新的 EC2 实例,并按照 有关如何使用 PyTorch 运行大型语言模型 (LLM) 聊天机器人 的说明进行操作。准备就绪后,运行模型并输入与上面相同的系统提示和用户查询。您将获得显著加速推理的结果:第一个 token 少于 1 秒,速度约为每秒 25 token。
这将推理时间从 245 秒缩短到约 10 秒。这样可以减少服务器的功耗,因为它有更多的时间处于空闲状态,而不是运行耗能的推理。在其他条件相同的情况下,这是一个比非优化应用更环保的解决方案。接下来的两种方法超越了模型推理优化,它们修改了解决方案架构以进一步降低计算负载。
优化 2:使用 FAISS 匹配数据库以处理常见问题
正如引言中所述,模型推理通常比其他搜索技术计算成本更高。如果您可以自动回复常见的用户查询而无需执行 LLM 推理怎么办?使用查询/回复数据库是一种绕过 LLM 推理并高效响应的选项。对于这个交互式故事应用程序,您可以想象关于特定角色、世界本身以及关于聊天机器人能做什么/不能做什么的规则的常见问题,这些问题可以有预先生成的答案。
然而,传统的精确匹配数据库不够充分,因为用户可以用多种方式表达相同的查询。询问聊天机器人的能力都可能得到相同的答案,但表达方式不同
- “你能做什么?”
- “告诉我你能做些什么。”
- “我如何与你互动?”
实现语义搜索解决了这个问题,通过理解用户的意图,将用户的查询与最相关的预生成答案匹配。FAISS 库 是实现语义搜索的一个绝佳选项。
这种方法的计算节省取决于三个因素
- 可以通过语义搜索而非 LLM 服务的用户查询百分比。
- 运行 LLM 推理的计算成本。
- 运行语义搜索的计算成本。
节省的计算方式如下
Computational_savings = (% of queries) * (LLM_cost – search_cost).
这种类型的架构在几种情况下是合理的。一种是如果您的系统有许多重复问题的常见查询。另一种是具有数十万个传入查询的大规模系统,其中小百分比的节省会累积成有意义的变化。最后,如果您的 LLM 推理相对于搜索成本非常昂贵,特别是对于更大参数的模型。
最后的优化方法是从服务器过渡到无服务器。
优化 3:无服务器方法
使用无服务器架构流行有很多原因,其中之一是只为活动计算时间付费,并消除了闲置服务器的成本。闲置服务器需要相当大的功率来保持开启,在等待时浪费能源。
这种成本效益转化为一种本质上更环保的架构,因为它减少了浪费的能源消耗。此外,多个应用程序共享底层物理基础设施,提高了资源效率。
要设置您自己的无服务器聊天机器人,您需要首先使用包含 Lambda 入口函数 lambda_handler
的 python 脚本对经过量化的 Llama-3.1-8b 以及 TorchChat、TorchAO 和 Arm KleidiAI 优化进行容器化。一种部署选项是将您的容器上传到 AWS ECR,并将容器附加到您的 Lambda 函数。然后设置一个 API Gateway WebSocket 或类似的东西,以便通过 API 与您的 Lambda 交互。
使用无服务器架构托管您的 LLM 有两个显著限制,第一个是 token 生成速度。回想一下,基于服务器的方法在使用 KleidiAI 优化时提供了约每秒 25 个 token 的速度。无服务器方法的速度慢了一个数量级,我们测量约为每秒 2.5 个 token。此限制主要源于 Lambda 函数部署在 Graviton2 服务器上。当部署迁移到具有更多 SIMD 通道(如 Graviton3 和 Graviton4)的 CPU 上时,每秒 token 数应会随着时间增加。通过 此处 了解更多关于 Graviton3 中引入的架构优化信息 Arm Neoverse-V1 CPU。
这种较慢的速度限制了无服务器 LLM 架构的可行用例,但在某些情况下,这可以被视为一个优势。在我们交互式故事讲述的用例中,缓慢地揭示信息会营造出沉浸感,构建期待并模仿实时叙事。其他用例包括
- 缓慢、舒缓地传递文字的引导式冥想应用
- 进行深思熟虑的对话或治疗性对话的虚拟朋友。
- 诗歌生成或互动艺术,通过缓慢的传递营造沉思的美学。
在合适的应用中,用户可能会因为较慢的 token 生成而获得更好的体验。当优先考虑更可持续的解决方案时,限制最终会变成优势。打个比方,对如今现代电影的一个常见批评是它们过度依赖视觉效果导致故事情节不如老电影引人入胜。VFX 的成本限制意味着老电影必须精心设计引人入胜的对话,利用巧妙的镜头角度和人物站位来充分吸引观众。同样,深思熟虑地专注于可持续的 AI 架构可以在恰当完成时带来更具吸引力、更身临其境的体验。
LLM 推理的第二个无服务器限制是冷启动时间约为 50 秒。如果实施不当,用户等待 50 秒而没有其他选择可能会离开应用。您可以通过几种设计技巧将此限制转化为基于邪恶之地体验的特色
- 创建一个“序幕体验”,通过硬编码的问题和答案引导用户,为他们即将进入翡翠城做准备,并收集输入以塑造他们即将到来的体验。
- 将等待时间设置为倒计时计时器,揭示硬编码的故事片段或世界构建内容。像巫师这样的角色可以向用户传递零散的台词,以营造悬念,并让用户进入合适的心态。
- 创建一个包含电影或音乐剧中音乐的音频介绍,并配有轮播视觉效果,将用户吸引到邪恶世界的氛围中。
跳出思维定势
实施以可持续性为导向的解决方案架构包括且超越了优化您的 AI 推理。了解用户将如何与您的系统交互,并相应地调整您的实现。始终优化以获得每秒快速的 token 或第一个 token 的时间将隐藏引人入胜的功能的机会。
话虽如此,您应在可能的情况下利用简单直接的优化。使用 TorchAO 和 Arm KleidiAI 微内核是加速您的 LLM 聊天机器人的好方法。通过结合创新的解决方案架构并在可能的情况下进行优化,您可以构建更可持续的基于 LLM 的应用程序。祝您编码愉快!