Исследователи из лаборатории Zlab Принстонского университета представили LLM-Pruning Collection. Это репозиторий на основе фреймворка JAX, который консолидирует основные алгоритмы прунинга (сжатия) больших языковых моделей в единую воспроизводимую среду. Цель проекта — упростить сравнение методов прунинга на уровне блоков, слоев и весов с использованием согласованного стека для обучения и оценки на GPU и TPU.
Репозиторий содержит реализации нескольких ключевых методов сжатия. В их числе Minitron от NVIDIA, ShortGPT, Wanda, SparseGPT, Magnitude, Sheared Llama и LLM-Pruner. Для обучения предоставлена интеграция с FMS-FSDP для GPU и MaxText для TPU. Оценка моделей проводится с помощью JAX-совместимых скриптов на основе lm-eval-harness, что дает ускорение в 2-4 раза при использовании MaxText.
Актуальность проекта обусловлена растущей потребностью в оптимизации огромных LLM для их более эффективного развертывания. Разные методы прунинга предлагают различные подходы: от посттренировочного удаления весов (Wanda, SparseGPT) до структурного удаления целых слоев или внимания (ShortGPT, LLM-Pruner). Наличие единой платформы позволяет объективно сравнивать их эффективность, что ускоряет исследования в области эффективных AI.
LLM-Pruning Collection, распространяемый под лицензией Apache-2.0, представляет собой практический инструмент для исследователей и инженеров, работающих над уменьшением размера и стоимости эксплуатации современных языковых моделей без значительной потери качества.
