torch.compile

原创 已于 2025-08-20 18:19:28 修改 · 447 阅读 · 0 ·
CC 4.0 BY-SA版权
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
文章标签：

#pytorch #ai编程

于 2025-08-15 18:07:35 首次发布

简介

在Pytorch 2.x中，引入了torch.compile特性，主要包含如下4个部分：

TorchDynamo：基于Python Frame Evaluation Hook技术，实现安全的Pytorch的计算图捕获。

AOTAutograd： AOT生成计算图的反向图。

PrimTorch：规范化2000+ PyTorch Operators为250+ Primitive Operators, 极大降低了开发Pytorch后端的难度。

TorchInductor：一个Deep Learning Compiler，为多种加速器生成高性能代码。对NVIDIA和AMD GPUs, 使用OpenAI Triton编译器作为Backend。

torch.compile编译过程如下：

在图编译视角下，Pytorch的软件栈如下，Triton是Inductor的一个Codegen Backend：

示例

以如下pytorch代码为例：

import torch

def mul(x1, x2):
    return x1 * x2

x1 = torch.randn(4096, 390, device='cuda', requires_grad=True)
x2 = torch.randn(4096, 390, device='cuda', requires_grad=True)
fn = torch.compile(mul, backend='inductor')
out = fn(x1, x2)
out.sum().backward()

编译过程

在pytorch/torch/_inductor/codecache.py中的write函数打印调用栈，观察torch.compile的实际编译过程：

def write(
    content: Union[str, bytes],
    extension: str,
    extra: str = "",
    hash_type: str = "code",
    specified_dir: str = "",
) -> Tuple[str, str]:
    # use striped content to compute hash so we don't end up with different
    # hashes just because the content begins/ends with different number of
    # spaces.
    print('#######################################################')
    import traceback
    traceback.print_stack()
    ...

前向图编译

