AWS 发布 TorchNeuron:Trainium 的原生 PyTorch 后端
Introducing TorchNeuron: Native PyTorch backend for AWS Trainium
AWS 宣布 TorchNeuron,一个面向 Trainium 的开源原生 PyTorch 后端,10 月 20 日起以 Apache 2.0 公开 Beta 发布于 github.com/aws-neuron/torch-neuronx。
原文给出了 TorchNeuron 的执行模式、分布式支持与实测覆盖情况,读者可以据此评估现有 PyTorch 工作流迁移到 Trainium 的成本。
前沿模型涵盖语言、视频和音频,架构从稠密和混合专家 Transformer 到扩散、状态空间、世界和结构预测模型。PyTorch 是 AI 研究发生的地方,是定义模型的地方,也日益成为生产运行的地方。研究人员希望进行实验和创新,客户则希望获得 AWS Trainium 提供的性能,以实现其性能和规模目标。训练和推理库是针对 PyTorch 自身的 API 编写的,新的库和新版本持续发布。过去几年中,PyTorch 使这些 API 与加速器无关,并提供了用于注册自定义硬件的官方扩展点。
今天,为满足开发者和研究社区对 Trainium 日益增长的需求,我们很高兴在 10 月 20 日于北美 PyTorch Conference 开源发布之前,分享一种更简单的方法,将芯片作为标准 PyTorch 设备进行编程。AWS Trainium 是一系列 AI 芯片,专为下一代云规模 AI 训练和推理从头设计。我们宣布推出 TorchNeuron,这是一个面向 Trainium 的开源原生 PyTorch 后端,使模型代码、库以及你已在使用的工具都能在 Trainium 上运行。这一集成直接建立在 PyTorch 社区创建的加速器无关扩展点之上,包括 PrivateUse1 设备后端和 torch.accelerator API。
TorchNeuron 在 Trainium2 和 Trainium3 上支持用于研究和调试的 eager 执行(包括性能改进的全新自适应 eager 模式)、用于生产性能的 torch.compile,以及用于自定义内核的 Neuron Kernel Interface (NKI)。此外,借助 Neuron Explorer,性能工程师可以使用相同的 torch.profiler 追踪从框架到硬件的执行过程,将每个事件链接回其源代码行,从而精确定位瓶颈并高效优化。TorchNeuron 将以 Apache 2.0 许可证在 github.com/aws-neuron/torch-neuronx 开源发布,并于 10 月 20 日作为公开 Beta 版上线。
“PyTorch 的愿景很明确:开发者应在不同硬件平台上获得一致的体验。对 AWS Trainium 的原生支持扩展了这一选择,使研究人员和工程师能够使用他们已熟悉的 eager 执行和分布式 API 进行实验和迭代。团队可以利用 Trainium 的能力,而无需重新构建工作流或离开 PyTorch 生态系统。”——Meta PyTorch 工程总监 Jana van Greunen
快速开始:轻松切换设备
快速开始非常简单。使用指针在 Trainium 上运行你的代码:.to('neuron')。
# Run a training loop on Trainium
model = MyModel().to('neuron') # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
optimizer.zero_grad()
loss = criterion(model(batch), targets)
loss.backward()
optimizer.step()
针对 PyTorch 推荐的设备无关 API torch.accelerator 编写的现代代码无需修改即可运行。编译模型需要指定一次 Neuron 后端,使用 torch.compiler.set_default_backend("neuron"),之后由 torch.compile(model) 接管,或按调用使用 torch.compile(model, backend="neuron")。
# Run a compiled training loop on Trainium
torch.compiler.set_default_backend("neuron") # call once, before any compile
model = MyModel().to('neuron')
model.compile() # no backend argument
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
optimizer.zero_grad()
loss = criterion(model(batch), targets)
loss.backward()
optimizer.step()
功能
TorchNeuron 通过 PrivateUse1(PyTorch 用于自定义硬件的官方扩展点)注册 Trainium,因此标准 API 可按原样工作:torch.nn、torch.tensor、torch.autocast、torch.distributed 和 torch.profiler 均在 neuron 设备上运行。PyTorch 的两种执行模式 eager 和 torch.compile 均可原生工作。

