PyTorch近日推出了一款名为torchao的新库，旨在通过采用低精度数据类型、量化以及稀疏性技术，降低模型的计算成本和内存消耗，从而实现更高效的模型执行。
低精度数据类型：加速模型训练与推理
torchao支持包括float8、int4在内的低精度数据类型，这些类型能够有效地减少计算资源的需求和内存使用量。例如，在LLaMA 3 70B模型的预训练过程中，使用torchao提供的float8训练流程，可以将模型的训练速度提高1.5倍。开发人员只需调用convert_to_float8_training函数，就能将模型训练运算转换为float8格式，实现高效的训练过程。
量化技术：优化推理性能
为了进一步提升模型的推理性能，torchao提供了多种量化的解决方案，包括权重量化（Weight-Only Quantization）和动态激活量化（Dynamic Activation Quantization）。用户可以根据实际需求选择合适的量化策略，以达到最优的推理效果。
对于LLaMA 3 8B模型，使用int4权重量化方法可以将推理速度提升至原来的近两倍。同时，通过量化键值缓存，可以在保持128K上下文长度的情况下，减少73%的内存使用量。此外，torchao的autoquant功能能够自动为模型的每一层选择最适合的量化方式，减轻了开发者的手动调优负担。
稀疏性技术：减少不必要的计算开销
torchao通过引入稀疏性技术来优化模型参数的计算，减少不必要的运算消耗。例如，在ViT-H模型的推理过程中，结合稀疏性和量化技术可以实现5%的速度提升。而对于LLaMA 3.1 8B模型，在处理完整的128K上下文长度时，通过int4权重量化和int8键值缓存量化，仅需占用18.9GB的显存。
应用案例：显著提升多个模型的性能
torchao不仅显著提升了LLaMA 3 8B模型的基准测试性能，还在Diffusion模型Flux.1-Dev上实现了53.88%的推理速度提升。
未来发展：持续改进与扩展支持
PyTorch官方表示，未来将不断更新和完善torchao，计划支持精度低于int4的数据类型，开发更高吞吐量的推理引擎，并扩展对更多硬件后端的支持。
0条评论