피드로 돌아가기
Universal Assisted Generation: Faster Decoding with Any Assistant Model
Hugging Face BlogHugging Face Blog
AI/ML

Intel Labs와 Hugging Face가 토크나이저 불일치 문제를 2-way 토크나이저 변환으로 해결해 서로 다른 모델 계열 간 Assisted Generation 적용 가능

Universal Assisted Generation: Faster Decoding with Any Assistant Model

2024년 10월 29일9intermediate

Context

Assisted Generation은 LLM 추론 속도를 1.5x-3x 향상시키는 검증된 기법이지만, 타겟 모델과 어시스턴트 모델이 동일한 토크나이저를 공유해야 한다는 제약이 있습니다. Gemma-2-9b나 CodeLlama-13b 같은 널리 사용되는 모델들은 충분히 작은(50-100배 이상 작은) 동일 계열 버전이 없어 기존 Assisted Generation 적용이 불가능했습니다.

Technical Solution

  • 2-way 토크나이저 변환 메커니즘 도입: 어시스턴트 모델의 토큰을 텍스트로 변환한 후 타겟 모델의 토크나이저로 재인코딩하고, 검증 후 다시 어시스턴트 토크나이저 형식으로 변환
  • 컨텍스트 윈도우 기반 토큰 정렬: 새로 생성된 어시스턴트 토큰 재인코딩 시 이전 토큰들의 컨텍스트 윈도우를 포함해 정확한 토큰 위치 파악
  • Key-Value 캐시 정합성 관리: 타겟에서 어시스턴트로 변환 시 불일치 토큰을 KV 캐시에서 제거해 데이터 무결성 보장
  • Hugging Face Transformers 4.46.0에 통합: generate() 함수에 tokenizer와 assistant_tokenizer 파라미터 추가로 사용자가 임의의 모델 조합 지정 가능
  • 초기 구현에서 다항식 샘플링(multinomial sampling) 지원: do_sample=True 사용 시 타겟 모델이 어시스턴트와 다른 토큰을 자동 거부

Impact

  • CodeLlama-13b-Instruct + tiny_starcoder_py 조합: 1.90x 추론 속도 향상
  • Mixtral-8x22B-Instruct-v0.1 + vicuna-68m 조합: 1.52x-1.78x 추론 속도 향상
  • Gemma-2-9b + vicuna-68m 조합: 1.76x 추론 속도 향상
  • Llama-3.1-70B + Qwen2-0.5B-Instruct 조합: 1.78x 추론 속도 향상
  • Phi-3-medium-128k-instruct + Qwen2-0.5B-Instruct 조합: 1.91x 추론 속도 향상

Key Takeaway

토크나이저 변환을 통해 모델 계열 간 호환성 문제를 해결함으로써, 기존 Assisted Generation의 적용 범위를 140,000개 Hugging Face Hub 모델로 확대했습니다. 이는 추론 최적화에서 모델 선택의 자유도를 크게 높이는 아키텍처 패턴으로, 프로덕션 환경에서 기존 모델 쌍의 제약 없이 최적의 속도-정확도 조합을 구성할 수 있게 합니다.


대규모 언어 모델을 배포하는 엔지니어는 Universal Assisted Generation을 통해, 기존에 호환되지 않던 작은 모델들(vicuna-68m, Qwen2-

0.5B 등)을 어시스턴트로 선택해

1.5x-

1.9x 추론 속도 향상을 달성할 수 있습니다. 이를 위해 Transformers

4.

4

6.0 이상에서 generate() 함수에 assistant_tokenizer 파라미터만 추가하면 되므로, 기존 코드 수정 최소화로 즉각적인 성능 개선이 가능합니다.

원문 읽기