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 的执行模式、分布式 API 支持和 100k 模型测试结果,读者可据此评估 Trainium 上现有 PyTorch 代码的迁移成本。
前沿模型涵盖语言、视频和音频,架构从稠密和专家混合 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 进行实验和迭代。团队无需重建工作流或离开 PyTorch 生态,即可利用 Trainium 的能力。”——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 开源发布,覆盖主流 LLM 和多模态场景中的稠密与 MoE 模型,例如 FlashAttention 和 MoE 分块内核,开发者可以用它们在 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 首席技术官
“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 x 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,以及 PyTorch 上游测试套件在 eager 和 torch.compile 两种模式下的验证。
模型覆盖
盲扫(大规模广泛测试)。我们使用 Hugging Face 模型宇宙对开源模型进行广泛扫描。我们持续运行 10 万个 Hugging Face 模型,覆盖 99.9% 的实际使用场景(基于下载量)。我们在 Trainium 上以精简但有代表性的配置运行每个模型,并记录发生的情况:它能否加载,前向和后向传播能否运行,能否编译,以及哪些算子在该设备上执行、哪些回退到 CPU。
下图显示了截至 2026 年 9 月,10 万个模型在 eager 和 torch.compile 模式下的运行结果。我们正在努力弥合剩余的差距,目标是在未来版本中弥合大部分(如果不是全部)差距。
PyTorch 上游测试
我们运行 PyTorch 上游测试套件:ATen、Dynamo、distributed 和 profiler,并报告相对于已启用测试的通过率。2026 年 8 月的运行产生了约 15.4 万个测试结果,其中约 6.9 万个因有效原因被跳过:其中约 3 万个由 PyTorch 自身的资格检查跳过(例如仅 CUDA 和仅 CPU 测试),并非 Neuron 特有。其余的约 3.9 万个是 Neuron 有意跳过的,分为三组
- 目前尚无原生 Neuron lowering 的算子(约 2.2 万个,是最大的一组;与上文提到的算子覆盖相关,我们计划在后续版本中逐步弥合)——elementwise、reduction、foreach、advanced-indexing 和 advanced-linalg 算子的边界情况,以及融合的 LSTM/GRU 单元。
- CUDA 特有的测试(即上游资格检查未标记)——CUDA graphs、NCCL 特有的集合通信、CUDA 特有的 Dynamo 路径测试在资格检查下被跳过。
- TorchNeuron 在首个版本中明确选择不支持的功能,因为客户未提出需求:复合算子合规性、稀疏张量和稀疏算子、NestedTensor、前向模式自动微分(JVP)、原生多头注意力中的任意注意力掩码(支持无掩码和因果)、逐通道量化、负/共轭视图位、torch.signal 窗口、双线性/双三次上采样,以及动态和符号形状。
作为这一过程的一部分,我们还通过提出多项 RFC 和 PR 来改进测试质量并移除硬编码的 CUDA 假设,从而回馈 PyTorch(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 列表
- 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
文章 推出 TorchNeuron:AWS Trainium 的原生 PyTorch 后端 首次出现在 PyTorch 上。
来源:PyTorch:Blog · pytorch.org


