TorchNeuron 性能分析（第1部分）：AWS Trainium 上的系统级性能分析
huggingface ·
本文介绍了在 AWS Trainium 上对 TorchNeuron 进行系统级性能剖析，旨在帮助机器学习从业者分析推理性能瓶颈。文章详细说明了如何配置 NeuronConfig 和 NeuronProfiler，以利用 PyTorch 的 PrivateUse1 扩展捕获设备和运行时追踪数据。指南展示了如何使用 Neuron Explorer 可视化从框架分发到硬件执行的操作生命周期。研究对比了三种 RMSNorm 实现：原始的 PyTorch 操作、内置的 torch.rms_norm，以及使用 torch.compile 编译的原始函数。结果表明，通过 torch.rms_norm 进行融合和使用编译技术，均能显著减少 Neuron 设备时间，因为它们消除了中间张量的物化以及操作之间的间隙。融合方法所需的设备时间约为原始实现的六分之一。作者总结认为，识别 key_averages 表中的热点并在追踪数据中加以验证，是一种有效的优化工作流程。
huggingface