到
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 }
