КПД

Flash Attention 3 вышел! [Статья][Блог] Не прошло и года с выхода Flash Attention 2, как вышло продолжение, доведя серию до трилогии. Метод Основная идея первого Flash Attention в уменьшении передачи памяти между HBM и кэшом GPU, а второго - в отпимизации не matmul операций. Flash Attention 2 позволяет довольно эффективно использовать вычислительные возможности A100, но для более современной H100 все еще имеет место сильная недоутилизация - всего 35%. В Flash Attention 3 отпимизируют Attention с учетом новых архитектурных особенностей Hopper+ - новых тензорных ядер с большим throughput, тензорным ускорителем памяти (Tensor Memory Accelerator) и поддержкой fp8. Основных источника ускорения два: 1️⃣ Паралеллизация вычислений между GEMM (General Matrix Multiply) и иными операциями (softmax, RMSNorm). Хоть по flops softmax и нормализации кажутся мизерными, они выполняются в 256 раз медленее на H100, потому могут занимать до половины времени матричного умножения. За их выполнение отвечают другие компоненты GPU, чем тензорные ядра для перемножения матриц, потому их можно обрабатывать параллельно. 2️⃣ Использование fp8. fp8 позволяет почти удвоить скорость. Однако есть нюанс - выбросы, приводящие к большим ошибкам квантования. Чтобы избавиться от них используют incoherence processing из QuIP. Благодаря этому ошибка квантизации уменьшается до 2.6 раз. Результаты С fp16 удается достичь ускорения до 1.6-1.8 раз по сравнению с Flash Attention 2. И под 1.2 PFLOPs c fp8 на больших последовательностях и с большими головами трансформера. Обещают в будущем накатить в торч. Вывод Сильно! Ждем через год Flash Attention 4️⃣!

Fast as CHITA: Neural Network Pruning with Combinatorial Optimization [Статья][Без Кода] CHITA - это гепард, а не административный центр Забайкальского края. Статья 2023 года от Google Research про прунинг. Введение Optimal 🧠 Surgeon, использующий приближение 2️⃣ порядка для определения оптимальной сжатой конфигурации, лежит в основе многих методов прунинга и квантизации. Алгоритм Iterative Hard Thresholding (IHT), в котором на каждом шаге делается шаг оптимизации и прунинг самых маленьких весов, лежит в основе некоторых методов sparse training - в частности, AC/DC. В данной статье решили поженить OBS и IHT и назвали полученную сущность CHITA ( Combinatorial Hessian- free Iterative Thresholding Algorithm). Метод Основная проблема с Гессианом нейронной сети, что его хрен посчитаешь честно. Существуют различные приближения; в данной статье опираются на Фишеровское (оценка Гессиана сумой внешних произведений градиентов). Обычно градиентов (n) много меньше, чем параметров (d), потому полученная матрица низкоранговая. Можно ее даже не материализовывать, так как в итоговом алгоритме потребуются только матрично-векторные произведения, а хранить лишь n градиентов в матрице A \in R^{n x d}. Фишеровская матрица выражается как A^T A. CHITA на каждом шаге делает шаг градиентного спуска для квадратичного разложения в окрестности оптимума (т.е L(w) = L_0 + g^ T + 1/2 w^T H w) с последующим прореживанием как в IHT. Метод требует времени и памяти линейной по числу параметров и сохраненных градиентов. Оптимальный шаг оптимизации (learning rate в IHT) находят с помощью умного 🧠 алгоритма поиска на прямой. Эксперименты Метод валидируют на небольших CNNках на CIFAR-10 и ImageNet. Метод быстрее M-FAC и при этом достигает лучшего качества. Однако честность сравнения в one-shot pruning, учитывая, что метод итеративный, в отличие от M-FAC, вызывает вопросы. Метод можно применять итеративно, постепенно повышая степень сжатия и дообучая сеть с фиксированной маской. При итеративном сжатии удается добиться 1% просадки на MobileNet-v1 при 75% cжатии и 5% просадки при 89%. В этом режиме почему-то забывают в сравнениях про M-FAC 😂, который достигает близких результатов по качеству. Метод довольно неплох и на ResNet50 - с умеренными просадками при 90-95% sparsity. Выводы Метод, неплох и не слишком сложен в реализации. Основная проблема - необходимость хранить много градиентов, что не позволяет масштабировать метод на LLMки и интересные в 2024 году модели. У вашего покорного слуги есть одна идейка, как это обойти. Но это пока секрет)

