大规模 AI 批量推理：向量生成速度提升 9 倍

向量生成是现代 AI 应用的核心，从推荐系统到检索增强生成（RAG）无所不在。然而，在百万/万亿规模数据集上运行向量生成的批量推理绝非易事。

经过数日开发精心调优的批量推理脚本后，获取数百或数千个 GPU 并在这些 GPU 上并行运行任务仍然是件非常痛苦的事。

在本文中，我们分享了如何利用 SkyPilot 在约 3000 万条记录上使用最先进的 LLM 向量模型进行大规模向量生成，并与典型的自动扩缩容 K8s 设置（AWS EKS）相比，将工作负载速度提升了 9 倍。

我们在此处开源了运行大规模向量生成的代码。

TL;DR：“被遗忘的”区域带来 9 倍加速

通过利用多个区域，我们成功将向量生成速度提高了 9 倍，并将时间从 20 小时缩短到 2.3 小时。此外，我们还节省了 61% 的成本。

下图显示了任务在全球范围内的分布，包括欧洲、北美、亚洲、澳大利亚等。

总任务数

406

使用的区域

最活跃区域

ap-northeast-1

可视化我们的执行轨迹。利用率最高的三大区域：ap-northeast-1、ap-southeast-2 和 eu-west-3。

我们通过使用 SkyPilot 实现了性能提升，它能自动跨区域搜索资源并管理数百个任务的生命周期。

一个关键洞察：像 US-East-1 这样的热门区域容量有限，但全球范围内的“被遗忘的区域”拥有大量可用容量。通过将工作负载分散到这些区域，您可以获得显著更多的计算资源。

现在，让我们深入了解工作负载的细节。

向量：AI 搜索的神经基础

向量构成了当代 AI 搜索和检索系统的基石。与传统的基于关键词的方法（Google 搜索等）不同，向量使 AI 能够理解诸如“解除租赁合同”和“终止租房协议”这类短语之间的概念相似性，即使它们没有任何共同关键词。

科技巨头如何使用向量

亚马逊

使用 Amazon Titan Text Embeddings 支持产品推荐。

谷歌

利用向量提高搜索相关性，增强查询理解。

Facebook

利用社交上下文进行基于向量的检索，增强搜索功能。

微软

在 Azure AI Search 中利用向量改进搜索和电子取证。

Netflix

大规模使用向量进行电影和节目推荐。

Spotify

应用神经向量进行个性化音乐推荐。

Airbnb

利用向量改进搜索排名和列表推荐。

Uber

使用双塔向量进行 Uber Eats 推荐。

为大型数据集生成向量需要巨大的计算资源。考虑这些例子：典型的 Nvidia L4 GPU 处理一个具有 70 亿参数的模型，大约每秒可以处理 2000 个文本 token。计算十亿条项目的向量在单台机器上需要超过 5.8 天。像 falcon-refinedweb 数据集这样包含 6000 亿文本 token 的更大规模数据集，处理起来将耗时超过 9.5 年！

真实世界案例研究：为亚马逊评论生成向量

我们基于亚马逊评论的真实世界案例研究开发了一个大规模向量生成流水线，并在此处开源了代码。

我们使用了 Amazon reviews 2023 数据集的书籍部分，该部分包含约 3000 万条亚马逊书籍评论，并使用最先进的专业向量 LLM Alibaba-NLP/gte-Qwen2-7B-instruct（MTEB 排行榜上的顶级向量模型之一）为这些评论生成向量。

我们使用 AWS 上的单个 L4 GPU（g6.xlarge 实例）来处理每个向量生成任务。

现有方法：遭遇可用性瓶颈

采用传统方法时，单区域的向量生成扩缩容能力有限

Kubernetes 集群（EKS、GKE 等）：Kubernetes 集群必须部署在单个区域，即使启用了自动扩缩容，也会受到区域 GPU 可用性的限制。
云批处理服务（AWS Batch、GCP Batch 等）：云批处理服务同样受限于您在单个区域可配置的 GPU 数量。
Slurm：Slurm 是 HPC 常用的作业调度器，但它并非为云环境设计，扩缩容性不佳，尤其是在需要跨区域时。
手动配置 GPU：手动配置 100 多个 GPU 并进行管理是一场噩梦，且不可扩展。

现有方法局限于一个区域，这意味着 GPU 可用性有限。

通过在 AWS 区域 ap-northeast-2 使用自动扩缩容集群，由于可用性有限，我们在 2 小时后最多只能获得 47 个 L4 GPU（使用 g6.xlarge 实例）。

在持续尝试配置更多 GPU 的过程中，我们遭遇了 InsufficientInstanceCapacity 错误 2,955 次。

