跳到主要内容

AWS Trainium / Inferentia

厂商: Amazon AWS

分类: ASIC 专用加速卡

架构: NeuronCore

简介​

AWS 自研 AI 芯片系列：Trainium（训练专用）和 Inferentia（推理专用）。Trainium2 提供极致训练性能，Inferentia2 提供低延迟高吞吐推理，通过 Amazon EC2 Trn1/Inf2 实例提供。

规格参数​

型号算力显存/内存接口TDP制程
Trainium2676 TFLOPS (FP8)192GB HBM3eNeuronLink500W5nm
Inferentia2220 TOPS (FP16)64GB HBM2eNeuronLink175W5nm

官方网站​

访问官方网站

驱动下载​

Linux​

相关文档​

操作系统支持​

WindowsLinuxmacOSAndroid
❌✅ (AWS EC2)❌❌

版本历史​

版本发布时间说明
Neuron SDK 3.02025-Q1Trainium2 全面支持
Neuron SDK 2.x2024Inferentia2 支持

性能基准​

型号任务性能指标
Trainium2 × 16Llama 3 405B 训练~2 天 (AWS 数据)
Inferentia2Llama 2 13B 推理~500 tok/s
Inferentia2Stable Diffusion XL~5s/img (批量)

定价信息​

型号参考价格备注
Trainium2 (Trn2)~$5.00/chip/hEC2 实例按需
Inferentia2 (Inf2)~$0.96/chip/h按需计费
Inf2 (xlarge)~$1.16/h最小实例

快速安装​

AWS (通过 EC2)​

# 1. 启动 Trn2 实例
aws ec2 run-instances --instance-type trn2.48xlarge --region us-east-1

# 2. 连接后安装 Neuron SDK
pip install torch-neuron neuronx-cc

# 3. 验证
neuron-ls

Trainium/Inferentia 仅在 AWS EC2 可用，无法单独购买裸卡。

代码示例​

Python (Neuron SDK)​

import torch
import torch_neuronx

# 将 PyTorch 模型编译为 Neuron 格式
model = torch.nn.Linear(1024, 1024)
model = model.to("xla")  # XLA 后端

# 编译
neuron_model = torch_neuronx.trace(model, example_input)

# 推理
output = neuron_model(input_data)

架构特色​

  • NeuronCore: AWS 自研 AI 加速器架构，Trainium 面向训练，Inferentia 面向推理
  • NeuronLink: 芯片间互联，支持大规模分布式训练
  • Trn2 UltraServers: 每节点 16 个 Trainium2 芯片，通过 NeuronLink Connect 互联

模型兼容性​

模型/框架支持情况备注
PyTorch✅ Neuron SDKtorch-neuronx
TensorFlow✅neuron-cc
JAX✅jax-neuronx
Llama / QLM 等 LLM✅Inf2 推理最优
HuggingFace✅transformers-neuronx

相关产品​

如果你在评估替代方案，以下产品可能也适合你的场景：