跳到主要内容

TorchNeuron 开源:AWS 用一行 .to("neuron") 拆 PyTorch 的 CUDA 偏好

· 阅读需 3 分钟
Industry Research Team

PyTorch Conference North America 2026 上演了教科书式的「独立趋同」:10 月 19 日 AWS 开源 TorchNeuron(Trainium 原生 PyTorch 后端),10 月 20 日谷歌开源 TorchTPU。两家在 48 小时内各自解决了同一个问题——PyTorch 开发者不再需要重写代码就能用上非英伟达芯片。

一、TorchNeuron 是什么​

要素内容
定位Trainium 原生 PyTorch 后端(Apache 2.0 公测)
切换方式.to('neuron') 一行(与 TorchTPU 的 .to('tpu') 对称)
覆盖官方内部扫描称覆盖按下载量计 99.9% 的真实使用场景
集成vLLM、SGLang、Ray、Hugging Face Transformers、Neuron 2.33
早期客户Deepgram(语音 AI 大规模负载)
硬件Trainium2 / Trainium3

它取代的是旧的 Hugging Face 适配路径——此前 Trainium 支持需要 HF 专属改造而非一等设备接入。AWS 还宣布将持续向 PyTorch 上游贡献代码。

二、两个后端的技术取向对比​

TorchTPU(谷歌+Meta)TorchNeuron(AWS)
已知缺口FlexAttention 未支持、动态形状进行中稀疏张量/NestedTensor/动态形状未支持、FP8 私测
标杆数据98% HF 模型无崩溃;InferenceX 2.01 美分/Mtok99.9% 使用场景覆盖(内部口径)
共同点都基于 PrivateUse1 后端 + torch.accelerator API,都集成 vLLM/SGLang/Ray/HF

两者仍是 beta——注意力变体、内存捷径等尚未全覆盖,性能数字均为厂商自报,尚无独立交叉基准。但对 CUDA 生态锁定而言,这是结构性松动:硬件迁移成本从「重写」降到「改一行 + 调优」。

三、llm-d 1.0 补上最后一块​

同场发布的 llm-d 1.0(谷歌云 + Red Hat + IBM)提供了多机推理编排层:把众多 vLLM/SGLang 服务器组成一个池,按缓存亲和路由请求、KV Cache 分层下沉、Prefill/Decode 分离。谷歌报告中位首 Token 时延下降超 85%、Ironwood TPU 编码负载吞吐 2.4 倍,且同时支持 NVIDIA/AMD/Intel GPU 与 TPU。框架层(TorchTPU/TorchNeuron)+ 编排层(llm-d)合起来,ASIC 阵营的软件栈第一次完整了。

四、对选型的直接影响​

  1. POC 成本骤降:把现有 PyTorch 推理服务搬到 Trainium/TPU 做对比测试,从「数周工程」变成「数天」
  2. 采购谈判更硬:CUDA 锁定作为溢价理由的说服力下降,直接用 TCO 计算器比每 Token 成本即可
  3. 训练侧仍需谨慎:两个后端都优先覆盖推理与中小规模训练场景,前沿训练负载的迁移还在早期

延伸阅读:TorchTPU 开源详解 · 推理成本经济学 · Trainium 2 卡片