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




NLP工程化

1.本公众号以对话系统为中心,专注于Python/C++/CUDA、ML/DL/RL和NLP/KG/DS/LLM领域的技术分享。
2.本公众号Roadmap可查看飞书文档:https://z0yrmerhgi8.feishu.cn/wiki/Zpewwe2T2iCQfwkSyMOcgwdInhf

NLP工程化(公众号)

NLP工程化(星球号)

posted on 2024-01-14 23:57  扫地升  阅读(156)  评论(0编辑  收藏  举报