PyTorch Functorch项目：AOT Autograd优化技术详解-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00603/article/details/149010919

PyTorch Functorch项目：AOT Autograd优化技术详解

引言

在深度学习模型训练过程中，计算图的优化对于提升训练效率至关重要。PyTorch Functorch项目中的AOT Autograd技术提供了一种创新的方法来加速模型训练。本文将深入解析AOT Autograd的工作原理、使用方法以及它如何通过编译器和重计算技术来优化训练过程。

AOT Autograd技术概述

AOT (Ahead-Of-Time) Autograd是PyTorch Functorch中的一个重要组件，它通过提前追踪前向和反向计算图来实现训练优化。这项技术具有两大核心优势：

提前图追踪：在训练开始前就捕获完整的前向和反向计算图，这使得联合图优化（如重计算/激活检查点）成为可能
编译支持：为提取的计算图提供简单的编译接口，支持多种深度学习编译器（如NVFuser、NNC、TVM等）

基础使用示例

让我们从一个简单的例子开始，了解如何使用AOT Autograd：

import torch

def fn(a, b, c, d):
    x = a + b + c + d
    return x.cos().cos()

这个函数执行了简单的张量相加和两次余弦运算。我们将使用AOT Autograd来优化它的训练过程。

计算图可视化

AOT Autograd使用__torch_dispatch__机制来追踪计算图，并将结果存储在torch.Fx的GraphModule容器中。我们可以编写一个简单的编译器函数来打印计算图：

from functorch.compile import aot_function

def compiler_fn(fx_module: torch.fx.GraphModule, _):
    print(fx_module.code)
    return fx_module

aot_print_fn = aot_function(fn, fw_compiler=compiler_fn, bw_compiler=compiler_fn)

运行后会打印出前向和反向计算图。值得注意的是，前向图除了输出最终结果外，还会输出一些中间张量，这些张量将在反向传播时用于梯度计算。

算子融合优化

算子融合是深度学习优化的关键技术之一，它通过减少内存访问次数来提升性能。AOT Autograd可以与多种编译器配合实现这一优化：

from functorch.compile import ts_compile

aot_nnc_fn = aot_function(fn, fw_compiler=ts_compile, bw_compiler=ts_compile)

在我们的基准测试中，使用NNC编译器后，前向和反向传播都获得了明显的加速。特别是反向传播，由于不受中间结果保存的限制，获得了更大的性能提升。

重计算（激活检查点）技术

重计算（又称激活检查点）是一种内存优化技术，它通过在前向传播中少保存一些激活值，在反向传播时重新计算这些值来节省内存。传统实现会带来性能开销，但结合融合编译器，我们可以实现内存和性能的双重优化。

AOT Autograd提供了min_cut_rematerialization_partition分区器来实现这一优化：

from functorch.compile import min_cut_rematerialization_partition

aot_fn = aot_function(fn, fw_compiler=compiler_fn, 
                     bw_compiler=compiler_fn, 
                     partition_fn=min_cut_rematerialization_partition)

从打印的计算图可以看出，使用重计算后，前向图输出的张量减少了，部分计算被移到了反向传播阶段。结合NNC编译器，我们获得了进一步的性能提升。