torch.compile

torch-neuronx：AWS Trainium 的原生 PyTorch 后端，支持 eager 执行、torch.compile、分布式训练与 NKI 自定义内核

aws-neuron/torch-neuronx · A native PyTorch backend for AWS Trainium with support for eager execution, torch.compile, distributed training, and custom NKI kernels

torch-neuronx 是 AWS Trainium 的原生 PyTorch 后端，通过 PrivateUse1 机制在 PyTorch 中新增 neuron 设备，让用户可以用熟悉的 PyTorch API 把模型和输入迁移到 neuron 设备上训练或推理。项目默认以 eager 模式逐算子执行，也可通过 torch.compile 的 neuron 后端编译模型或模型区域以提升性能；支持 torch.distributed、DTensor 与函数式集合通信，覆盖数据并行、张量并行和完全分片并行；并支持在 eager 模式和 torch.compile 编译图内调用 NKI 自定义内核。项目当前为 Beta，需要 Trn2 或 Trn3 实例及对应 Python 版本，并列出静态形状、连续张量、float64 降级等已知限制。