首页
/ PyTorch AO v0.9.0发布:稀疏化技术成熟与量化API重大升级

PyTorch AO v0.9.0发布:稀疏化技术成熟与量化API重大升级

2025-06-24 08:25:00作者:江焘钦

PyTorch AO(torchao)是PyTorch生态系统中专注于模型优化的重要组件,主要提供模型量化和稀疏化等前沿优化技术。近日发布的v0.9.0版本标志着该项目在稀疏化技术成熟度和量化API设计方面取得了重大进展。

稀疏化技术正式进入生产阶段

本次版本最引人注目的变化是将块稀疏(Block Sparsity)技术从原型状态提升为正式功能。块稀疏是一种高效的内存压缩技术,通过将权重矩阵划分为固定大小的块并选择性保留重要块来实现模型压缩和加速。

新版本提供了简洁易用的API接口,开发者只需几行代码即可实现模型稀疏化:

from torchao.sparsity import sparsify, block_sparse_weight
sparsify_(model, block_sparse_weight(blocksize=64))

性能测试数据显示,在Meta-Llama-3.1-8B模型上,使用64x64块大小和90%稀疏率时,解码速度达到262.94 tokens/s,同时模型大小压缩至4.88GB,相比原始模型实现了显著的性能提升和存储节省。

超级掩码技术提升稀疏模型精度

针对稀疏化可能带来的精度损失问题,v0.9.0引入了Supermask技术。这是一种训练阶段的技术,通过在训练过程中学习最优的块稀疏掩码,有效缓解了稀疏化带来的精度下降问题。

使用方式上,开发者可以先使用Supermask进行训练,再转换为常规稀疏格式进行推理:

from torchao.sparsity import SupermaskLinear
sparsify_(model, lambda x: SupermaskLinear.from_linear(x, block_size=64, sparsity_level=0.9))
# 训练代码...
sparsify_(model, lambda x: SupermaskLinear.to_linear(x, sparsity_level=0.9))
sparsify_(model, block_sparse_weight(blocksize=64))

量化API架构重构

v0.9.0对量化API进行了重大重构,从原先基于可调用对象的配置方式转向更直观的配置对象模式。这一变化使API设计更加一致,便于配置检查和后续扩展。

新旧API对比示例:

# 旧版API (v0.8.0及之前)
quantize_(model, int8_weight_only(group_size=128))

# 新版API (v0.9.0)
from torchao.quantization import Int8WeightOnlyConfig
quantize_(model, Int8WeightOnlyConfig(group_size=128))

尽管旧版API仍可继续使用,但建议开发者逐步迁移到新的配置对象模式。这一变化为未来更多训练优化技术的引入奠定了良好基础。

4位动态量化新内核

针对资源受限场景,v0.9.0新增了基于CUTLASS的4位动态激活+4位权重量化内核,进一步降低模型内存占用和计算需求:

from torchao.quantization import int4_dynamic_activation_int4_weight
quantize_(model, int4_dynamic_activation_int4_weight)

面向NVIDIA Blackwell的MXFP支持

作为早期原型特性,v0.9.0开始支持NVIDIA Blackwell GPU的MXFP8和MXFP4训练与推理。MX格式是一种新兴的微缩放浮点格式,可在保持精度的同时显著提升计算效率。

训练示例:

from torchao.prototype.mx_formats import swap_linear_with_mx_linear
from torchao.prototype.mx_formats.config import MXLinearConfig

config = MXLinearConfig(elem_dtype=torch.float8_e4m3fn, block_size=32)
swap_linear_with_mx_linear(model, config=config)

需要注意的是,当前MX支持仍处于早期阶段,性能和API稳定性将在后续版本中持续改进。

其他重要改进

  1. 浮点8训练支持2的幂次缩放因子,并使用e4m3格式提升前向和反向传播效率
  2. 优化了自动量化(autoquant)对不同数据类型的支持
  3. 修复了多项与分布式训练、CPU卸载优化器等相关的关键问题
  4. 大幅扩充了文档体系,新增快速入门指南和API迁移教程

总结

PyTorch AO v0.9.0通过将块稀疏技术正式化、重构量化API架构、引入4位量化内核等重大改进,为开发者提供了更强大、更易用的模型优化工具。特别是稀疏化技术的成熟,使得在保持模型精度的同时实现显著加速成为可能。随着MXFP等新兴格式支持的逐步完善,PyTorch AO正在为下一代AI硬件的性能优化奠定基础。

登录后查看全文
热门项目推荐

项目优选

收起
RuoYi-Vue3RuoYi-Vue3
🎉 (RuoYi)官方仓库 基于SpringBoot,Spring Security,JWT,Vue3 & Vite、Element Plus 的前后端分离权限管理系统
Vue
703
459
ohos_react_nativeohos_react_native
React Native鸿蒙化仓库
C++
141
224
openGauss-serveropenGauss-server
openGauss kernel ~ openGauss is an open source relational database management system
C++
102
159
leetcodeleetcode
🔥LeetCode solutions in any programming language | 多种编程语言实现 LeetCode、《剑指 Offer(第 2 版)》、《程序员面试金典(第 6 版)》题解
Java
53
15
openHiTLSopenHiTLS
旨在打造算法先进、性能卓越、高效敏捷、安全可靠的密码套件,通过轻量级、可剪裁的软件技术架构满足各行业不同场景的多样化要求,让密码技术应用更简单,同时探索后量子等先进算法创新实践,构建密码前沿技术底座!
C
114
255
MateChatMateChat
前端智能化场景解决方案UI库,轻松构建你的AI应用,我们将持续完善更新,欢迎你的使用与建议。 官网地址:https://matechat.gitcode.com
706
97
SnowySnowy
💖国内首个国密前后分离快速开发平台💖《免费商用》,基于开源技术栈精心打造,融合Vue3+AntDesignVue4+Vite5+SpringBoot3+Mp+HuTool+Sa-Token。平台内置国密加解密功能,保障前后端数据传输安全;全面支持国产化环境,适配多种机型、中间件及数据库。特别推荐:插件提供工作流、多租户、多数据源、即时通讯等高级插件,灵活接入,让您的项目开发如虎添翼。
Java
179
23
HarmonyOS-ExamplesHarmonyOS-Examples
本仓将收集和展示仓颉鸿蒙应用示例代码,欢迎大家投稿,在仓颉鸿蒙社区展现你的妙趣设计!
Cangjie
363
355
RuoYi-Cloud-Vue3RuoYi-Cloud-Vue3
🎉 基于Spring Boot、Spring Cloud & Alibaba、Vue3 & Vite、Element Plus的分布式前后端分离微服务架构权限管理系统
Vue
122
85
cherry-studiocherry-studio
🍒 Cherry Studio 是一款支持多个 LLM 提供商的桌面客户端
TypeScript
530
45