AWS Trainium & Inferentia 文档
引言
并获得增强的文档体验
开始使用
简介
在当今世界，几乎每位 AI 工程师都熟悉通过简单的 API 调用来运行推理，但后端是如何以最优方式处理这些请求的呢？您所使用的模型提供商或服务是如何确保满足延迟和吞吐量要求的呢？
在这篇博客中，我们将介绍如何使用 Optimum Neuron 在 AWS Inferentia2 上通过 HuggingFace TGI 容器来部署模型。我还会深入探讨如何优化延迟和吞吐量，以及我们可以做出哪些决策来影响我们的优先级。
了解工具
- Inferentia2 芯片：Inferentia2 是 AWS 专为机器学习推理打造的第二代加速器。
- Optimum Neuron：连接 🤗 Transformers 库与 AWS 加速器（包括 AWS Trainium 和 AWS Inferentia）的接口。
- 文本生成推理 (TGI) 容器：文本生成推理 (TGI) 是一个用于部署和提供大型语言模型 (LLM) 服务的工具包。
- GuideLLM：一个用于评估和优化大型语言模型 (LLM) 部署的工具。
本次实验我使用的实例是 inf2.48xlarge。我可以运行 neuron-ls 来检查实例类型并查看每个设备，该命令会给出以下输出：
instance-type: inf2.48xlarge
+--------+--------+--------+-----------+---------+
| NEURON | NEURON | NEURON | CONNECTED | PCI |
| DEVICE | CORES | MEMORY | DEVICES | BDF |
+--------+--------+--------+-----------+---------+
| 0 | 2 | 32 GB | 11, 1 | 80:1e.0 |
| 1 | 2 | 32 GB | 0, 2 | 90:1e.0 |
| 2 | 2 | 32 GB | 1, 3 | 80:1d.0 |
| 3 | 2 | 32 GB | 2, 4 | 90:1f.0 |
| 4 | 2 | 32 GB | 3, 5 | 80:1f.0 |
| 5 | 2 | 32 GB | 4, 6 | 90:1d.0 |
| 6 | 2 | 32 GB | 5, 7 | 20:1e.0 |
| 7 | 2 | 32 GB | 6, 8 | 20:1f.0 |
| 8 | 2 | 32 GB | 7, 9 | 10:1e.0 |
| 9 | 2 | 32 GB | 8, 10 | 10:1f.0 |
| 10 | 2 | 32 GB | 9, 11 | 10:1d.0 |
| 11 | 2 | 32 GB | 10, 0 | 20:1d.0 |
+--------+--------+--------+-----------+---------+设置和安装
首先，我运行以下命令来安装必要的依赖项，并拉取编译模型以及为基准测试提供已编译模型服务所需的容器。
!pip install hftransfer guidellm==0.1.0 !git clone https://github.com/huggingface/optimum-neuron.git !docker pull ghcr.io/huggingface/text-generation-inference:latest-neuron
根据模型的不同，可选择性地配置你的 HF_TOKEN，如下所示：
!export HF_TOKEN=你的HF_TOKEN
模型编译与部署
在我的用例中，我需要用独特的特定参数来编译模型。值得一提的是，编译并非总是必需的。例如，如果已经缓存的配置对我有效，optimum 会默认使用它。
摘自文档：“Neuron 模型缓存是一个远程缓存，用于存储 neff 格式的已编译 Neuron 模型。它已集成到 NeuronTrainer 和 NeuronModelForCausalLM 类中，以便从缓存加载预训练模型，而不是在本地进行编译。”
现在我使用以下命令编译我选择的模型 meta-llama-3.1-8b-instruct：
!docker run -p 8080:80 -e HF_TOKEN=YOUR_TOKEN \
-v $(pwd):/data \
--device=/dev/neuron0 \
--device=/dev/neuron1 \
--device=/dev/neuron2 \
--device=/dev/neuron3 \
--device=/dev/neuron4 \
--device=/dev/neuron5 \
--device=/dev/neuron6 \
--device=/dev/neuron7 \
--device=/dev/neuron8 \
--device=/dev/neuron9 \
--device=/dev/neuron10 \
--device=/dev/neuron11 \
-ti \
--entrypoint "optimum-cli" ghcr.io/huggingface/text-generation-inference:latest-neuron \
export neuron --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
--sequence_length 16512 \
--batch_size 8 \
--num_cores 8 \
/data/exportedmodel/请注意，在我的用例中，我决定使用大小为 8 的批量大小（batch size）和 8 的张量并行度（tensor parallel degree）。由于一个 inf2.48xlarge 实例有 24 个核心，我可以使用 3 的数据并行度（data parallel），这意味着我将在该实例上拥有 3 个模型的副本。`
优化批量大小以实现最大吞吐量
当为了成本效益而优化硬件利用率时，特别是对于按需每小时 12.98 美元的 inf2.48xlarge 实例，屋顶线模型（roofline model）是一个有价值的框架。
屋顶线模型定义了理论性能的边界。在一个极端，受内存限制的工作负载受限于内存容量，需要频繁的读/写操作。在另一个极端，受计算限制的工作负载充分利用加速器的计算能力，最大化设备上的数据处理。批量大小是控制这种平衡的关键杠杆。较大的批量大小倾向于将工作负载推向计算密集型，而较小的批量大小可能导致更多的内存密集型操作。尽管如此，最大化批量大小并不总是可行的。在指定的延迟预算（我们希望返回响应所花费的时间）内，牢记最大批量大小至关重要。这最直接地由批量大小控制。有关此主题的更多信息，请查看此资源： https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-features/neuroncore-batching.html
创建用于服务的文件
需要几个文件来确保我们的配置设置正确，并确保使用的是我编译的模型而不是缓存的配置。
首先，我需要创建我的 .env 文件，其中指定了我的批量大小、精度等。需要注意的是，由于我编译了我的模型，我需要将 model_id 从通常的 Hugging Face 仓库标识符更改为我在编译命令中指定的容器卷位置。
MODEL_ID='/data/exportedmodel'
HF_AUTO_CAST_TYPE='bf16'
MAX_BATCH_SIZE=8
MAX_INPUT_TOKENS=16000
MAX_TOTAL_TOKENS=16512接下来，我根据我期望的设置创建 benchmark.sh 脚本。
#!/bin/bash
model=${1:-meta-llama/Meta-Llama-3.1-8B-Instruct}
date_str=$(date '+%Y-%m-%d-%H-%M-%S')
output_path="${model//\//_}#${date_str}_guidellm_report.json"
export HF_TOKEN=YOUR_TOKEN
export GUIDELLM__NUM_SWEEP_PROFILES=1
export GUIDELLM__MAX_CONCURRENCY=128
export GUIDELLM__REQUEST_TIMEOUT=60
guidellm \
--target "https://:8080/v1" \
--model ${model} \
--data-type emulated \
--data "prompt_tokens=15900,prompt_tokens_variance=100,generated_tokens=450,generated_tokens_variance=50" \
--output-path ${output_path} \请注意通过 --data 标志传递的参数。由于我的用例是长提示和长生成，我已经相应地设置了 `prompt_tokens` 和 `generated_tokens`。请记住根据你的用例和你预期的输入/输出词元负载来设置这些值。根据这些数字，GuideLLM 将生成大小在 15900 个词元左右的正态分布中的随机提示，并要求生成在 450 个词元左右的正态分布中的随机数量的词元。
docker-compose 文件对于定义你的数据并行非常重要，它通过指定我希望分配给每个容器的设备数量来实现。这也是我指定负载均衡器的地方。
version: '3.7'
services:
tgi-1:
image: ghcr.io/huggingface/text-generation-inference:latest-neuron
ports:
- "8081:8081"
volumes:
- $PWD:/data
environment:
- PORT=8081
- MODEL_ID=${MODEL_ID}
- HF_AUTO_CAST_TYPE=${HF_AUTO_CAST_TYPE}
- HF_NUM_CORES=8
- MAX_BATCH_SIZE=${MAX_BATCH_SIZE}
- HF_TOKEN=YOUR_TOKEN
- MAX_INPUT_TOKENS=${MAX_INPUT_TOKENS}
- MAX_TOTAL_TOKENS=${MAX_TOTAL_TOKENS}
- MAX_CONCURRENT_REQUESTS=512
devices:
- "/dev/neuron0"
- "/dev/neuron1"
- "/dev/neuron2"
- "/dev/neuron3"
tgi-2:
image: ghcr.io/huggingface/text-generation-inference:latest-neuron
ports:
- "8082:8082"
volumes:
- $PWD:/data
environment:
- PORT=8082
- MODEL_ID=${MODEL_ID}
- HF_AUTO_CAST_TYPE=${HF_AUTO_CAST_TYPE}
- HF_NUM_CORES=8
- MAX_BATCH_SIZE=${MAX_BATCH_SIZE}
- HF_TOKEN=YOUR_TOKEN
- MAX_INPUT_TOKENS=${MAX_INPUT_TOKENS}
- MAX_TOTAL_TOKENS=${MAX_TOTAL_TOKENS}
- MAX_CONCURRENT_REQUESTS=512
devices:
- "/dev/neuron4"
- "/dev/neuron5"
- "/dev/neuron6"
- "/dev/neuron7"
tgi-3:
image: ghcr.io/huggingface/text-generation-inference:latest-neuron
ports:
- "8083:8083"
volumes:
- $PWD:/data
environment:
- PORT=8083
- MODEL_ID=${MODEL_ID}
- HF_AUTO_CAST_TYPE=${HF_AUTO_CAST_TYPE}
- HF_NUM_CORES=8
- MAX_BATCH_SIZE=${MAX_BATCH_SIZE}
- HF_TOKEN=YOUR_TOKEN
- MAX_INPUT_TOKENS=${MAX_INPUT_TOKENS}
- MAX_TOTAL_TOKENS=${MAX_TOTAL_TOKENS}
- MAX_CONCURRENT_REQUESTS=512
devices:
- "/dev/neuron8"
- "/dev/neuron9"
- "/dev/neuron10"
- "/dev/neuron11"
loadbalancer:
image: nginx:alpine
ports:
- "8080:80"
volumes:
- ./nginx.conf:/etc/nginx/nginx.conf:ro
depends_on:
- tgi-1
- tgi-2
- tgi-3
deploy:
placement:
constraints: [node.role == manager]最后，我为负载均衡器定义 nginx.conf 文件
### Nginx TGI Load Balancer
events {}
http {
upstream tgicluster {
server tgi-1:8081;
server tgi-2:8082;
server tgi-3:8083;
}
server {
listen 80;
location / {
proxy_pass http://tgicluster;
}
}
}使用 GuideLLM 进行基准测试
现在我已经定义了必要的文件，我开始用 TGI 后端来为我的 optimum-neuron 模型提供服务。
!docker compose -f docker-compose.yaml --env-file .env up
为了确保一切正常，我可以观察上述命令的输出来确认每个容器以及负载均衡器是否都正确启动。一旦我成功启动了容器，我就可以使用先前定义的基准测试脚本开始进行基准测试。
!benchmark.sh "meta-llama/Meta-Llama-3.1-8B-Instruct"
当 guidellm 开始测试你的模型服务设置时，终端将开始填充彩色的标准输出。
性能分析
大约 15-20 分钟后，基准测试完成，并在终端显示以下详细分解：
╭─ Benchmarks ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ [15:02:17] 100% synchronous (0.10 req/sec avg)│
│ [15:04:17] 100% throughput (0.85 req/sec avg)│
│ [15:05:25] 100% constant@0.85 req/s (0.77 req/sec avg) │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
Generating report... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ (3/3) [ 0:05:04 < 0:00:00 ]
╭─ GuideLLM Benchmarks Report (meta-llama_Meta-Llama-3.1-8B-Instruct#2025-05-27-15-02-11_guidellm_report.json) ──────────────────────────────────────────────────────────────────────────────────╮
│ ╭─ Benchmark Report 1 ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ │
│ │ Backend(type=openai_server, target=http://:8080/v1, model=meta-llama/Meta-Llama-3.1-8B-Instruct) │ │
│ │ Data(type=emulated, source=prompt_tokens=15900,prompt_tokens_variance=100,generated_tokens=450,generated_tokens_variance=50, tokenizer=meta-llama/Meta-Llama-3.1-8B-Instruct) │ │
│ │ Rate(type=sweep, rate=None) │ │
│ │ Limits(max_number=None requests, max_duration=120 sec) │ │
│ │ │ │
│ │ │ │
│ │ Requests Data by Benchmark │ │
│ │ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━┓ │ │
│ │ ┃ Benchmark ┃ Requests Completed ┃ Request Failed ┃ Duration ┃ Start Time ┃ End Time ┃ │ │
│ │ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━┩ │ │
│ │ │ synchronous │ 11/11 │ 0/11 │ 113.56 sec │ 15:02:17 │ 15:04:11 │ │ │
│ │ │ asynchronous@0.85 req/sec │ 88/88 │ 0/88 │ 114.59 sec │ 15:05:25 │ 15:07:19 │ │ │
│ │ │ throughput │ 55/55 │ 0/55 │ 64.83 sec │ 15:04:17 │ 15:05:22 │ │ │
│ │ └───────────────────────────┴────────────────────┴────────────────┴────────────┴────────────┴──────────┘ │ │
│ │ │ │
│ │ Tokens Data by Benchmark │ │
│ │ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ │ │
│ │ ┃ Benchmark ┃ Prompt ┃ Prompt (1%, 5%, 50%, 95%, 99%) ┃ Output ┃ Output (1%, 5%, 50%, 95%, 99%) ┃ │ │
│ │ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ │
│ │ │ synchronous │ 15902.82 │ 15896.0, 15896.0, 15902.0, 15913.0, 15914.6 │ 293.09 │ 70.3, 119.5, 315.0, 423.5, 443.1 │ │ │
│ │ │ asynchronous@0.85 req/sec │ 15899.06 │ 15877.4, 15879.4, 15898.5, 15918.0, 15919.8 │ 288.75 │ 24.6, 74.1, 298.5, 452.6, 459.1 │ │ │
│ │ │ throughput │ 15899.22 │ 15879.5, 15883.7, 15898.0, 15914.6, 15920.5 │ 294.24 │ 59.1, 114.9, 285.0, 452.9, 456.4 │ │ │
│ │ └───────────────────────────┴──────────┴─────────────────────────────────────────────┴────────┴──────────────────────────────────┘ │ │
│ │ │ │
│ │ Performance Stats by Benchmark │ │
│ │ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ │ │
│ │ ┃ ┃ Request Latency [1%, 5%, 10%, 50%, 90%, 95%, 99%] ┃ Time to First Token [1%, 5%, 10%, 50%, 90%, 95%, ┃ Inter Token Latency [1%, 5%, 10%, 50%, 90% 95%, ┃ │ │
│ │ ┃ Benchmark ┃ (sec) ┃ 99%] (ms) ┃ 99%] (ms) ┃ │ │
│ │ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ │
│ │ │ synchronous │ 3.68, 5.13, 6.94, 10.91, 13.51, 14.26, 14.87 │ 1563.3, 1569.2, 1576.5, 1589.4, 1594.0, 1595.3, │ 23.2, 28.2, 29.4, 29.8, 30.3, 31.7, 36.5 │ │ │
│ │ │ │ │ 1596.4 │ │ │ │
│ │ │ asynchronous@0.85 req/sec │ 2.62, 6.55, 9.40, 20.66, 30.60, 32.78, 35.07 │ 1594.1, 1602.5, 1605.7, 1629.7, 4650.1, 4924.1, │ 0.2, 0.2, 0.2, 34.3, 44.9, 54.5, 1613.9 │ │ │
│ │ │ │ │ 5345.6 │ │ │ │
│ │ │ throughput │ 18.29, 21.24, 23.81, 44.60, 61.50, 62.80, 63.72 │ 2157.6, 9185.1, 12220.5, 23333.5, 44214.1, │ 28.2, 31.5, 33.1, 39.1, 59.0, 65.2, 1604.6 │ │ │
│ │ │ │ │ 45329.8, 51276.9 │ │ │ │
│ │ └───────────────────────────┴───────────────────────────────────────────────────┴───────────────────────────────────────────────────┴────────────────────────────────────────────────────┘ │ │
│ │ │ │
│ │ Performance Summary by Benchmark │ │
│ │ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┓ │ │
│ │ ┃ Benchmark ┃ Requests per Second ┃ Request Latency ┃ Time to First Token ┃ Inter Token Latency ┃ Output Token Throughput ┃ │ │
│ │ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ │
│ │ │ synchronous │ 0.10 req/sec │ 10.32 sec │ 1585.08 ms │ 29.81 ms │ 28.39 tokens/sec │ │ │
│ │ │ asynchronous@0.85 req/sec │ 0.77 req/sec │ 20.77 sec │ 2401.32 ms │ 63.69 ms │ 221.75 tokens/sec │ │ │
│ │ │ throughput │ 0.85 req/sec │ 43.78 sec │ 24624.46 ms │ 65.18 ms │ 249.64 tokens/sec │ │ │
│ │ └───────────────────────────┴─────────────────────┴─────────────────┴─────────────────────┴─────────────────────┴─────────────────────────┘ │ │
│ ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯解开结果，我们得到了相当多有用的数据点供我们使用。在底层，guidellm 运行三个独立的“负载”来对系统进行基准测试。
- 同步 - 一次只处理一个请求
- 异步 - 以固定的每秒请求数（本例中为 0.85）同时处理多个请求
- 吞吐量 - 服务系统能够承受的最大请求数
通过这些测试，我们得到了每个测试的多个指标，例如成功执行的请求数与失败的请求数、首个 token 的生成时间、提示输入和输出的大小等。在我的实验中，我可以看到在最大负载下，我每秒最多可以服务 0.85 个请求，每个请求的最大延迟略低于 44 秒。根据我的延迟预算，下一步是如果我能容忍更长的响应时间并期望更高的吞吐量，就增加我的批量大小。或者，我可以降低我的批量大小以减少延迟，但代价是可能降低吞吐量。
最后，我的工作负载所需的大量输入和输出词元直接影响了基准测试结果，特别是编码输入上下文所需的时间，这构成了大部分基准测试时间。
结论
在这篇博客文章中，我带您了解了如何编译和加载 Optimum Neuron 模型，如何使用 HuggingFace 文本生成推理容器来提供服务，以及如何对您的设置进行基准测试，以便为您的工作负载进行优化。
参考文献
https://huggingface.co/docs/optimum-neuron/en/guides/cache_system https://github.com/huggingface/optimum-neuron/tree/main/benchmark/text-generation-inference/performance https://github.com/vllm-project/guidellm