最近PyTorch生态在性能优化、硬件支持、模型部署等方面有诸多能力更新,主要体现在PyTorch 2.7版本发布以及高性能通用AI算子库FlagGems加入生态等,具体如下:
- PyTorch 2.7版本更新:
- 深度集成TensorRT-LLM:显著加强了与NVIDIA TensorRT-LLM的兼容性和易用性,用户可在不改动模型代码的情况下,通过TensorRT获得最高数倍的推理加速,利于企业部署大型模型。
- 全面支持CUDA 12.x:针对CUDA 12.1/12.2引擎进行优化,在多流并发、异步内存拷贝、异构计算调度等方面带来显著性能提升,使训练复杂模型更为高效可靠。
- 升级Triton Inference Server集成体验:优化了模型的导出、序列化和格式兼容性,能更顺畅地对接NVIDIA Triton Inference Server,降低线上推理部署和维护的复杂度。
- 增强多实例GPU(MIG)支持:在调度引擎中引入更灵活的资源管理能力,支持A100、H100上基于MIG的细粒度资源划分,可在同一块GPU上同时部署多种不同工作负载。
- 默认启用torch.compile:将torch.compile(model)作为默认推荐方式,自动应用图优化、内核融合、低级别内存调优等,普遍带来20%-50%的性能加速。
- 强化低精度计算(FP8)支持:加强了对FP8数据格式的支持,配合NVIDIA Transformer Engine,可将训练吞吐量提高1.5倍以上,且进一步降低显存占用。
- FlagGems加入PyTorch生态:2025年,高性能通用AI算子库FlagGems正式加入PyTorch生态项目体系。它是唯一支持多种AI芯片架构的算子库,能加速大语言模型的训练和推理过程。通过在PyTorch的ATen后端进行注册,用户无需修改模型代码即可切换到Triton函数库。其平均性能超过各个芯片厂商的原生算子性能,比PyTorch ATen的CUDA算子库平均提速30%,可帮助PyTorch适配多个芯片硬件后端生态,降低算法在不同硬件平台之间的迁移难度。
- 其他更新:PyTorch 2.7版本还增加了对NVIDIA Blackwell架构的支持,提升了Intel GPU加速能力,优化了FlexAttention大型语言模型在X86 CPU上的处理性能,新增了torch.fx.passes.graph_transform_observer.GraphTransformObserver接口用于更细粒度的节点级别追踪等。