Sparse maximal update parameterization: A holistic approach to sparse training dynamics [Папира][Кода нет] Введение Подбор шага обучения - занятие довольно потное 🥲 и утомительное 😞. Learning rate, хороший для меньших моделей, зачастую приводит к расходимости для больших. Потому большие модели обычно обучают с меньшим learning rate. Но подбор по сетке (или с помощью байесовской оптимизации) может оказаться слишком накладным на практике, а иметь хорошие значения гиперпараметров все-таки хочется. В Tensor Programs V предложили параметризацию весов, градиентов и learning rate, такую, что оптимальные гиперпараметры обучения переносятся с маленьких моделей на большие. Динамика обучения в разреженных сетях существенно отличается от плотных. Потому естественно предположить, что обобщаемая параметризация на разные степени прореживания потребует дополнительных телодвижений. И в данной статье выводят такую параметризацию. Метод Параметризация, при которой learning rate и прочие гиперпараметры обучения будут переноситься на разные размеры моделей и степени прореживания, должна удовлетворять следующим свойствам: 1️⃣ Нормы весов не зависят от ширины и sparsity 2️⃣ Нормы градиентов по весам не зависят от ширины и sparsity 3️⃣ Изменения весов (для произвольного алгоритма оптимизации) не зависят от ширины и sparsity Из нехитрой математики следует, что variance весов при иниациализации и learning rate следует масштабировать как 1 / (ширину_слоя * (1 - sparsity)). Чем шире сеть - тем меньше разброс параметров и learning rate, чем более разреженная сеть - тем, наоборот, больше разброс параметров и learning rate. При sparsity = 0, предложенная SμPar вырождается в μP. Эксперименты Метод валидируют на LMках c SwiGLU и Alibi, обучаемых на токенах из Slim Pajama. В качестве маленькой прокси модели для тюнинга гиперпараметрво берут LMку с 40M параметрами, а основную серию экспериментов проводят на 610M модели. Замеряют loss на обучении (датасет большой, потому переобучения нет). При стандартной параметризации и μP оптимальный learning rate приходится подбирать для каждой степени сжатия свой, в то время, как для SμPar оптимальные learning rate зафиксирован. Перенесенные с меньшей модели параметры на большую для разных степеней прореживания позволяют достичь лучшего качества по сравнению с попытками затюнить стандартную параметризацию и μP. Авторы - ребята из Cerebras, потому умеют вполне эффективно эксплуатировать разреженные операции 😉. Вывод Полезная идея при обучении разреженных сетей. Ранее в Sparsity Scaling Laws было показано, что при очень большом числе данных sparse модели более compute оптимальны, чем плотные. Если в будущем будут активно учить LLMки с использованием специализированного железа SμPar будет весьма кстати для подбора гиперпараметров.

Sparse-Marlin [Код] Некоторое время назад я затрагивал Marlin - быстрый кернел для батчового инференса INT4xFP16. Пару месяцев назад коллеги из IST выпустили SparseMarlin - где дополнительно к квантизации весов добавляется 2:4 sparsity, за счет чего достигается еще большее ускорение по сравнению с fp16. Как я понял, основные инженерные наработки следующие: 1️⃣ Эффективная обработка 2:4 метаданных - позиций нулевых и ненулевых весов 2️⃣ Использование Sparse Tensor Cores (SPTCs) 3️⃣ Умное расположение квантованных весов, метаданных для sparsity, и статистик квантования Итоговое ускорение до 5.3x 🏃‍♂️ на RTX3090 (у Marlin ~4x) при перемножении больших матриц. Интересно🥴 , какое end-2-end ускорение может быть достигнуто на эффективном движке для инференса типа vLLM? На Hopper и более новые архитектуры пока не завезли Таким образом, 2:4 может давать определенный выигрыш по скорости. Основная проблема в том, что на текущий момент 2:4 прунинг сильно просаживает качество LLM ☹️.

