КПД

Liger Kernel: Efficient Triton Kernels for LLM Training
[Техрепорт]
[Репозиторий]

Ребята из LinkedIn написали кернелы для разных операций в LLMках на Тритоне (когда нибудь я научусь писать на нем, а не про него), которые ускоряют процедуру обучения и снижают расход памяти по сравнению с ванильной торчовой реализации.

Метод

В частности, Liger Kernel предлагает следующее:

⭐️ Зафьюженные RMSNorm и LayerNorm на прямом и обратном проходе
⭐️ Зафьюженные SwiGLU и GeGLU на прямом и обратном проходе
⭐️ Оптимизированный RoPE
⭐️ И самое интересное - оптимизация вычисления кросс-энтропии

По поводу последнего:

Словари нынче у моделек перевалили за 100к, и на сколь-либо длинных последовательностях матрица логитов будет весить десятки гигов. Потому авторы реализовали прямой проход, вычисление кросс-энтропии и обратный проход в одном kernel-е без необходимости материализации одновременно матрицы логитов и градиентов по логитам. Кроме того, считать логиты можно не разом для всей последовательности, а чанками.

Бенчмарки

Замеряют на A100.

Предложенные нормализации дают хорошее ускорение по сравнению с торчом (не хватает сравнения с Apex). RoPE прямо знатно ускорился. SwiGLU и GeGLU по скорости такие же, как в торче, но снижают расход памяти в 1.6 раз.

В end-2-end сценарии гоняют finetune на 4A100 на Alpaca на Llama-3-8B, Qwen-2-7b, Gemma-1-7b. Скорость обучения возрастает от 10 до 40 %, и пиковый расход памяти уменьшается в среднем на 50%.

Метод также валидируют в связке с Medusa, методом где LLM учится предсказывать несколько следующих токенов за раз, и на каждый n+1 токен своя обучаемая голова. Liger Kernel снижает заметно расход памяти как с замороженной, так и обучаемой тушкой трансформера, благодаря трюкам с вычислениями логитов.

Использование

Liger Kernel можно вызвать по щелчку пальца через AutoLigerKernelForCausalLM или патчинг модели из 🤗. А можно импортировать модули и из них собрать модельку.

Liger Kernel уже интегрирован в transformers, TRL и axolotl.

👍10

www.tgoop.com/quant_prune_distill/354

2.25K viewsNov 8, 2024 at 20:49

Liger Kernel: Efficient Triton Kernels for LLM Training
[Техрепорт]
[Репозиторий]

Ребята из LinkedIn написали кернелы для разных операций в LLMках на Тритоне (когда нибудь я научусь писать на нем, а не про него), которые ускоряют процедуру обучения и снижают расход памяти по сравнению с ванильной торчовой реализации.

Метод

В частности, Liger Kernel предлагает следующее:

⭐️ Зафьюженные RMSNorm и LayerNorm на прямом и обратном проходе
⭐️ Зафьюженные SwiGLU и GeGLU на прямом и обратном проходе
⭐️ Оптимизированный RoPE
⭐️ И самое интересное - оптимизация вычисления кросс-энтропии

По поводу последнего:

Словари нынче у моделек перевалили за 100к, и на сколь-либо длинных последовательностях матрица логитов будет весить десятки гигов. Потому авторы реализовали прямой проход, вычисление кросс-энтропии и обратный проход в одном kernel-е без необходимости материализации одновременно матрицы логитов и градиентов по логитам. Кроме того, считать логиты можно не разом для всей последовательности, а чанками.

Бенчмарки

Замеряют на A100.

Предложенные нормализации дают хорошее ускорение по сравнению с торчом (не хватает сравнения с Apex). RoPE прямо знатно ускорился. SwiGLU и GeGLU по скорости такие же, как в торче, но снижают расход памяти в 1.6 раз.

В end-2-end сценарии гоняют finetune на 4A100 на Alpaca на Llama-3-8B, Qwen-2-7b, Gemma-1-7b. Скорость обучения возрастает от 10 до 40 %, и пиковый расход памяти уменьшается в среднем на 50%.

Метод также валидируют в связке с Medusa, методом где LLM учится предсказывать несколько следующих токенов за раз, и на каждый n+1 токен своя обучаемая голова. Liger Kernel снижает заметно расход памяти как с замороженной, так и обучаемой тушкой трансформера, благодаря трюкам с вычислениями логитов.

Использование

Liger Kernel можно вызвать по щелчку пальца через AutoLigerKernelForCausalLM или патчинг модели из 🤗. А можно импортировать модули и из них собрать модельку.

Liger Kernel уже интегрирован в transformers, TRL и axolotl.

BY КПД