Ай Дайджест - категория kernel

SageAttention2: Технический отчет

Хотя квантование для линейных слоев широко используется, его применение для ускорения процесса внимания остается ограниченным. SageAttention использует умножение матриц на 8 бит, умножение матриц на 16 бит с аккумулятором на 16 бит и методы повышения точности, реализуя точное и ускоренное на 2 раза ядро по сравнению с FlashAttention2. Для дальнейшего улучшения эффективности вычислений внимания при сохранении точности мы предлагаем SageAttention2, который использует значительно более быстрое умножение матриц на 4 бита (Matmul) вместе с дополнительными методами повышения точности. Во-первых, мы предлагаем квантование матриц (Q, K) до INT4 на уровне warp и квантование матриц (widetilde P, V) до FP8. Во-вторых, мы предлагаем метод для сглаживания Q и V, повышая точность внимания с INT4 QK и FP8 PV. В-третьих, мы анализируем точность квантования по временным шагам и слоям, затем предлагаем адаптивный метод квантования для обеспечения метрик от начала до конца по различным моделям. Операции в секунду (OPS) SageAttention2 превосходят FlashAttention2 и xformers примерно на 3 и 5 раз на RTX4090 соответственно. Всеобъемлющие эксперименты подтверждают, что наш подход приводит к незначительным потерям метрик от начала до конца на разнообразных моделях, включая те, что используются для обработки больших языковых моделей, генерации изображений и видео. Код доступен по адресу https://github.com/thu-ml/SageAttention.

2024-11-21attention adaptive kernel

Эффективное использование памяти при обучении больших языковых моделей

По мере того как языковые модели становятся всё больше, растут и их словари. Это привело к непропорциональному увеличению объема памяти, используемого LLMs во время обучения, особенно в одном слое: при вычислении кросс-энтропии в функции потерь. Кросс-энтропия создает матрицу логитов с элементами для каждой пары входных токенов и элементов словаря и, для небольших моделей, потребляет на порядок больше памяти, чем остальная часть LLM вместе взятая. Мы предлагаем метод Cut Cross-Entropy (CCE), который вычисляет потери кросс-энтропии без материализации логитов для всех токенов в глобальной памяти. Вместо этого, CCE вычисляет логит только для правильного токена и оценивает log-sum-exp по всем логитам на лету. Мы реализовали специальное ядро, которое выполняет матричные умножения и редукцию log-sum-exp по словарю во флэш-памяти, делая потребление глобальной памяти для вычислений кросс-энтропии незначительным. Это имеет драматический эффект. Например, в случае модели Gemma 2 (2B), CCE снижает объем памяти, необходимый для вычисления потерь, с 24 ГБ до 1 МБ, а общее потребление памяти при обучении головы классификатора - с 28 ГБ до 1 ГБ. Для улучшения пропускной способности CCE мы используем врожденную разреженность softmax и предлагаем пропускать элементы вычисления градиента, которые имеют пренебрежимо малый вклад в градиент (то есть ниже численной точности). Эксперименты показывают, что драматическое сокращение потребления памяти достигается без ущерба для скорости обучения или сходимости.

2024-11-15kernel logits cross-entropy