pytorch 源码阅读(1)——torch.complie

0 torch.complie(model, fullgraph, dynamic, backend, mode, options, disable),使用 TorchDynamo 和指定的后端优化给定的模型/函数

对于在编译区域内执行的每一帧,都会尝试编译,并将编译结果缓存到代码对象中,供将来使用。 如果之前的编译结果不适用于后续调用那么单帧可能会被编译多次,这被称为 guard_failure,此时可以使用 TORCH_LOGS=guards 来进行调试。

TorchDynamo 是一个设计用于加速未修改的 PyTorch 程序的 Python 级即时(JIT)编译器。它通过 Python Frame Evaluation Hooks(Python 框架评估钩子)来实现这一目标,以便在运行时动态地生成和优化代码。这使得 TorchDynamo 可以有效地处理各种 Python 代码,包括包含控制流(如循环和条件语句)的代码,而无需进行任何修改。

1 参数说明

torch.complie(model, fullgraph, dynamic, backend, mode, options, disable)

  • model(callable):需要进行优化的模型或者函数
  • fullgraph(bool):是否可以将模型分解为多个子图
  • dynamic (bool or None): 使用动态 shape 进行 tracing
    • True:尝试生成尽可能动态的 kernel,但这个也不总是有效的,因为某些特化优化会使动态 kernel 失效
    • False:将不生成任何动态的 kernel,只进行特化优化
    • None:默认值,会自动检测是否发生了动态变化,并在重新编译时编译一个更动态的动态内核。
  • backend (str or Callable):使用的后端
    • inductor,默认的后端,在性能和开销之间有较好的平衡
    • 可以使用 torch._dynamo.list_backends() 查看非实验性树内后端
    • 可使用 torch._dynamo.list_backends(None) 查看实验性或调试性的树内后端
    • 注册树外后端的方式:https://pytorch.org/docs/main/compile/custom-backends.html
  • mode(str),可以配置的值:"default", "reduce-overhead", "max-autotune" or "max-autotune-no-cudagraphs"
    • default:默认模式,在性能和开销之间进行很好的平衡
    • reduce-overhead:降低开销模式,小 batch 的场景比较有用,降低 python 使用 CUDA garphs 的开销。降低开销的代价是使用更多的内存,因为会缓存调用所需的工作区内存,这样就在后续运行中无需重新分配。减少开销并不能保证一定有效,只是减少了不改变输入的 CUDA graphs 的开销。在 CUDA graphs 不适用的其他情况下,可以使用 TORCH_LOG=perf_hints 来进行 debug。
    • max-autotune:利用基于 Triton 的矩阵乘法和卷积的模,默认开启 CUDA graphs。
    • max-autotune-no-cudagraphs:与 max-autotune 模式类似,但是不开启 CUDA graphs。
    • 调用 torch._inductor.list_mode_options() 可以查看每个模式的具体配置。
  • options (dict): 传递给后端的一些字典选项,一些值得尝试的配置:
    • epilogue_fusion,将点对点融合设置到模板中,需要同时设置max_autotune
    • max_autotune,通过配置文件选择最佳的 matmul 配置
    • fallback_random,在调试精度问题的时候比较有用
    • shape_padding,填充矩阵形状,以便更好地调整 GPU 上的负载,尤其是张量内核的负载
    • triton.cudagraphs,可减少 Python 在使用 CUDA graphs 时的开销
    • trace.enabled,最好用的 debug 选项,建议打开
    • trace.graph_diagram,将显示融合后的 graph 的图片
    • 可以通过调用 torch._inductor.list_options() 查看支持的全部配置列表
  • disable (bool): 将 torch.compile 设置为 no-op 的测试模式

2 额外说明

  • 这个函数一开始有这么一行:_C._log_api_usage_once("torch.compile"),这个是用来进行调试的。

    • _C._log_api_usage_once("torch.compile"),这是严格意义上的调试,很可能大部分使用都是在 facebook 内部进行的。如果设置环境变量PYTORCH_API_USAGE_STDERR=1,就能看到 PyTorch 记录了 "涉及 PyTorch 的哪些部分",这很可能会在出现意外情况时大致知道该去哪里找(以及找谁)。这个函数出现的原因是 pytorch 现在的规模已经很多大了,往往会出现很多意想不到的交互调用,这个函数可以帮助缩小 debug 的范围。
  • 这个函数的作用是使用 TorchDynamo 和指定的后端优化给定的模型/函数,进行包装和参数配置的设置,最终会调用torch._dynamo.optimize进行优化处理。

相关推荐
CountingStars61913 分钟前
目标检测常用评估指标(metrics)
人工智能·目标检测·目标跟踪
tangjunjun-owen21 分钟前
第四节:GLM-4v-9b模型的tokenizer源码解读
人工智能·glm-4v-9b·多模态大模型教程
冰蓝蓝25 分钟前
深度学习中的注意力机制:解锁智能模型的新视角
人工智能·深度学习
橙子小哥的代码世界33 分钟前
【计算机视觉基础CV-图像分类】01- 从历史源头到深度时代:一文读懂计算机视觉的进化脉络、核心任务与产业蓝图
人工智能·计算机视觉
黄公子学安全43 分钟前
Java的基础概念(一)
java·开发语言·python
新加坡内哥谈技术1 小时前
苏黎世联邦理工学院与加州大学伯克利分校推出MaxInfoRL:平衡内在与外在探索的全新强化学习框架
大数据·人工智能·语言模型
程序员一诺1 小时前
【Python使用】嘿马python高级进阶全体系教程第10篇:静态Web服务器-返回固定页面数据,1. 开发自己的静态Web服务器【附代码文档】
后端·python
小木_.2 小时前
【Python 图片下载器】一款专门为爬虫制作的图片下载器,多线程下载,速度快,支持续传/图片缩放/图片压缩/图片转换
爬虫·python·学习·分享·批量下载·图片下载器
fanstuck2 小时前
Prompt提示工程上手指南(七)Prompt编写实战-基于智能客服问答系统下的Prompt编写
人工智能·数据挖掘·openai
lovelin+v175030409662 小时前
安全性升级:API接口在零信任架构下的安全防护策略
大数据·数据库·人工智能·爬虫·数据分析