Los investigadores de Zlab Princeton han lanzado LLM-Pruning Collection, un repositorio basado en JAX que consolida los principales algoritmos de poda para modelos de lenguaje grandes en un marco único y reproducible. Tiene como objetivo un objetivo concreto: facilitar la comparación de métodos de poda a nivel de bloque, nivel de capa y nivel de peso bajo una pila de capacitación y evaluación consistente tanto en GPU como en TPU.
¿Qué contiene la colección LLM-Pruning?
Se describe como un repositorio basado en JAX para poda LLM. Está organizado en tres directorios principales:
La poda contiene implementaciones para varios métodos de poda: Minitron, ShortGPT, Wanda, SparseGPT, Magnitude, Sheared Llama y LLM-Pruner. La capacitación proporciona integración con FMS-FSDP para capacitación de GPU y MaxText para capacitación de TPU. eval expone scripts de evaluación compatibles con JAX creados alrededor de lm-eval-harness, con soporte basado en aceleración para MaxText que proporciona una aceleración de 2 a 4 veces.
Métodos de poda cubiertos
LLM-Pruning Collection abarca varias familias de algoritmos de poda con diferentes niveles de granularidad:
Minitrón
Minitron es una práctica receta de poda y destilación desarrollada por NVIDIA que comprime Llama 3.1 8B y Mistral NeMo 12B a 4B y 8B preservando el rendimiento. Explora la poda en profundidad y la poda del ancho de las juntas de tamaños ocultos, atención y MLP, seguidas de destilación.
En LLM-Pruning Collection, la carpeta pruning/minitron proporciona scripts como prune_llama3.1-8b.sh que ejecutan poda estilo Minitron en Llama 3.1 8B.
GPT corto
ShortGPT se basa en la observación de que muchas capas de Transformer son redundantes. El método define Block Influence, una métrica que mide la contribución de cada capa y luego elimina las capas de baja influencia mediante la eliminación directa de capas. Los experimentos muestran que ShortGPT supera a los métodos de poda anteriores para tareas generativas y de opción múltiple.
En la colección, ShortGPT se implementa a través de la carpeta Minitron con un script dedicado prune_llama2-7b.sh.
Wanda, SparseGPT, Magnitud
Wanda es un método de poda posterior al entrenamiento que califica los pesos según el producto de la magnitud del peso y la activación de entrada correspondiente por salida. Elimina las puntuaciones más pequeñas, no requiere reentrenamiento e induce una escasez que funciona bien incluso en una escala de mil millones de parámetros.
SparseGPT es otro método posterior al entrenamiento que utiliza un paso de reconstrucción inspirado en segundo orden para podar modelos grandes de estilo GPT con altos índices de dispersión. La poda de magnitud es la línea de base clásica que elimina pesos con valor absoluto pequeño.
En LLM-Pruning Collection, los tres viven bajo pruning/wanda con una ruta de instalación compartida. El archivo README incluye una tabla densa de resultados de Llama 2 7B que compara Wanda, SparseGPT y Magnitude en BoolQ, RTE, HellaSwag, Winogrande, ARC E, ARC C y OBQA, bajo patrones de dispersión estructurados y no estructurados, como 4:8 y 2:4.
Llama esquilada
Sheared LLaMA es un método de poda estructurado que aprende máscaras para capas, cabezas de atención y dimensiones ocultas y luego vuelve a entrenar la arquitectura podada. La versión original proporciona modelos en múltiples escalas, incluidas 2.7B y 1.3B.
El directorio pruning/llmshearing en LLM-Pruning Collection integra esta receta. Utiliza un subconjunto de RedPajama para la calibración, al que se accede a través de Hugging Face, y scripts de ayuda para convertir entre los formatos Hugging Face y MosaicML Composer.
LLM-Podadora
LLM-Pruner es un marco para la poda estructural de modelos de lenguaje grandes. Elimina estructuras acopladas no críticas, como cabezas de atención o canales MLP, utilizando puntuaciones de importancia basadas en gradientes y luego recupera el rendimiento con una breve etapa de ajuste LoRA que utiliza alrededor de 50.000 muestras. La colección incluye LLM-Pruner bajo poda/LLM-Pruner con scripts para LLaMA, LLaMA 2 y Llama 3.1 8B.
Conclusiones clave
LLM-Pruning Collection es un repositorio Apache-2.0 basado en JAX de zlab-princeton que unifica los métodos modernos de poda de LLM con procesos compartidos de poda, capacitación y evaluación para GPU y TPU. El código base implementa enfoques de poda de bloques, capas y niveles de peso, incluidos Minitron, ShortGPT, Wanda, SparseGPT, Sheared LLaMA, Magnitude pruning y LLM-Pruner, con scripts de métodos específicos para los modelos de la familia Llama. La capacitación integra FMS-FSDP en GPU y MaxText en TPU con scripts de evaluación compatibles con JAX creados en lm-eval-harness, lo que brinda una evaluación aproximadamente de 2 a 4 veces más rápida para los puntos de control de MaxText mediante aceleración. El repositorio reproduce resultados clave de trabajos de poda anteriores, publicando tablas en paralelo “en papel versus reproducidas” para métodos como Wanda, SparseGPT, Sheared LLaMA y LLM-Pruner para que los ingenieros puedan verificar sus ejecuciones con respecto a líneas de base conocidas.
Consulte el repositorio de GitHub. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 100.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
Shobha es un analista de datos con una trayectoria comprobada en el desarrollo de soluciones innovadoras de aprendizaje automático que impulsan el valor empresarial.