ao: 一个用于从训练到推理的模型量化与稀疏化的 PyTorch 原生架构优化库

它解决了什么问题

TorchAO 提供了一个原生的 PyTorch 库，用于优化 AI 模型，使其运行更快且更节省内存。它解决了模型大小与准确度之间的常见权衡问题，允许用户在不显著损失质量的情况下，减少大型模型（如 LLMs 和 diffusion models）的内存占用，并加速训练和推理过程。

工作原理

TorchAO 实现了多种架构优化技术：

• Quantization: 它将模型权重和激活值转换为低精度格式（例如 int4, int8, 和 float8），从而减少内存使用并提高吞吐量。
• Quantization-Aware Training (QAT): 为了防止量化过程中的精度损失，它允许模型在训练过程中适应低精度。
• Sparsity: 它使用半结构化 2:4 sparsity 来移除冗余权重，从而进一步提升速度。
• uma-native integration: 它与 torch.compile() 和 FSDP2 无缝集成，可在各种硬件（CUDA, XPU, CPU, 和 ARM）上实现高性能执行。

适用人群

该库专为 AI 研究人员和工程师设计，他们需要在有限的硬件上部署大规模模型、加速海量模型的预训练，或者通过 ExecuTorch 为边缘设备优化模型。

亮点

• Training Speedups: 使用 float8 训练，预训练 Llama-3.1-70B 的速度最高可提升 1.5 倍。
• Inference Gains: 将 Llama-3-8B 量化为 int4，推理速度可提升 1.89 倍，内存占用减少 58%。
• Broad Integration: 内置支持 Hugging Face Transformers, Diffusers, vLLM, 和 SGLang。
• Memory Efficiency: 包括量化优化器（AdamW 4/8-bit）和 CPU offloading，可将 VRAM 需求降低高达 60%。