に
v0.6.1
はじめに |推論 |トレーニング |構成可能性 |カスタム カーネル |アルファ版の機能 |インストール |統合 |動画 |ライセンス |引用
torchao: カスタム データ型と最適化のための PyTorch ライブラリ。推論とトレーニングのために重み、勾配、オプティマイザー、アクティベーションを量子化およびスパース化します。
Fast シリーズを提供したチームから
torchao は、Huggingface 上のほとんどの PyTorch モデル上で、すぐにtorch.compile()とFSDP2で動作します。
モデルの量子化とスパース化は、お気に入りの HuggingFace モデルを含む、 nn.Linearを持つあらゆるモデルで機能する 1 つのライナーです。より包括的な使用手順はこちら、スパース性はこちら、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 の場合、bs=1 ではint4_weight_only()が最適なオプションです。tok /s が 2 倍になり、torch.compiled ベースラインよりも VRAM 要件が約 65% 削減されます。
モデル全体を GPU で量子化するのに十分な VRAM がなく、CPU の量子化が遅すぎる場合は、デバイス引数を so quantize_(model, int8_weight_only(), device="cuda")のように使用できます。各レイヤーを GPU に個別に量子化します。
これらの手法のいずれかで速度の低下が見られる場合、またはどのオプションを使用するかわからない場合は、レイヤーを自動的にプロファイリングし、各レイヤーを量子化する最適な方法を選択する autoquant の使用を検討してください。
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))また、独自の量子化アルゴリズムを実装できる開発者向け API も提供しています。優れた HQQ アルゴリズムを動機付けの例として使用してください。
長いコンテキスト長の (必然的にメモリ効率の高い) 推論を可能にするために、kv キャッシュ量子化とその他の機能を追加しました。
実際には、これらの機能と int4 重みのみの量子化を併用することで、ピーク メモリを最大 55% 削減できます。これは、わずか 18.9 GB のピーク メモリで 130k コンテキスト長の Llama3.1-8B 推論ができることを意味します。詳細はこちらからご覧いただけます
トレーニング後の量子化により、高速でコンパクトなモデルが得られますが、精度の低下につながる可能性もあります。この制限を克服するには、量子化対応トレーニング (QAT) を検討することをお勧めします。トーチチューンと協力して、従来の PTQ に比べて大幅な精度の向上を示す QAT レシピを開発しました。トレーニング後の量子化 (PTQ) と比較して、Llama3 のヘラスワグでの精度の低下の 96%、ウィキテキストでのパープレキシティの低下の 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 は、https://arxiv.org/abs/2209.05433 に記載されているように、スケーリングされた float8 dtype を使用したトレーニング レシピを実装します。
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 をチェックしてください。
ViT-L でのエンドツーエンドの 6% の高速化を実現する半構造化 2:4 スパース性のサポートを追加しました。ブログ全文はこちら
コード変更は 1 ライナーであり、完全な例はここから入手できます。
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM はモデル パラメータの 2 倍のメモリを必要とするため、オプティマイザの状態を 8 ビットまたは 4 ビットに量子化し、オプティマイザの VRAM 要件を fp16 ベースラインと比べてそれぞれ 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 : 私たちが提供する新しい dtype やレイアウトはコンパイラで動作する必要があるため、私たちにとって重要な設計原則はコンポーザビリティです。カーネルが純粋な PyTorch、CUDA、C++、または Triton で書かれているかどうかは問題ではありません。問題なく動作するはずです。そこで、純粋な PyTorch で dtype、レイアウト、またはビット パッキング ロジックを記述し、効率的なカーネルをコード生成します。下位ビット dtype のコンポーザビリティとコンパイルおよび fsdp を組み合わせた最良の例は、QLoRA アルゴリズムの実装に使用した NF4 です。したがって、この地域の交差点で調査を行っている場合は、ぜひご連絡ください。
torch.compile()を使用してグラフ ブレークを行わないカスタム操作のオーサリングとリリースのサポートを追加しました。そのため、カーネルを書くのは好きだが、すべてのオペレーティング システムと cuda バージョンで動作するようにカーネルをパッケージ化するのが嫌いな場合は、次のようなコントリビュートを喜んで受け入れます。あなたのカスタムオペレーション。参考にしていただける例をいくつかご紹介します
quantize_(model, fpx_weight_only(3, 2))を備え、fp16 よりも 2 倍高速な推論を実現します。もっと詳しく調べるべき他の CUDA カーネルがあると思われる場合は、この問題についてコメントを残してください。
楽しみだけどオーブンで調理するのにもっと時間が必要なもの
quantize_(model, int8_weight_only_quantized_training())で簡単に使用できます。メモリ ベンチマークがまだ説得力を持っていないため、この作業はプロトタイプです。torchao Pytorch のいくつかの新機能を最大限に活用しているため、PyTorch の現在の Nightly バージョンまたは最新の安定したバージョンで使用することをお勧めします。
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 }
