Сворачивание весов, стримы CUDA и баг, из-за которого моя модель заговорила задом наперед — Филип Макрадули

РусскийالعربيةEnglishFrançaisहिन्दीBahasa Indonesia日本語한국어Português中文
AAI Engineer
Computing/SoftwareInternet Technology

Transcript

00:00:00.
00:00:13Всем привет, спасибо, что пришли,
00:00:18и я начинаю свой доклад.
00:00:22Этот доклад посвящен статье,
00:00:27которую я написал, и она очень простая.
00:00:33Основная идея предельно понятна.
00:00:36По сути, это всего пару строк
00:00:39алгебры, которые делают слой RMSNorm
00:00:42в трансформерах менее ресурсоемким,
00:00:45более быстрым и в целом улучшают его
00:00:48как слой в архитектуре
00:00:51трансформера. Примерно так же,
00:00:54как LayerNorm когда-то был стандартом,
00:00:57а потом уступил место RMSNorm.
00:01:00И эта работа идет в том же направлении.
00:01:04Мне посчастливилось познакомиться
00:01:09с разработчиками из open-source сообщества,
00:01:11и я стал соавтором этой статьи
00:01:14вместе с Нильсом Графом, который был
00:01:17главным создателем этого подхода.
00:01:21Отсюда и развилось наше исследование.
00:01:23Оно опубликовано на arXiv.
00:01:26Можете взглянуть, почитать, протестировать.
00:01:29Также есть репозиторий.
00:01:31И концепцию, сам подход,
00:01:36проще всего объяснить,
00:01:38пожалуй, на примере FlashAttention.
00:01:40Подобно тому, как FlashAttention
00:01:45откладывает вычисления до умножения
00:01:48и стремится сократить число обращений
00:01:52к памяти, чтобы весь процесс
00:01:54шел быстрее,
00:01:55здесь используется сходная логика.
00:01:58Мы вносим определенные улучшения,
00:02:01которые ускоряют обработку RMSNorm
00:02:08и в конечном счете оптимизируют весь трансформер.
00:02:15Возникает вопрос: почему именно RMSNorm?
00:02:19Ведь этот слой почти не выполняет математических вычислений.
00:02:24И это действительно так.
00:02:26Доля именно вычислительной нагрузки
00:02:29здесь довольно мала.
00:02:31Однако реальное время выполнения,
00:02:34Wall time, весьма ощутимо.
00:02:36Например, за один шаг декодирования,
00:02:40то есть непосредственно при инференсе,
00:02:43RMSNorm может вызываться около 33 раз.
00:02:47Конечно, это зависит от конкретной модели и других факторов.
00:02:50В статье приведены детали по моделям
00:02:52и описана методика тестирования.
00:02:54Главный вопрос — как это оптимизировать
00:02:59и избежать лишнего ожидания весов
00:03:03при матричном умножении.
00:03:06Причина медленной работы вовсе не в том,
00:03:09что GPU плохо или медленно считают математику,
00:03:15а в накладных расходах вокруг самих вычислений.
00:03:20То есть на этапе запуска самой работы.
00:03:24Например, запуск процесса 33 раза подряд,
00:03:30как это происходило в некоторых экспериментах,
00:03:32занимает много времени.
00:03:35И, к примеру, объединение каждой нормализации
00:03:40с матричным умножением помогает этого избежать.
00:03:44Также сворачивание весов (weight folding)
00:03:46помогает сократить перемещение данных между блоками памяти,
00:03:50а этот процесс тоже является узким местом для GPU.
00:03:54И ещё один момент — ожидание.
00:03:56Например, откладывание деления,
00:04:00которое выполняется в слое RMSNorm,
00:04:02тоже позволяет устранить эту фазу ожидания.
00:04:06Таким образом, суть данной статьи
00:04:09заключается в улучшении всех трех аспектов
00:04:13с помощью нескольких алгебраических трюков
00:04:16при вычислении RMSNorm.
00:04:19Вот и всё.
00:04:21С точки зрения математики, вот эти трюки.
00:04:25В основном всё строится на первых двух положениях.
00:04:29Первое — нормализация без весов.
00:04:31Вы можете видеть это здесь.
00:04:33И отложенная нормализация.
00:04:35Это второе положение.
00:04:37А в более современных архитектурах
00:04:39бывают ситуации, когда RMS встречается дважды.
00:04:44Например, так происходит в Gemma 4.
00:04:48Поэтому сокращение предварительной нормализации тоже работает.
00:04:52И всё это алгебраически доказано в статье.
00:04:58Первое положение состоит в том,
00:05:00что масштабирование (gain) и сворачивание весов
00:05:04объединяются в одну матрицу.
00:05:06Матрица W отмечена здесь звездочкой.
00:05:09Она вычисляется оффлайн,
00:05:12подобно тому, как во FlashAttention
00:05:14некоторые значения рассчитываются заранее,
00:05:16чтобы избежать постоянного обмена
00:05:18данными с памятью.
00:05:20Так что это первый выполненный шаг —
00:05:24сворачивание весов.
00:05:26А второй шаг — это откладывание скалярного деления
00:05:32при матричном умножении для их параллельного выполнения.
00:05:35 В обычном случае вам пришлось бы посчитать один раз,
00:05:38затем подождать и посчитать снова.
00:05:41Здесь же идея состоит в том, чтобы разделить операции
00:05:44для их распараллеливания.
00:05:48Третий момент, являющийся развитием этой идеи,
00:05:51заключается в том, что при наличии двух операций,
00:05:56из-за инвариантности к масштабу,
00:05:58одну из них можно опустить, и всё продолжит работать.
00:06:01Это применимо к более новым моделям,
00:06:04поддерживающим такую архитектуру и реализацию.
00:06:10Чтобы воплотить это на практике,
00:06:13особенно второе положение...
00:06:16Вот с этим, например, всё просто.
00:06:20Есть репозиторий transformer-tricks,
00:06:22можно просто применить это к любой модели, и всё работает.
00:06:25Но для второго требуется работа с кернелами.
00:06:29И реализация здесь не столь очевидна.
00:06:31И вот в процессе реализации
00:06:35я однажды наткнулся на такой результат эксперимента.
00:06:42В целом всё выглядело нормально, вроде:
00:06:46"Архитектура трансформеров
00:06:48произвела революцию в NLP, потому что...",
00:06:51а дальше идет ожидаемый сгенерированный текст.
00:06:54Но в итоговом выводе
00:06:56я заметил повторение и задержку на один шаг.
00:06:59Видите, слово "потому что" появляется снова.
00:07:01Проблема была связана с потоками GPU (GPU streams),
00:07:07и я пытался разобраться, в чем дело.
00:07:10Я получал этот лаг в один шаг и результаты,
00:07:15которые словно подтягивались из прошлого.
00:07:18И в ходе отладки
00:07:21я понял, что в процессе создания подобной системы...
00:07:26Как я уже объяснял во втором положении про откладывание двух операций,
00:07:32в CUDA можно задействовать два ресурса.
00:07:35Тензорные ядра, выполняющие часть матричного умножения,
00:07:39и ядра CUDA, исполняющие поэлементные операции,
00:07:44редукции, вычисление квадратных корней и так далее.
00:07:47Идея заключалась в их параллельном запуске для получения преимущества,
00:07:52описанного в статье, чтобы протестировать эту концепцию.
00:07:58Вот как это должно было выглядеть.
00:08:00Если делать всё последовательно,
00:08:03возникает время простоя, пока векторный блок вычисляет RMS и масштабирование,
00:08:10а затем происходит умножение матриц.
00:08:12И идея Flash Norm — метода из статьи —
00:08:16заключалась в том, чтобы выполнять эти шаги параллельно.
00:08:19Матричный блок считает matmul, а векторный — RMS.
00:08:23Таким образом экономится время.
00:08:26Однако это нельзя сделать просто на Python.
00:08:28Нужно опуститься немного ниже.
00:08:30И я сделал это с помощью вот такого CUDA-кода.
00:08:35И в целом на тот момент всё выглядело нормально.
00:08:42Однако потом я понял, что сделал кое-что неправильно.
00:08:47А именно то, что объединение двух потоков в конце в моём случае получилось неявным.
00:08:57И когда я это протестировал, юнит-тест прошёл успешно.
00:09:01Качество казалось схожим, как при проверке перплексии и прочем, так как генерация была похожей.
00:09:08Но при длинной генерации мне удалось увидеть эту проблему.
00:09:11Я понятия не имел, в чём дело.
00:09:14А причина была в том, что из-за неявного объединения один из потоков не успевал завершить работу.
00:09:24Из-за этого возникало состояние гонки и считывание старых данных из незавершённого умножения матриц.
00:09:31Для исправления возникла идея о том, что объединение должно быть явным.
00:09:40И нужно ждать завершения одной из операций, чтобы при объединении не считывать устаревшие данные.
00:09:48Вот к какому выводу я пришёл в процессе изучения CUDA-потоков.
00:09:54И вот как у меня всё было устроено.
00:09:57Объединение происходило неявно.
00:10:00Поэтому post scale считывал старое значение буфера.
00:10:06Исправляется это так: нужно явно отметить конец умножения матриц,
00:10:14затем отметить конец вычисления RMS,
00:10:17после чего в post scale дождаться первого потока,
00:10:21а затем дождаться второго потока.
00:10:24Это исправило баг, заставило подход из статьи работать, и модель заговорила нормально, а не задом наперёд.
00:10:33Это была интересная академическая перспектива.
00:10:38Но мне также хотелось опробовать это на практике,
00:10:40задеплоить, протестировать и посмотреть, как заставить это работать в продуктовой среде.
00:10:47Вы также можете прочитать статью и посмотреть все тесты.
00:10:50Большинство из них сделаны на моделях Llama, но метод работает и с другими архитектурами.
00:10:56Что касается этой статьи, то, например, свёртку весов из утверждения 1 можно сделать с помощью кода из репозитория.
00:11:10Вы вызываете «flashify», и код делает это сам.
00:11:13Однако для второй вещи, о которой я упоминал, придётся немного поработать с кернелами,
00:11:19как я и объяснял на своём примере.
00:11:22Вот некоторые результаты на базе моделей Llama, здесь есть разные детали, с которыми можно ознакомиться.
00:11:29Например, что происходит, если использовать только отложенную нормализацию, или если применить полностью слитный кернел.
00:11:36Было проведено много экспериментов с опусканием на более низкий уровень для проверки всех гипотез.
00:11:41И вот наши результаты на разных уровнях детализации и проверки.
00:11:48Но даже простой вариант со свёрткой весов показывает некоторое улучшение.
00:11:54И это работает с повседневными инструментами для работы с моделями.
00:11:59Так что вам не придётся изобретать велосипед или делать всё с нуля.
00:12:05Это совместимо с torch.compile, так как это просто новый чекпоинт, и всё.
00:12:13Flash Attention использует похожие трюки на другом уровне.
00:12:16И это также работает с квантованными моделями.
00:12:18Так что это отлично подходит для применения, и вы можете получить модель с этим классным новым слоем нормализации.
00:12:27Подробности и код для запуска можно найти в репозитории transformer-tricks.
00:12:34Там есть разные алгебраические трюки, о которых я рассказывал, а также упомянутая статья.
00:12:41Кроме того, есть репозиторий моделей на Hugging Face, где я выложил некоторые из них.
00:12:50Вы можете перейти по ссылке на Hugging Face к этой модели и протестировать её.
00:12:54А ещё эти модели с Hugging Face можно развернуть в продакшене.
00:13:00Размышляя об этом, я понял: когда научная часть готова и есть ссылка на модель в Hugging Face,
00:13:11движок инференса Superlink оказывается отличным способом развернуть любую модель с Hugging Face.
00:13:19Мы делали это на хакатонах, куда участники приносили кастомные модели или чекпоинты с Hugging Face со своими файнтюнами.
00:13:27И даже если у вас есть своя версия таких алгебраических трюков для улучшения модели и проверки исследовательской идеи,
00:13:37вы можете опробовать её и получить развёрнутую на кластере версию модели, не беспокоясь об обвязочном коде для деплоя.
00:13:48И это очень здорово. Главное, что если весь кластер и инференс моделей имеют открытый исходный код,
00:13:58вы можете тестировать подобные более новаторские исследования: если вы хотите модифицировать кернелы или Flash Norm и тому подобное,
00:14:12это гораздо сложнее сделать на арендованном эндпоинте, где инференс вам не принадлежит.
00:14:18Вам нужно что-то портативное и гибкое, позволяющее делать подобные вещи,
00:14:23но при этом достаточно готовое к продакшену, чтобы тестировать решения на масштабе.
00:14:27И вы можете, например, использовать Cy для объединения этого с другими моделями.
00:14:33Как видно вверху слева, вы можете связывать эти оптимизированные модели с другими для выполнения агентных задач,
00:14:43и реализовывать сквозные масштабные сценарии.
00:14:46Принцип работы Cy в том, что этот продуктовый кластер помогает вам деплоить модели.
00:14:52Вы можете заглянуть в репозиторий Cy, чтобы узнать больше подробностей.
00:14:57Также там есть более умный механизм очередей, который помогает, особенно при работе с небольшими моделями.
00:15:03Ведь работая над Flash Norm, я использовал небольшие модели Llama, а также компактных агентов с Hugging Face.
00:15:11Возможность разворачивать небольшие модели на одном и том же GPU, чтобы не тратить лишние деньги на инфраструктуру,
00:15:24и при этом легко переключаться между ними, оказалась очень полезной.
00:15:30Вы также можете управлять конфигурациями моделей и кластера через API,
00:15:35что крайне удобно, когда у вас нет отдельного инженера для поддержки open-source исследований.
00:15:40Так что это тоже здорово.
00:15:43Вы полностью контролируете своё облако, что полезно для работы с открытыми весами и open-source моделями.
00:15:50У Cy также есть каталог различных моделей, не только тех, о которых я упомянул,
00:15:57вы можете его посмотреть.
00:15:59Там есть и модели реранкинга эмбеддингов, если вы разрабатываете что-то подобное.
00:16:03На этом я подхожу к концу своего рассказа о научном пути, где я стал соавтором статьи о методе улучшения трансформеров,
00:16:15а также нашёл способ довести это до продакшена, протестировать и поэкспериментировать с открытыми моделями.
00:16:24Не стесняйтесь писать мне в LinkedIn, если у вас есть вопросы или идеи для сотрудничества.
00:16:30Многое из того, что я упомянул, опубликовано в виде PR в VLLM или на Hugging Face.
00:16:36Вы легко сможете их найти.
00:16:38Также вы можете ознакомиться со статьёй.
00:16:40Ссылка на arXiv указана там.
00:16:43Там же есть ссылка на репозиторий Cy и мой LinkedIn.
00:16:47Большое спасибо за внимание.
00:16:58Вы можете задать мне вопросы.
00:16:59Мы будем здесь,
00:17:00неподалёку.
00:17:13Спасибо.

