认真扩展 AI 模型
2024年4月4日 • 作者：Sean Sheng
如果您正在阅读本文，那么您可能已经对部署开源模型的优势了如指掌。在过去几年中，我们看到开源模型在数量和质量上都取得了令人难以置信的增长。
- Hugging Face 等平台普及了对各种模型的访问，包括大型语言模型 (LLM) 和扩散模型，使开发者能够自由高效地进行创新。
- 开发者享有更大的自主权，可以随意微调和组合不同的模型，从而产生创新方法，如检索增强生成 (RAG) 和高级代理的创建。
- 从经济角度来看，开源模型可大幅节省成本，与 GPT-4 等通用模型相比，使用更小、更专业的模型更经济实惠。
开源模型提供了一个有吸引力的解决方案，但下一个障碍是什么？与使用像 OpenAI 这样的模型端点不同（其中模型是 API 背后的可扩展黑箱），部署您自己的开源模型会带来扩展挑战。确保您的模型能够有效地随着生产流量进行扩展并在流量高峰期间保持无缝体验至关重要。此外，高效管理成本也很重要，这样您只需按使用量付费，避免月底出现意外的财务支出。
北极星方向：无服务器函数及更多
有趣的是，这听起来像是现代无服务器架构（如 AWS Lambda）已经解决的挑战——一个已经存在了近十年的解决方案。然而，在 AI 模型部署方面，情况并非如此。
无服务器函数在 AI 部署方面的局限性是多方面的。
- 不支持 GPU。像 AWS Lambda 这样的平台不支持 GPU。这不仅仅是技术上的疏忽；它植根于架构和实际考虑。
- GPU 不易共享。 GPU 虽然作为设备功能强大且高度并行化，但在同时处理不同模型上的多个推理任务方面不如 CPU 灵活。
- GPU 非常昂贵。 它们是计算世界中的高性能跑车，对于特定任务来说非常出色，但维护成本很高，特别是在不持续使用的情况下。
接下来，让我们回顾一下我们的扩展历程以及在此过程中学到的重要经验。
冷启动问题
在我们开始着手扩展之前，就遇到了臭名昭著的“冷启动”问题。这个问题体现在三个不同的阶段：
- 云资源调配：此阶段涉及云提供商分配实例并将其集成到集群所需的时间。此过程变化很大，从最短 30 秒到几分钟不等，在某些情况下甚至长达数小时，尤其对于 Nvidia A100 和 H100 GPU 等高需求实例更是如此。
- 容器镜像拉取：与简单的 Python 作业镜像不同，AI 模型服务镜像由于其所需的依赖项和自定义库而非常复杂。尽管云提供商声称拥有多千兆位的网络带宽，但我们的经验经常看到下载速度远低于此，镜像拉取时间长达三到五分钟。
- 模型加载。此处所需的时间很大程度上取决于模型的大小，像 LLM 和扩散模型这样较大的模型由于其数十亿个参数而需要更长的时间。我们发现的一个有用经验法则是，通过将模型的参数数量乘以二（磁盘上通常使用的半精度存储格式）来估计以千兆字节为单位的模型下载大小。例如，使用 1Gbps Internet 带宽加载像 Stable Diffusion 2 这样的 5GB 模型可能需要大约 1.3 分钟，而像 Llama 13B 和 Mixtral 8x7B 这样更大的模型可能分别需要 3.5 分钟和 12.5 分钟。
冷启动问题的每个阶段都需要特定的策略来最大限度地减少延迟。在接下来的部分中，我们将更详细地探讨每个阶段，分享我们在实现 AI 模型可扩展、高效 GPU 部署方面的策略和解决方案。
云资源调配
与无服务器 CPU 的同构环境不同，处理 GPU 时管理各种计算实例类型至关重要，每种类型都针对特定用例量身定制。例如，IO 密集型 LLM 需要高 GPU 内存带宽和容量，而生成模型需要更强大的 GPU 计算能力。
通过维护所有 GPU 实例类型来确保高峰流量期间的可用性可能会导致高得令人望而却步的成本。为了避免空闲实例造成的财务负担，我们实施了“备用实例”机制。我们没有为最大潜在负载做准备，而是维护了与我们的增量扩展步骤相匹配的计算数量的备用实例。例如，如果我们一次扩展两个 GPU，我们就需要准备两个备用实例。这使我们能够在需求激增时快速将资源添加到我们的服务队列中，显着缩短等待时间，同时保持成本可控。
在多租户环境中，多个团队或（在我们的案例中）多个组织共享一个公共资源池，我们可以实现更高效的利用率。这种共享环境使我们能够平衡不同的资源需求，从而提高成本效率。然而，管理多租户会带来挑战，例如强制执行配额和确保网络隔离，这可能会增加集群的复杂性。
容器镜像拉取
无服务器 CPU 工作负载通常使用轻量级镜像，例如 Python slim 镜像（约 154 MB）。相比之下，为服务 LLM 而构建的容器镜像要大得多（即使压缩后也有 6.7 GB）；大部分大小来自于运行 AI 模型所需的各种依赖项。
尽管云提供商宣传了高带宽网络，但实际情况往往不尽如人意，实际下载速度仅为承诺速率的一小部分。
实际上，很大一部分文件从未被使用。一种方法是优化容器镜像本身，但这很快被证明是难以管理的。相反，我们将重点转移到按需文件拉取方法。具体来说，我们首先仅下载镜像元数据，实际的远程文件在稍后需要时才获取。此外，我们利用集群内的点对点网络显着提高了拉取效率。
通过这些优化，我们将镜像拉取时间从几分钟缩短到了几秒钟。然而，我们都知道这个测量结果有点“作弊”，因为实际文件并未在此阶段拉取。真正的文件拉取发生在服务运行时。因此，拥有一个允许您在不同生命周期阶段（例如初始化和服务）定义行为的服务框架至关重要。通过在初始化期间完成大部分（如果不是全部）引导过程，我们可以确保所有文件依赖项都已拉取。这样，在服务运行时，就不会因文件拉取而产生延迟。
在上面的示例中，模型加载是在初始化生命周期内的 __init__ 方法中完成的，而服务发生在命名为 txt2img 的 @bentoml.api 中。
模型加载
最初，最直接的模型加载方法是直接从 Hugging Face 等远程存储中获取。使用内容分发网络 (CDN)、NVMe SSD 和共享内存，我们可以绕过一些瓶颈。尽管这种方法有效，但远非最优。
为了改进此过程，我们考虑了我们的区域内网络带宽。我们在分布式文件系统中播种模型，并将其分解成更小的块，以便进行并行下载。这极大地提高了性能，但我们仍然遇到了云提供商的网络瓶颈。
作为响应，我们通过使用点对点共享和利用本地缓存，进一步优化以利用集群内网络带宽。尽管这些改进是显着的，但它们增加了流程的复杂性，我们需要将其抽象出来，以便开发者不必关心。
即使采用了上述实践，我们仍然面临一个顺序瓶颈：需要在完成每个步骤后才能继续进行下一个步骤。模型必须完全下载到持久化存储才能加载到 CPU 内存，然后再加载到 GPU。
我们转向了一种基于流的方法来加载模型权重，使用了我们现有的分布式文件缓存系统。该系统允许程序像所有文件都在磁盘上逻辑可用一样运行。实际上，所需数据是按需从远程存储中获取的，因此绕过了磁盘写入。通过利用像 Safetensors 这样的格式，我们可以通过内存映射 (mmap) 将模型权重高效地加载到主内存中，然后再以流式方式加载到 GPU 内存。
此外，我们采用了异步写入磁盘的方式。通过这样做，我们在本地磁盘上创建了一个更快的访问缓存层。因此，仅进行代码更改的新部署可以绕过较慢的远程存储获取阶段，直接从本地缓存读取模型权重。
总而言之，我们成功地在一定程度上优化了冷启动时间，并且对结果感到满意
- 通过备用实例实现无云资源调配延迟。
- 通过按需和点对点流式传输实现更快的容器镜像拉取。
- 通过分布式文件系统、点对点缓存和流式加载到 GPU 内存，实现加速模型加载时间。
- 服务框架支持的镜像拉取和模型加载并行化。
扩展指标
接下来，我们需要确定用于扩展 GPU 上 AI 模型部署的最具指示性的信号。
资源利用率指标
最初，我们考虑了 CPU 利用率。它很简单，并且有一个直观的默认阈值，例如 80%。然而，显而易见的缺点是 CPU 指标无法反映 GPU 利用率。此外，Python 中的全局解释器锁 (GIL) 限制了并行性，阻止了多核实例上的高 CPU 利用率，这使得 CPU 利用率成为一个不太可行的指标。
我们还探索了 GPU 利用率作为衡量模型工作负载的更直接指标。然而，我们遇到了一个问题：像 nvml 这样的工具报告的 GPU 利用率并不能准确代表 GPU 的实际利用率。该指标在一段时间内对内核使用情况进行采样，如果至少有一个内核正在执行，则认为 GPU 被利用。这与我们的观察结果一致，即即使 GPU 设备已被报告为高利用率，通常也可以通过改进批处理来获得更好的性能。
注意：根据 NVIDIA 文档，utilization.gpu 表示“在过去的采样周期内，一个或多个内核在 GPU 上执行的时间百分比。采样周期可能在 1 秒到 1/6 秒之间，具体取决于产品”。
基于资源的指标本质上是回顾性的，因为它们只反映资源消耗后的使用情况。它们也被限制在 100%，这带来了一个问题：基于这些指标进行扩展时，调整的最大比例通常是当前利用率与期望阈值之比（参见下面的扩展公式）。这导致了一种保守的横向扩展行为，不一定与实际生产流量需求相匹配。
desiredReplicas = ceil[currentReplicas * ( currentMetricValue / desiredMetricValue )]
基于请求的指标
我们转向基于请求的指标，以获得更主动的信号，这些信号不受 100% 的限制。
QPS（每秒查询次数）因其简单性而成为一个广泛认可的指标。然而，它在生成式 AI（例如 LLM）中的应用仍然存在疑问。它不容易配置，并且由于每个请求的成本是可变的（取决于处理和生成的令牌数量），因此使用 QPS 作为扩展指标可能会导致不准确。
另一方面，并发性已被证明是反映系统实际负载的理想指标。它表示处于排队或正在处理状态的活动请求数量。此指标
- 精确反映系统负载。Little 定律指出 QPS 乘以平均延迟等于并发性，提供了一种理解 QPS 和并发性之间关系的优雅方式。实际上，在模型服务中，每个请求的平均延迟是未知数。然而，通过衡量并发性，我们不需要计算平均延迟。
- 使用扩展公式准确计算所需的副本数。允许部署直接扩展到理想大小，无需中间步骤。
- 易于根据批处理大小进行配置。对于不可批处理的模型，它就是 GPU 的数量，因为每个 GPU 一次只能处理一个生成任务。对于支持批处理的模型，批处理大小决定了并发级别。
为了使并发性起作用，我们需要服务框架的支持，以便自动将并发性作为指标进行检测，并将其作为部署平台的扩展信号。我们还必须制定正确的扩展策略，以防止在流量高峰期间过度扩展或在流量稀少时过早缩减。
请求队列
我们与并发性集成的另一个重要机制是请求队列。它充当缓冲区和协调器，确保有效地处理传入请求，并且不会使任何单个服务器副本过载。
在没有请求队列的情况下，所有传入请求都会直接分派到服务器（下图中的 6 个请求）。如果多个请求同时到达，并且只有一个活动服务器副本，它就会成为瓶颈。服务器会尝试以先到先服务的方式处理每个请求，这通常会导致超时和糟糕的客户端体验。
相反，在有了请求队列的情况下，服务器以最优速率消费请求，处理速率基于为服务定义的并发性。当额外的服务器副本扩展时，它们也会从队列中拉取请求。这种机制可以防止任何单个服务器过载，并允许在可用基础设施之间更平滑、更易于管理的请求分发。
结论
我们探索 AI 模型扩展解决方案的旅程是一次冒险，最终使我们创建了 BentoCloud 上的扩展体验——这是一个囊括了我们所有经验的平台。
为了避免给人留下宣传的印象，我们将用一幅胜过千言万语的图片来说明我们的观点。下面的监控仪表板展示了传入请求与服务器实例扩展之间的关联。
与扩展同样重要的是缩减能力。随着请求减少到零，部署会相应地减少活动实例的数量。这种能力确保不会因未使用的资源而产生不必要的成本，使支出与实际使用情况保持一致。
我们希望本文的核心是：模型部署的扩展应该被视为生产应用程序的一个重要方面。与扩展 CPU 工作负载不同，在 GPU 上扩展模型部署面临独特的挑战，包括冷启动时间、配置扩展指标和编排请求。在评估部署平台时，应彻底评估它们对这些挑战的解决方案。