AWS Trainium & Inferentia 文档

EC2 设置

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

EC2 设置

本指南将帮助您启动并运行 Optimum Neuron。主要有两种方法

  1. 🚀 推荐:使用深度学习 AMI 的 AWS EC2 - 使用预配置环境的最简单入门方式
  2. ⚙️ 手动安装 - 在现有基础设施上安装 Optimum Neuron

推荐:使用深度学习 AMI 的 AWS EC2

在 Amazon EC2 上使用 AWS Trainium 或 Inferentia 和 Optimum Neuron 最简单的方法是使用 Hugging Face Neuron 深度学习 AMI (DLAMI)。DLAMI 预装了所有必需的库,包括 Optimum Neuron、Neuron 驱动程序、Transformers、Datasets 和 Accelerate。HF DLAMI 对 Amazon EC2 用户是免费提供的。

Optimum Neuron 支持 Inf1、Inf2、Trn1 和 Trn2,所有这些都可以在 Amazon EC2 上访问。您可以在这里的“加速计算”部分找到 Trn 和 Inf 实例的所有规格。

在本节中,我们将向您展示

  1. 如何使用 HF DLAMI 在 Amazon EC2 上创建 AWS Trainium 或 Inferentia 实例
    1. 查找支持的区域
    2. 增加服务配额
    3. 使用 HF DLAMI 启动 Amazon EC2 实例
    4. 通过 SSH 连接
  2. 如何设置您的远程开发环境
    1. 通过 Jupyter Notebook 访问
    2. 通过 VS Code 远程服务器访问

使用 HF DLAMI 在 Amazon EC2 上创建 AWS Trainium 或 Inferentia 实例

在创建 EC2 实例之前,请确保您位于所选实例支持的区域,并且您的 AWS 账户中有配额。

查找支持的区域

截至 2025 年 2 月,支持至少一种 Trainium 或 Inferentia2 实例类型的区域列表如下:

  • us-east-1: 美国东部(弗吉尼亚北部)
  • us-east-2: 美国东部(俄亥俄)
  • us-west-2: 美国西部(俄勒冈)
  • ap-south-1: 亚太地区(孟买)
  • ap-northeast-1: 亚太地区(东京)
  • ap-southeast-1: 亚太地区(新加坡)
  • ap-southeast-2: 亚太地区(悉尼)
  • ap-southeast-4: 亚太地区(墨尔本)
  • eu-north-1: 欧洲(斯德哥尔摩)
  • eu-west-3: 欧洲(巴黎)
  • eu-west-2: 欧洲(伦敦)
  • eu-west-1: 欧洲(爱尔兰)
  • eu-central-1: 欧洲(法兰克福)
  • sa-east-1: 南美洲(圣保罗)

这里有一个 Python 脚本,可以帮助您在每个已启用的区域中获取最新的受支持实例类型

import boto3  
from datetime import datetime

ec2 = boto3.client('ec2')  

regions = [region['RegionName'] for region in ec2.describe_regions()['Regions']]  


#Edit this line to change the instance types displayed  
instance_types = ['trn1.32xlarge', 'trn1.2xlarge', 'inf2.48xlarge', 'inf2.24xlarge', 'inf2.8xlarge', 'inf2.xlarge', 'trn2.48xlarge']  
  
supported_regions = {}  
  
for region in regions:  
   ec2_region = boto3.client('ec2', region_name=region)  
   response = ec2_region.describe_instance_type_offerings(  
      #LocationType='availability-zone',  
      Filters=[  
        {'Name': 'instance-type', 'Values': instance_types},  
      ]  
   )
   if response['InstanceTypeOfferings']:  
      supported_regions[region] = [offer['InstanceType'] for offer in response['InstanceTypeOfferings']]  

print('# Supported Regions as of',datetime.now().strftime('%B %d, %Y'))
print('================')  


client = boto3.client('ssm')

for region, instance_types in supported_regions.items():
    try:
        response = client.get_parameter(Name=f'/aws/service/global-infrastructure/regions/{region}/longName')
        region_long_name = response['Parameter']['Value']
    except (client.exceptions.ParameterNotFound, KeyError):
        region_long_name = region
    print(f' * {region}: {region_long_name}')
    for instance_type in instance_types:
      print(f'  - {instance_type}')
    print('\n') 

增加服务配额

现在您已经选择了区域并切换到该区域,您可以通过 AWS 控制台请求增加服务配额。导航到服务配额,在左侧面板中选择 AWS 服务,然后搜索 Amazon EC2,再搜索 “trn” 或 “inf”。您可以分别为按需实例和竞价实例请求增加配额。

默认情况下,Inferentia 和 Trainium 的所有配额均为 0。增加配额不收取任何费用。Inferentia 和 Trainium 有单独的配额,竞价实例和按需实例也有单独的配额。配额指的是分配给每种实例类型的 vCPU 总数的最大值。

例如,192 的配额可以让您运行单个 inf2.48xlarge、两个 inf2.24xlarge、六个 inf2.8xlarge 或四十八个 inf2.xlarge。它还允许您运行 inf1 实例类型。同样,对于 Trainium,128 的配额可以让您运行单个 trn1n.32xlarge 或 trn1.32xlarge,但它也允许您运行十六个 trn1.2xlarge。

使用 HF DLAMI 启动 Amazon EC2 实例

让我们通过 EC2 控制台在 us-east-1 区域(弗吉尼亚北部)部署一个 trn1.2xlarge 实例。

首先,单击 启动实例 并为实例定义一个名称 (trainium-huggingface-demo)。

name instance

接下来,在 Amazon Marketplace 中搜索 Hugging Face AMI。在“应用程序和操作系统映像”的搜索栏中输入“Hugging Face”并按“enter”。

search ami