Real-Time Video Generation with Pyramid Attention Broadcast [Cтатьи Нет][Код есть] В ряде отечественных тг каналов (пост на эйай ньюз, пост на Machine Learning) упомянули проект Real-Time Video Generation with Pyramid Attention Broadcast. Самой папиры еще нет и потому какие-то нюансы могут быть неизвестны. От себя добавлю несколько деталей. Суть подхода заключается в следующем. В видео диффузии есть 3 вида attention операций: 1️⃣ Пространственное 2️⃣ Временное 3️⃣ Перекрестное внимание на condition Ранее в Cache Me If You Can было замечено, что карты attention между соседними шагами диффузии мало меняются между соседними шагами на большей части процесса (за исключением начала и конца) при картиночной генерации. В данном проекте авторы замечают, что для скорости изменения attention карт справедливо следующее неравенство: v_cross < v_time < v_spatial И соотвественно, чем медленее меняется attention, тем чаще он переиспользуется. За счет переиспользования карт attention можно параллелить эффективно между разными GPU разные шаги генерации по времени (уменьшении оверхеда на 50% без переиспользования карт). Переиспользование карт дает ускорение ~30%. А распаралелливание на 8 GPU дает почти линейное ускорение (8.4x-10.6x) по сравнению с генерацией на одной GPU наивным способом. То есть ускорение достигается в первую очередь за счет эффективного параллелизма. Тем не менее, достойный инженерный результат.

Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity [Блог] Торч недавно выкатил блог, где показывает, что 2:4 sparsity паттерн может быть полезен не только для инференса, но и на стадии обучения модели. Напомню, что 2:4 sparsity - это когда из 4 подряд идущих весов 2-нулевые. Начиная с Ampere нвидиевские GPU имеют поддержку такого формата. Веса хранятся в виде сжатого вдвое тензора и еще есть тензор индексов элементов. Ранее было показано, что такой паттерн выдает ускорение 10-30% по сравнению с fp16 на инференсе, но чтобы добиться эффективного обучения необходимо учесть еще ряд нюансов. Обучение с фиксированной маской работает не очень, и нужно ее обновлять по ходу дела чтобы не просесть по качеству. Наивный kernel слишком медленный и по порядку величины близок к времени forward пасса. И заслуга авторов в разработке эффективной процедуры прореживания. 1️⃣ На обратном проходе при подсчете градиента участвует не W, W.T, которая вообще говоря не 2:4. Авторы ограничивают паттерн прореживания так, что 2:4 паттерн имеет место как для W, так и для W.T (прореживая блоки 4x4). Еще есть нюанс в том, что на уровне железа реализовано умножение sparse матриц на dense, но не наоборот. Потому предлагается считать транспонированное умножение, а затем его транспонировать. В cuSPARSELt есть опция выбирать индексирование в памяти как по колонкам, так и по столбцам, что позволяет иметь непрерывное расположение элементов в обоих случаях. 2️⃣ Потоки CUDA считывают за раз по 128 байт, и чтобы отпимизировать операции чтения/записи каждый поток берет на себя по 4 блока 4x4 (4x4x4x2 (fp16) = 128 байт). 3️⃣ if/else операции, когда разные потоки в группе (warp) попадают в разные условные ветви - работают сильно неэффективно на GPU. Используется sorting network для определения важных/неважных элементов без условных операторов. 4️⃣ Сжатые матрицы и метаданные (индексы ненулевых элементов) хранятся в column-major для оптимизации операций записи. Эксперименты Обучают DINOv2 (ViT-L) c разными пропорциями sparse/dense обучения. Сначала идет sparse обучение, а затем dense. Для оценки качества смотрят на точность лог-регрессии на ImageNet-1k поверх признаков обученной модели. Обучение в sparse режиме на протяжении 40-70% обучения не просаживается по качеству по сути по сравнению с dense обучением. Всегда sparse - теряет полпроцента качества. Ускорение времени обучения при 70% sparse обучения порядка 6% (на А100). Вывод 2:4 паттерн не дает какого-то впечатляющего ускорения, но может слегка повысить эффективность. Интересно 🤔, насколько полученные результаты масштабируются на LLMки с миллиардами параметров. Будет ли выигрыш от 2:4 компенсироваться более медленной сходимостью?

