【ITBEAR】PyTorch官方近日宣布，旗下全新架构优化库torchao已正式上线。该库专注于模型的量化和稀疏性优化，致力于在保障性能的同时，降低AI模型的计算成本及RAM占用，从而提升其运行效率。
torchao提供了一整套优化工具，特别针对如LLaMA 3等热门AI模型进行性能提升。通过支持float8、int4等低精度数据类型，该库能有效减轻硬件负担，减少内存使用。
在LLaMA 3 70B模型的预训练过程中，torchao展现了显著效果。其float8训练流程能提升模型计算速度达1.5倍，开发者仅需调用convert_to_float8_training函数，便可轻松实现模型的高效训练转换。
在推理方面，torchao提供了包括权重量化和动态激活量化在内的多种量化策略，供用户灵活选择，以达到最佳的模型推理效果。
在稀疏性优化领域，torchao同样表现出色。它不仅能提升模型参数的计算效率，还能显著降低显存占用。例如，在ViT-H模型的推理中，torchao可提升速度达5%。同时，通过将LLaMA 3.1 8B模型的权重量化为int4，并将键值缓存量化为int8，torchao使得该模型在完整的128K上下文长度下，显存占用仅需18.9GB。
