LLM-Pruning Collection: JAX Репозиторий для Сжатия Моделей

Обзор

Исследователи из Zlab Princeton выпустили LLM-Pruning Collection — репозиторий на основе JAX, который объединяет основные алгоритмы обрезки для больших языковых моделей в единую, воспроизводимую структуру. Его цель — облегчить сравнение методов обрезки по блокам, слоям и уровням весов в рамках единого стека обучения и оценки как на GPU, так и на TPU.

Что содержит LLM-Pruning Collection?

Этот репозиторий организован в три основные директории:

pruning: Содержит реализации нескольких методов обрезки: Minitron, ShortGPT, Wanda, SparseGPT, Magnitude, Sheared Llama и LLM-Pruner.
training: Обеспечивает интеграцию с FMS-FSDP для обучения на GPU и MaxText для TPU.
eval: Открывает совместимые с JAX скрипты оценки, основанные на lm-eval-harness, с поддержкой MaxText, что обеспечивает прирост скорости примерно в 2-4 раза.

Методы Oбрезки

LLM-Pruning Collection охватывает несколько семейств алгоритмов обрезки с различными уровнями гранулярности:

Minitron

Minitron — это метод обрезки и дистилляции, разработанный NVIDIA, который сжимает Llama 3.1 8B и Mistral NeMo 12B до 4B и 8B, сохраняя производительность. Он исследует обрезку по глубине и совместную обрезку по ширине скрытых размеров, внимания и MLP, за которой следует дистилляция. В каталоге pruning/minitron находятся скрипты, такие как prune_llama3.1-8b.sh, позволяющие выполнять обрезку стиля Minitron на Llama 3.1 8B.

ShortGPT

ShortGPT определяет избыточные слои трансформатора, используя метрику Block Influence, измеряющую вклад каждого слоя, что позволяет удалять слои с низким вкладом. Он превосходит предыдущие методы обрезки в заданиях с выбором и генерацией. В коллекции ShortGPT реализован через папку Minitron со специальным скриптом prune_llama2-7b.sh.

Wanda, SparseGPT, Magnitude

Wanda — это метод обрезки после обучения, который оценивает веса по произведению величины веса и соответствующей входной активации, обрезая самые маленькие оценки без необходимости повторного обучения. SparseGPT использует второй порядок, чтобы обрезать большие модели стиля GPT с высокими коэффициентами разреженности, в то время как обрезка по величине является классическим базовым методом, который удаляет веса с малыми абсолютными значениями. В LLM-Pruning Collection все три доступны в каталоге pruning/wanda.

Sheared Llama

Sheared LLaMA применяет структурированную обрезку, обучая маски для слоев, голов внимания и скрытых измерений, в то время как реконструируемая архитектура проходит повторное обучение. Этот каталог предлагает модели на нескольких масштабах, включая 2.7B и 1.3B.

LLM-Pruner

LLM-Pruner предназначен для структурной обрезки больших языковых моделей, используя градиентные оценки важности для удаления некритических структур и используя короткий этап LoRA для восстановления производительности. Он включает скрипты для LLaMA, LLaMA 2 и Llama 3.1 8B в каталоге pruning/LLM-Pruner.

Основные моменты

LLM-Pruning Collection — это репозиторий с лицензией Apache-2.0 от zlab-princeton, который объединяет современные методы обрезки LLM, поддерживая общие процессы обрезки, обучения и оценки для GPU и TPU.
Он включает методы обрезки по блокам, слоям и уровням весов, включая Minitron, ShortGPT, Wanda, SparseGPT и другие, с шаблонами для моделей семейства Llama.
Обучающая платформа интегрирует FMS-FSDP на GPU и MaxText на TPU, предлагая совместимые с JAX скрипты оценки с значительным ускорением.
Репозиторий воспроизводит ключевые результаты из предыдущих исследований по обрезке, публикуя соседние “статья против воспроизведена” таблицы для проверки.

Посмотрите GitHub Repo для получения дополнительных сведений.