Эксперименты Метод валидируют на 🦙-2, 3, и в качестве бейзлайнов выступает QuIP# и AQLM первой версии. С PV-tuning не сравниваются, шельмецы 👺, хоть и лежит он на архиве 3 недели с их публикации) Квантуют группами 16x16 (число весов по входной и выходной размерности) с L=16, и Q=9 (гиперпараметр в обучаемом гибридном коде). Разные битности отличаются числом ребер k из каждой вершины. Метод демонстрирует выдающееся качество даже без какого-то дообучения, в 4 бита почти без просадки, а в 2 бита почти на уровне AQLM и QuIP# c end-to-end finetuning (жулики в FT репортят цифры AQLM только с блочным файтьюном 😈 ). С полноценным файнтьюном (дообучение знаков в Адамаровых матрицах и параметров HYB преобразования) метод вырывается далеко вперед при низких битностях от QuIP#/AQLM. Неупомянутый PV-tuning работает несколько лучше на меньших моделях - 7(8B), 13B, и сравнивается на 70B (используя, правда, больше данных и вычислений). На 🦙-3 просадки больше, но данной семейство модели, согласно наблюдениям извне и личному опыту, тяжелее квантизуются. Скорость инференса (замеряют на RTX4090) на уровне QuIP#. Сравнение с AQLM полная шляпа 🎩, чуваки, небось коряво поставили либу для инференса, либо что еще наворотили, ибо tokens/s в нашей статье (на более медленной RTX3090) куда больше. Выводы Несмотря на корявое сравнение с AQLM, в общем и целом, конкуренты проделали отличную работу. Здорово, когда красивые идеи из математики находят применение на практике. Кажется, что стоит реально засесть за ботанье 👨‍🏫 матана и чтение статей из 80-х)

