ao
v0.6.1
Introdução | Inferência | Treinamento | Composição | Núcleos personalizados | Recursos alfa | Instalação | Integrações | Vídeos | Licença | Citação
torchao: biblioteca PyTorch para tipos de dados personalizados e otimizações. Quantize e esparsifique pesos, gradientes, otimizadores e ativações para inferência e treinamento.
Da equipe que trouxe para você a série rápida
torchao apenas funciona com torch.compile() e FSDP2 na maioria dos modelos PyTorch no Huggingface prontos para uso.
Quantizar e esparsificar seus modelos é um liner que deve funcionar em qualquer modelo com um nn.Linear incluindo seu modelo HuggingFace favorito. Você pode encontrar instruções de uso mais abrangentes aqui, esparsidade aqui e um exemplo de inferência HuggingFace aqui
Para inferência, temos a opção de
from torchao . quantization . quant_api import (
quantize_ ,
int8_dynamic_activation_int8_weight ,
int4_weight_only ,
int8_weight_only
)
quantize_ ( m , int4_weight_only ()) Para gpt-fast int4_weight_only() é a melhor opção em bs=1, pois 2x o tok/s e reduz os requisitos de VRAM em cerca de 65% em uma linha de base torch.compiled.
Se você não tiver VRAM suficiente para quantizar todo o seu modelo na GPU e achar que a quantização da CPU é muito lenta, você pode usar o argumento do dispositivo como quantize_(model, int8_weight_only(), device="cuda") que enviará e quantize cada camada individualmente para sua GPU.
Se você notar lentidão com qualquer uma dessas técnicas ou não tiver certeza de qual opção usar, considere usar o autoquant, que criará automaticamente o perfil das camadas e escolherá a melhor maneira de quantizar cada camada.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))Também fornecemos uma API voltada para o desenvolvedor para que você possa implementar seus próprios algoritmos de quantização. Use o excelente algoritmo HQQ como exemplo motivador.
Adicionamos quantização de cache kv e outros recursos para permitir inferência de longo comprimento de contexto (e necessariamente eficiente em termos de memória).
Na prática, esses recursos junto com a quantização apenas do peso int4 nos permitem reduzir o pico de memória em ~ 55% , o que significa que podemos inferir Llama3.1-8B com um comprimento de contexto de 130k com apenas 18,9 GB de memória de pico. Mais detalhes podem ser encontrados aqui
A quantização pós-treinamento pode resultar em um modelo rápido e compacto, mas também pode levar à degradação da precisão. Recomendamos explorar o Quantization Aware Training (QAT) para superar essa limitação. Em colaboração com Torchtune, desenvolvemos uma receita QAT que demonstra melhorias significativas na precisão em relação ao PTQ tradicional, recuperando 96% da degradação da precisão no hellaswag e 68% da degradação da perplexidade no wikitexto para Llama3 em comparação com a quantização pós-treinamento (PTQ) . E nós fornecemos uma receita completa aqui
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 implementa receitas de treinamento com os dtypes float8 em escala, conforme descrito em https://arxiv.org/abs/2209.05433.
Com torch.compile ativado, os resultados atuais mostram velocidades de transferência de até 1,5x em trabalhos de pré-treinamento de 128 GPU H100 LLaMa 3 70B (detalhes)
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)E para uma receita de treinamento mínimo de pré-treinamento com float8, você pode conferir torchtitan
Adicionamos suporte para esparsidade 2:4 semiestruturada com 6% de aceleração ponta a ponta no ViT-L . Blog completo aqui
A mudança de código é de 1 linha com o exemplo completo disponível aqui
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM usa 2x mais memória que os parâmetros do modelo, então podemos quantizar o estado do otimizador para 8 ou 4 bits, reduzindo efetivamente os requisitos de VRAM do otimizador em 2x ou 4x, respectivamente, em uma linha de base fp16
from torchao . prototype . low_bit_optim import AdamW8bit , AdamW4bit , AdamWFp8
optim = AdamW8bit ( model . parameters ()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versionsNa prática, somos um pouco mais lentos do que kernels escritos por especialistas, mas as implementações para esses otimizadores foram escritas em algumas centenas de linhas de código PyTorch e compiladas, portanto, use-as ou copie-as e cole-as para seus otimizadores quantizados. Referências aqui
Também temos suporte para descarregamento de CPU de GPU única, onde tanto os gradientes (mesmo tamanho dos pesos) quanto os otimizadores serão enviados com eficiência para a CPU. Só isso pode reduzir seus requisitos de VRAM em 60%
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile : Um princípio chave de design para nós é a capacidade de composição, pois qualquer novo dtype ou layout que fornecemos precisa funcionar com nosso compilador. Não deveria importar se os kernels são escritos em PyTorch, CUDA, C++ ou Triton puro - as coisas deveriam simplesmente funcionar! Portanto, escrevemos a lógica dtype, layout ou bit pack em PyTorch puro e geramos kernels eficientes para geração de código.O melhor exemplo que temos combinando a capacidade de composição do tipo de bit inferior com compilação e fsdp é o NF4, que usamos para implementar o algoritmo QLoRA. Então, se você estiver fazendo pesquisas na intersecção desta área, adoraríamos ouvir sua opinião.
Adicionamos suporte para criação e lançamento de operações personalizadas que não quebram gráficos com torch.compile() então, se você adora escrever kernels, mas odeia empacotá-los para que funcionem em todos os sistemas operacionais e versões cuda, adoraríamos aceitar contribuições para suas operações personalizadas. Temos alguns exemplos que você pode seguir
quantize_(model, fpx_weight_only(3, 2))Se você acredita que há outros kernels CUDA que deveríamos examinar mais de perto, por favor deixe um comentário sobre este assunto
Coisas que nos entusiasmam, mas que precisam de mais tempo para cozinhar no forno
quantize_(model, int8_weight_only_quantized_training()) . Este trabalho é um protótipo, pois os benchmarks de memória ainda não são convincentes. torchao faz uso liberal de vários novos recursos do Pytorch, é recomendado usá-lo com a versão noturna atual ou a versão estável mais recente do PyTorch.
Versão estável do Pypi cujo padrão será CUDA 12.1
pip install torchaoVersão estável do índice PyTorch
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124Lançamento noturno
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124Para a maioria dos desenvolvedores, você provavelmente desejará pular a construção de extensões C++/CUDA personalizadas para uma iteração mais rápida
USE_CPP=0 pip install -e . Também temos a sorte de estar integrados em algumas das principais bibliotecas de código aberto, incluindo
torchao é lançado sob a licença BSD 3.
Se você achar a biblioteca torchao útil, cite-a em seu trabalho conforme abaixo.
@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 }
