PyTorch中的AOTAutograd、PrimTorch和TorchInductor

一.PyTorch 中的 AOTAutograd

  在 PyTorch 中，AOTAutograd 是一个先进的功能，它是 Ahead-Of-Time (AOT) Autograd 的缩写。这是一种新的自动微分方法，它在编译时而不是运行时计算梯度。这使得它可以更有效地处理大型计算图。以下是一些关于 AOTAutograd 的论文：

1. “TorchProbe: Fuzzing Dynamic Deep Learning Compilers” by Q Su, C Geng, G Pekhimenko, X Si. 这篇文章讨论了如何在 PyTorch 的最新主分支上从源代码构建 PyTorch，并提到了由于 AOT Autograd 错误而导致的崩溃。可以在这里[1]找到这篇文章。
2. “μ-TWO: 3× Faster Multi-Model Training with Orchestration and Memory Optimization” by S Purandare, A Wasay, S Idreos. 这篇文章提到了现代 GPU 在开发神经网络时的严重未利用情况，并提到了一些优化技术。可以在这里[2]找到这篇文章。

  在 PyTorch 中，AOTAutograd 的实现涉及到一些复杂的概念和技术。以下是一些关于 AOTAutograd 实现的资源：

1. PyTorch AOTAutograd: A pythonic extension point[3]：这是一个 YouTube 视频，讲解了如何使用这个扩展点来加速 PyTorch 模型。
2. TorchDynamo Update 6: Training support with AOTAutograd[4]：这是一个关于 TorchDynamo 项目的更新，该项目是一个 Python 级别的 JIT 编译器，设计用来使未修改的 PyTorch 模型更快。在这个更新中，他们讨论了如何使用 AOTAutograd 来支持训练。
3. AOT Autograd - How to use and optimize?[5]：这是一个教程，讲解了如何使用 AOT Autograd 来加速深度学习模型的训练。
4. PyTorch 2.0[6]：在这个页面中，他们讨论了 TorchInductor 的核心循环级别 IR 只包含大约 50 个操作符，并且它是用 Python 实现的，使得它容易被修改和扩展。他们还提到了如何使用 AOTAutograd 来重用现有的 autograd 图。
5. How does torch.compile work with autograd?[7]：在这个讨论中，他们提到了 AOTAutograd 是处理 torch.compile 运行时反向传播的主要组件。

二.PyTorch 中的 PrimTorch

  在 PyTorch 中，PrimTorch 是一个重要的组件，它主要用于将 PyTorch 的操作符进行规范化。根据我找到的一篇文章，PrimTorch 可以将大约 2000 个 PyTorch 操作符规范化为一个封闭的集合。这样做的好处是可以简化 PyTorch 的操作符集合，使得在进行计算时更加高效。以下是一篇关于 PrimTorch 的论文：

• “MAP: Memory-aware Automated Intra-op Parallel Training For Foundation Models” by Y Liu, S Li, J Fang, Y Shao, B Yao, Y You. 这篇文章讨论了如何使用 PrimTorch 来提高模型训练的效率。可以在这里[8]找到这篇文章。

三.PyTorch 中的 TorchInductor

  TorchInductor 是 PyTorch 2.0 中引入的一种优化编译器。它的设计目标是为了提高 PyTorch 在动态形状机器学习负载上的性能。以下是一些相关的学术论文和研究，可能会对你有所帮助：

1. TorchProbe: Fuzzing Dynamic Deep Learning Compilers[9]: 这篇论文讨论了 TorchInductor 在 CPU 后端上的表现，并提供了一些示例。
2. BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach[10]: 这篇论文提到了 TorchInductor 在 HuggingFace 的 TensorFlow 和 PyTorch 实现中的应用。
3. PyTorch 2.0: The Journey to Bringing Compiler Technologies to the Core of PyTorch (Keynote)[11]: 这是一篇关于 PyTorch 2.0 的演讲稿，其中讨论了 TorchDynamo 和 TorchInductor 的设计。
4. A journey to enable generative AI on a new hardware platform with PyTorch 2.0[12]: 这篇文章分享了在新硬件平台上使用 PyTorch 2.0 启用生成 AI 的经验，其中讨论了 PyTorch 运行时和 TorchInductor 的一些特性。

参考文献：
[1] https://link.springer.com/chapter/10.1007/978-981-99-8311-7_15
[2] https://proceedings.mlsys.org/paper_files/paper/2023/hash/a72071d84c001596e97a2c7e1e880559-Abstract-mlsys2023.html
[3] PyTorch AOTAutograd: A pythonic extension point：https://www.youtube.com/watch?v=KpH0qPMGGTI
[4] TorchDynamo Update 6: Training support with AOTAutograd：https://dev-discuss.pytorch.org/t/torchdynamo-update-6-training-support-with-aotautograd/570
[5] AOT Autograd - How to use and optimize?：https://pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html
[6] PyTorch 2.0：https://pytorch.org/get-started/pytorch-2.0/
[7] How does torch.compile work with autograd?：https://dev-discuss.pytorch.org/t/how-does-torch-compile-work-with-autograd/1621
[8] https://arxiv.org/abs/2302.02599
[9] TorchProbe: Fuzzing Dynamic Deep Learning Compilers: https://link.springer.com/chapter/10.1007/978-981-99-8311-7_15
[10] BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach：https://dl.acm.org/doi/abs/10.1145/3617327
[11] PyTorch 2.0: The Journey to Bringing Compiler Technologies to the Core of PyTorch (Keynote): https://dl.acm.org/doi/abs/10.1145/3579990.3583093
[12] A journey to enable generative AI on a new hardware platform with PyTorch 2.0: https://research.ibm.com/publications/a-journey-to-enable-generative-ai-on-a-new-hardware-platform-with-pytorch-20