Pendahuluan | Inferensi | Pelatihan | Komposabilitas | Kernel Khusus | Fitur Alfa | Instalasi | Integrasi | Video | Lisensi | Kutipan
torchao: Pustaka PyTorch untuk tipe data & pengoptimalan khusus. Mengukur dan memperkecil bobot, gradien, pengoptimal & aktivasi untuk inferensi dan pelatihan.
Dari tim yang membawakan Anda seri cepat
torchao hanya berfungsi dengan torch.compile()
dan FSDP2
pada sebagian besar model PyTorch di Huggingface.
Mengkuantisasi dan Sparsifikasi model Anda adalah 1 liner yang dapat digunakan pada model apa pun dengan nn.Linear
termasuk model HuggingFace favorit Anda. Anda dapat menemukan petunjuk penggunaan yang lebih komprehensif di sini, ketersebaran di sini, dan contoh inferensi HuggingFace di sini
Sebagai kesimpulan, kita mempunyai pilihan
from torchao . quantization . quant_api import (
quantize_ ,
int8_dynamic_activation_int8_weight ,
int4_weight_only ,
int8_weight_only
)
quantize_ ( m , int4_weight_only ())
Untuk gpt-fast int4_weight_only()
adalah opsi terbaik di bs=1 karena 2x tok/s dan mengurangi persyaratan VRAM sekitar 65% dibandingkan baseline yang dikompilasi torch.
Jika Anda tidak memiliki cukup VRAM untuk mengkuantisasi seluruh model pada GPU dan Anda merasa kuantisasi CPU terlalu lambat maka Anda dapat menggunakan argumen perangkat seperti quantize_(model, int8_weight_only(), device="cuda")
yang akan mengirimkan dan kuantisasi setiap lapisan satu per satu ke GPU Anda.
Jika Anda melihat perlambatan pada salah satu teknik ini atau Anda tidak yakin opsi mana yang harus digunakan, pertimbangkan untuk menggunakan autoquant yang secara otomatis akan membuat profil lapisan dan memilih cara terbaik untuk mengkuantisasi setiap lapisan.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))
Kami juga menyediakan API yang menghadap pengembang sehingga Anda dapat mengimplementasikan algoritme kuantisasi Anda sendiri, jadi harap gunakan algoritme HQQ yang sangat baik sebagai contoh yang memotivasi.
Kami telah menambahkan kuantisasi cache kv dan fitur lainnya untuk mengaktifkan inferensi dengan panjang konteks yang panjang (dan tentu saja hemat memori).
Dalam praktiknya, fitur-fitur ini bersama dengan kuantisasi hanya bobot int4 memungkinkan kami mengurangi memori puncak sebesar ~55% , yang berarti kami dapat melakukan inferensi Llama3.1-8B dengan panjang konteks 130k dengan memori puncak hanya 18,9 GB. Detail lebih lanjut dapat ditemukan di sini
Kuantisasi pasca pelatihan dapat menghasilkan model yang cepat dan ringkas, namun juga dapat menyebabkan penurunan akurasi. Kami merekomendasikan untuk mempelajari Quantization Aware Training (QAT) untuk mengatasi keterbatasan ini. Bekerja sama dengan Torchtune, kami telah mengembangkan resep QAT yang menunjukkan peningkatan akurasi yang signifikan dibandingkan PTQ tradisional, memulihkan 96% penurunan akurasi pada hellaswag dan 68% penurunan kebingungan pada teks wiki untuk Llama3 dibandingkan dengan kuantisasi pasca-pelatihan (PTQ) . Dan kami telah menyediakan resep lengkapnya di sini
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 mengimplementasikan resep pelatihan dengan tipe float8 berskala, seperti yang dijelaskan di https://arxiv.org/abs/2209.05433.
Dengan torch.compile
aktif, hasil saat ini menunjukkan peningkatan throughput hingga 1,5x pada pekerjaan pra-pelatihan 128 GPU H100 LLaMa 3 70B (detail)
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)
Dan untuk resep pelatihan awal hingga minimal dengan float8, Anda dapat melihat torchtitan
Kami telah menambahkan dukungan untuk ketersebaran 2:4 semi-terstruktur dengan percepatan end-to-end sebesar 6% pada ViT-L . Blog selengkapnya di sini
Perubahan kodenya adalah 1 liner dengan contoh lengkap tersedia di sini
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })
ADAM membutuhkan memori 2x lebih banyak dibandingkan parameter model sehingga kita dapat mengkuantisasi status pengoptimal menjadi 8 atau 4 bit yang secara efektif mengurangi persyaratan VRAM pengoptimal masing-masing sebesar 2x atau 4x pada garis dasar fp16
from torchao . prototype . low_bit_optim import AdamW8bit , AdamW4bit , AdamWFp8
optim = AdamW8bit ( model . parameters ()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versions
Dalam praktiknya, kami sedikit lebih lambat dibandingkan kernel yang ditulis secara ahli, tetapi implementasi untuk pengoptimal ini ditulis dalam beberapa ratus baris kode PyTorch dan dikompilasi, jadi silakan gunakan atau salin-tempel untuk pengoptimal terkuantisasi Anda. Tolok ukur di sini
Kami juga memiliki dukungan untuk pembongkaran CPU GPU tunggal di mana gradien (ukuran sama dengan bobot) dan pengoptimal akan dikirim secara efisien ke CPU. Ini saja dapat mengurangi kebutuhan VRAM Anda hingga 60%
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])
torch.compile
: Prinsip desain utama bagi kami adalah komposisi karena setiap dtype atau tata letak baru yang kami sediakan harus berfungsi dengan kompiler kami. Tidak masalah jika kernel ditulis dalam PyTorch murni, CUDA, C++, atau Triton - semuanya akan berfungsi! Jadi kami menulis logika dtype, layout, atau bit packing di PyTorch murni dan kernel efisien yang menghasilkan kode.Contoh terbaik yang kami miliki menggabungkan komposisi dtype bit rendah dengan kompilasi dan fsdp adalah NF4 yang kami gunakan untuk mengimplementasikan algoritma QLoRA. Jadi jika Anda melakukan penelitian di persimpangan area ini, kami ingin mendengar pendapat Anda.
Kami telah menambahkan dukungan untuk membuat dan merilis operasi khusus yang tidak merusak grafik dengan torch.compile()
jadi jika Anda suka menulis kernel tetapi tidak suka mengemasnya sehingga dapat berfungsi di semua sistem operasi dan versi cuda, kami dengan senang hati menerima kontribusi untuk operasi khusus Anda. Kami memiliki beberapa contoh yang dapat Anda ikuti
quantize_(model, fpx_weight_only(3, 2))
yang mudah digunakanJika Anda yakin ada kernel CUDA lain yang harus kita lihat lebih dekat, silakan tinggalkan komentar tentang masalah ini
Hal-hal yang kami sukai tetapi membutuhkan lebih banyak waktu untuk dimasak di oven
quantize_(model, int8_weight_only_quantized_training())
. Karya ini merupakan prototipe karena tolok ukur memorinya belum meyakinkan. torchao
memanfaatkan beberapa fitur baru di Pytorch secara bebas, disarankan untuk menggunakannya dengan PyTorch versi stabil nightly atau terbaru.
Rilis stabil dari Pypi yang akan default ke CUDA 12.1
pip install torchao
Rilis Stabil dari indeks PyTorch
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124
Rilis Malam Hari
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124
Bagi sebagian besar pengembang, Anda mungkin ingin melewatkan pembuatan ekstensi C++/CUDA khusus untuk iterasi yang lebih cepat
USE_CPP=0 pip install -e .
Kami juga beruntung dapat diintegrasikan ke dalam beberapa perpustakaan sumber terbuka terkemuka termasuk
torchao
dirilis di bawah lisensi BSD 3.
Jika menurut Anda perpustakaan torchao bermanfaat, silakan kutip dalam pekerjaan Anda seperti di bawah ini.
@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 }