au
v0.6.1
Présentation | Inférence | Formation | Composabilité | Noyaux personnalisés | Fonctionnalités alpha | Installation | Intégrations | Vidéos | Licence | Citation
torchao : bibliothèque PyTorch pour les types de données personnalisés et les optimisations. Quantifiez et rationalisez les poids, les gradients, les optimiseurs et les activations pour l'inférence et la formation.
De l'équipe qui vous a apporté la série rapide
torchao fonctionne simplement avec torch.compile() et FSDP2 sur la plupart des modèles PyTorch sur Huggingface, prêts à l'emploi.
La quantification et la sparsification de vos modèles sont une ligne unique qui devrait fonctionner sur n'importe quel modèle avec un nn.Linear y compris votre modèle HuggingFace préféré. Vous pouvez trouver des instructions d'utilisation plus complètes ici, la parcimonie ici et un exemple d'inférence HuggingFace ici
Pour déduire, nous avons la possibilité de
from torchao . quantization . quant_api import (
quantize_ ,
int8_dynamic_activation_int8_weight ,
int4_weight_only ,
int8_weight_only
)
quantize_ ( m , int4_weight_only ()) Pour gpt-fast int4_weight_only() est la meilleure option à bs=1 car elle 2x le tok/s et réduit les exigences de VRAM d'environ 65 % par rapport à une ligne de base torch.compilée.
Si vous n'avez pas assez de VRAM pour quantifier l'intégralité de votre modèle sur GPU et que vous trouvez que la quantification du CPU est trop lente, vous pouvez utiliser l'argument de périphérique comme suit quantize_(model, int8_weight_only(), device="cuda") qui enverra et quantifiez chaque couche individuellement sur votre GPU.
Si vous constatez des ralentissements avec l'une de ces techniques ou si vous ne savez pas quelle option utiliser, envisagez d'utiliser la quantification automatique qui profilera automatiquement les couches et choisira la meilleure façon de quantifier chaque couche.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))Nous fournissons également une API destinée aux développeurs afin que vous puissiez implémenter vos propres algorithmes de quantification. Veuillez donc utiliser l'excellent algorithme HQQ comme exemple motivant.
Nous avons ajouté la quantification du cache kv et d'autres fonctionnalités afin de permettre une inférence de longue durée de contexte (et nécessairement efficace en termes de mémoire).
En pratique, ces fonctionnalités ainsi que la quantification par poids int4 uniquement nous permettent de réduire la mémoire maximale d'environ 55 % , ce qui signifie que nous pouvons inférer Llama3.1-8B avec une longueur de contexte de 130 000 avec seulement 18,9 Go de mémoire maximale. Plus de détails peuvent être trouvés ici
La quantification post-entraînement peut aboutir à un modèle rapide et compact, mais peut également conduire à une dégradation de la précision. Nous vous recommandons d’explorer la Quantization Aware Training (QAT) pour surmonter cette limitation. En collaboration avec Torchtune, nous avons développé une recette QAT qui démontre des améliorations significatives de la précision par rapport au PTQ traditionnel, récupérant 96 % de la dégradation de la précision sur hellaswag et 68 % de la dégradation de la perplexité sur le wikitext pour Llama3 par rapport à la quantification post-entraînement (PTQ). . Et nous avons fourni une recette complète ici
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 implémente des recettes de formation avec les types float8 mis à l'échelle, comme indiqué dans https://arxiv.org/abs/2209.05433.
Avec torch.compile activé, les résultats actuels montrent des accélérations de débit allant jusqu'à 1,5x sur 128 tâches de pré-entraînement H100 GPU LLaMa 3 70B (détails)
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)Et pour une recette d'entraînement minimale de pré-entraînement avec float8, vous pouvez consulter torchtitan
Nous avons ajouté la prise en charge de la parcimonie 2:4 semi-structurée avec des accélérations de bout en bout de 6 % sur ViT-L . Blog complet ici
Le changement de code est une ligne avec l'exemple complet disponible ici
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM prend 2 fois plus de mémoire que les paramètres du modèle afin que nous puissions quantifier l'état de l'optimiseur sur 8 ou 4 bits, réduisant efficacement les exigences de VRAM de l'optimiseur de 2x ou 4x respectivement sur une base de référence fp16.
from torchao . prototype . low_bit_optim import AdamW8bit , AdamW4bit , AdamWFp8
optim = AdamW8bit ( model . parameters ()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versionsEn pratique, nous sommes un peu plus lents que les noyaux écrits par des experts, mais les implémentations de ces optimiseurs ont été écrites en quelques centaines de lignes de code PyTorch et compilées, veuillez donc les utiliser ou les copier-coller pour vos optimiseurs quantifiés. Des repères ici
Nous prenons également en charge le déchargement du processeur sur un seul GPU, où les dégradés (de la même taille que les poids) et les optimiseurs seront efficacement envoyés au processeur. Cela seul peut réduire vos besoins en VRAM de 60 %
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile : Un principe de conception clé pour nous est la composabilité, car tout nouveau type ou mise en page que nous fournissons doit fonctionner avec notre compilateur. Cela ne devrait pas avoir d'importance si les noyaux sont écrits en PyTorch pur, CUDA, C++ ou Triton - les choses devraient fonctionner ! Nous écrivons donc la logique de type, de disposition ou de compression de bits dans des noyaux PyTorch purs et efficaces pour générer du code.Le meilleur exemple que nous ayons combinant la composabilité du type de bit inférieur avec compile et fsdp est NF4 que nous avons utilisé pour implémenter l'algorithme QLoRA. Donc, si vous effectuez des recherches à l’intersection de ce domaine, nous serions ravis d’avoir votre avis.
Nous avons ajouté la prise en charge de la création et de la publication d'opérations personnalisées qui ne rompent pas le graphique avec torch.compile() donc si vous aimez écrire des noyaux mais détestez les empaqueter pour qu'ils fonctionnent avec tous les systèmes d'exploitation et versions de cuda, nous serions ravis d'accepter des contributions pour vos opérations personnalisées. Nous avons quelques exemples que vous pouvez suivre
quantize_(model, fpx_weight_only(3, 2))Si vous pensez qu'il existe d'autres noyaux CUDA que nous devrions examiner de plus près, veuillez laisser un commentaire sur ce problème.
Des choses qui nous passionnent mais qui ont besoin de plus de temps pour cuisiner au four
quantize_(model, int8_weight_only_quantized_training()) . Ce travail est un prototype car les tests de mémoire ne sont pas encore convaincants. torchao utilise généreusement plusieurs nouvelles fonctionnalités de Pytorch, il est recommandé de l'utiliser avec la version nocturne actuelle ou la dernière version stable de PyTorch.
Version stable de Pypi qui sera par défaut CUDA 12.1
pip install torchaoVersion stable de l'index PyTorch
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124Sortie nocturne
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124Pour la plupart des développeurs, vous souhaiterez probablement ignorer la création d'extensions C++/CUDA personnalisées pour une itération plus rapide.
USE_CPP=0 pip install -e . Nous avons également la chance d'être intégrés à certaines des principales bibliothèques open source, notamment
torchao est publié sous la licence BSD 3.
Si vous trouvez la bibliothèque torchao utile, veuillez la citer dans votre travail comme ci-dessous.
@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 }