Key Takeaway

Сворачивание весов оффлайн и параллельное выполнение RMSNorm с матричным умножением через явную синхронизацию CUDA-потоков устраняет задержки памяти трансформеров без потери точности генерации.

Highlights

  • Слой RMSNorm вызывает задержки выполнения (Wall time) из-за накладных расходов на запуск операций и частые обращения к памяти — до 33 вызовов за один шаг декодирования.

  • Метод Flash Norm использует сворачивание весов (weight folding) оффлайн и отложенное скалярное деление для параллельного выполнения операций на CUDA-ядрах и тензорных ядрах.

  • Неявная синхронизация потоков CUDA при оптимизации RMSNorm приводит к состоянию гонки (race condition) и считыванию устаревших данных из буфера памяти.

  • Явная синхронизация CUDA-потоков с фиксацией точек завершения матричного умножения и расчета RMS полностью устраняет баги генерации текста.

  • Оптимизированные слои нормализации полностью совместимы с квантованными моделями и компилятором torch.compile.

Timeline

Узкие места слоя RMSNorm в архитектуре трансформеров

  • Слой RMSNorm выполняет минимальный объем математических вычислений, но создает высокую задержку реального времени (Wall time).
  • За один шаг декодирования при инференсе модель может вызывать RMSNorm около 33 раз.
  • Основное время расходуется на накладные расходы при запуске ядер GPU и постоянное перемещение данных между блоками памяти.

