Hugging Face 博客发布 TorchNeuron 系统级性能分析教程：在 AWS Trainium 上定位推理耗时

OmniTools 10月22日消息，Hugging Face 博客消息，AWS 团队发布文章《Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.》，介绍如何从零开始系统性分析模型在 AWS Trainium 上的推理性能并定位耗时环节。TorchNeuron 是 Trainium 的开源原生 PyTorch 后端，支持 eager 执行、torch.compile 和标准 PyTorch API，并通过 PyTorch 的 PrivateUse1 扩展点注册为原生后端。

文章给出了配置 torch.profiler 的具体方法：用 NeuronConfig 指定要采集的数据流与 NeuronCore，把 correlate_device_time 设为 True，将设备时间写入 key_averages() 表；再把 NeuronConfig 作为 experimental_config 传入，由 NeuronProfiler 的 export_trace 回调导出 trace.json 与 Neuron 系统 trace。结果可在基于浏览器的 Neuron Explorer 中查看，追踪算子从框架调度、运行时调度（nrta_execute_schedule）到 NeuronCore 硬件执行（nc_exec_running）的完整链路。

文中以 RMSNorm 为例对比三种实现：朴素 PyTorch 算子版本包含 6 次算子调用，Neuron 设备时间 2,355 µs；torch.rms_norm 被降级为单一融合内核，设备时间 403 µs；对同一朴素实现使用 torch.compile(backend="neuron") 后编译为单个 NEFF 产物，设备时间 376 µs。后两种实现均只需一次设备调度。