访问官网
Github
论文torchao: PyTorch 架构优化
简介
torchao 是一个专为 PyTorch 设计的库，主要用于自定义数据类型和优化。它支持推理和训练过程中的权重、梯度、优化器和激活函数的量化和稀疏化。这个库由开发了快速系列的团队打造，其代表性成果包括：
- 使用 sam-fast 系列为图像分割模型提速 9.5 倍
- 使用 gpt-fast 系列为语言模型提速 10 倍
- 使用 sd-fast 系列为扩散模型提速 3 倍
torchao 可以与 torch.compile() 和 FSDP2 一起使用，支持大多数 Huggingface 上的 PyTorch 模型。
推理
后训练量化
torchao 提供了一键量化和稀疏化模型的功能，能够应用于含有 nn.Linear 的任何模型，甚至包括您喜爱的 HuggingFace 模型。支持以下选择：
- 只量化权重：适用于内存受限模型
- 量化权重和激活函数：适用于计算受限模型
- 量化激活和权重并稀疏化权重
如果 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 还在开发和测试中的一些新功能包括：
- MX 训练和推理支持
- 完整的 Int8 量化训练
- IntX 算法，支持所有整数类型
- 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 账户获取。希望此工具能够为研究人员和开发者带来更加高效的工作体验。
