Saltar capa

Esta base de código es la implementación de LayerSkip: Habilitación de inferencia de salida temprana y decodificación autoespeculativa.

Empezando

Clonar repositorio:

$ git clone [email protected]:facebookresearch/LayerSkip.git
$ cd LayerSkip

Entorno de configuración:

$ conda create --name layer_skip python=3.10
$ conda activate layer_skip

$ pip install -r requirements.txt

Acceda a los modelos: para observar la aceleración, debe acceder a los LLM que hayan sido entrenados utilizando la receta LayerSkip. Proporcionamos 6 puntos de control en HuggingFace de diferentes modelos de Llama entrenados previamente usando la receta LayerSkip:
- facebook/layerskip-llama2-7B
- facebook/layerskip-llama2-13B
- facebook/layerskip-codellama-7B
- facebook/layerskip-codellama-34B
- facebook/layerskip-llama3-8B
- facebook/layerskip-llama3.2-1B

Para acceder a cada modelo:

Visite el enlace correspondiente al modelo arriba, asegúrese de haber iniciado sesión en el sitio web de HuggingFace con su cuenta.
Complete el formulario de solicitud y envíelo. La aprobación puede tardar un poco y debería recibir una notificación por correo electrónico para informarle que se ha concedido el permiso al modelo.
Siga los pasos aquí para obtener un token de acceso de usuario.
En la línea de comando, ejecute huggingface-cli login y se le pedirá que proporcione el token que obtuvo en el Paso 3.

Una vez que ejecute esos pasos, los siguientes comandos para ejecutar los puntos de control de LayerSkip deberían funcionar.

Generar

Para ejecutar uno de nuestros modelos en modo interactivo utilizando decodificación autorregresiva normal:

$ torchrun generate.py --model facebook/layerskip-llama2-7B 
    --sample True 
    --max_steps 512

Para observar la aceleración, debe utilizar la decodificación autoespeculativa para generar tokens y especificar --exit_layer , la capa por la que saldrá la etapa de borrador, y --num_speculations , la cantidad de tokens de borrador:

$ torchrun generate.py --model facebook/layerskip-llama2-7B 
    --sample True 
    --max_steps 512 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6

Consejos:

Puede cambiar --model a cualquier modelo de HuggingFace, pero para observar la aceleración con la decodificación autoespeculativa, use un modelo entrenado usando la receta LayerSkip, como los que tenemos de código abierto en HuggingFace.
Por defecto habilitamos el muestreo. Puede cambiar el comportamiento de muestreo utilizando los argumentos --sample , --temperature , --top_p y --top_k .
Puede ejecutar python generate.py --help para obtener detalles sobre diferentes argumentos de la línea de comandos.

Punto de referencia

Para comparar un conjunto de datos:

$ torchrun benchmark.py --model facebook/layerskip-llama2-7B 
    --dataset cnn_dm_summarization 
    --num_samples 100 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6 
    --output_dir ./logs

Consejos:

Puede especificar diferentes tareas modificando el argumento --dataset :
- cnn_dm_summarization : Resumen CNN/DM
- xsum_summarization : Resumen XSUM
- cnn_dm_lm : Modelado de lenguaje CNN/DM (dadas las primeras palabras de un artículo, genera el artículo restante)
- human_eval : Codificación HumanEval
De forma predeterminada, las tareas se ejecutan como 0-shot. Puede cambiar a cualquier n -shot especificado especificando el argumento --n_shot .
De forma predeterminada, habilitamos el muestreo, mientras que los resultados informados en el artículo fueron una decodificación codiciosa sin muestreo. Puede cambiar el comportamiento de muestreo utilizando los argumentos --sample , --temperature , --top_p y --top_k .
Puede ejecutar python benchmark.py --help para obtener detalles sobre diferentes argumentos de la línea de comandos.

Evaluar

Hemos integrado nuestros scripts de generación con Eleuther Language Model Assessment Harness para permitir una gran cantidad de tareas y postprocesar adecuadamente el texto generado.

$ torchrun eval.py --model facebook/layerskip-llama2-7B 
    --tasks gsm8k 
    --limit 10 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6 
    --output_dir ./logs

Consejos:

Tenga en cuenta que con la decodificación especulativa solo podemos obtener aceleraciones a partir de tareas de generación (por ejemplo, gsm8k o cnn_dailymail ), mientras que las tareas de clasificación, es decir, tareas de preguntas de opción múltiple (por ejemplo, piqa , social_iqa ) o tareas de preguntas de Verdadero/Falso (por ejemplo, boolq ) no conduce a una aceleración.
Puede especificar un número arbitrario de tareas admitidas por Eleuther Assessment Harness utilizando el argumento --tasks . Para obtener una lista de todas las tareas posibles, consulte este enlace.
De manera similar a los scripts generate.py y benchmark.py , puede especificar diferentes modelos, conjuntos de datos y parámetros de muestreo.
Puede ejecutar python benchmark.py --help para obtener detalles sobre diferentes argumentos de la línea de comandos.