现在应该会打开带有搜索结果的“选择亚马逊机器映像”视图。您可以导航到“AWS Marketplace AMI”,找到 Hugging Face Neuron 深度学习 AMI 并点击选择。

select ami

如果您尚未订阅,系统会要求您订阅。AMI 完全免费,您只需支付 EC2 的计算费用。

然后您需要定义一个密钥对,它将用于通过 ssh 连接到实例。如果您没有密钥对,可以就地创建一个。

select ssh key

之后,创建或选择一个允许 ssh 流量的安全组

select security group

您现在可以启动实例了。因此,请点击右侧的“启动实例”。

select ssh key

AWS 现在将使用 Hugging Face Neuron 深度学习 AMI 配置实例。

通过 SSH 连接

实例准备就绪后,您可以查看并复制公共 IPv4 地址以通过 ssh 连接到机器。

select public dns

将下面代码段中的空字符串 "" 替换为您的实例的 IP 地址以及您在启动实例时创建/选择的密钥对的路径。

PUBLIC_DNS="" # IP address
KEY_PATH="" # local path to key pair

ssh -i $KEY_PATH ubuntu@$PUBLIC_DNS

连接成功后,您可以运行 neuron-ls 来确保您有权访问 Trainium 加速器。您应该会看到类似下面的输出。

ubuntu@ip-172-31-79-164:~$ neuron-ls
instance-type: trn1.2xlarge
instance-id: i-0570615e41700a481
+--------+--------+--------+---------+
| NEURON | NEURON | NEURON |   PCI   |
| DEVICE | CORES  | MEMORY |   BDF   |
+--------+--------+--------+---------+
| 0      | 2      | 32 GB  | 00:1e.0 |
+--------+--------+--------+---------+

设置您的远程开发环境

我们将逐步介绍如何在 Amazon EC2 实例上设置 Jupyter Notebooks 或 VS Code 远程服务器。

这两种方法都需要某种形式的 SSH 连接。这些说明是为 Mac 编写的,但也应该适用于 Linux 系统。PC 可能需要使用 Putty。

您应该有一个在部署实例时创建的 .pem 文件,或者来自之前的部署。您可以使用以下命令连接到您的系统

ssh -i "/path/to/sshkey.pem" ubuntu@instance_ip_address

通过 Jupyter Notebook 访问

该方法涉及在 Neuron 实例上运行 Jupyter notebook 服务器,本地映射一个端口,然后使用桌面上的浏览器访问 notebook 服务器。

首先,将本地机器上的一个端口映射到 Neuron 实例。在您的系统终端中运行

ssh -i "/path/to/sshkey.pem" -N -f -L localhost:8888:localhost:8888 ubuntu@instance_ip_address

然后从您的计算机使用 SSH 连接到您的 Amazon EC2 实例。连接后,在命令提示符下运行

nohup jupyter notebook --no-browser --port=8888

几秒钟后,检查 nohup.out 文件以找到您的服务器令牌

cat nohup.out | grep localhost

复制连接字符串并将其粘贴到浏览器中。几秒钟后,您应该会看到 Jupyter Notebook 浏览器。它看起来应该像 https://:8888/tree?token=337fc8de2aenot_a_real_tokene952c43946e4fb57131

这个过程之所以有效,是因为您已将本地机器上的 8888 端口映射到 Neuron 实例上的 8888 端口,因此当您连接到 localhost:8888 时,您最终会访问 Neuron 实例上的 Jupyter 服务器。

如果遇到问题,请确保初始端口映射成功。如果您的机器上已经有东西在 8888 端口上运行,这可能会给您一个错误。如果需要,您可以随时在所有说明中更改端口(例如 8885)。

通过 VS Code 远程服务器访问

在本地计算机上安装 Visual Studio Code 后,您可以使用 Remote-SSH 命令编辑和运行存储在 Neuron 实例上的文件。有关更多详细信息,请参阅 VS Code 文章。

  1. 从命令面板(F1, ⇧⌘P)中选择“远程 SSH: 连接到主机…”
  2. 输入上面 ssh 部分的完整连接字符串:ssh -i “/path/to/sshkey.pem” ubuntu@instance_ip_address
  3. VS Code 应该会连接并自动设置 VS Code 服务器。
  4. 最终,您应该会被提示输入一个基本目录。您可以浏览到 Neuron 实例上的一个目录。
  5. 如果您发现菜单中的某些命令似乎变灰了,但键盘命令仍然有效(⌘S 保存或 ^⇧` 打开终端),您可能需要重启 VS Code。

替代方案:手动安装

手动安装在以下几种情况下非常有用:

  • 使用较新版本:安装最新的 Optimum Neuron 版本,该版本可能尚未在 DLAMI 中提供
  • 自定义 AMI 要求:使用您组织的标准 AMI 或安全加固的映像
  • 现有基础设施:向预配置的环境或 Docker 容器添加 Neuron 支持
  • 开发设置:安装预发布或开发版本进行测试
  • 最小化安装:创建仅包含所需依赖项的轻量级环境

如果选择手动安装,您需要确保首先正确安装了 Neuron 驱动程序和工具。

在安装 optimum-neuron 之前,请确保您已安装 Neuron 驱动程序和工具,请查看此处的更详细指南

添加 pip 包 URL

指向 AWS Neuron 存储库

python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

为 AWS Trainium (trn1) 或 AWS inferentia2 (inf2) 安装 optimum-neuron

python -m pip install optimum-neuron[neuronx]

为 AWS inferentia (inf1) 安装 optimum-neuron

python -m pip install optimum-neuron[neuron]

下一步是什么?

现在您已经设置好了 Optimum Neuron,请查看快速入门指南,学习使用 Optimum Neuron 进行训练和推理的基础知识。