TorchNeuron 的高层架构图(详情)。
Eager
Eager 模式会立即执行操作,因此 Python 控制流、打印和断点都能正常工作,无需编译或跟踪步骤。这是用于研究和调试的模式。Eager 模式会逐个分派操作。每个操作直接执行时都会编译自己的 Neuron 可执行文件格式(NEFF)文件,在设备上启动它,并通过 HBM 移动其中间结果。我们还引入了自适应 eager 模式(实验性),以降低启动开销,并提供更好的操作融合机会,在不修改代码的情况下接近编译图的性能,同时保留交互式体验。
torch.compile
torch.compile(model, backend="neuron") 将模型捕获为图,由 Neuron 编译器进行优化,是用于稳定、高性能模型投入生产的模式。torch.compile 可用于优化整个模型或部分模型,并与 eager 模式无缝协作。原生支持 FlexAttention。TorchDynamo 将模型捕获为 FX 图,AOTAutograd 提前跟踪反向图,Neuron 后端通过 Torch-MLIR 将结果降为 StableHLO,Neuron 编译器生成在硬件上执行的 NEFF 文件。
Neuron 推理栈基于 TorchNeuron 中相同的 torch.compile 支持构建。vLLM Neuron 是 vLLM 的插件,用于服务 LLM;vLLM Omni Neuron 是 vLLM Omni 的插件,用于服务图像和视频生成的扩散模型。
算子与数据类型
广泛覆盖 aten 算子,涵盖文本、多模态与视觉语言、图像与视频生成和理解、状态空间模型以及世界模型中出现的算子。模型可以使用 BF16、FP32 和 FP16,以及整数和布尔类型,torch.autocast 和混合精度代码按原样运行。TorchAO 和 FP8 目前处于私有 Beta 阶段,将在下一版本中正式可用。
分布式 API
分布式推理和训练可以使用标准 PyTorch API,高级并行模块可按原样与你的模型组合:DDP、FSDP2、SimpleFSDP、张量并行和流水线并行,以及 FSDP2 与张量并行、ZeRO 等常见组合,这些对训练大型模型至关重要。支持 DTensor 和 DeviceMesh 来处理分片和多维进程组。支持集合通信 API(例如:all_gather、reduce_scatter、all_reduce)和 P2P API(send、recv)。集合通信在节点内通过 Trainium 专用的集合通信核心执行,在节点间通过 Elastic Fabric Adapter(EFA)执行,因此通信传输原生适配硬件,无需更改 PyTorch API。
流与事件
Torch 流和事件原生受支持。torch_neuronx.Stream 和 torch_neuronx.Event 将 PyTorch 的流语义映射到 Trainium 的硬件队列上。在同一个流内,数据移动和计算会自动保持顺序。不同的流之间没有隐式依赖关系,因此跨流时,你需要通过在生产者流上记录事件,并在消费者流上等待该事件,然后再使用结果,来显式地建立这种顺序。与设备无关的 torch.Stream、torch.Event 和 torch.accelerator 入口点会到达相同的实现,因此针对它们编写的代码无需修改即可运行。错误会在下一个同步点暴露,而不是在导致该错误的操作处暴露;TorchNeuron 会跟踪最早的异常,并在该同步点以原始上下文重新抛出它。在执行时,非默认流会多路复用到单个硬件队列上,因此它们会保留顺序、优先级和事件语义,但不会各自获得专用引擎;并且当主机集合通信被禁用时,在非默认流上发出的集合通信会保留其在该流执行顺序中的位置,但会在默认硬件队列上运行。
Neuron 内核接口(NKI)
Neuron 内核接口(NKI)让研究人员和开发者能够为 eager 和 torch.compile 开发、优化并执行自定义算子。研究人员可以针对需要底层硬件控制的想法进行原型验证,包括新的模型架构。开发者可以通过在模型代码中直接调用性能关键的 NKI 内核来替代一系列标准 PyTorch 操作(ATen 操作),从而获得超越 Adaptive Eager 和 torch.compile 的更精细的性能控制。
TorchNeuron 还提供由 NKI 实现、针对 Trainium 优化的高性能算子,并以开源形式发布为 NKI-lib,覆盖稠密模型和 MoE 模型的主流 LLM 与多模态场景,例如 FlashAttention 和 MoE blockwise 内核,开发者可以使用它们来提升在 Trainium 上的性能。以下示例使用 PyTorch 推荐的 register_autograd 风格定义并注册了一个自定义 NKI 算子:
import torch
import nki
import nki.language as nl
from torch_neuronx import nki_op
# Declaring and implementing NKI kernels
# Note in NKI kernels, the output tensor should be defined as parameter to the kernel
@nki.jit
def sin_kernel(in_ptr0, out_ptr):
input_tile = nl.load(in_ptr0[0:128])
output_tile = nl.sin(input_tile)
nl.store(out_ptr[0:128], value=output_tile)
return out_ptr
# defining cos since it will be used in backward path of sin
@nki.jit
def cos_kernel(in_ptr0, out_ptr):
input_tile = nl.load(in_ptr0[0:128])
output_tile = nl.cos(input_tile)
nl.store(out_ptr[0:128], value=output_tile)
return out_ptr
# after this line, there is no NKI code, just native PyTorch, registering custom ops
@nki_op("mylib::pt_nki_sin", mutates_args={})
def pt_nki_sin(x: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
sin_kernel(x, out)
return out
@nki_op("mylib::pt_nki_cos", mutates_args={})
def pt_nki_cos(x: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
cos_kernel(x, out)
return out
def pt_nki_sin_backward(ctx, grad):
x, = ctx.saved_tensors
return grad * pt_nki_cos(x) # cos is derivative of sin
def pt_sin_setup_context(ctx, inputs, output):
x, = inputs
ctx.save_for_backward(x)
# Register backward pass and context function used during autograd for nki_sin
pt_nki_sin.register_autograd(pt_nki_sin_backward, setup_context=pt_sin_setup_context)
if __name__ == "__main__":
x = torch.randn(128, 1, device="neuron", requires_grad=True)
# Forward pass uses sin_kernel
y = pt_nki_sin(x)
# Backward pass, autograd calls uses nki_sin_backward() -> cos_kernel()
loss = y.sum()
loss.backward()
使用 TorchNeuron 进行性能分析
TorchNeuron 与标准 torch.profiler.profile API 集成,并扩展了一个 NeuronConfig 对象,用于控制 Neuron 特定的数据收集。这使用户能够捕获并关联来自框架级 CPU 操作(PyTorch 调度、Python 栈)以及 Neuron 运行时和 Trainium 活动的跟踪数据(由 NeuronConfig 控制),从而实现端到端性能分析。
捕获的性能分析数据可以在 Neuron Explorer UI 中可视化,用户可以跨框架、运行时和设备跟踪执行过程,从任意事件跳转到确切的源代码行(源代码查看器),通过调用层级浏览工作负载(层级查看器),并沿着依赖链查看每个停顿在等待什么(依赖链查看器)。
Neuron Explorer 还会输出一个 Parquet 输出,其中包含每个算子和每个引擎的指标,因此 UI 背后的同一份数据可以直接查询,或由代理用于自动化瓶颈分析。对于程序化访问,用户可以使用内置的 torch.profiler 工具。例如,key_averages 表让用户可以深入查看特定算子,并使用 sort_by 参数显示每个算子的设备侧开销,这对于识别热点很有用。
以下代码片段展示了一个示例(详情):
import torch
from torch.profiler import profile, ProfilerActivity
from torch_neuronx.profiling import NeuronConfig, ProfileMode, NeuronProfiler
# Configure Neuron profiling and pass appropriate path as OUTPUT_DIR
nc = NeuronConfig(
modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
profile_output_dir=OUTPUT_DIR,
capture_enabled_for_nc='0',
correlate_device_time=True,
)
# NeuronProfiler.export_trace() places the framework trace
# in the same directory as Neuron system traces
exporter = NeuronProfiler(nc)
device = torch.device('neuron')
x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)
def rms_norm_fn(t):
x_sq = t * t
mean_sq = x_sq.mean(dim=-1, keepdim=True)
rms = torch.rsqrt(mean_sq + 1e-6)
return t * rms * weight
# Profile your workload
with torch.profiler.profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
experimental_config=nc,
record_shapes=False,
on_trace_ready=exporter.export_trace,
schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
for step in range(4):
with torch.no_grad():
_ = rms_norm_fn(x)
torch.neuron.synchronize()
prof.step()
# Print the per-op summary, sorted by device time to surface hotspots
print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))


