레이어 건너뛰기

이 코드 베이스는 LayerSkip: 조기 종료 추론 및 자체 추측 디코딩 활성화의 구현입니다.

시작하기

복제 저장소:

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

설치 환경:

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

$ pip install -r requirements.txt

모델 액세스: 속도 향상을 관찰하려면 LayerSkip 레시피를 사용하여 훈련된 LLM에 액세스해야 합니다. 우리는 LayerSkip 레시피를 사용하여 지속적으로 사전 훈련된 다양한 Llama 모델의 HuggingFace에 대해 6개의 체크포인트를 제공합니다.
- 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

각 모델에 액세스하려면:

위의 모델에 해당하는 링크를 방문하여 HuggingFace 웹사이트에 귀하의 계정으로 로그인되어 있는지 확인하세요.
요청 양식을 작성하여 제출하세요. 승인에는 다소 시간이 걸릴 수 있으며 모델에 대한 권한이 부여되었음을 알리는 이메일 알림을 받게 됩니다.
사용자 액세스 토큰을 얻으려면 여기의 단계를 따르세요.
명령줄에서 huggingface-cli login 실행하면 3단계에서 얻은 토큰을 제공하라는 메시지가 표시됩니다.

해당 단계를 실행하면 LayerSkip 체크포인트를 실행하는 아래 명령이 작동합니다.

생성하다

일반 자동 회귀 디코딩을 사용하여 대화형 모드에서 모델 중 하나를 실행하려면 다음을 수행하세요.

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

속도 향상을 관찰하려면 자체 추측 디코딩을 사용하여 토큰을 생성하고 --exit_layer , 초안 단계에서 종료할 계층 및 --num_speculations , 초안 토큰 수를 지정해야 합니다.

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

팁:

--model HuggingFace 모델로 변경할 수 있지만 자체 추측 디코딩으로 속도 향상을 관찰하려면 HuggingFace에서 오픈 소스로 제공되는 것과 같은 LayerSkip 레시피를 사용하여 훈련된 모델을 사용하세요.
기본적으로 샘플링을 활성화합니다. --sample , --temperature , --top_p 및 --top_k 인수를 사용하여 샘플링 동작을 변경할 수 있습니다.
다양한 명령줄 인수에 대한 자세한 내용을 보려면 python generate.py --help 실행하세요.

기준

데이터세트를 벤치마킹하려면 다음 안내를 따르세요.

$ 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

팁:

--dataset 인수를 수정하여 다양한 작업을 지정할 수 있습니다.
- cnn_dm_summarization : CNN/DM 요약
- xsum_summarization : XSUM 요약
- cnn_dm_lm : CNN/DM 언어 모델링(기사의 처음 몇 단어가 주어지면 나머지 기사 생성)
- human_eval : HumanEval 코딩
기본적으로 작업은 0-shot으로 실행됩니다. --n_shot 인수를 지정하여 지정된 n -shot으로 변경할 수 있습니다.
기본적으로 우리는 샘플링을 활성화하지만, 논문에 보고된 결과는 샘플링 없이 탐욕스러운 디코딩이었습니다. --sample , --temperature , --top_p 및 --top_k 인수를 사용하여 샘플링 동작을 변경할 수 있습니다.
다양한 명령줄 인수에 대한 자세한 내용을 보려면 python benchmark.py --help 실행하세요.

평가하다

우리는 생성 스크립트를 Eleuther 언어 모델 평가 하네스와 통합하여 많은 작업을 가능하게 하고 생성된 텍스트를 적절하게 사후 처리합니다.

$ 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

팁:

추론적 디코딩을 사용하면 생성 작업(예: gsm8k 또는 cnn_dailymail )에서만 속도 향상을 얻을 수 있는 반면 분류 작업, 즉 객관식 질문 작업(예: piqa , social_iqa ) 또는 True/False 질문 작업(예: boolq )에서는 속도 향상을 얻을 수 있습니다. 속도 향상으로 이어지지 않습니다.
--tasks 인수를 사용하여 Eleuther Evaluation Harness에서 지원하는 작업 수를 임의대로 지정할 수 있습니다. 가능한 모든 작업 목록을 얻으려면 이 링크를 확인하세요.
generate.py 및 benchmark.py 스크립트와 유사하게 다양한 모델, 데이터세트 및 샘플링 매개변수를 지정할 수 있습니다.
다양한 명령줄 인수에 대한 자세한 내용을 보려면 python benchmark.py --help 실행하세요.

스위프

추론 하이퍼파라미터인 exit_layer 및 num_speculations 추론 중 속도 향상을 결정합니다.

exit_layer :
- 작다는 것은 더 빠르지만 덜 정확한 초안 단계를 의미합니다.
- 클수록 더 정확하지만 드래프트 단계가 느려짐을 의미합니다.
num_speculations :
- 작을수록 승인률이 높지만 검증 단계에서는 초안 단계보다 적게 상각됩니다.
- 크기가 클수록 검증 단계에서 초안 단계를 더 잘 상각하지만 승인률은 감소함을 의미합니다.

exit_layer 와 num_speculations 의 최적 조합은 모델, 데이터세트 및 샘플링 매개변수에 따라 변경될 수 있습니다. 따라서 우리는 다양한 exit_layer 및 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

--outpu_dir 인수에 지정된 디렉터리에 CSV 파일이 생성됩니다.

팁:

generate.py 및 benchmark.py 스크립트와 유사하게 다양한 모델, 데이터세트 및 샘플링 매개변수를 지정할 수 있습니다.
다양한 명령줄 인수에 대한 자세한 내용을 보려면 python sweep.py --help 실행하세요.

단정

자체 추측 디코딩 알고리즘에서 생성된 토큰이 올바른지 확인하기 위해 자동 회귀 디코딩의 출력을 자체 추측 디코딩과 비교하는 스크립트를 만들었습니다. 샘플링이 없는 경우에만(예: --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

도커 사용

docker를 사용하여 프로젝트를 설정하려면 DOCKER.md를 확인하세요.

기타 구현

또한 LayerSkip 추론의 다른 구현도 있습니다.

gpt-fast: gpt-fast는 간단하고 효율적인 pytorch 기본 변환기 텍스트 생성입니다. 우리는 torch.compile() , 양자화 및 텐서 병렬 처리와 같은 다른 최적화와 함께 합성할 수 있도록 gpt-fast 코드베이스에 LayerSkip을 구현했습니다.
기본 HuggingFace: 각 HuggingFace 모델의 모델 카드에는 가중치를 복제하지 않고 기본 모델의 이전 레이어를 복제하는 간단한 트릭을 사용하여 HuggingFace 추측 디코딩 기능을 활용하는 간단한 코드 조각을 제공했습니다. 이 구현은 간단하고 다른 기능을 구현하거나 다른 라이브러리를 가져올 필요가 없지만 초안 단계와 검증 단계 간에 KV 캐시 또는 실행을 공유하지 않습니다.

훈련

우리의 교육 구현은 진행 중입니다. 자세한 내용과 토론은 이 풀 요청을 확인하세요.

특허

LayerSkip은 CC-by-NC 라이선스에 따라 라이선스가 부여됩니다. 최상위 디렉터리의 LICENSE 파일을 참조하세요.

기여

LayerSkip에 대한 기여를 환영합니다. 기여에 관심이 있다면 이 문서를 참조하세요.

소환

연구에 LayerSkip을 사용하는 경우 다음 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 " ,
}