======= Job Status Summary =======
Total jobs: 128
Successfully provisioned or in progress: 47
Failed to provision: 81

======= Error Type Summary =======
InsufficientInstanceCapacity errors: 2955
  - On-demand instances: 2955
  Affected zones:
    - ap-northeast-2: 2955 occurrences

使用传统云服务面临的挑战

可用性：我们在特定区域 ap-northeast-2 的配额允许使用 128 个 GPU，但由于可用性限制，我们只能获得 47 个 GPU。
处理时间长：处理完整的亚马逊评论数据集耗时：20+ 小时
成本高昂：按需定价导致预估成本高昂：约 $710

跨越单个区域

为了缓解可用性限制，我们使用 SkyPilot 将工作负载分布到多个区域以获得更多可用资源，这带来了多达 9 倍的资源回报。

下图显示了亚马逊评论数据集向量生成在单区域和多区域方法之间的吞吐量比较

配置：gte-Qwen2-7B-instruct 向量模型，亚马逊评论数据集，L4 GPU，运行 2 小时。

请注意，出于预算限制，我们将图表截断在 2 小时处，但与单区域方法（如自动扩缩容的 EKS 集群）相比，我们已经可以看到 9 倍的加速。

通过跨越多个区域，尽管我们在不同区域也遇到了可用性错误，但 SkyPilot 会自动在全球范围内搜索区域并重试配置 GPU，最终为我们获取了 406 个 L4 GPU，token 吞吐量高达 797.7k tokens/s！

======= Job Status Summary =======
Total jobs: 406
Successfully provisioned or in progress: 406
Failed to provision: 0
Spot instances: 396
  - Successfully provisioned or in progress: 396
  - Failed to provision: 0
On-demand instances: 10
  - Successfully provisioned or in progress: 10
  - Failed to provision: 0

查看完整的错误类型摘要

======= Error Type Summary =======
MaxSpotInstanceCountExceeded errors: 6673
  - Spot instances: 6673
  - On-demand instances: 0
  Affected zones:
    - ca-central-1b: 607 occurrences
    - ca-central-1a: 606 occurrences
    - ap-south-1a: 564 occurrences
    - eu-west-2b: 520 occurrences
    - eu-west-2a: 509 occurrences
    - us-east-1d: 450 occurrences
    - us-east-2b: 427 occurrences
    - ap-south-1b: 425 occurrences
    - us-east-1a: 414 occurrences
    - eu-central-1b: 390 occurrences
    - eu-central-1a: 374 occurrences
    - us-east-2a: 334 occurrences
    - us-west-2c: 319 occurrences
    - us-west-2d: 313 occurrences
    - us-east-1b: 211 occurrences
    - us-west-2b: 136 occurrences
    - us-east-2c: 74 occurrences
InsufficientInstanceCapacity errors: 6566
  - Spot instances: 6360
  - On-demand instances: 206
  Affected zones:
    - ap-northeast-2a: 600 occurrences
    - ap-northeast-2d: 580 occurrences
    - us-east-2c: 567 occurrences
    - eu-west-3c: 560 occurrences
    - eu-north-1a: 535 occurrences
    - ap-southeast-2c: 494 occurrences
    - eu-north-1b: 455 occurrences
    - ap-southeast-2a: 435 occurrences
    - eu-west-3b: 427 occurrences
    - us-east-1b: 412 occurrences
    - ap-northeast-1a: 299 occurrences
    - us-west-2b: 281 occurrences
    - ap-northeast-1c: 256 occurrences
    - us-west-2: 206 occurrences
    - us-east-2b: 189 occurrences
    - eu-west-2a: 84 occurrences
    - ap-south-1b: 84 occurrences
    - us-east-2a: 58 occurrences
    - us-west-2d: 16 occurrences
    - ca-central-1a: 12 occurrences
    - eu-west-2b: 6 occurrences
    - eu-central-1b: 3 occurrences
    - ca-central-1b: 3 occurrences
    - us-east-1a: 2 occurrences
    - ap-south-1a: 2 occurrences
VcpuLimitExceeded errors: 0
  - Spot instances: 0
  - On-demand instances: 0

Other errors: 0
  - Spot instances: 0
  - On-demand instances: 0

完整的事件和错误日志可以在此处找到。

结果总结

指标	单区域	多区域	提升
实际获得的 GPU 数量	47	406	资源增加 9 倍
处理时间	20 小时	2.3 小时	速度提升 9 倍
成本	$710	$277	节省 61%
区域	1	12	可用性更高
任务恢复	手动	自动	可靠性提高

跨越单个区域的好处

