hacia
v0.6.1
Introducción | Inferencia | Formación | Componibilidad | Núcleos personalizados | Funciones alfa | Instalación | Integraciones | Vídeos | Licencia | Citación
torchao: biblioteca PyTorch para optimizaciones y tipos de datos personalizados. Cuantifique y distribuya pesos, gradientes, optimizadores y activaciones para inferencia y entrenamiento.
Del equipo que te trajo la serie rápida.
torchao simplemente funciona con torch.compile() y FSDP2 en la mayoría de los modelos de PyTorch en Huggingface listos para usar.
Cuantificar y dispersar sus modelos es una línea 1 que debería funcionar en cualquier modelo con un nn.Linear incluido su modelo HuggingFace favorito. Puede encontrar instrucciones de uso más completas aquí, escasez aquí y un ejemplo de inferencia de HuggingFace aquí
Por inferencia, tenemos la opción 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() es la mejor opción en bs=1, ya que duplica el tok/s y reduce los requisitos de VRAM en aproximadamente un 65 % con respecto a una línea de base compilada con torch.
Si no tiene suficiente VRAM para cuantificar todo su modelo en la GPU y considera que la cuantificación de la CPU es demasiado lenta, puede usar el argumento del dispositivo como quantize_(model, int8_weight_only(), device="cuda") que enviará y cuantice cada capa individualmente en su GPU.
Si ve desaceleraciones con cualquiera de estas técnicas o no está seguro de qué opción usar, considere usar autoquant, que perfilará automáticamente las capas y elegirá la mejor manera de cuantizar cada capa.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))También proporcionamos una API orientada al desarrollador para que pueda implementar sus propios algoritmos de cuantificación, así que utilice el excelente algoritmo HQQ como ejemplo motivador.
Hemos agregado cuantificación de caché kv y otras características para permitir una inferencia de contexto largo (y necesariamente eficiente en memoria).
En la práctica, estas características junto con la cuantificación de solo peso int4 nos permiten reducir la memoria máxima en aproximadamente un 55 % , lo que significa que podemos realizar inferencias Llama3.1-8B con una longitud de contexto de 130 k y solo 18,9 GB de memoria máxima. Más detalles se pueden encontrar aquí.
La cuantificación posterior al entrenamiento puede dar como resultado un modelo rápido y compacto, pero también puede provocar una degradación de la precisión. Recomendamos explorar Quantization Aware Training (QAT) para superar esta limitación. En colaboración con Torchtune, hemos desarrollado una receta QAT que demuestra mejoras significativas en la precisión con respecto al PTQ tradicional, recuperando el 96 % de la degradación de la precisión en hellaswag y el 68 % de la degradación de la perplejidad en wikitext para Llama3 en comparación con la cuantificación posterior al entrenamiento (PTQ). . Y hemos proporcionado una receta completa aquí.
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 recetas de entrenamiento con los tipos de float8 escalados, como se establece en https://arxiv.org/abs/2209.05433.
Con torch.compile activado, los resultados actuales muestran velocidades de rendimiento de hasta 1,5 veces en 128 trabajos de preentrenamiento de GPU H100 LLaMa 3 70B (detalles)
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)Y para obtener una receta de entrenamiento previo al mínimo con float8, puedes consultar torchtitan
Agregamos soporte para escasez semiestructurada 2:4 con aceleraciones del 6 % de extremo a extremo en ViT-L . Blog completo aquí
El cambio de código es de 1 línea y el ejemplo completo está disponible aquí.
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM requiere el doble de memoria que los parámetros del modelo, por lo que podemos cuantificar el estado del optimizador a 8 o 4 bits, reduciendo efectivamente los requisitos de VRAM del optimizador en 2 o 4 veces respectivamente sobre una línea 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 versionsEn la práctica, somos un poquito más lentos que los kernels escritos por expertos, pero las implementaciones de estos optimizadores se escribieron en unos cientos de líneas de código PyTorch y se compilaron, así que úselas o cópielas y péguelas para sus optimizadores cuantificados. Puntos de referencia aquí
También tenemos soporte para la descarga de CPU de una sola GPU, donde tanto los gradientes (mismo tamaño que los pesos) como los optimizadores se enviarán de manera eficiente a la CPU. Esto por sí solo puede reducir sus requisitos de VRAM en un 60%
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile : Un principio de diseño clave para nosotros es la componibilidad, ya que cualquier tipo de diseño o diseño nuevo que proporcionemos debe funcionar con nuestro compilador. No debería importar si los núcleos están escritos en PyTorch, CUDA, C++ o Triton puro: ¡todo debería funcionar! Por lo tanto, escribimos la lógica de tipo d, diseño o empaquetado de bits en PyTorch puro y generamos código con núcleos eficientes.El mejor ejemplo que tenemos combinando la componibilidad de tipo d de bits inferiores con compilación y fsdp es NF4, que utilizamos para implementar el algoritmo QLoRA. Entonces, si está investigando en la intersección de esta área, nos encantaría saber de usted.
Hemos agregado soporte para crear y lanzar operaciones personalizadas que no se rompen en gráficos con torch.compile() por lo que si te encanta escribir kernels pero odias empaquetarlos para que funcionen en todos los sistemas operativos y versiones de cuda, nos encantaría aceptar contribuciones para tus operaciones personalizadas. Tenemos algunos ejemplos que puedes seguir.
quantize_(model, fpx_weight_only(3, 2))Si cree que hay otros núcleos CUDA que deberíamos analizar más de cerca, deje un comentario sobre este tema.
Cosas que nos entusiasman pero que necesitamos más tiempo para cocinar en el horno
quantize_(model, int8_weight_only_quantized_training()) . Este trabajo es un prototipo ya que los puntos de referencia de la memoria aún no son convincentes. torchao hace un uso liberal de varias funciones nuevas en Pytorch; se recomienda usarlo con la versión nocturna actual o la última versión estable de PyTorch.
Versión estable de Pypi que por defecto será CUDA 12.1
pip install torchaoLanzamiento estable del índice PyTorch
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124Lanzamiento nocturno
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124Para la mayoría de los desarrolladores, probablemente desee omitir la creación de extensiones C++/CUDA personalizadas para una iteración más rápida.
USE_CPP=0 pip install -e . También tenemos la suerte de estar integrados en algunas de las principales bibliotecas de código abierto, incluidas
torchao se lanza bajo la licencia BSD 3.
Si encuentra útil la biblioteca de torchao, cítela en su trabajo como se muestra a continuación.
@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 }
