Почему происходит ошибка нехватки памяти (OOM) в последнем слое LLM и как это исправить с помощью пользовательского ядра Triton.
Делиться
Если вы когда-либо обучали или совершенствовали модель LLM, вы, вероятно, сталкивались с препятствием на самом последнем этапе: функцией потерь кросс-энтропии .
Виновником является узкое место логистической регрессии . Для предсказания следующего токена мы проецируем скрытое состояние в огромное пространство словарей. Для Llama 3 (128 256 токенов) только матрица весов содержит более 525 миллионов параметров . Хотя это всего лишь ~1 ГБ в формате bfloat16, настоящая проблема заключается в промежуточном тензоре логистической регрессии. Для больших пакетов данных вычисление одной скалярной функции потерь может легко превысить 80 ГБ видеопамяти.
Оптимизация этого слоя — вот как такие библиотеки, как Unsloth и Liger-Kernel, достигают столь значительного сокращения потребления памяти. В этой статье мы с нуля создадим ядро, объединяющее линейную и кросс-энтропийную обработку, в Triton. Мы выведем математические формулы и реализуем пошаговый прямой и обратный проход, который сократит пиковое использование памяти на 84% .
Примечание о производительности: Данная реализация носит преимущественно образовательный характер . Мы отдаем приоритет математической ясности и читаемому коду Triton, используя глобальные атомарные операции. Хотя это решает проблему узкого места в памяти, для достижения скорости, характерной для промышленного использования, потребовались бы значительно более сложные реализации, которые выходят за рамки данной статьи.
Этот пост является частью моей серии статей о Triton. Мы будем использовать такие концепции, как тайлинг и онлайн-оптимизация Softmax, которые мы уже рассматривали ранее. Если они вам незнакомы, рекомендую сначала ознакомиться с ними!
Изучение Triton: ядро за ядром — умножение матриц.
Изучение Triton по одному ядру за раз: Softmax
Изучение Triton: ядро за ядром — сложение векторов.
Узкое место Logit
Для начала давайте добавим еще несколько цифр к логит-узкому месту. Рассмотрим входную матрицу X размером [NxD], матрицу весов W размером [DxV] и логит-матрицу Y=X@W размером [NxV]. В контексте LLM, N будет представлять собой длину последовательности, умноженную на размер пакета (т.е. общее количество токенов в пакете), D — размер скрытого состояния, а V — размер словаря.
Для модели Llama3 8B у нас будет контекстное окно размером 8192 токена, скрытое состояние с 4096 измерениями и размер словаря 128 256 токенов. Используя умеренный размер пакета данных, равный 8, мы получим N = 8192 x 8 = 65 536.
В результате матрица Y имеет форму [NxV]=[65 536 x 128 256], или примерно 8,4 миллиарда элементов. В формате bfloat16 это заняло бы 16,8 ГБ памяти. Однако, если следовать передовым методам и использовать float32 для расчета функции потерь, чтобы обеспечить численную стабильность, требования удваиваются до 33,6 ГБ .
Чтобы представить эту цифру в перспективе, нам также потребуется около 16 ГБ памяти для хранения весов Llama3 8B в памяти в формате bfloat16. На большинстве графических процессоров это не оставляет места для огромных накладных расходов состояний оптимизатора (например, моментов Адама) и других функций активации, что приводит к печально известной ошибке нехватки памяти (OOM) в PyTorch.
Как правило, эта проблема решается с помощью следующих методов:
- Накопление градиентов: Используйте меньший размер пакета и накапливайте градиенты по нескольким пакетам между каждым шагом оптимизатора, имитируя больший размер пакета, но при этом храня меньше данных в памяти.
- Сохранение промежуточных активаций: PyTorch сохраняет все промежуточные активации для повторного использования в обратном проходе, а сохранение промежуточных активаций очищает эти активации и пересчитывает их на лету во время обратного прохода. Это приводит к значительной экономии памяти, но увеличивает время обучения, поскольку количество необходимых прямых проходов удваивается.
- Микропакетная обработка функции потерь: вместо вычисления функции потерь по N измерениям за один раз, мы можем разделить ее на более мелкие фрагменты размером n < N и накапливать результат. Теперь в памяти хранится только фрагмент размером [n, V] за раз.
- Обучение с использованием смешанной точности: применение половинной точности во время обучения обеспечивает двукратное сокращение объема памяти и значительное ускорение работы на ядрах Tensor Core.
Хотя эти решения кажутся привлекательными, у всех них есть существенные недостатки: накопление градиента и контрольные точки активации замедляют обучение, смешанная точность может быть нестабильной, а микропакетная обработка требует (медленной) итерации на уровне PyTorch, и даже если n выбрано меньше N, размер словаря остается огромным по сравнению с ним.
Что еще более важно, эти решения не решают проблему, которую мы неоднократно рассматривали в этой серии: перемещение данных . Действительно, мы по-прежнему тратим время впустую, записывая миллиарды логитов в видеопамять только для того, чтобы прочитать их обратно через миллисекунды.
Решение на уровне ядра
Как мы увидим через минуту, прямой и обратный проходы функции потерь кросс-энтропии включают скалярное произведение, умножение матриц и функцию softmax. Как мы узнали из этой серии, все эти операции можно эффективно выполнять по частям. Другими словами, мы можем выполнять их итеративно, удерживая в памяти лишь небольшую часть входных данных в любой момент времени.
Кроме того, кросс-энтропии обычно предшествует матричное умножение: линейная проекция скрытого состояния в пространство словаря. Это отличная возможность для слияния операторов : объединения нескольких операций в рамках одного ядра, что приводит к значительному ускорению и потенциальному выигрышу в памяти.
В следующих разделах мы рассмотрим, как эффективно объединять прямые и обратные проходы через ядро, сочетающее линейный слой с кросс-энтропией.
Как упоминалось в предыдущей статье, ядра Triton изначально не регистрируются в функции автоградиента PyTorch. Поэтому нам нужно самостоятельно вычислить градиент, что является прекрасной возможностью освежить в памяти некоторые моменты математического анализа 😉
Математические основы объединенной линейной кросс-энтропии
Определение и передача вперед
В этом разделе мы выводим математическое выражение для нашего слоя объединенной линейной кросс-энтропии, чтобы увидеть, как он естественным образом подходит для мозаичного размещения.
Для двух дискретных распределений вероятностей p и q перекрестная энтропия определяется следующим образом:
В нашем контексте p — это one-hot вектор, представляющий целевой токен, а q — распределение модели по словарю. Мы получаем q, применяя функцию softmax к логитам l, которые сами являются выходами предыдущего линейного слоя.
Поскольку p положительно для единственного целевого токена y, суммирование схлопывается. Затем мы можем подставить численно устойчивую функцию softmax (как обсуждалось в предыдущей статье), чтобы получить окончательное выражение:
Заменив логиты l линейным слоем x . w, мы видим, что прямой проход сводится к трем основным величинам:
- Целевой логит x . w_y.
- Логарифм суммы-экспресса (LSE) всех скалярных произведений.
- Для оценки численной устойчивости используется глобальный максимальный логит.
Благодаря онлайн-алгоритму softmax мы можем вычислять эти величины, не материализуя весь словарь в памяти. Вместо узкого места памяти O(V) мы итерируем по скрытой размерности D и словарю V небольшими блоками (D_block и V_block). Это преобразует вычисление в задачу с регистрами O(1).
Для эффективного распараллеливания процесса мы запускаем по одной программе на графическом процессоре для каждой строки входной матрицы. Каждая программа независимо выполняет следующие шаги:
- Предварительно вычислите целевой логит: выполните скалярное произведение между текущей строкой X и столбцом W, соответствующим токену Y.
- Онлайн-сокращение: Пройдите по скрытым блокам и блокам лексики, чтобы:
1. Отслеживайте текущий максимум (м)
2. Обновите текущую сумму экспоненциальных функций (d), используя онлайн-формулу softmax:
Теперь, когда мы лучше понимаем прямой проход, давайте рассмотрим вывод обратного прохода.
Обратный проход
Обозначение
Для эффективного вычисления градиентов мы будем использовать обозначения Эйнштейна и дельту Кронекера .
В обозначениях Эйнштейна повторяющиеся индексы суммируются неявно. Например, стандартное умножение матриц Y = X@W упрощается от многословного суммирования до простого сопоставления индексов:
Дельта Кронекера (δ_ij) используется вместе с этим обозначением для обработки логики тождества. Она равна 1, если i=j, и 0 в противном случае. Как мы увидим, это особенно полезно для объединения индексов при дифференцировании.
Умножение матриц
В этом разделе мы выводим градиенты обратного распространения для умножения матриц. Мы предполагаем существование градиента ℓ, распространяющегося вверх по потоку.
Чтобы определить, как происходит обратное распространение ошибки при умножении матриц, мы применяем правило цепочки к входным данным x и матрице весов w. Здесь y представляет собой выходные данные умножения:
Начнём с вывода частных производных y по x, выполнив следующие шаги:
- Выразите y через x и w.
- Обратите внимание, что w является константой относительно производной x, поэтому мы можем вынести её за производную.
- Выразите тот факт, что частная производная x_ik по x_mn равна 1 только тогда, когда i=m и k=n, используя символ Кронекера.
- Обратите внимание, что ẟ_kn обеспечивает k=n, поэтому w_kj * ẟ_kn сводится к w_nj.
Затем мы рассматриваем полное выражение и получаем градиент. Последний шаг мы выводим, еще раз заметив, что 1/y_ij * ẟ_im сводится к 1/y_mj.
Однако матричная запись концептуально ближе к нашему ядру Тритона, поэтому мы перепишем это выражение как матричное умножение, используя тождество X_ij = [X^T]_ji:
Для вычисления градиента относительно W мы выполняем те же самые шаги:
Затем следует обратное распространение градиента:
Что эквивалентно матричной записи:
Перекрестная энтропия
В этом разделе мы сосредоточимся на кросс-энтропии, применяемой к дискретным вероятностным распределениям. Рассмотрим тензор из j логитов с меткой y; кросс-энтропия вычисляется следующим образом:
Где x_y соответствует логиту, связанному с меткой.
И снова нас интересует частная производная любого выходного значения i по любому входному значению k. Из-за нормирующего множителя каждый элемент i влияет на значение каждого другого элемента, поэтому частная производная получается путем определения функции, зависящей от значения i, по частям:
Суммируя оба случая, получаем градиент:
А в матричной записи:
Где y_{one hot} — это вектор из нулей, в котором элемент, соответствующий метке, равен единице. Этот результат показывает, что градиент — это просто разница между предсказанием и истинным значением .
Слитая линейная кросс-энтропия
Объединив линейную проекцию с кросс-энтропией в одном выражении, получаем:
Благодаря правилу цепочки, вычисление градиента этого выражения сводится к умножению градиентов, которые мы вычислили ранее:
Где x и y обозначают входные и выходные данные линейного слоя соответственно, а w — соответствующую матрицу весов.
Примечание: в пакетном режиме нам потребуется уменьшить градиенты W по размерности пакета. Обычно мы используем уменьшение суммы или среднего значения.
Реализация ядра
Разработав теорию, мы можем реализовать объединенное ядро в Triton. Поскольку кросс-энтропия обычно является последним слоем в языковой модели, мы можем объединить прямой и обратный проходы в одно ядро. Такое объединение дает два преимущества: минимизирует накладные расходы на запуск нескольких ядер и значительно улучшает локальность данных, сохраняя промежуточные значения на чипе.
Мы проанализируем ядро шаг за шагом с точки зрения отдельного экземпляра программы , который в нашей стратегии распараллеливания обрабатывает одну конкретную строку входной матрицы.
1. Настройка и предварительный расчет целевой логит-кривой.
Начальный этап включает стандартную настройку системы Triton:
- Идентификация программы: Мы используем tl.program_id для определения того, за какую строку входной матрицы отвечает текущая программа.
- Инициализация параметров: Мы определяем блоки, используя D_BLOCK и V_BLOCK, и инициализируем скользящий максимум (m) и сумму (d), необходимые для онлайн-алгоритма softmax.
- Арифметика указателей: Мы вычисляем базовые адреса памяти для наших тензоров. Указатели для X (входные данные) и dX (градиент) смещаются с использованием шага по строкам, так что каждая программа обращается к своему уникальному вектору токенов. И наоборот, указатель W (вес) остается по базовому адресу, поскольку каждая программа в конечном итоге должна пройтись по всему пространству словаря.
- Маскировка и досрочное завершение: Мы определяем ignore_index (по умолчанию -100). Если программа встречает эту метку (например, для токенов заполнения), она завершается досрочно с потерей 0 циклов для экономии ресурсов.
2. Вычисление целевой логиты
Перед основным циклом необходимо выделить целевой логит x . w_y. Мы итерируем по скрытому измерению D блоками D_BLOCK, выполняя скалярное произведение между входной строкой X и конкретным столбцом W, соответствующим метке истинности Y.
Поскольку W — это двумерная матрица, для вычисления указателей на эти конкретные ячейки столбцов требуется точное манипулирование шагом. Приведенная ниже иллюстрация помогает визуализировать, как мы «перемещаемся» по памяти, чтобы извлечь только необходимые веса для целевого токена.
После загрузки тайлов мы преобразуем их в тип float32 для обеспечения численной стабильности и добавляем их скалярное произведение к переменной-аккумулятору, прежде чем перейти к следующей итерации.
Вот код на данный момент:
Далее мы выполняем прямой проход, который обрабатывает пространство словаря в два вложенных этапа:
- Вычисление логитов для плиток: Мы вычисляем логиты для каждого блока V_BLOCK за раз. Это достигается путем итерации по размерности словаря V (внешний цикл) и скрытой размерности D (внутренний цикл). Внутри внутреннего цикла мы загружаем плитку X и блок W, накапливая их частичные скалярные произведения в регистре высокой точности.
- Обновление функции Softmax в режиме реального времени: После завершения вычисления скалярного произведения для логита мы не сохраняем его в видеопамять. Вместо этого мы немедленно обновляем нашу текущую статистику: максимальное значение m и текущую сумму экспоненциальных функций d, используя формулу Softmax в режиме реального времени. Выполняя это «на лету», мы гарантируем, что в регистрах графического процессора в любой момент времени хранится только небольшой V_BLOCK логитов.
После этих итераций окончательные значения m и d используются для восстановления LSE. Затем вычисляется окончательная скалярная функция потерь для строки путем вычитания целевого логита (x . w_y) из этого значения LSE.
Вот наглядное представление паса вперёд:
Вот код для передачи вперед:
Теперь мы подошли к последней части ядра: обратному проходу. Наша цель — вычислить градиенты относительно X и W, используя выведенное нами ранее выражение:
Для обеспечения эффективности использования памяти мы снова обрабатываем словарь по плиткам, используя двухэтапный подход:
- Пересчет нормализованных вероятностей (P): Поскольку мы не сохранили полную матрицу логита во время прямого прохода, нам необходимо пересчитать активации для каждого фрагмента. Используя логарифмическую сумму экспонент, вычисленную во время прямого прохода, мы можем нормализовать эти активации на лету. Вычитание истинной метки Y из целевого логита в этом фрагменте дает нам локальный фрагмент градиентного логита, P.
2. Накопление градиентов: Имея в наличии блок P, мы вычисляем частичные градиенты. Для dX мы выполняем скалярное произведение с блоками W^T; для dW мы умножаем на блоки X^T. Для безопасного агрегирования этих значений по всей партии мы используем функцию tl.atomic_add из пакета Triton.
Эта операция действует как потокобезопасное +=, гарантируя, что разные программы, обновляющие один и тот же градиент веса, не перезапишут друг друга.
Вот еще несколько подробностей о реализации:
- Изменение шага: При вычислении P . W_T нам фактически не нужно физически транспонировать массивную матрицу W в памяти. Вместо этого мы инвертируем формы и шаги в указателе блока W, чтобы считывать строки W как столбцы W^T. Это приводит к «бесплатной» транспонировке, которая экономит время и видеопамять.
- Точность численных значений: Стоит отметить, что хотя X и W могут быть в формате bfloat16, накопление dW и dX с помощью atomic_add обычно выполняется в формате float32, чтобы предотвратить накопление мелких ошибок округления в тысячах строк.
- Примечание о конкуренции: Хотя atomic_add необходим для dW (поскольку каждая программа обновляет одни и те же веса), dX является приватным для каждой программы, что означает отсутствие конкуренции между идентификаторами программ для конкретного тензора.
- Маскирование атомарного сложения: atomic_add не поддерживает указатели на блоки. Поэтому мы реализуем логику указателей и масок для dW явно.
На следующем рисунке представлен обратный проход для одной итерации внешнего цикла (т.е. одного блока вдоль V и всех блоков вдоль D):
Вот полный код для обратного прохода:
На этом завершается реализация нашего ядра! Полный код, включая ядро и скрипт для бенчмарка, доступен здесь.
Тест производительности памяти
Наконец, мы сравниваем наше ядро с базовым алгоритмом PyTorch, используя гиперпараметры, заимствованные из Llama3, и графический процессор A100. В частности, мы рассматриваем длину последовательности S=16 384, размер пакета B=1 и размерность встраивания D=4096; размер словаря установлен равным V=128 256.
Как и ожидалось, базовый алгоритм PyTorch выделяет огромный промежуточный тензор для хранения активаций, что приводит к пиковому использованию памяти в 36,02 ГБ . Для сравнения, наше ядро Triton снижает пиковое использование памяти на 84% , выделяя всего 5,04 ГБ при использовании D_BLOCK=64 и V_BLOCK=64!
Использование еще меньших размеров блоков позволило бы дополнительно увеличить объем памяти за счет повышения эффективности.
Атомные ограничения и масштабирование производства
В этой статье мы сосредоточились на технической и математической интуиции, лежащей в основе объединенных ядер линейной кросс-энтропии. Мы использовали атомарные операции, такие как tl.atomic_add, чтобы сделать код минимальным и читаемым. Однако, хотя наше ядро успешно сократило использование памяти на ошеломляющие 86% , ядро Triton значительно медленнее, чем нативный PyTorch.
К сожалению, те же атомарные операции, которые упрощают написание и понимание этого ядра, приводят к масштабной перегрузке сети, поскольку тысячи потоков пытаются одновременно изменить один и тот же адрес в памяти. Как правило, tl.atomic_add работает эффективно при низкой конкуренции за ресурсы. В нашей текущей реализации мы имеем:
- Высокая конкуренция: В случае градиента весов каждая программа в пакете (до 16 384 в нашем тесте) пытается одновременно обновить одни и те же блоки памяти. Аппаратное обеспечение должно сериализовать эти обновления, заставляя тысячи потоков ждать в очереди.
- Числовая неассоциативность: В компьютерах сложение чисел с плавающей запятой является неассоциативным . Ошибки округления могут накапливаться по-разному в зависимости от порядка операций, поэтому тесты на корректность могут пройти на T4, но провалиться на A100, поскольку последний имеет больше потоковых мультипроцессоров (SM), выполняющих больше параллельных, недетерминированных операций сложения.
Примечание о точности: на архитектурах Ampere и более новых форматах формат TF32 может дополнительно способствовать этим расхождениям. Для обеспечения строгой числовой четности следует установить allow_tf32=False или использовать типы с более высокой точностью на этапах накопления.
Путь к производству
Чтобы выйти за рамки этой образовательной реализации и перейти к готовому к использованию в производственной среде ядру (рекомендую ознакомиться с реализацией Liger-Kernel), можно реализовать несколько оптимизаций:
- Замена атомарных операций dX: Поскольку каждая программа «владеет» своей строкой X, мы можем использовать простое накопление регистров с последующим использованием tl.store, полностью исключая атомарные операции для входных градиентов.
- Специализированное ядро dW: Для оптимизации вычислений dW в производственных ядрах обычно используется другая стратегия сетки, где каждая программа обрабатывает блок W и итерирует по размерности пакета, накапливая градиенты локально перед одной глобальной записью.
- Микропакетная обработка: В более совершенных реализациях, таких как в библиотеке Liger-Kernel , последовательность обрабатывается блоками по N-мерному измерению, что делает масштабирование памяти постоянным по длине последовательности, а не линейным. Это позволяет использовать гораздо большие размеры пакетов при сниженных затратах памяти.
Заключение
На этом завершается наше подробное изучение объединенных линейных ядер кросс-энтропии. Спасибо, что дочитали до конца, и я надеюсь, что эта статья дала вам как интуитивное понимание, так и практическое представление, необходимые для дальнейшего развития этих идей и их изучения.
Если статья оказалась для вас полезной, пожалуйста, поделитесь ею; это действительно поможет поддержать время и усилия, затраченные на её создание. Как всегда, не стесняйтесь обращаться ко мне с вопросами, мыслями или идеями для дальнейших исследований. Если вы хотите поддержать мои независимые исследования и написание статей, вы также можете угостить меня кофе .
До новых встреч! 👋
Источники
- Представляем Meta Llama 3: самую функциональную из когда-либо доступных LLM-систем.
- LigerKernel (лекция)
- Реализация линейной кросс-энтропии LigerKernel
- Реализация Unsloth (только кросс-энтропия)
Источник: towardsdatascience.com
