到
v0.6.1
简介 |推论|培训|可组合性|定制内核 |阿尔法功能|安装|集成 |视频 |许可证 |引文
torchao：用于自定义数据类型和优化的 PyTorch 库。量化和稀疏权重、梯度、优化器和激活以进行推理和训练。
来自为您带来 Fast 系列的团队
torchao 可以在 Huggingface 上的大多数 PyTorch 模型上与torch.compile()和FSDP2一起使用，开箱即用。
量化和稀疏模型是一个 1 衬垫，应该适用于任何具有nn.Linear的模型，包括您最喜欢的 HuggingFace 模型。您可以在此处找到更全面的使用说明，在此处找到稀疏性，在此处找到 HuggingFace 推理示例
为了推理，我们可以选择
from torchao . quantization . quant_api import (
quantize_ ,
int8_dynamic_activation_int8_weight ,
int4_weight_only ,
int8_weight_only
)
quantize_ ( m , int4_weight_only ())对于 gpt-fast， int4_weight_only()是 bs=1 时的最佳选择，因为它是 tok/s 的 2 倍，并且比 torch.compiled 基线减少了约 65% 的 VRAM 要求。
如果您没有足够的 VRAM 来量化 GPU 上的整个模型，并且发现 CPU 量化太慢，那么您可以使用设备参数，如quantize_(model, int8_weight_only(), device="cuda")它将发送并将每一层单独量化到您的 GPU。
如果您发现这些技术中的任何一种都变慢，或者您不确定使用哪个选项，请考虑使用自动量化，它将自动分析图层并选择量化每个图层的最佳方法。
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))我们还提供面向开发人员的 API，以便您可以实现自己的量化算法，因此请使用优秀的 HQQ 算法作为激励示例。
我们添加了 kv 缓存量化和其他功能，以实现长上下文长度（并且必然具有内存效率）推理。
在实践中，这些功能与 int4 仅权重量化一起使我们能够将峰值内存减少约 55% ，这意味着我们可以仅使用 18.9 GB 峰值内存就 130k 上下文长度进行 Llama3.1-8B 推理。更多详细信息可以在这里找到
训练后量化可以产生快速且紧凑的模型，但也可能导致准确性下降。我们建议探索量化感知训练（QAT）来克服这一限制。我们与 Torchtune 合作开发了一种 QAT 配方，与训练后量化 (PTQ) 相比，它的准确性比传统 PTQ 有了显着提高，Llama3在 hellaswag 上恢复了 96% 的准确性下降，在 wikitext 上恢复了 68% 的困惑度下降。我们在这里提供了完整的食谱
from torchao . quantization . qat import Int8DynActInt4WeightQATQuantizer
qat_quantizer = Int8DynActInt4WeightQATQuantizer ()
# Insert "fake quantize" operations into linear layers.
# These operations simulate quantization numerics
model = qat_quantizer . prepare ( model )
# Run Training...
# Convert fake quantize to actual quantize operations
model = qat_quantizer . convert ( model )torchao.float8 使用缩放的 float8 dtypes 实现训练配方，如 https://arxiv.org/abs/2209.05433 中所述。
打开torch.compile后，当前结果显示128 H100 GPU LLaMa 3 70B 预训练作业的吞吐量加速高达 1.5 倍（详情）
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)对于使用 float8 进行预训练的端到端最小训练方法，您可以查看 torchtitan
我们添加了对半结构化 2:4 稀疏性的支持，在 ViT-L 上实现了 6% 的端到端加速。完整博客在这里
代码更改为 1 行，完整示例可在此处获取
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM 占用的内存是模型参数的 2 倍，因此我们可以将优化器状态量化为 8 或 4 位，从而在 fp16 基线上分别有效地将优化器 VRAM 需求减少 2 倍或 4 倍
from torchao . prototype . low_bit_optim import AdamW8bit , AdamW4bit , AdamWFp8
optim = AdamW8bit ( model . parameters ()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versions实际上，我们比专业编写的内核慢一点，但这些优化器的实现是用几百行 PyTorch 代码编写并编译的，因此请使用它们或将它们复制粘贴到您的量化优化器中。基准在这里
我们还支持单 GPU CPU 卸载，其中梯度（与权重大小相同）和优化器都将有效发送到 CPU。仅此一项就可以将您的 VRAM 需求降低 60%
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile ：我们的一个关键设计原则是可组合性，因为我们提供的任何新数据类型或布局都需要与编译器一起使用。内核是否是用纯 PyTorch、CUDA、C++ 或 Triton 编写的并不重要 - 一切都应该正常工作！因此，我们在纯 PyTorch 中编写数据类型、布局或位打包逻辑，并通过代码生成高效的内核。我们将低位 dtype 的可组合性与编译和 fsdp 相结合的最好例子是 NF4，我们用它来实现 QLoRA 算法。因此，如果您正在该领域的交叉领域进行研究，我们很乐意听取您的意见。
我们添加了对编写和发布自定义操作的支持，这些操作不会使用torch.compile()进行图表中断，因此，如果您喜欢编写内核但讨厌打包它们，以便它们适用于所有操作系统和 cuda 版本，我们很乐意接受以下贡献您的自定义操作。我们有一些您可以遵循的示例
quantize_(model, fpx_weight_only(3, 2))如果您认为还有其他 CUDA 内核，我们应该仔细研究一下，请对此问题发表评论
我们很兴奋但需要更多时间在烤箱里烹饪的东西
quantize_(model, int8_weight_only_quantized_training())一起使用。这项工作只是原型，因为内存基准测试尚未引人注目。torchao充分利用了 Pytorch 中的几个新功能，建议与当前的 Nightly 或最新稳定版本的 PyTorch 一起使用。
Pypi 的稳定版本，默认为 CUDA 12.1
pip install torchaoPyTorch 索引的稳定版本
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124每晚发布
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124对于大多数开发人员来说，您可能希望跳过构建自定义 C++/CUDA 扩展以加快迭代速度
USE_CPP=0 pip install -e . 我们也很幸运能够集成到一些领先的开源库中，包括
torchao是在 BSD 3 许可证下发布的。
如果您发现 torchao 库有用，请在您的工作中引用它，如下所示。
@software { torchao ,
title = { torchao: PyTorch native quantization and sparsity for training and inference } ,
author = { torchao maintainers and contributors } ,
url = { https//github.com/pytorch/torchao } ,
license = { BSD-3-Clause } ,
month = oct,
year = { 2024 }