Низкая вычислительная сложность RMSNorm скрывает проблему низкой эффективности на GPU. Основной задержкой становится не сама математика, а постоянный запуск мелких операций и ожидания при обмене данными с памятью. Для решения этого используется подход Flash Norm, снижающий число обращений к памяти аналогично FlashAttention.

Алгебраические методы оптимизации: свёртка весов и отложенная нормализация

  • Масштабирование (gain) и веса нормализации объединяются в единую матрицу, вычисляемую оффлайн.
  • Скалярное деление в RMSNorm откладывается для распараллеливания с матричным умножением.
  • В архитектурах с двойным RMSNorm одна из операций нормализации опускается благодаря инвариантности к масштабу.

Три математических трюка позволяют оптимизировать работу нормализации. Объединение параметров масштабирования с матрицей весов до запуска инференса (weight folding) полностью убирает этап предварительной нормализации. Разделение скалярного деления и умножения матриц позволяет запускать их одновременно на разных блоках GPU.

Управление CUDA-потоками и исправление состояния гонки

  • Тензорные ядра выполняют матричное умножение параллельно с CUDA-ядрами, рассчитывающими поэлементные операции и RMS.
  • Неявное объединение CUDA-потоков приводит к состоянию гонки и считыванию устаревших данных из прошлых шагов генерации.
  • Явное ожидание завершения обоих потоков перед операцией post scale устраняет ошибку повторяющегося текста.

Реализация отложенной нормализации требует низкоуровневой работы с CUDA-кернелами. При попытке параллельного запуска матричных операций на тензорных ядрах и векторных вычислений на CUDA-ядрах возникла ошибка: модель повторяла слова с задержкой в один шаг. Причиной стало неявное слияние потоков, при котором операция считывала данные из незавершенного буфера. Введение явных меток завершения процессов (cudaEventRecord) и точек ожидания (cudaStreamWaitEvent) полностью решило проблему.

Практическое применение и деплой оптимизированных моделей

  • Библиотека transformer-tricks позволяет применять свёртку весов функции flashify без переписывания архитектуры.
  • Предложенный подход полностью совместим с torch.compile, квантованием и библиотекой vLLM.
  • Платформа Cy и инференс-движок Superlink обеспечивают развертывание кастомных моделей с Hugging Face на GPU-кластерах.

Сворачивание весов подготавливает новый чекпоинт модели, сохраняя полную совместимость со стандартным стеком PyTorch. Оптимизированные модели легко интегрируются в продуктовые среды с помощью open-source инфраструктуры Cy, которая поддерживает динамическую смену моделей и умное управление очередями запросов для компактных LLM-агентов.

Community Posts

No posts yet. Be the first to write about this video!

Write about this video