重磅！PyTorch官宣2.0版本即将发布，最新torch.compile特性说明！

2022年的PyTorch Conference在新奥尔良市举办。刚刚会上的keynote官宣PyTorch2.0版本即将到来。PyTorch是目前最流行的深度学习框架之一，它的易用性被广大的用户所喜爱。关于PyTorch2.0，官方透露其最重要的就是支持动态图，极大地提升了现有模型的运行速度。

而这一切都是围绕torch.compile展开的。本文将简单介绍以下这个前瞻性的特性。

目前PyTorch最大的挑战是eager模式难以跟上不断增长的GPU带宽和更疯狂的模型架构，而保持PyTorch易于使用和调试也是很重要的方面。因此，解决这些问题需要在代码可及性与速度之间保持棘手的平衡。

为了提升PyTorch处理大模型的速度，PyTorch2.0中有一个新特性，即torch.compile。这是一个新增的特性，可以极大地提升模型运行速度。由于是是additive且optional，所以不会对现有版本的PyTorch代码产生任何影响。

经过测试，torch.compile可以极大地提升模型的运行速度。下图是使用torch.compile之后运行GPT-2的测试结果，可以看到，动态图模式下的torch.compile性能提升明显，且seq长度越长，提升越明显。

Torch.compile的基础是新技术--TorchDynamo、AOTAutograd、PrimTorch和TorchInductor。

TorchDynamo使用Python框架评估钩子安全（Python Frame Evaluation Hooks）地捕获PyTorch程序，这是一项重大创新，是我们5年来在安全图形捕获方面的研发成果。

AOTAutograd重载了PyTorch的autograd引擎，作为一个追踪的autodiff，用于生成超前的反向追踪。

PrimTorch将约2000多个PyTorch运算符归纳为约250个原始运算符的封闭集，开发人员可以针对这些运算符构建一个完整的PyTorch后端。这大大降低了编写PyTorch功能或后端的障碍。

TorchInductor是一个深度学习编译器，可以为多个加速器和后端生成快速代码。对于英伟达GPU，它使用OpenAI Triton作为关键构建模块。

TorchDynamo、AOTAutograd、PrimTorch和TorchInductor是用Python编写的，支持动态形状（即能够在不引起重新编译的情况下送入不同大小的Tensors），这使得它们很灵活，容易被破解，并降低了开发者和供应商的准入门槛。

首先PyTorch2.0目前最主要关注的就是性能和可访问性，而剩下的部分完全后向兼容。因此，torch.compile使用是向后兼容的！

不过，如果想使用这个特性，那么需要使用torch.compilewrap后调用现有模型。下图是一个样例：

@torch.compile
def train_fn(input, target):
  optimizer.zero_grad()
  loss(model(input), target).backward()
  optimize.step()

使用很简单。在GitHub的7000个项目上测试结果显示，目前已经支持93%的项目都可以直接使用。而且在A100测试结果显示评价每个项目都有43%的性能提升！

从今天开始，我们既可以在nightly版本中测试这个新特性。即在nightly版本中会引入torch.compile的开发者测试版本，目前这是一个早期开发版本，可能有一些问题。而2023年三月发布的稳定版本的PyTorch2.0也将正式发布这个特性。

不过需要注意的是，PyTorch2.0默认不会开启，需要手动指定开启。

官方也做了一个说明，其实这个版本按照正常的节奏应该是1.14，不过官方认为这个新特性十分重要，所以直接就是2.0版本了！