国庆啥也没干，等了好久，PyTorch 2.1终于正式发布（发布说明），有些重磅功能终于完善的差不多了。
在PyTorch 2.1中，torch.compile支持automatic dynamic shape，还有torch.distributed.checkpoint功能，它能够并行地在多个级别上保存和加载分布式训练任务，以及torch.compile对NumPy API的全面支持。
值得注意的是，这个新版本还带来了一系列性能的提升，例如CPU inductor的优化、对AVX512的支持、以及对scaled-dot-product-attention的支持等。同时，我们还首次推出了torch.export的原型版本，这是一个完善的全图捕获机制，并引入了基于torch.export的量化方式。
简要概述下：
- torch.compile 可以自动检测并最小化由于张量形状改变导致的重复编译，这一切都得益于它现在支持自动动态形状。
- 通过torch.distributed.checkpoint，您可以并行地从多个级别保存和加载模型，同时还支持因集群拓扑改变而进行的resharding操作。
- torch.compile 可以编译NumPy操作，并将其转换为相应的PyTorch操作。
- torch.compile 为Python 3.11带来了更好的支持和兼容性。
- CPU性能方面进行了多项优化和改进，包括对inductor的优化（比如支持bfloat16和动态形状）、引入了支持AVX512的内核，以及推出了scaled-dot-product-attention内核。
- 推出了torch.export prototype 功能，这是一个稳定的全图捕获机制，同时还有基于torch.export的量化方式。
- torch.sparse现在支持在NVIDIA® GPU上使用semi-structured (2:4) sparsity ，目前这一功能还处于prototype阶段。
细节说明
BETA FEATURES
(Beta) Automatic Dynamic Shapes
Dynamic Shapes是内嵌在torch.compile中的功能，通过跟踪和基于张量的符号形状（而不是静态形状，例如*[B, 128, 4]而不是[64, 128, 4]）生成代码来减少重编译。这允许torch.compile生成一个可以适用于多个大小的单一内核，只需以轻微的效率成本为代价。在PyTorch 2.1中，动态形状已得到极大稳定，并且现在如果torch.compile*注意到由于输入形状变化而重新编译，将自动启用。可以通过将dynamic=False传递给torch.compile，或通过设置torch._dynamo.config.automatic_dynamic_shapes = False来禁用自动动态。
在PyTorch 2.1中，在包括大型语言模型在内的各种模型类型上显示出启用Dynamic Shapes的良好性能，包括CUDA和CPU。
有关动态形状的更多信息，请查看此文档。
[Beta]torch.distributed.checkpoint
torch.distributed.checkpoint使并行保存和加载来自多个rank的模型成为可能。此外，检查点自动处理模型和优化器之间的完全限定名（FQN）映射，从而在不同集群拓扑中进行加载时重分片。
更多信息，请参阅torch.distributed.checkpoint的文档和教程。
[Beta]torch.compile + NumPy
torch.compile现在了解如何通过将numpy的op翻译成PyTorch等效操作来编译NumPy操作。因为这个集成以设备无关的方式运作，您现在可以仅通过使用torch.compile来加速NumPy程序，或甚至混合NumPy/PyTorch程序。
有关torch.compile + NumPy interaction的更多信息，请参阅torch.compile FAQ中的此部分，并关注PyTorch Blog以获取有关此功能的即将发布的博客。
[Beta]torch.compile + Python 3.11
torch.compile之前只支持Python 3.8-3.10版本。用户现在可以使用Python 3.11中的torch.compile来优化模型。
[Beta]torch.compile + autograd.Function
torch.compile现在可以跟踪和优化用户定义的autograd Functions的反向函数，这解锁了对使用扩展机制的模型的训练优化。
[Beta] 改进的第三方设备支持：PrivateUse1
现在可以使用privateuse1调度键将第三方设备类型注册到PyTorch。这允许设备扩展向PyTorch注册新内核，并将它们与新键关联，允许用户代码与内置设备类型等效地工作。例如，要注册*“my_hardware_device*”, 一个人可以执行以下操作：
torch.rename_privateuse1_backend("my_hardware_device")
torch.utils.generate_methods_for_privateuse1_backend()
x = torch.randn((2, 3), device='my_hardware_device')
y = x + x # 在'my_hardware_device'上运行add kernel
为了验证这个特性，Ascend NPU的OSS团队已成功将torch_npu集成到pytorch作为通过PrivateUse1功能的插件。
更多信息，请参见PrivateUse1教程这里。
原型特性
[Prototype] torch.export()
*torch.export()*提供了一个基于PT2.0提供的新技术的稳定追踪机制，以从PyTorch程序中捕获完整图表。
用户可以提取PyTorch程序的清晰表示形式（Export IR），它是以数据流图的形式，主要由对PyTorch操作的直线调用组成。然后可以转换，序列化，保存到文件，转移，加载回执行环境，无论是有Python还是无Python。
更多信息，请参阅教程这里。
[Prototype] 基于torch.export的量化
torch.ao.quantization现在支持在PyTorch 2 torch.export-based flows上的量化。这包括对内置XNNPACK和X64Inductor Quantizer的支持，以及指定自己Quantizer的能力。
有关使用torch.export的后训练静态量化的解释，请参阅此教程，有关使用torch.export的量化感知训练的静态量化，请参阅此教程。
关于如何编写自己的Quantizer的解释，请参见此教程。
[Prototype] NVIDIA® GPUs的semi-structured (2:4) sparsit
torch.sparse现在支持在半结构化稀疏（2:4）张量上创建和加速计算。有关该格式的更多信息，请参阅NVIDIA的此博客。引入半结构化稀疏性的一个简单示例如下：
from torch.sparse import to_sparse_semi_structured
x = torch.rand(64, 64).half().cuda()
mask = torch.tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
linear = nn.Linear(64, 64).half().cuda()
linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
linear(x)
[Prototype] cpp_wrapper for torchinductor
cpp_wrapper可以通过在C++中生成内核包装器代码来减少在torchinductor中调用内核的Python开销。此功能仍处于原型阶段; 它不支持今天在PT2中成功编译的所有程序。
如果发现限制，请报告问题以帮助Pytorch官方确定优先级。
启用此功能的API是：
import torch
import torch._inductor.config as config
config.cpp_wrapper = True
有关更多信息，请参阅教程。
性能改进
AVX512内核支持
在PyTorch 2.0中，即使CPU支持AVX512指令，也会使用AVX2内核。现在，如果CPU支持这些指令，PyTorch默认使用AVX512 CPU内核，相当于在先前版本中设置ATEN_CPU_CAPABILITY=avx512。可以通过设置ATEN_CPU_CAPABILITY=avx2来启用先前的行为。
针对scaled-dot-product-attention (SDPA)的CPU优化
PyTorch先前版本通过torch.nn.functiona.scaled_dot_product_attention为转换器基元提供了优化的CUDA实现。 PyTorch 2.1包括基于FlashAttention的优化CPU例程。
请参阅文档这里。
bfloat16的CPU优化
PyTorch 2.1包括针对bfloat16的CPU优化，包括改进的矢量化支持和torchinductor代码生成。