Barrer

Nuestros hiperparámetros de inferencia, exit_layer y num_speculations determinan la aceleración durante la inferencia:

exit_layer :
- Más pequeño significa una etapa de borrador más rápida pero menos precisa.
- más grande significa una etapa de borrador más precisa pero más lenta
num_speculations :
- Más pequeño significa una tasa de aceptación más alta, pero la etapa de verificación amortizará menos que la etapa de borrador.
- Más grande significa que la etapa de verificación amortizará mejor la etapa de borrador, pero la tasa de aceptación disminuye

La combinación óptima de exit_layer y num_speculations puede cambiar con el modelo, el conjunto de datos y los parámetros de muestreo. Por lo tanto, proporcionamos un script para recorrer una cuadrícula de diferentes exit_layer y num_speculations :

$ torchrun sweep.py --model facebook/layerskip-llama2-7B 
    --dataset human_eval 
    --generation_strategy self_speculative 
    --num_samples 150 
    --max_steps 256 
    --output_dir ./logs/ 
    --sample False

Esto creará un archivo CSV en el directorio especificado en el argumento --outpu_dir .

Consejos:

De manera similar a los scripts generate.py y benchmark.py , puede especificar diferentes modelos, conjuntos de datos y parámetros de muestreo.
Puede ejecutar python sweep.py --help para obtener detalles sobre diferentes argumentos de la línea de comandos.

Exactitud

Para verificar que los tokens generados por nuestro algoritmo de decodificación autoespeculativa sean correctos, hemos creado un script para comparar los resultados de la decodificación autorregresiva con la decodificación autoespeculativa. Tenga en cuenta que las salidas solo podemos garantizar la equivalencia cuando no hay muestreo (es decir, --sample False ):

$ torchrun correctness.py --model facebook/layerskip-llama2-7B 
    --dataset human_eval 
    --generation_strategy self_speculative 
    --num_speculations 6 
    --exit_layer 4 
    --num_samples 10 
    --sample False 
    --output_dir ./logs

Usando ventana acoplable

Por favor revise DOCKER.md para configurar el proyecto usando Docker

Otras implementaciones

También tenemos otras implementaciones de inferencia LayerSkip:

gpt-fast: gpt-fast es una generación de texto transformador nativa de pytorch simple y eficiente. Hemos implementado LayerSkip en el código base de gpt-fast para permitir combinarlo con otras optimizaciones como torch.compile() , cuantificación y paralelismo tensorial.
HuggingFace nativo: en la tarjeta de modelo de cada uno de nuestros modelos HuggingFace, proporcionamos fragmentos de código simples que aprovechan las capacidades de decodificación especulativa de HuggingFace utilizando un truco simple para clonar las capas anteriores del modelo principal sin clonar sus pesos. Aunque esta implementación es simple y no requiere implementar otras funciones ni importar otras bibliotecas, no comparte el caché KV ni la ejecución entre las etapas de borrador y verificación.

Capacitación

Nuestra implementación de capacitación es un trabajo en progreso. Puede consultar esta solicitud de extracción para obtener detalles y discusiones.

Licencia

LayerSkip tiene una licencia CC-by-NC. Consulte el archivo LICENCIA en el directorio de nivel superior.

Contribuyendo

Agradecemos las contribuciones a LayerSkip. Si está interesado en contribuir, consulte este documento.

Citación

Si utiliza LayerSkip en su investigación, utilice la siguiente entrada BibTex:

 @misc { layerskip ,
    title = { LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding } ,
    author = { Mostafa Elhoushi and Akshat Shrivastava and Diana Liskovich and Basil Hosmer and Bram Wasti and Liangzhen Lai and Anas Mahmoud and Bilge Acun and Saurabh Agarwal and Ahmed Roman and Ahmed A Aly and Beidi Chen and Carole-Jean Wu } ,
    booktitle = " Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) " ,
    month = aug,
    year = " 2024 " ,
    address = " Bangkok, Thailand " ,
    publisher = " Association for Computational Linguistics " ,
    url = " https://aclanthology.org/2024.acl-long.681 " ,
    doi = " 10.18653/v1/2024.acl-long.681 " ,
    pages = " 12622--12642 " ,
}