к
v0.6.1
Введение | Вывод | Обучение | Компонуемость | Пользовательские ядра | Альфа-функции | Установка | Интеграции | Видео | Лицензия | Цитирование
torchao: библиотека PyTorch для пользовательских типов данных и оптимизации. Квантуйте и разрежайте веса, градиенты, оптимизаторы и активации для вывода и обучения.
От команды, подарившей вам быструю серию
torchao просто работает с torch.compile() и FSDP2 с большинством моделей PyTorch на Huggingface из коробки.
Квантование и разрежение ваших моделей — это 1-лайнер, который должен работать с любой моделью с nn.Linear , включая вашу любимую модель HuggingFace. Более подробные инструкции по использованию можно найти здесь, разреженность здесь и пример вывода 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 int4_weight_only() является лучшим вариантом при bs=1, поскольку он в 2 раза увеличивает ток/с и снижает требования к видеопамяти примерно на 65 % по сравнению с базовым показателем, скомпилированным torch.
Если у вас недостаточно видеопамяти для квантования всей вашей модели на графическом процессоре, и вы обнаружите, что квантование процессора слишком медленное, вы можете использовать аргумент устройства, например, quantize_(model, int8_weight_only(), device="cuda") который отправит и квантовать каждый слой индивидуально на вашем графическом процессоре.
Если вы заметили замедление при использовании любого из этих методов или не уверены, какой вариант использовать, рассмотрите возможность использования автокванта, который автоматически профилирует слои и выбирает лучший способ квантования каждого слоя.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))Мы также предоставляем API, ориентированный на разработчиков, чтобы вы могли реализовать свои собственные алгоритмы квантования, поэтому используйте отличный алгоритм HQQ в качестве мотивирующего примера.
Мы добавили квантование кэша kv и другие функции, чтобы обеспечить вывод длинного контекста (и обязательно эффективного использования памяти).
На практике эти функции наряду с квантованием только веса int4 позволяют нам уменьшить пиковую память на ~ 55% , что означает, что мы можем делать выводы Llama3.1-8B с длиной контекста 130 КБ, имея только 18,9 ГБ пиковой памяти. Более подробную информацию можно найти здесь
Квантование после обучения может привести к созданию быстрой и компактной модели, но также может привести к снижению точности. Мы рекомендуем изучить обучение с учетом квантования (QAT), чтобы преодолеть это ограничение. В сотрудничестве с Torchtune мы разработали рецепт QAT, который демонстрирует значительные улучшения точности по сравнению с традиционным PTQ, восстанавливая 96% снижения точности в hellaswag и 68% снижения недоумения в викитексте для Llama3 по сравнению с квантованием после обучения (PTQ). . А полный рецепт мы предоставили здесь.
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 реализует рецепты обучения с масштабированными типами float8, как описано в https://arxiv.org/abs/2209.05433.
При включенном torch.compile текущие результаты показывают увеличение пропускной способности до 1,5 раз на 128 заданиях предварительной подготовки графического процессора H100 LLaMa 3 70B (подробности)
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)А рецепт предварительной тренировки от конца до минимального с помощью float8 можно найти на torchtitan.
Мы добавили поддержку полуструктурированной разреженности 2:4 с 6% сквозным ускорением на ViT-L . Полный блог здесь
Изменение кода занимает одну строку, полный пример доступен здесь.
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM занимает в 2 раза больше памяти, чем параметры модели, поэтому мы можем квантовать состояние оптимизатора до 8 или 4 бит, эффективно снижая требования к VRAM оптимизатору в 2 или 4 раза соответственно по сравнению с базовым уровнем 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На практике мы немного медленнее, чем профессионально написанные ядра, но реализации этих оптимизаторов были написаны в несколько сотен строк кода PyTorch и скомпилированы, поэтому, пожалуйста, используйте их или скопируйте и вставьте для своих квантованных оптимизаторов. Тесты здесь
У нас также есть поддержка разгрузки ЦП одного графического процессора, при которой как градиенты (того же размера, что и веса), так и оптимизаторы будут эффективно отправляться в ЦП. Одно это может снизить требования к VRAM на 60 %.
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile : Ключевым принципом проектирования для нас является возможность компоновки, поскольку любой новый dtype или макет, который мы предоставляем, должен работать с нашим компилятором. Не имеет значения, написаны ли ядра на чистом PyTorch, CUDA, C++ или Triton — все должно работать! Поэтому мы пишем dtype, макет или логику упаковки битов в чистом PyTorch и эффективно генерируем код ядра.Лучшим примером сочетания возможности компоновки младшего бита dtype с компиляцией и fsdp является NF4, который мы использовали для реализации алгоритма QLoRA. Так что, если вы проводите исследования на стыке этой области, мы будем рады услышать ваше мнение.
Мы добавили поддержку создания и выпуска пользовательских операций, которые не разбивают графики с помощью torch.compile() поэтому, если вы любите писать ядра, но ненавидите их упаковывать, чтобы они работали во всех операционных системах и версиях cuda, мы будем рады принять вклады для ваши индивидуальные операции. У нас есть несколько примеров, которым вы можете следовать
quantize_(model, fpx_weight_only(3, 2))Если вы считаете, что существуют другие ядра CUDA, к которым нам следует присмотреться, оставьте комментарий по этому вопросу.
Вещи, которые нам очень нравятся, но нужно больше времени, чтобы приготовить их в духовке
quantize_(model, int8_weight_only_quantized_training()) . Эта работа является прототипом, поскольку тесты памяти еще не убедительны. torchao широко использует несколько новых функций Pytorch, рекомендуется использовать его с текущей ночной или последней стабильной версией PyTorch.
Стабильная версия Pypi, которая по умолчанию будет CUDA 12.1.
pip install torchaoСтабильный выпуск из индекса PyTorch
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 }