更多资源：通过利用跨区域资源，我们同时获得了 406 个 GPU，token 吞吐量达到 364.4k token/s！
更快：处理相同数据集的时间：从 20 小时缩短到仅 2.3 小时（比单区域方法快 9 倍）
更便宜：Spot 实例的使用将成本从 $710 降低到 $277（节省 61%）

使用 SkyPilot 运行大规模任务

我们使用 SkyPilot 将向量工作负载分布到多个区域，并同时利用按需实例和 Spot 实例，显著增加了可用资源，同时降低了成本。

使用 SkyPilot 运行和扩缩容向量生成的批量推理任务十分简单

首先，我们为向量任务定义 SkyPilot YAML 配置。

# compute_text_embeddings.yaml
workdir: .

resources:
  cpus: 4
  accelerators: L4
  cloud: aws
  # Notably, we don't specify a region to allow multi-region deployment
  any_of:
    - use_spot: true
    - use_spot: false

envs: # Will be overridden by batch launcher script
  START_IDX: 0
  END_IDX: 10000  

run: python compute_text_embeddings.py

为了均匀分配工作负载，我们使用一个批处理启动脚本，该脚本在多个区域创建数百个独立的 worker 任务

# batch_compute_embeddings.py
def main():
    # ... initialization code ...
    task = sky.Task.from_yaml('compute_text_embeddings.yaml')
    
    # Launch jobs for each partition
    for job_rank in range(args.num_jobs):
        # Update environment variables based on partition method
        env_vars = {
            'START_IDX': str(args.start_idx),
            'END_IDX': str(args.end_idx),
        }
           
        task_copy = task.update_envs(env_vars)
        print(f"Launching job {job_rank}...")
        sky.jobs.launch(
            task_copy,
        )

运行大规模向量生成的完整代码已在此处开源。

提示： 由于 huggingface 数据集仅支持迭代器访问而不支持索引，我们实现了跨步分区，以确保每个任务处理的数据集文档均匀分布（任务 0 处理条目 {0, N, 2N, ...}；任务 1 处理条目 {1, N+1, 2N+1, ...}）。

任务队列

所有任务都可以在 UI (sky jobs dashboard) 或通过 CLI 显示

Job queue showing distributed jobs across regions

向量生成进度仪表板

我们还构建了一个自定义仪表板来可视化向量生成的进度

Dashboard showing global job distribution

关于出口流量费用？

当使用多区域方法时，云服务提供商对区域间数据传输收取的出口流量费用是一个自然而然的担忧。然而，对于向量生成工作负载，由于以下几个原因，这些成本是可控的甚至是微不足道的：

向量是紧凑的表示（通常每个向量有 768-3072 个维度）
- 数据流主要是单向的（从计算节点到存储）
- 我们集中的 S3 存储桶方法最大限度地减少了跨区域传输
在我们的实验中，亚马逊评论数据集（约 3000 万条评论）生成了约 120GB 的向量数据，所有区域的出口流量费用不到 $50。计算资源节省的数百美元成本远超这些传输费用。

结论

为了运行大规模批量推理任务，SkyPilot 提供了一种简单且可扩展的方式，将工作负载分布到多个区域（甚至跨云提供商）。利用“被遗忘的区域”和 Spot 实例，可用的资源显著增加了 9 倍，成本降低了 61%。

9 倍加速：原本需要 20 多小时的处理，仅需 2 小时完成
成本降低 61%：总支出从 $710 降至 $277
突破规模限制：绕过单区域可用性限制，在 12 个区域同时访问 406 个 GPU（相比单区域约 40 个）
可靠性增强：自动从 Spot 实例中断中恢复，保持处理连续性

下一步

在亚马逊评论数据集上运行向量生成的代码：此处
使用 SkyPilot运行大量任务
为您的团队部署 SkyPilot：SkyPilot 团队部署

TL;DR：“被遗忘的”区域带来 9 倍加速#

向量：AI 搜索的神经基础#

科技巨头如何使用向量

亚马逊

谷歌

Facebook

微软

Netflix

Spotify

Airbnb

Uber

真实世界案例研究：为亚马逊评论生成向量#

现有方法：遭遇可用性瓶颈#

跨越单个区域#

使用 SkyPilot 运行大规模任务#

任务队列#

向量生成进度仪表板#

关于出口流量费用？#

结论#

下一步#

TL;DR：“被遗忘的”区域带来 9 倍加速

向量：AI 搜索的神经基础

真实世界案例研究：为亚马逊评论生成向量

现有方法：遭遇可用性瓶颈

跨越单个区域

使用 SkyPilot 运行大规模任务

任务队列

向量生成进度仪表板

关于出口流量费用？

结论

下一步