通过分析调用栈，可以看到：

  • JIT编译：前向图的编译在调用编译的模型时触发：out = fn(x1, x2)，并不是在fn = torch.compile(mul, backend='inductor')时触发。
  • TorchDynamo：从out = fn(x1, x2)到return compile_fx(model_, inputs_, config_patches=self.config)完成了Python Bytecode捕获，并转换为FX Graph。
  • AOT Autograd：从return compile_fx(model_, inputs_, config_patches=self.config)到return self.compiler_fn(gm, example_inputs)，完成了反向图FX Graph的生成和Autograd注册，为后续反向图编译打下了基础。
  • Inductor：从return inner_compile到return self._compile_to_module实现了FX Graph到Triton Kernel的编译。
  File "/home/vincent/mul.py", line 10, in <module>
    out = fn(x1, x2)
  File "/home/vincent/pytorch/torch/_dynamo/eval_frame.py", line 574, in _fn
    return fn(*args, **kwargs)
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 1380, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 1164, in __call__
    result = self._inner_convert(
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 547, in __call__
    return _compile(
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/home/vincent/pytorch/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/vincent/pytorch/torch/_dynamo/bytecode_transformation.py", line 1361, in transform_code_object
    transformations(instructions, code_options)
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "/home/vincent/pytorch/torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "/home/vincent/pytorch/torch/_dynamo/symbolic_convert.py", line 2868, in run
    super().run()
  File "/home/vincent/pytorch/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
  File "/home/vincent/pytorch/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/vincent/pytorch/torch/_dynamo/symbolic_convert.py", line 3048, in RETURN_VALUE
    self._return(inst)
  File "/home/vincent/pytorch/torch/_dynamo/symbolic_convert.py", line 3033, in _return
    self.output.compile_subgraph(
  File "/home/vincent/pytorch/torch/_dynamo/output_graph.py", line 1101, in compile_subgraph
    self.compile_and_call_fx_graph(
  File "/home/vincent/pytorch/torch/_dynamo/output_graph.py", line 1382, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/vincent/pytorch/torch/_dynamo/output_graph.py", line 1432, in call_user_compiler
    return self._call_user_compiler(gm)
  File "/home/vincent/pytorch/torch/_dynamo/output_graph.py", line 1462, in _call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/vincent/pytorch/torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/vincent/pytorch/torch/__init__.py", line 2340, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/home/vincent/pytorch/torch/_inductor/compile_fx.py", line 1863, in compile_fx
    return aot_autograd(
  File "/home/vincent/pytorch/torch/_dynamo/backends/common.py", line 83, in __call__
    cg = aot_module_simplified(gm, example_inputs, **s
最低0.47元/天 解锁文章
VLLM专题（三十八）—vLLM 的 torch.compile 集成
03-19 783
在vLLM的V1架构中，默认启用，并且是框架的关键部分。本文档通过一个简单的示例来展示如何理解的使用。在整个示例中，我们将使用v1运行一个常见的Llama模型，并开启调试级别的日志记录以显示所有细节。
Triton - Introduction
weixin_42654107的博客
07-26 787
Triton是一种为并行编程设计的语言和编译器，旨在提供一个Python-based变成环境，提升DNN compute kernels的开发效率的同时，也能最大化利用现代GPU硬件的计算吞吐能力。原生的Triton编译器的编译管线与GPU深度绑定，微软在原生Triton的基础上，在社区贡献了triton-shared项目，对接到了MLIR的linalg Dialect，方便支持其他类型的AI加速器。
PyTorch编译优化：torch.compile与图编译技术
gitblog_00707的博客
09-19 1001
在深度学习模型训练与推理过程中，计算效率始终是核心挑战。随着模型规模的指数级增长（如GPT-4的万亿参数规模），传统即时执行（Eager Execution）模式的性能瓶颈日益凸显。PyTorch作为主流框架，自2.0版本起引入`torch.compile` API，标志着其向静态图编译优化范式的重大演进。本文将系统剖析PyTorch的编译优化技术栈，通过原理分析、性能对比和实战案例，展示如何通过...
推理加速 --- torch.compile
热门推荐
在人工智能的道路上爬行
08-01 1万+
PyTorch 2.0 官宣了一个重要特性 —— torch.compile，这一特性将 PyTorch 的性能推向了新的高度，并将 PyTorch 的部分内容从 C++ 移回 Python。torch.compile 是一个完全附加的（可选的）特性，因此 PyTorch 2.0 是 100% 向后兼容的。支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。
Triton Pipeines的实现方式及对比
Untitled的专栏
11-02 2908
部署yolov5 Triton Pipelines中，简单介绍了BLS和Ensemble这两种实现Triton Pipelines的方式，本文进行了对比
torch.compile用法
zhuikefeng的博客
02-21 5089
torch.compile 通过 JIT 将 PyTorch 代码编译成优化的内核，使 PyTorch 代码运行得更快。加速主要来自减少了 Python 开销和 GPU 读/写，因此观察到的加速可能因模型架构和批量大小等因素而异。例如，如果一个模型的架构很简单并且数据量很大，那么瓶颈将是 GPU 计算并且观察到的加速可能不那么显着。要求torch>2.0，还需要安装torchtriton。
【编译系列】Torch.compile()流程解析——2. TorchDynamo
StarCap的博客
11-27 1769
在上一篇【编译系列】Torch.compile()流程解析——1. torch.compile介绍，我们解释了torch.compile出现的背景并初步了解了其使用和基础组成(感兴趣的小伙伴可以去翻一翻哦~~~)。本章我们将解释四个基础组件中的TorchDynamo。
【编译系列】Torch.compile()训练编译——算子融合逻辑 & 工程化
StarCap的博客
02-02 1509
从2.4版本开始，PyTorch开始支持通过torch.library的方式往torch.compile中注册新operator，对环境有一定要求torch>=2.4且gcc>=9(训练机器最高为8.2，需要手动安装)。<2.4版本也支持，但使用起来会比较繁琐。在训练场景下，要求实现operator对应的前向函数和反传函数。
进一步认识torch.compile的decomp对编译、执行与inplace操作的影响
weixin_38718148的博客
01-26 673
在torch.compile模式下，当我们企图实现某个计算语义，可能存在两种选择：1、将该计算语义定义为一个算子，如"reuse_input_tensor_op(Tensor[] output_tensor_list, Tensor input_tensor, str tag, int[] ranks, int group_size) -> Tensor[]"，此时该计算语义具备被定义好的输入、输出。
【编译系列】Torch.compile()流程解析——1. torch.compile介绍
StarCap的博客
11-27 2518
本系列主要分享训练编译的相关知识，第一章主要解析torch.compile的工作流程以及各个组件的作用和源码调用关系，配合代码示例进行解析。主要分为两大块：前端静态图捕获和后端编译。a. 前端对应TorchDynamo，主要介绍python编译执行的基本流程、PEP 523、TorchDynamo模拟执行字节码、FX Graph的构建过程；
pytorch 源码阅读(1)——torch.complie
哦豁灬
06-25 7163
如上所述，为了更快地运行工作流，torch.compile 通过 TorchDynamo 需要一个后端，该后端将捕获的图形转换为快速的机器代码。对于在编译区域内执行的每一帧，都会尝试编译，并将编译结果缓存到代码对象中，供将来使用。中引入的 PyTorch 函数，旨在解决 PyTorch 中的准确图形捕获问题，并最终使软件工程师能够更快地运行其 PyTorch 程序。这个函数的作用是使用 TorchDynamo 和指定的后端优化给定的模型/函数，进行包装和参数配置的设置，最终会调用。，这个是用来进行调试的。
深入解析torch.compile：提升PyTorch模型性能、高效解决常见问题
deephub
05-17 8363
PyTorch 2.0的torch.compile功能显著提升了深度学习模型的性能。本文从实用角度出发，介绍了torch.compile的核心技巧，帮助用户提升开发效率。文章首先将模型划分为三种复杂度类别，分别讨论了直接适配型、需调整适配型和高复杂度调整型的应用场景。接着，分析了训练工作流中可编译的组件，如模型定义、优化器流程、自动微分系统和日志记录功能，并指出了当前尚不完全支持的编译场景。文章还提供了系统化的调试策略，包括跟踪分析与可视化、分层消融测试、问题最小化复现和独立复现环境构建。
高性能 ：OpenAI Triton＆ torch.compile(add,backend=“inductor“) & torch生成Triton 内核代码+使用用户编写的Triton 内核
ResumeProject的博客
02-11 1919
的行为、查看中间代码、诊断错误或性能问题。进行编译时，可以启用调试模式，以便分析。是一个环境变量，在使用。
【编译系列】Torch.compile()流程解析——3. AOTAutograd
StarCap的博客
11-30 1724
在上一节TorchDynamo的介绍中，我们解析了torch.compile()是如何捕获计算图并保存为GraphModule的，但在这个过程中只是对整个Python字节码进行了模拟执行、解析并构建FX Graph，相当于只是初步构建了前向计算图，没有捕获训练场景下的反向计算图。
Pytorch实用教程（一）：torch.compile计算提速
所有专家都曾始于新手
04-22 1484
开源仓库：TingsongYu/PyTorch-Tutorial-2nd: 《Pytorch实用教程》（第二版）无论是零基础入门，还是CV、NLP、LLM项目应用，或是进阶工程化部署落地，在这里都有。相信在本书的帮助下，读者将能够轻松掌握 PyTorch 的使用，成为一名优秀的深度学习工程师。在线阅读：简介 · PyTorch实用教程（第二版） (tingsongyu.github.io)
torch.compile模型编译加速
weixin_40777649的博客
07-15 1906
定义接口介绍案例。
炸裂提速！PyTorch 2.2 Torch.compile优化器实战：从原理到模型训练加速300%全攻略
weixin_40593051的博客
05-19 2294
深度学习训练性能在传统PyTorch模式下存在计算效率低、内存开销大和部署复杂等问题。PyTorch 2.2引入的Torch.compile通过多级编译架构和关键优化技术（如算子融合、内存布局优化和动态形状支持），显著提升了训练性能。实验表明，Torch.compile将单步迭代时间从800ms降至200ms，GPU利用率从45%提升至90%，显存占用减少50%，训练总时长缩短至24小时。与传统JIT编译相比，Torch.compile具有全局优化、硬件适配和易用性等优势。
torch.compile() 加速推理速度
Drug discovery
05-13 1589
需安装2.0及以上版本的torch即可。
使用torch.compile进行CPU优化
AI知识搬运工
02-25 1025
在PyTorch中，使用torch.compile可以自动地将模型转换成优化的执行代码，这对于提升模型在CPU上的运行效率尤其有用。torch.compile是基于TorchDynamo实现的，它可以将Python代码转换为高效的TorchScript代码。这对于那些在CPU上运行的大型模型尤其有益，因为它可以减少运行时开销并提高整体性能。
torch.compile是什么用法
最新发布
12-10
<think>我们被要求用中文回答，并在最后提出3-5个相关问题，以“§§相关问题§§”开头，然后列出问题。 用户的问题是：torch.compile是什么用法 首先，回答torch.compile的用法： torch.compile 是 PyTorch 2.0 引入的一个功能，用于将 PyTorch 模型编译成更高效的表示形式，从而加速模型运行。它通过使用即时（JIT）编译技术，将模型图优化并转换为底层硬件特定的代码，以提高执行速度。 具体用法： 1. 基本用法：使用 torch.compile 包裹一个模型（nn.Module 实例）或一个函数。 2. 参数说明： - model_or_function: 要编译的模型或函数。 - mode: 编译模式，可选 "default"（平衡优化）、"reduce-overhead"（减少框架开销，适用于小模型）和 "max-autotune"（最大程度优化，编译时间长）。 - dynamic: 是否开启动态形状支持（适用于输入形状变化的场景）。 示例代码： ```python import torch # 定义一个简单的模型 class MyModel(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(10, 10) def forward(self, x): return self.linear(x) model = MyModel() # 编译模型 compiled_model = torch.compile(model) # 使用编译后的模型 input_tensor = torch.randn(1, 10) output = compiled_model(input_tensor) ``` 3. 注意事项： - 第一次调用编译后的函数或模型时，会进行编译，因此会较慢，后续调用会使用缓存。 - 对于动态图（每次执行图结构变化的情况）可能不适用。 4. 高级用法：使用 `torch.compile` 也可以编译任意函数，例如： ```python def my_function(x, y): return x @ y compiled_function = torch.compile(my_function) result = compiled_function(torch.randn(2,2), torch.randn(2,2)) ``` 5. 性能建议： - 对于循环中的小操作，编译整个循环而不是内部操作。 - 对于大模型，使用 `mode="max-autotune"` 可能获得最佳性能，但需要更长的编译时间。 接下来，根据问题和回答，提出3-5个相关问题。 相关问题可围绕以下方面： 1. torch.compile 的工作原理 2. 不同模式的区别和选择 3. 实际应用中的性能提升 4. 使用限制和注意事项 因此，我们可以提出以下问题： §§相关问题§§ 1. torch.compile 在底层使用了哪些编译技术来加速 PyTorch 模型？ 2. 如何根据模型的大小和结构选择合适的 torch.compile 模式（如 "default"、"reduce-overhead"、"max-autotune"）？ 3. 在实际项目中，使用 torch.compile 通常能带来多少性能提升？请举例说明。 4. 在哪些情况下不适合使用 torch.compile？ 5. 如何诊断 torch.compile 编译后的模型性能瓶颈？ </think> ### `torch.compile` 的用法详解 `tororch.compile` 是 PyTorch 2.0 引入的核心特性，用于**动态编译模型图**，显著提升计算效率。其核心机制是通过即时编译（JIT）将 Python 模型转换为优化的底层代码（如 C++/CUDA）。以下是关键用法： #### 1. 基础语法 ```python compiled_fn = torch.compile( fn, # 待编译的函数或 nn.Module *, mode="default", # 编译模式 dynamic=False, # 动态形状支持 fullgraph=False, # 强制完整编译 backend="inductor", # 默认后端 options=None # 后端配置 ) ``` #### 2. 核心参数解析 | 参数 | 可选值 | 作用 | |------|--------|------| | **`mode`** | `"default"` (平衡优化)<br>`"reduce-overhead"` (减少小模型开销)<br>`"max-autotune"` (极致性能) | 优化策略选择 | | **`dynamic`** | `True`/`False` | 启用动态形状支持（处理可变输入尺寸） | | **`fullgraph`** | `True`/`False` | 强制整个模型可编译（否则回退 eager 模式） | | **`backend`** | `"inductor"` (默认)<br>`"aot_eager"`<br>`"nvfuser"` | 编译引擎选择 | #### 3. 典型使用场景 **场景 1：加速自定义函数** ```python def cosine_sim(x, y): return x @ y.T / (x.norm(dim=1) * y.norm(dim=1)) # 编译函数（减少 3-5 倍耗时） compiled_cos = torch.compile(cosine_sim) result = compiled_cos(tensor_a, tensor_b) ``` **场景 2：优化神经网络（官网示例应用）** ```python from e3nn import o3 # 原始张量积操作 tp = o3.FullTensorProduct(irreps_in1="1o", irreps_in2="1o") # 编译为高效内核（首次运行慢，后续加速 2-10 倍） tp_compiled = torch.compile(tp, fullgraph=True, mode="max-autotune") z = tp_compiled(x, x) # 编译发生在首次调用 ``` #### 4. 性能对比建议 ```python import timeit # 基准测试（未编译） t1 = timeit.timeit(lambda: tp(x,x), number=100) # 基准测试（已编译） t2 = timeit.timeit(lambda: tp_compiled(x,x), number=100) print(f"加速比: {t1/t2:.1f}x") # 典型输出: 加速比: 4.8x ``` #### 5. 注意事项 - **预热阶段**：首次调用会触发编译（较慢），后续调用高速运行 - **动态形状限制**：默认仅支持静态形状，启用 `dynamic=True` 处理变长输入 - **算子覆盖**：部分自定义算子需注册才能编译（[参见文档](https://pytorch.org/docs/stable/torch.compiler.html)） > **适用场景**：循环计算/科学计算/3D 处理（如 e3nn）/大模型推理
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
表情包 代码片
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值