龙岗网站建设网站制作,seo优化推广多少钱,网站建设 知乎,wordpress js加速【新智元导读】新的一年#xff0c;PyTorch也迎来了重大更新#xff0c;PyTorch 2.2集成了FlashAttention-2和AOTInductor等新特性#xff0c;计算性能翻倍。
新的一年#xff0c;PyTorch也迎来了重大更新#xff01;
继去年十月份的PyTorch大会发布了2.1版本之后#…
【新智元导读】新的一年PyTorch也迎来了重大更新PyTorch 2.2集成了FlashAttention-2和AOTInductor等新特性计算性能翻倍。
新的一年PyTorch也迎来了重大更新
继去年十月份的PyTorch大会发布了2.1版本之后全世界各地的521位开发者贡献了3628个提交由此形成了最新的PyTorch 2.2版本。 新的版本集成了FlashAttention-2使得scaled_dot_product_attention SDPA相较于之前的版本有了约2倍的性能提升。
PyTorch 2.2还引入了一个新的TorchInductor提前扩展称为 AOTInductor旨在为非python服务器端编译和部署PyTorch程序。
PyTorch中的torch.distributed支持了一个叫做device_mesh的新抽象用于初始化和表示ProcessGroups。 另外PyTorch 2.2提供了一个标准化的、可配置的日志记录机制——TORCH_LOGS。
PyTorch 2.2还对torch.compile做了许多改进包括改进了对编译优化器的支持以及TorchInductor融合和布局优化。 最后值得注意的是PyTorch将放弃对macOS x86的支持PyTorch 2.2.x是支持macOS x64的最后一个版本。
PyTorch 2.2新特性
首先请注意如果从源代码构建PyTorch 2.2需要GCC 9.4或更高版本PyTorch 代码库已从C 14迁移到C 17。 FlashAttention-2
FlashAttention-2通过优化GPU上不同线程块和warps之间的工作分区来解决占用率低或不必要的共享内存读写。 FlashAttention-2调整了算法以减少非matmul的计算量同时提升了Attention计算的并行性即使是单个头也可以跨不同的线程块以增加占用率在每个线程块中优化warps之间的工作分配以减少通过共享内存的通信。
PyTorch 2.2将FlashAttention内核更新到了v2版本不过需要注意的是之前的Flash Attention内核具有Windows实现Windows用户可以强制使用sdp_kernel仅启用Flash Attention的上下文管理器。 而在2.2中如果必须使用 sdp_kernel 上下文管理器请使用memory efficient或math内核在Windows上。 在FlashAttention-2的加持之下torch.nn.functional.scaled_dot_product_attention的速度提升了大约2倍在A100 GPU上达到了理论计算峰值的50%-73%。
AOTInductor
AOTInductor是TorchInductor的扩展用于处理导出的PyTorch模型对其进行优化并生成共享库以及其他相关工件。
这些编译的工件可以部署在非Python环境中经常用于服务器端的推理。
下面的示例演示了如何调用 aot_compile 将模型转换为共享库。 AOTInductor支持与Inductor相同的后端包括CUDA、ROCm和CPU。
TORCH_LOGS
PyTorch 2.2提供了一个标准化的、可配置的日志记录机制可用于分析各种子系统的状态例如编译和分布式操作
可以通过TORCH_LOGS环境变量启用日志。比如通过在命令行中修改环境变量 将TorchDynamo的日志级别设置为logging.ERROR将TorchInductor的日志级别设置为logging.DEBUG。
当然也可以在代码中以API的形式使用 torch.distributed.device_mesh
PyTorch 2.2引入了一个新的抽象用于表示分布式并行中涉及的 ProcessGroups称为torch.distributed.device_mesh。
为分布式训练设置分布式通信器NCCL是一件麻烦的事情。用户需要编写不同并行度的工作负载并为每个并行度手动设置和管理NCCL通信器ProcessGroup 。
这个过程可能很复杂容易出错。而DeviceMesh 可以简化此过程使其更易于管理。
DeviceMesh 是管理 ProcessGroup 的更高级别的抽象。它允许用户毫不费力地创建节点间和节点内进程组而不必担心如何为不同的子进程组正确设置等级。
例如数组的其中一个维度可以表示FSDP中的数据并行data parallelism而另一个维度可以表示FSDP中的张量并行tensor parallelism。
用户还可以通过 DeviceMesh 轻松管理底层process_groups以实现多维并行。 DeviceMesh在处理多维并行性如3D并行时很有用。如上图所示当你的并行解决方案需要跨主机和每个主机内部进行通信时可以创建一个2D网格用于连接每个主机中的设备并以同构设置将每个设备与其他主机上的对应设备连接起来。
借助 init_device_mesh() 我们可以在短短两行内完成上面这个2D设置 而如果不使用DeviceMesh我们大概需要自己写下面这一堆代码 当然如果需要我们仍然可以访问底层 ProcessGroup 优化器的改进
大概有以下几点 编译优化器在所有基准测试中都提高了性能HuggingFace 18%、TorchBench 19%、TIMM 8% E2E 编译的优化器增加对cudagraphs的支持 对测试套件中所有模型进行平均每个测试套件的基准测试平均编译时间增加约40秒正在进行的优化可能会将其降低到30秒以下。 用于多张量优化器编译的inductor中缺少的主要功能是foreach算子的高效编码生成。
在调度器内部将所有在下放过程中注册的缓冲区列表凝聚到ForeachKernelSchedulerNodes中FusedSchedulerNode的子类。
为了检查融合是否合法每个内部 SchedulerNode 执行的写操作必须与消费SchedulerNode在同一列表索引处的读操作相匹配。 此外正常的垂直融合规则必须允许在消费者和生产者SchedulerNode列表的每个索引处进行融合。
如果满足了这些条件ForeachKernelSchedulerNode将垂直融合成一个 ForeachKernelSchedulerNode其中每个列表上的相应点操作都将被融合。
通过实现这种融合可以将一系列 foreach 运算融合到单个内核中从而实现多张量优化器的完全融合。
性能改进
TorchInductor中添加了许多性能优化包括对torch.concat的水平融合支持、改进的卷积布局优化、以及改进scaled_dot_product_attention模式匹配。 PyTorch 2.2还包括aarch64的许多性能增强包括对mkldnn权重预打包的支持、改进的ideep基元缓存以及通过对OneDNN的固定格式内核改进来提高推理速度。
参考资料
https://pytorch.org/blog/pytorch2-2/
