ao

torchao: PyTorch 架构优化

简介

torchao 是一个专为 PyTorch 设计的库，主要用于自定义数据类型和优化。它支持推理和训练过程中的权重、梯度、优化器和激活函数的量化和稀疏化。这个库由开发了快速系列的团队打造，其代表性成果包括：

• 使用 sam-fast 系列为图像分割模型提速 9.5 倍
• 使用 gpt-fast 系列为语言模型提速 10 倍
• 使用 sd-fast 系列为扩散模型提速 3 倍

torchao 可以与 torch.compile() 和 FSDP2 一起使用，支持大多数 Huggingface 上的 PyTorch 模型。

推理

后训练量化

torchao 提供了一键量化和稀疏化模型的功能，能够应用于含有 nn.Linear 的任何模型，甚至包括您喜爱的 HuggingFace 模型。支持以下选择：

1. 只量化权重：适用于内存受限模型
2. 量化权重和激活函数：适用于计算受限模型
3. 量化激活和权重并稀疏化权重

如果 GPU 内存不足以在 GPU 上量化整个模型，又无法忍受 CPU 量化的缓慢，可以使用 device 参数，例如：quantize_(model, int8_weight_only(), device="cuda")，这会逐层将模型发送到 GPU 进行量化。

KV 缓存量化

通过 kv 缓存量化和其他特性，实现长上下文长度推理（即内存高效），例如可以用仅 18.9 GB 的峰值内存在 130k 的上下文长度下对 Llama3.1-8B 进行推理。

量化感知训练（QAT）

QAT 能够克服后训练量化可能带来的精度损失。例如，我们开发的 QAT 配方使得在 hellaswag 上恢复了 96% 的精度损失。

训练

Float8

torchao 提供了 Float8 数据类型的训练配方，可以在 LLaMa 3 70B 的训练任务上提高 1.5 倍的吞吐量。

稀疏训练

支持半结构化的 2:4 稀疏性，能够在 ViT-L 模型上提升 6% 的性能。

内存高效优化器

torchao 支持 ADAM 优化器状态的量化，从而有效地减少优化器对 VRAM 的需求。更为紧凑的优化器实现提高了内存利用率。

组合能力

torch.compile 支持不同数据类型和布局的组合，这使得不同内核（即便是由 PyTorch、CUDA、C++ 或 Triton 编写的）都能很好地协作运行。

自定义内核

torchao 支持编写和发布自定义操作，以便与 torch.compile() 无缝结合。这使得开发人员可以编写更快速、更有效率的内核。

Alpha 特性

torchao 还在开发和测试中的一些新功能包括：

1. MX 训练和推理支持
2. 完整的 Int8 量化训练
3. IntX 算法，支持所有整数类型
4. Bitnet，用于更好的硬件和内核支持

安装

推荐使用 PyTorch 的当前夜间版或最新稳定版本来确保 torchao 的最佳运行。可以通过 pip 安装不同版本的 torchao，例如：

• 使用 CUDA 12.1 的稳定版本：pip install torchao
• 使用夜间版本：pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121

开源整合

torchao 与多种领先的开源库进行了集成，如 Hugging Face transformers 和 diffusers 等，提高了模型推理和训练的效率。

此项目的详细文档、代码示例和更多信息可以通过相关的 PyTorch 和 torchao 账户获取。希望此工具能够为研究人员和开发者带来更加高效的工作体验。