训练与微调AI 资源
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
ELI5：它是什么，解决什么问题
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 降级等已知限制。
核心能力与快速开始
★ 119 · Python
- 能做什么
- 为 AWS Trainium 提供原生 PyTorch 后端，新增 neuron 设备；默认 eager 逐算子执行，无预先追踪的计算图；支持 torch.compile 的 neuron 后端进行全图编译；支持 torch.distributed、DTensor 与函数式集合通信，实现数据并行、张量并行和完全分片并行；支持在 PyTorch 代码及 torch.compile 图中调用 NKI 自定义内核，并可结合 torch.library.register_autograd 注册反向传播；提供 API 兼容性指南和环境变量等参考文档。
- 适用场景
- 在 Trn2/Trn3 实例上使用 PyTorch 训练或推理模型；以 eager 模式快速运行和调试模型；通过 torch.compile 编译模型或模型区域以获得更好性能；进行多节点、张量并行与 FSDP 等分布式训练；在 PyTorch 工作流中集成 NKI 自定义内核；参考 Llama、Whisper、Qwen2、GPT-2、MoE、Conformer ASR 等示例开展训练与推理。
- 项目形态
- 开发框架
- 部署方式
- 需使用 Trn2 或 Trn3 实例；推荐使用已预装 torch-neuronx 的 Neuron DLAMI 或 Deep Learning Container；其他 Linux 实例需先安装 Neuron 驱动和运行时，再安装 torch-neuronx；多节点训练还需安装 EFA 驱动。
- 主要依赖
- 需要 Trn2 或 Trn3 实例，Python 3.12、3.13 或 3.14；需安装 torch==2.14.0（CPU 索引源）、torch-neuronx==2.14.*、neuronx-cc==2.* 与 nki>=0.6，并保留版本固定；安装后可用 python -m torch_neuronx.collect_env 检查 neuron 设备是否可用。
- 使用限制
- 项目处于 Beta。首次调用算子或编译图包含编译时间，大模型可能较明显；torch.compile 仅支持静态形状（dynamic=False），每个新输入形状都会重新编译，仅支持连续张量；FlexAttention 支持有限；float64 以 float32 运行，int64 支持取决于算子；本版本不提供 TorchAO 集成（如 MXFP8 等复合数据类型）；adaptive eager execution 为实验性功能且默认关闭。
- 许可证
- Apache-2.0
- 最近推送
- 2026-10-21 01:16
资料检查：2026-10-24 07:00 · 卡片整理：2026-10-24 07:00
卡片由 AI 根据仓库资料整理，未进行安装或运行验证
讨论 0
登录后参与讨论，阅读资讯无需注册。
还没有评论，来分享第一个看法。