[Техрепорт]
[Репозиторий]
Ребята из 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.