今天看啥  ›  专栏  ›  吃果冻不吐果冻皮

一文搞懂 TorchDynamo 原理

吃果冻不吐果冻皮  · 公众号  · 科技媒体  · 2024-08-12 18:05
    

主要观点总结

文章介绍了PyTorch 2.0的TorchDynamo模块,该模块的主要功能是从PyTorch应用中捕获计算图,相比于TorchScript和TorchFX,TorchDynamo更加灵活和可靠。TorchDynamo通过翻译Python字节码来构建FX Graph,并使用Guard来检查输入数据的属性是否有变化,以决定是否能重用已编译的计算图。当遇到不支持的算子时,TorchDynamo会将计算图切分为子图,由Python解释器执行。循环在TorchDynamo中被展开,它会尝试内联被调函数,并在DistributedDataParallel的bucket边界引入graph break,以确保allreduce能与反向传播同时执行。

关键观点总结

关键观点1: TorchDynamo的功能

TorchDynamo从PyTorch应用中捕获计算图,相比其他模块更灵活和可靠。

关键观点2: TorchDynamo的工作原理

TorchDynamo通过翻译Python字节码来构建FX Graph,并使用Guard来检查输入数据的属性。

关键观点3: TorchDynamo处理不支持算子的方式

TorchDynamo通过graph break将计算图切分为子图,由Python解释器执行不支持的算子。

关键观点4: TorchDynamo处理循环的方式

TorchDynamo在捕获计算图时将循环展开,并尝试内联被调函数。

关键观点5: TorchDynamo在DistributedDataParallel中的作用

TorchDynamo在DDP的bucket边界引入graph break,确保allreduce能与反向传播同时执行。


文章预览

【点击】 加入大模型技术交流群 原文:https://fkong.tech/posts/2023-05-20-dynamo/ 简介 PyTorch 2.0 的使命是更快、更 Pythonic 以及一如既往地支持动态特性。为了达到这个目的,PyTorch 2.0 引入了  torch.compile ,在解决 PyTorch 固有的性能问题的同时,把部分用 C++ 实现的东西引入 Python 中。PyTorch 2.0 利用了 4 个组件: TorchDynamo,AOTAutograd,PrimTorch 和 TorchInductor。本文以几个简单的案例讲解 TorchDynamo 的使用方法和实现原理。 PyTorch 2.0 TorchDynamo 的作用是从 PyTorch 应用中抓取计算图 ,相比于 TorchScript 和 TorchFX,TorchDynamo 更加灵活、可靠性更高。用过 TorchScript 的朋友知道,通过  jit.trace  或者  jit.script  把模型转化为 TorchScript 的过程困难重重,往往需要修改大量源代码。而 TorchFX 在捕获计算图时,遇到不支持的算子会直接报错,最常见的就是  if  语句。TorchDynamo 克服 ………………………………

原文地址:访问原文地址
快照地址: 访问文章快照
总结与预览地址:访问总结与预览