跳至内容

AWS Neuron

AWS Neuron 是用于在基于AWS Inferentia和AWS Trainium的Amazon EC2实例及UltraServers（Inf1、Inf2、Trn1、Trn2和Trn2 UltraServer）上运行深度学习和生成式AI工作负载的软件开发工具包（SDK）。Trainium和Inferentia均由完全独立的异构计算单元NeuronCore驱动。本文介绍如何配置环境以在Neuron上运行vLLM。

警告

该设备没有预构建的wheel包或镜像，因此您必须从源代码构建vLLM。

要求

  • 操作系统: Linux
  • Python: 3.9或更新版本
  • PyTorch 2.5/2.6
  • 加速器：NeuronCore-v2（用于trn1/inf2芯片）或NeuronCore-v3（用于trn2芯片）
  • AWS Neuron SDK 2.23

配置新环境

启动Trn1/Trn2/Inf2实例并验证Neuron依赖项

启动预装Neuron依赖的Trainium或Inferentia实例的最简单方法是按照这份快速入门指南，使用Neuron深度学习AMI（亚马逊机器镜像）。

  • 启动实例后，按照连接到您的实例中的说明连接到实例
  • 进入实例后，运行以下命令激活预装的虚拟环境以进行推理
source /opt/aws_neuronx_venv_pytorch_2_6_nxd_inference/bin/activate

请参考NxD推理设置指南获取其他设置说明，包括使用Docker和手动安装依赖项。

注意

NxD Inference 是默认推荐的用于在 Neuron 上运行推理的后端。如果您希望使用传统的 transformers-neuronx 库，请参考 Transformers NeuronX 设置。

使用Python设置

预构建的wheel包

目前还没有预构建的Neuron wheel包。

从源码构建wheel

要从源代码构建并安装vLLM，请运行：

git clone https://github.com/vllm-project/vllm.git
cd vllm
pip install -U -r requirements/neuron.txt
VLLM_TARGET_DEVICE="neuron" pip install -e .

AWS Neuron 维护了一个vLLM的Github分支在https://github.com/aws-neuron/upstreaming-to-vllm/tree/neuron-2.23-vllm-v0.7.2，该分支除了包含vLLM V0版本的功能外，还具备多项额外特性。请使用AWS分支以获得以下功能：

  • Llama-3.2 多模态支持
  • 多节点分布式推理

更多详情和使用示例请参阅vLLM用户指南中的NxD推理部分。

要安装AWS Neuron分支版本，请运行以下命令：

git clone -b neuron-2.23-vllm-v0.7.2 https://github.com/aws-neuron/upstreaming-to-vllm.git
cd upstreaming-to-vllm
pip install -r requirements/neuron.txt
VLLM_TARGET_DEVICE="neuron" pip install -e .

请注意，AWS Neuron分支仅旨在支持Neuron硬件；未测试与其他硬件的兼容性。

使用Docker进行设置

预构建镜像

目前还没有预构建的Neuron镜像。

从源码构建镜像

查看deployment-docker-build-image-from-source获取构建Docker镜像的说明。

请确保使用 docker/Dockerfile.neuron替代默认的Dockerfile。

额外信息

通过NxD推理后端支持的功能特性

当前的vLLM与Neuron集成依赖于neuronx-distributed-inference（首选）或transformers-neuronx后端来执行大部分繁重工作，包括PyTorch模型初始化、编译和运行时执行。因此，Neuron支持的大多数功能也可以通过vLLM集成获得。

要通过vLLM入口点配置NxD推理功能，请使用override_neuron_config设置。以字典形式（或从CLI启动vLLM时使用JSON对象）提供您想要覆盖的配置。例如，要禁用自动分桶功能，请包含

override_neuron_config={
    "enable_bucketing":False,
}

或者在通过命令行界面启动vLLM时，传递

--override-neuron-config "{\"enable_bucketing\":false}"

或者，用户可以直接调用NxDI库来追踪和编译模型，然后通过NEURON_COMPILED_ARTIFACTS环境变量在vLLM中加载预编译产物以运行推理工作负载。

已知限制

  • EAGLE推测解码：NxD推理要求EAGLE草稿检查点包含目标模型的LM头权重。请参阅此指南了解如何将预训练的EAGLE模型检查点转换为与NxDI兼容。
  • 量化：vLLM中原生的量化流程在NxD Inference上支持不佳。建议按照此Neuron量化指南使用NxD Inference对模型进行量化和编译，然后将编译后的产物加载到vLLM中。
  • 多LoRA服务：NxD推理仅支持在服务器启动时加载LoRA适配器。目前不支持在运行时动态加载LoRA适配器。请参考multi-lora example
  • 多模态支持：多模态支持仅通过AWS Neuron分支提供。该功能尚未合并到vLLM主分支，因为NxD Inference目前依赖对核心vLLM逻辑的特定修改来实现此功能。
  • 多节点支持：仅在AWS Neuron分支上支持跨多个Trainium/Inferentia实例的分布式推理。请参考此多节点示例运行。请注意，张量并行（跨NeuronCore的分布式推理）在vLLM主分支中可用。
  • Known edge case bug in speculative decoding: An edge case failure may occur in speculative decoding when sequence length approaches max model length (e.g. when requesting max tokens up to the max model length and ignoring eos). In this scenario, vLLM may attempt to allocate an additional block to ensure there is enough memory for number of lookahead slots, but since we do not have good support for paged attention, there isn't another Neuron block for vLLM to allocate. A workaround fix (to terminate 1 iteration early) is implemented in the AWS Neuron fork but is not upstreamed to vLLM main as it modifies core vLLM logic.

环境变量

  • NEURON_COMPILED_ARTIFACTS：设置此环境变量指向您预编译的模型工件目录，以避免服务器初始化时的编译时间。如果未设置此变量，Neuron模块将执行编译并将工件保存在模型路径下的neuron-compiled-artifacts/{unique_hash}/子目录中。如果设置了此环境变量但目录不存在或内容无效，Neuron也将回退到新编译并将工件存储在此指定路径下。
  • NEURON_CONTEXT_LENGTH_BUCKETS: 用于上下文编码的分桶大小。(仅适用于transformers-neuronx后端)。
  • NEURON_TOKEN_GEN_BUCKETS: 用于令牌生成的桶大小。(仅适用于transformers-neuronx后端)。
优云智算