有关更多能力和功能,请参阅 Neuron Explorer 文档。
基于 TorchNeuron 构建的客户与合作伙伴
TorchNeuron 已被初创公司、企业、机器学习行业合作伙伴和学术机构广泛采用:涵盖自动驾驶、机器人、视频生成、语音、电信、医疗健康和地理空间建模。以下是各团队当前构建成果的概览。
“Cartesia 正在构建能够像人类一样从世界中学习并与之交互的智能,起点是对话。我们在 Mamba 等状态空间架构上的开创性工作支撑了 Sonic(文本转语音)和 Ink(语音转文本),这些语音 AI 模型驱动着真正像人类一样的对话体验。真实对话同时需要三样东西:质量(自然度和准确度)、效率(大规模下的低延迟)和可控性。大多数解决方案只能提供其中一两个。我们从第一性原理出发设计研究和产品,以同时实现这三点,这意味着我们需要能够在不重写技术栈的情况下,从每一层榨取最大性能的基础设施。借助 AWS Trainium 上的 PyTorch 原生支持,我们以对 PyTorch 代码的最小改动运行了标准推理流水线。小而快的算子通过原生路径高效映射,因此我们无需为它们手写内核。这让我们团队能够专注于为被反复调用、决定端到端延迟的层编写高性能 NKI 内核。对于实时语音,每一毫秒的网络距离都至关重要。在 Trainium 上运行有助于将 Sonic 和 Ink 部署到全球客户附近。”
Arjun Desai,Cartesia 联合创始人
“语音正在成为 AI 的一等模态,Deepgram 每年为超过 200,000 名开发者和 1,400 家组织处理数千万亿的音频 token。为满足这一需求,我们需要高效扩展,同时持续交付更好的模型。在为期数月的合作中,AWS 生成式 AI 创新中心帮助我们将其模型移植到 AWS Trainium 上的 PyTorch 原生支持。在我们评估的基准测试中,Trainium 每美元处理的输入音频秒数比我们现有的 GPU 基线多约 60%,这有助于我们以更低价格向客户提供更好的模型。我们期待在解决方案中支持未来的 AWS AI 芯片。”
Adam Sypniewski,Deepgram CTO
“Hugging Face 是领先的 AI 构建者开放平台,通过 Transformers 和 Diffusers 等由 PyTorch 驱动的开源库,免费提供对超过 600 万个开放模型、数据集和应用的访问。此前,要利用 AWS Trainium 需要用 Neuron 专用软件和配置来适配 Hugging Face 模型;TorchNeuron 使 Trainium 成为 Transformers 和 Diffusers 的原生后端,因此开发者可以通过完全开源的栈使用熟悉的模型定义——这切实体现了我们对芯片中立性和兼容性的承诺。现在,在 AWS Trainium 上部署最先进的开放 LLM 以及最新的开放图像和视频生成模型比以往任何时候都更容易,TRL 支持微调和强化学习,Neuron 内核通过我们的库自动提供优化性能。”
Jeff Boudier,Hugging Face 产品副总裁
在这些客户之外,我们还一直在与众多合作伙伴合作,他们正在基于 TorchNeuron 进行构建,以增强 Trainium 生态系统。这些合作伙伴包括 Anyscale、Decart、Kernelize、Loka、Luminal、Makora、Modular、Nota AI、Radix、Wafer 和 Yasp。
TorchNeuron 如何实现更广泛的兼容性
我们对原生 PyTorch 支持的目标之一,是让开发者能够利用 Trainium 的性能和成本优势,同时继续使用他们已经熟悉的模型、库和工作流。这有助于保护他们在 PyTorch 上的投资,包括其模型代码、推理和训练基础设施、调试实践,以及他们所依赖的社区项目。
运行 AI 工作负载的成本包括使其正常工作、优化其性能,以及随着模型和库的演进而维护它们所需的工程时间。我们希望开发者花更多时间试验架构和改进模型质量,减少将软件适配底层硬件的时间。他们应当能够从熟悉的模型或训练配方开始,将其适配到自己的工作负载,并在 Trainium 上运行,只需做极少的 Neuron 特定改动。
TorchNeuron 通过与更广泛的 PyTorch 兼容,为这一体验提供了基础。开发者可以使用 TorchTitan 等库——它汇集了 PyTorch 的分布式和编译能力以支持大规模训练——以及 Hugging Face Transformers,后者提供对社区模型实现、预训练检查点以及用于将模型适配到其应用的 API 的访问。
TorchTitan 示例
TorchTitan 是 PyTorch 用于快速实验和大规模训练的原生库。其模块化设计使自定义训练、测试新想法和扩展实验变得更加容易。TorchTitan 目前已在 Trainium 上训练。我们持续向上游贡献到主线 TorchTitan,torchtitan-neuron 会在这些改进合并之前提供给你,并在每项改进落地时与主线保持一致。Trainium 特定的产物也随 torchtitan-neuron 发布,并通过 TorchTitan 自身的扩展点注册;这些包括 NKI 内核、为 Trainium 调优的模型配置和训练配方,以及 Neuron Explorer 的性能分析集成。
# Install from source
git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan
pip install -r requirements.txt
# Download the llama3-8B tokenizer
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...
# Train Llama3-8B with Simple FSDP + torch.compile on a single trn2.48xlarge instance:
export NEURON_RT_NUM_CORES=64
torchrun --nnodes 1 --nproc_per_node $NEURON_RT_NUM_CORES \
--local-ranks-filter 0 --role rank --tee 3 \
--rdzv_backend c10d --rdzv_endpoint "localhost:0" \
-m torchtitan.train --module torchtitan_neuron.experiments.graph_trainer.llama3 --config simple_fsdp_llama3_8b \
--parallelism.tensor-parallel-degree 4 \
--parallelism.data-parallel-shard-degree -1 \
--training.global-batch-size 16 \
--training.local-batch-size 1 \
--compile.enable --compile.backend neuron
以下是使用 TorchTitan 在 8 台 trn2.48xlarge 上以张量并行(TP)和混合分片数据并行(HSDP)训练 Llama 3 8B 的示例损失曲线
Hugging Face Transformers v5
Transformers v5 为互操作性而构建:一个模型实现可在任何平台和任何硬件上运行,并可直接被下游训练和服务栈使用。
Trainium 是 Hugging Face 库上游支持的设备,包括 Transformers、Diffusers、trl 和 kernels。Transformers v5 支持文本和音频模型的训练与推理,Diffusers 支持图像和视频生成,trl 支持大型语言模型的强化学习,NKI 内核可通过 Hugging Face kernels 库获得。AWS 和 Hugging Face 为每个库添加了支持,因此现有的 Hugging Face 工作流可在 Trainium 上运行,无需设备特定代码。
- 训练。 Transformers v5 增强了完整预训练和微调的完备性,具有更好的初始化,并与并行和工具生态系统兼容——TorchTitan 用于扩展,Unsloth 和 TRL 用于后训练。
- 推理。 Transformers v5 在各模型实现之间标准化了注意力和内核。
- Diffusion。 Diffusers 中的张量并行已在 Neuron 上游合并,因此图像和视频生成可通过 diffusers API 原样运行。
- 兼容性。 Transformers v5 是 TorchTitan、vLLM 和 SGLang 等运行时库的原生后端,将其模型覆盖范围扩展到 Hugging Face 上可用的 300 万个公开模型中的大多数。我们正在努力在 Neuron 上启用并验证这一集成,将这一扩展的模型支持带到 Trainium 实例。
将模型加载到 Trainium 上
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
# Neuron hardware is automatically detected
device_map="neuron"
)
TorchNeuron 通过 PyTorch 的 PrivateUse1 扩展点注册 Trainium,从而将 Neuron 纳入 PyTorch 的设备枚举。device_map=”auto” 无需配置即可找到它,device_map=”neuron” 则显式指定它。无论哪种方式,权重都会加载到设备上用于训练或推理,脚本的其余部分保持不变。
测试方法与结果
我们的测试方法包括在 HuggingFace 和 TorchTitan 模型的广泛覆盖范围内验证 TorchNeuron,以及在 eager 和 torch.compile 两种模式下运行 PyTorch 上游测试套件。
模型覆盖
盲扫(大规模广泛测试)。 我们利用 Hugging Face 模型宇宙对开源模型进行广泛扫描。我们持续运行 10 万个 Hugging Face 模型,覆盖 99.9% 的实际使用场景(基于下载量)。我们在 Trainium 上以缩减但具有代表性的配置运行每个模型,并记录发生的情况:是否加载、前向和反向传播是否运行、是否编译,以及哪些算子在该设备上执行、哪些回退到 CPU。
下图展示了截至 2026 年 9 月,在 eager 和 torch.compile 模式下运行 10 万个模型的结果。我们正在努力缩小剩余差距,目标是在未来版本中缩小大部分(如果不是全部)差距。
PyTorch 上游测试
我们运行 PyTorch 上游测试套件:ATen、Dynamo、分布式和 profiler,并报告相对于已启用测试的通过率。2026 年 8 月的运行产生了约 15.4 万个测试结果,其中约 6.9 万个因有效原因被跳过:其中约 3 万个被 PyTorch 自身的资格检查跳过(例如仅 CUDA 和仅 CPU 测试),并非 Neuron 特有。其余约 3.9 万个是刻意跳过 Neuron 的,分为三组
- 目前尚无原生 Neuron lowering 的算子(约 2.2 万个,是最大的一组;与上述算子覆盖相关,我们计划在后续版本中逐步缩小)——逐元素、归约、foreach、高级索引和高级线性代数算子的边缘情况,以及融合 LSTM/GRU 单元。
- CUDA 特定测试(即上游资格检查未标记的) —— CUDA graphs、NCCL 特定集合通信、CUDA 特定 Dynamo 路径测试因资格检查而被跳过。
- TorchNeuron 在首个版本中明确选择不支持的功能,因为客户没有要求: 复合算子合规性、稀疏张量和稀疏算子、NestedTensor、前向模式自动微分(JVP)、原生多头注意力中的任意注意力掩码(支持无掩码和因果)、逐通道量化、负/共轭视图位、torch.signal 窗口、双线性/双三次上采样,以及动态和符号形状。
作为这一过程的一部分,我们还通过提出多项 RFC 和 PR 来回馈 PyTorch,以提升测试质量并移除硬编码的 CUDA 假设(193439、193441、193442,完整 PR 列表见附录 A)。
如何开始
从 2026 年 10 月 20 日起,TorchNeuron 将以公开测试版形式提供。研究人员和开发者可以通过 AWS Neuron pip 仓库执行 pip install torch-neuronx 进行安装,详情。
资源
以下资源将于 10 月 20 日上线:
- TorchNeuron:https://github.com/aws-neuron/torch-neuronx
- AWS Trainium:https://aws.amazon.com/machine-learning/trainium/
- AWS Neuron SDK:https://aws.amazon.com/machine-learning/neuron/
- Neuron 文档:https://awsdocs-neuron.readthedocs-hosted.com/
快速参考
有关入门、执行模式、分布式训练和开发路线图的更多指导,请参阅 Neuron 文档中的常见问题解答(FAQ)部分。
欢迎参加在北美举办的 PyTorch Conference
- 在 GitHub 上以 Apache 2.0 许可证开源 TorchNeuron
- 发布可通过 pip 安装的软件包,供立即使用
- 举办关于在 Trainium 上原生运行 PyTorch 的技术深度探讨会
- 进行现场演示
致谢
我们感谢 PyTorch 社区和 Meta 的 PyTorch 团队在 PyTorch 可扩展机制方面所做的工作,特别是 PrivateUse1 设备后端和 torch.accelerator API,正是它们使这一集成成为可能。我们也感谢 PyTorch 基金会在 CI/CD 基础设施方面给予的支持。
AWS 致力于支持 PyTorch 生态系统,为 PyTorch CI/CD 基础设施提供 Trainium 资源,并为开源 PyTorch 项目做出贡献。
特别感谢 Anthropic、Anyscale、Amazon AGI、Amazon Rufus、ByteDance、Cartesia、Clay Foundation、Databricks、Decart、Deepgram、Hugging Face、Karakuri、Lieber Institute for Brain Development、Meta、MIT、Neura Robotics、NinjaTech AI、Nota AI、Odyssey、PDT Partners、poolside、Qube Research and Technologies、RadixArk、Ricoh、Rivian、Salesforce、SemiAnalysis、UC Berkeley、Stanford、CMU、UCLA、Zyphra 以及许多其他机构,感谢他们在塑造 TorchNeuron 的设计和能力方面提供的反馈。
贡献者:Apoorv Gupta、Arash Jamalian、Rui Silva、Peter Sam Raj、Kunind Sahu、Kamran Khan、Shreya Srivastava、Charan Shettyhalli Guruswamy、Alex Hoffman、Abigail Fernandes、Sam Jin
附录:完整 PR 列表
- https://github.com/pytorch/pytorch/pull/179063
- https://github.com/pytorch/pytorch/pull/179511
- https://github.com/pytorch/pytorch/pull/187736
- https://github.com/pytorch/pytorch/pull/188839
- https://github.com/pytorch/pytorch/pull/192365
- https://github.com/pytorch/pytorch/pull/192401
- https://github.com/pytorch/pytorch/pull/192404
- https://github.com/pytorch/pytorch/pull/193428
- https://github.com/pytorch/pytorch/pull/193430
- https://github.com/pytorch/pytorch/pull/193433
- https://github.com/pytorch/pytorch/pull/193436
- https://github.com/pytorch/pytorch/pull/193438
- https://github.com/pytorch/pytorch/pull/193439
- https://github.com/pytorch/pytorch/pull/193441
- https://github.com/pytorch/pytorch/pull/193442
- https://github.com/pytorch/pytorch/pull/193444
- https://github.com/pytorch/pytorch/pull/193445
- https://github.com/pytorch/pytorch/pull/193450
- https://github.com/pytorch/pytorch/pull/193456
- https://github.com/pytorch/pytorch/pull/193457
- https://github.com/pytorch/pytorch/pull/193686
- https://github.com/pytorch/torchtitan/pull/2382
- https://github.com/pytorch/torchtitan/pull/2440
- https://github.com/pytorch/torchtitan/pull/2445
- https://github.com/pytorch/torchtitan/pull/2450
- https://github.com/pytorch/torchtitan/pull/2668
文章 Introducing TorchNeuron: Native PyTorch backend for AWS Trainium 首次发表于 PyTorch。
来源:PyTorch:Blog · pytorch.org


