LayerSkip

Esta base de código é a implementação de LayerSkip: Habilitando Inferência de Saída Antecipada e Decodificação Autoespeculativa.

Começando

Clonar repositório:

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

Ambiente de configuração:

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

$ pip install -r requirements.txt

Modelos de acesso: Para observar a aceleração, você precisa acessar LLMs que foram treinados usando a receita LayerSkip. Fornecemos 6 pontos de verificação no HuggingFace de diferentes modelos de Llama continuamente pré-treinados usando a receita 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 acessar cada modelo:

Visite o link correspondente do modelo acima, certifique-se de estar logado no site HuggingFace com sua conta.
Preencha o formulário de solicitação e envie-o. A aprovação pode demorar um pouco e você deverá receber uma notificação por e-mail informando que a permissão para o modelo foi concedida.
Siga as etapas aqui para obter um token de acesso do usuário.
Na linha de comando, execute huggingface-cli login e você será solicitado a fornecer o token obtido na Etapa 3.

Depois de executar essas etapas, os comandos abaixo para executar os pontos de verificação do LayerSkip deverão funcionar.

Gerar

Para executar um de nossos modelos em modo interativo usando decodificação autorregressiva regular:

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

Para observar a aceleração, você precisa usar a decodificação autoespeculativa para gerar tokens e especificar --exit_layer , a camada do estágio de rascunho para sair, e --num_speculations , o número de tokens de rascunho:

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

Pontas:

Você pode alterar --model para qualquer modelo HuggingFace, mas para observar a aceleração com a decodificação autoespeculativa, use um modelo treinado usando a receita LayerSkip, como aqueles que temos código aberto no HuggingFace.
Por padrão, habilitamos a amostragem. Você pode alterar o comportamento da amostragem usando os argumentos --sample , --temperature , --top_p e --top_k .
Você pode executar python generate.py --help para obter detalhes sobre diferentes argumentos de linha de comando.

Referência

Para comparar um conjunto de dados:

$ 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

Pontas:

Você pode especificar tarefas diferentes modificando o argumento --dataset :
- cnn_dm_summarization : Resumo CNN/DM
- xsum_summarization : Resumo XSUM
- cnn_dm_lm : Modelagem de linguagem CNN/DM (dadas as primeiras palavras de um artigo, gere o artigo restante)
- human_eval : codificação HumanEval
Por padrão, as tarefas são executadas como 0-shot. Você pode mudar para qualquer n -shot especificado especificando o argumento --n_shot .
Por padrão, habilitamos a amostragem, enquanto os resultados relatados no artigo foram uma decodificação gananciosa sem amostragem. Você pode alterar o comportamento da amostragem usando os argumentos --sample , --temperature , --top_p e --top_k .
Você pode executar python benchmark.py --help para obter detalhes sobre diferentes argumentos de linha de comando.

Avaliar

Integramos nossos scripts de geração com o Eleuther Language Model Evaluation Harness para permitir um grande número de tarefas e pós-processar adequadamente o texto gerado.

$ 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

Pontas:

Observe que com a decodificação especulativa só podemos obter acelerações de tarefas de geração (por exemplo, gsm8k ou cnn_dailymail ), enquanto tarefas de classificação, ou seja, tarefas de perguntas de múltipla escolha (por exemplo, piqa , social_iqa ) ou tarefas de perguntas Verdadeiro/Falso (por exemplo, boolq ) irão não leva à aceleração.
Você pode especificar um número arbitrário de tarefas suportadas pelo Eleuther Evaluation Harness usando o argumento --tasks . Para obter uma lista de todas as tarefas possíveis, verifique este link.
Semelhante aos scripts generate.py e benchmark.py , você pode especificar diferentes modelos, conjuntos de dados e parâmetros de amostragem
Você pode executar python benchmark.py --help para obter detalhes sobre diferentes argumentos de linha de comando.

Varrer

Nossos hiperparâmetros de inferência, exit_layer e num_speculations determinam a aceleração durante a inferência:

exit_layer :
- menor significa um estágio de rascunho mais rápido, mas menos preciso
- maior significa um estágio de rascunho mais preciso, mas mais lento
num_speculations :
- menor significa maior taxa de aceitação, mas a fase de verificação amortizará menos a fase de rascunho
- maior significa que o estágio de verificação amortizará melhor o estágio de rascunho, mas a taxa de aceitação diminui

A combinação ideal de exit_layer e num_speculations pode mudar com o modelo, conjunto de dados e parâmetros de amostragem. Portanto, fornecemos um script para varrer uma grade de diferentes exit_layer e 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

Isso criará um arquivo CSV no diretório especificado no argumento --outpu_dir .

Pontas:

Semelhante aos scripts generate.py e benchmark.py , você pode especificar diferentes modelos, conjuntos de dados e parâmetros de amostragem
Você pode executar python sweep.py --help para obter detalhes sobre diferentes argumentos de linha de comando.

Correção

Para verificar se os tokens gerados pelo nosso algoritmo de decodificação autoespeculativa estão corretos, criamos um script para comparar as saídas da decodificação autorregressiva com a decodificação autoespeculativa. Observe que as saídas só podemos garantir equivalência quando não há amostragem (ou seja, --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 Docker

Verifique DOCKER.md para configurar o projeto usando docker

Outras Implementações

Também temos outras implementações de inferência do LayerSkip:

gpt-fast: gpt-fast é uma geração de texto de transformador nativo de pytorch simples e eficiente. Implementamos LayerSkip na base de código gpt-fast para permitir combiná-lo com outras otimizações, como torch.compile() , quantização e paralelismo de tensor.
HuggingFace nativo: no cartão de modelo de cada um de nossos modelos HuggingFace, fornecemos trechos de código simples que aproveitam os recursos de decodificação especulativa do HuggingFace usando um truque simples para clonar as camadas anteriores do modelo principal sem clonar seus pesos. Embora esta implementação seja simples e não exija a implementação de outras funções ou a importação de outras bibliotecas, ela não compartilha o cache KV ou a execução entre as etapas de rascunho e verificação.

Treinamento

Nossa implementação de treinamento está em andamento. Você pode verificar esta solicitação pull para obter detalhes e discussões.

Licença

LayerSkip é licenciado sob licença CC-by-NC. Consulte o arquivo LICENSE no diretório de nível superior.

Contribuindo

Aceitamos contribuições para o LayerSkip. Se você estiver interessado em contribuir, consulte este documento.

Citação

Se você usar LayerSkip em sua pesquisa, use a seguinte 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 " ,
}