범용 보조 생성: 모든 어시스턴트 모델을 활용한 빠른 디코딩
TL;DR
Hugging Face와 Intel Labs는 Universal Assisted Generation (UAG)이라는 기술을 개발했습니다. 이 기술은 대상 모델과 토크나이저가 동일하지 않더라도 작은 언어 모델을 어시스턴트로 사용할 수 있게 해 주어, 기존에 호환 가능한 소형 변형이 없던 모델들의 추론 속도를 1.5배에서 2.0배까지 높일 수 있습니다.
보조 생성에서 토크나이저 의존성 문제
보조 생성(또는 speculative decoding)은 작은 어시스턴트 모델이 토큰 시퀀스를 예측하고, 큰 대상 모델이 이를 한 번의 포워드 패스로 검증함으로써 LLM 추론을 가속화합니다. 이는 지연 시간을 크게 줄이지만, 표준 보조 생성은 대상 모델과 어시스턴트 모델이 동일한 토크나이저(즉, 같은 모델 패밀리) 를 공유해야 한다는 제약이 있습니다.
이 의존성은 많은 고성능 모델이 충분히 작은 버전(대상 모델보다 보통 50~100배 작음)을 갖고 있지 않아 의미 있는 속도 향상을 제공하지 못한다는 병목을 만들게 됩니다. 예를 들어 gemma-2-9b는 2B 변형만 존재하고, CodeLlama-13b는 이 목적을 위한 더 작은 버전이 전혀 없습니다.
범용 보조 생성 (UAG) 메커니즘
범용 보조 생성은 두 방향 토크나이저 변환 과정을 구현함으로써 토크나이저 제약을 없앱니다. 이를 통해 vicuna-68m 같은 초소형 모델을 gemma-2-9b와 같은 대상 모델에 연결해 가속화할 수 있습니다.
양방향 토크나이저 변환
- Assistant to Target: 어시스턴트 모델이 토큰 시퀀스를 생성하면, 해당 토큰을 텍스트로 변환한 뒤 대상 모델의 토크나이저로 다시 토크나이징합니다.
- Target to Assistant: 대상 모델이 토큰을 검증한 후, 결과 토큰을 어시스턴트 모델의 토큰 형식으로 다시 변환하고 다음 반복을 위해 어시스턴트 모델의 컨텍스트에 추가합니다.
어휘 차이 처리
재인코딩 정확성을 보장하기 위해 UAG는 새로 생성된 시퀀스 앞에 이전 토큰들의 컨텍스트 윈도우를 앞에 붙입니다. 이 전체 시퀀스를 대상 형식으로 다시 인코딩하고 가장 최근의 대상 토큰과 정렬하여 새 토큰이 삽입될 정확한 위치를 결정합니다. 또한 대상 → 어시스턴트 변환 시, 어시스턴트 모델의 KV 캐시에서 일치하지 않는 토큰은 삭제해 데이터 무결성을 유지합니다.
성능 벤치마크
UAG는 동일 패밀리 어시스턴트가 없는 모델에 대해 큰 지연 시간 감소를 제공합니다. 다양한 작업에서 관찰된 속도 향상은 다음과 같습니다:
| 대상 모델 | 보조 모델 | 작업 | 속도 향상 |
|---|---|---|---|
CodeLlama-13b-Instruct-hf |
tiny_starcoder_py |
코드 생성 | 1.90x |
Mixtral-8x22B-Instruct-v0.1 |
vicuna-68m |
요약 | 1.52x |
gemma-2-9b |
vicuna-68m |
요약 | 1.76x |
Mixtral-8x22B-Instruct-v0.1 |
Qwen2-0.5B-Instruct |
장문 요약 | 1.78x |
Llama-3.1-70B |
Qwen2-0.5B-Instruct |
장문 요약 | 1.78x |
Phi-3-medium-128k-instruct |
Qwen2-0.5B-Instruct |
장문 요약 | 1.91x |
실험은 모델 크기에 따라 단일 A6000 GPU부터 4개의 A100 GPU까지 다양한 하드웨어에서 무작위로 선택된 100개의 예시를 사용해 수행되었습니다.
구현 및 사용법
범용 보조 생성은 🤗 Transformers 라이브러리 4.46.0 버전부터 통합되었습니다. 사용자는 generate() 메서드에 tokenizer와 assistant_tokenizer를 모두 전달함으로써 UAG를 활용할 수 있습니다.
from transformers import AutoModelForCausalLM, AutoTokenizer
prompt = "Alice and Bob"
checkpoint = "google/gemma-2-9b"
assistant_checkpoint = "double7/vicuna-68m"
assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt")
model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)
outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
현재 제한 사항 및 향후 작업
UAG는 현재 do_sample=True인 경우에만 다항 샘플링을 지원합니다. speculative 샘플링과 달리, 다항 샘플링은 대상 모델이 어시스턴트와 동일한 토큰을 샘플링하지 않을 경우 해당 토큰을 자동으로 거부하므로, 동일 토크나이저를 공유하는 어시스턴트에 비해 처리량이 낮아질 수 있습니다. 향후 업데이트에서는 speculative 샘플링 지원을 추가하고, 🤗 Transformers 파이프라인에 UAG를 직접 통합해 보다 간편한 사용자 경험을 제공할 예정입니다.