亚马逊发布开源PyTorch后端TorchNeuron，让PyTorch代码在Trainium上运行
亚马逊称，已在其他加速器上运行的PyTorch代码只需极少改动甚至无需改动即可在Trainium上运行；该后端提供自适应eager模式和编译缓存，并依据 841 个代表性模型和数万种张量尺寸进行调优与基准测试。
AI解读：亚马逊本周宣布开源PyTorch后端TorchNeuron，用于把PyTorch代码转换为可在AWS Trainium系统上执行的操作，目标是让原本运行于其他加速器的代码“极少或无需修改”即可迁移。
TorchNeuron增加可选的“自适应eager”模式：CPU上的Python解释器累积算子指令并组成小子图，交给TorchNeuron编译、缓存和排队执行，以接近编译模式的性能并隐藏编译时间。
TorchNeuron把编译结果转换为NEFF可执行文件，并在内存和磁盘多级缓存；在跨多台服务器的大型模型训练中，共享编译服务让其他加速器复用编译结果，亚马逊称测试中首次训练步耗时减少约三分之一。
评估环节由AI agent扫描Hugging Face仓库，筛出可执行模型并按架构特征分类，最终选 841 个代表性模型覆盖各类设计特征，并从中归纳数万种实际张量尺寸，形成用于快速单元测试和基准测试的shape registry。
上述数字与效果来自亚马逊科学博客，作者为AWS副总裁兼杰出工程师Robert Hundt；属于亚马逊单方描述，来源未给出第三方复现或完整公开基准结果。
适用边界在来源中也有说明：上述“约三分之一”为特定大语言模型跨加速器测试的结果，不能直接外推到其他模型或训练任务。
亚马逊本周宣布开源TorchNeuron，这是一个把PyTorch代码转换为可在AWS Trainium上执行操作的PyTorch后端。亚马逊科学博客称，已在其他加速器上运行的PyTorch代码在Trainium上“极少或无需修改”即可运行。
该文由AWS副总裁兼杰出工程师Robert Hundt署名，介绍TorchNeuron的运行时优化与评估方法；文中性能数字均为亚马逊方面给出的测试结果。
自适应eager与编译缓存
标准PyTorch有eager与编译两种执行模式。TorchNeuron增加一个可选的“自适应eager”模式：CPU上的Python解释器从PyTorch算子分发中累积指令，构建小型子图，再交给TorchNeuron编译、缓存并排队执行。
亚马逊称该模式同时接近eager模式的响应性和编译模式的性能，并保留eager模式下算子与源码行的对应关系，便于调试。自适应算法按子图和输入张量尺寸选择子图，使编译时间与已编译片段的执行重叠，从而在多种模型上隐藏编译时间。
编译产物是名为NEFF的可执行文件。TorchNeuron在内存和磁盘多级缓存NEFF，避免重复编译；一个算子、完整图或部分图都可对应NEFF。
在由多台服务器、数千颗加速器组成的大规模部署中，各加速器运行同一模型的不同部分，需要相同的NEFF。TorchNeuron运行时用共享编译服务集中编译，其他加速器复用结果。亚马逊称，在一个跨加速器分布的大语言模型测试中，这一做法把首次训练步的耗时减少约三分之一。
用 841 个模型和数万种张量尺寸做评估
亚马逊称，其团队用AI agent扫描Hugging Face仓库，排除配置损坏、无法在任何加速器上运行的模型，并筛选出数十万个可执行模型。
另一个agent按架构特征对模型分类，例如大语言模型的前端编码器、包含Transformer块与注意力的主干，以及输出后端。最终确定 841 个代表性模型，亚马逊称其覆盖了TorchNeuron可能遇到的设计特征。
为加快回归检测，团队记录每次评估中Trainium处理的张量运算及尺寸，形成包含数万种运算与张量维度的“shape registry”，并据此生成单元测试。新版本先跑快速单元测试，再跑 841 个模型，最后跑完整模型分布。
亚马逊表示，这套分层测试加快了开发迭代，并让其对TorchNeuron在Trainium上准确、高效执行PyTorch模型有信心。来源未公开完整基准结果或第三方复现数据。