AWS 发布 TorchNeuron：Trainium 的原生 PyTorch 后端

PyTorch：Blog· Nafea Bshara, Robert Hundt, Ron Diamant, Colin Brace, Rahul Solanki, Maen Suleiman, Khaled ElGalaind, Temesghen Kahsai, Cyril Bortolato, Indu Thangakrishnan, Shauheen Zahirazami, Sarang Ozarde, Adarsh Singh, Hanbo Wang, Esha Lakhotia, Jianying Lang·2026-10-19 20:45· 28 分钟前精选AI 评分66

Introducing TorchNeuron: Native PyTorch backend for AWS Trainium

AI 导读

AWS 宣布 TorchNeuron，一个面向 Trainium 芯片的开源原生 PyTorch 后端，10 月 20 日起以 Apache 2.0 协议公开 Beta，代码位于 github.com/aws-neuron/torch-neuronx。

推荐理由

原文给出执行模式、分布式 API 和 100k 模型兼容性测试等细节，开发者可据此评估现有 PyTorch 代码迁移到 Trainium 的改动成本。

前沿模型涵盖语言、视频和音频，架构从稠密和混合专家 transformer 到扩散、状态空间、世界和结构预测模型。PyTorch 是 AI 研究发生的地方，是模型被定义的地方，也日益成为生产发生的地方。研究人员希望进行实验和创新，客户则希望获得 AWS Trainium 提供的性能，以实现其性能和规模目标。训练和推理库是针对 PyTorch 自身的 API 编写的，新的库和新版本持续发布。过去几年中，PyTorch 使这些 API 与加速器无关，并提供了用于注册自定义硬件的官方扩展点。

今天，为满足开发者和研究社区对 Trainium 日益增长的需求，我们很高兴分享一种更简单的方法，将芯片作为标准 PyTorch 设备进行编程，这早于其于 10 月 20 日在北美 PyTorch Conference 上的开源发布。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()

针对 torch.accelerator（PyTorch 推荐的设备无关 API）编写的现代代码无需修改即可运行。编译模型需要指定一次 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 CI/CD 基础设施提供 Trainium 资源并贡献开源 PyTorch 项目来支持 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 列表

文章 Introducing TorchNeuron: Native PyTorch backend for AWS Trainium 首次发表于 PyTorch。

来源：PyTorch：Blog · pytorch.org