QTIP: Quantization with Trellises and Incoherence Processing [Статья][Кода нет, но обещают, что будетт] Было очевидно, что после выхода обновленного AQLM и PV-tuning, что ответочка от конкурентов из Cornell University - это лишь вопрос времени. Тем более, что в issue в PEFT первый автор QuIP# дал ясно понять, что новый, более совершенный, метод квантизации скоро выйдет на свет. И тихо, незаметно, так что сам я обнаружил лишь по чистой случайности, на архиве появилась статья и классическая репа с ридмишкой. Метод Векторная квантизация при той же битности точнее (если реализована не криворуко), чем скалярная. Раньше я думал, что данное явление обусловлено некоторыми корреляциями между измерениями, но на самом деле векторная квантизация работает точнее даже для последовательностей i.i.d (хорошее обьяснение здесь) . Чем больше группа - тем более высокой точности при заданном числе бит можно добиться потенциально. Однако, большие кодбуки становятся довольно обьемными и вносят нетривиальный вклад в общую битность (в AQLM 1x16 кодбуки накидывают ~0.3 дополнительных бит на параметр) и не влезают в кэш GPU. В идеале бы хотелось иметь компактный кодбук и большие группы. И авторы QTIP (названного в честь рэпера?) прибегают к красивой идее из теории кодирования - trellis coding. Обычно при кодировке последовательностей в L бит, каждый элемент может принимать любое из 2^L значений, и последовательность длины кодируется TL битами. Если представить последовательность в виде блуждания по графу, то граф полносвязный - из любой вершины можно прийти в любую другую. В trellis coding (названном в честь архитектурной решетки), каждый узел графа (соответствующий одному из 2^L) значений соединен только с 2^k (где k<L) ребрами с другими вершинами. Последовательность кодируется как индекс начального узла и индексы ребер при блуждании по графу. Итого расход памяти L + (T-1)k, что может быть заметно меньше TL. Однако, наивная реализация будет чрезмерно тяжеловесной для произвольного trellis при достаточно большом L и k, так как где-то надо хранить граф связей. Кроме того, декодирование для произвольного trellis - много последовательных операций, что непозволительно медленно на практике, и QTIP предлагает несколько хитрых идей, чтобы это работало эффективно. Как и в QuIP/QuIP# квантуется не исходный вес, а обработанный случайным ортогональным преобразованием (Адамаровыми матрицами), так что элементы матриц близки к i.i.d нормальному распределению. Далее, чтобы не хранить полный произвольный граф, авторы рассматривают специальный анзатц - bit-shift trellis, где связаны только вершины с индексами, отличающимися на битовый сдвиг. Однако, такой выбор слишком ограничителен, и авторы предлагают “перемешивать” связи некоторой хитрой случайной перестановкой и предлагают три методы: 1️⃣ 1MAD 2️⃣ 3INST 3️⃣ HYB В первых двух случаях в качестве индексов в случайной перестановке выступает некоторое фиксированное преобразование использующее линейный конгруэнтный генератор случайных чисел и побитовые операции. А в третьем случае - некоторое обучаемое преобразование. Для определения оптимальной конфигурации минимизируется ошибка на выходе слоя с помощью алгоритма Витерби (динамическое программирование). Полученный метод выдает MSE ошибку близкую к минимально теоретически возможной.

Вывод Трансформеры, вероятно, являясь в каком-то смысле ансамблем, устойчивы, как печень, к ампутации крупных кусков. Однако просадки (там где результатам можно доверять - первая и последняя работа), все же значительны при степени сжатия более 20%. Как ни крути, scaling laws существуют не просто так, и представления, выученные разными блоками, хоть имеют и разную полезность, но все же содержат в себе релевантное знание, хоть для какого-то входа. Отсюда и успех контекстуальных методов сжатия, вроде Deja Vu и PowerInfer.

Your Transformer is Secretly Linear [Статья][Пост в телеге][Пост на Хабре][Код] Интересное релевантное исследование от коллег из AIRI/Сколтеха. Кратко, ибо все хорошо и доступно изложено самим автором в постах в телеге и на Хабре. Метод Оказывается, что активации трансформера в соседних блоках связаны почти что линейным преобразованием. В качестве меры линейности используется linearity score, который по существу является r2 (коэффициентом детерминации) . Сам по себе факт не столь удивителен, ибо норма skip-connection обыкновенно значительно больше преобразования в attn/mlp. Но даже если вычесть skip connection выделяется явная линейная и нелинейная компонента. Проанализировав ряд промежуточных чекпоинтов открытых моделей, авторы замечают что норма линейной компоненты уменьшается в процессе предобучения, но возрастает в ходе файнтьюна . И предложенная регуляризация на нелинейность улучшает качество. Исходя из наблюдений, предлагается прунить самые “линейные” слои, а для восстановления качества заместо запруненного блока учится линейный слой. И такой адаптер неплохо восстанавливает качество. Эксперименты Рассматривают большое число открытых моделей - от GPT-2 (когда-то OpenAI был действительно открытым), Bloom до свежих Phi-3. Рост линейности имеет место для всех моделей и задач. Прирост качества от регуляризации хоть и невелик, но все же заметен. Интересно, насколько он будет иметь место на большом масштабе? Несколько нетипичный сетап замера перплексии - короткие последовательности вместо обычно используемого 2к/4к контекста (как на обучении), но общая суть от сего не меняется.