MLP Draft Models

MLP Draft Models (MLP 드래프트 모델)

MLP 스펙큘레이터(speculator)는 컨텍스트 벡터와 샘플링된 토큰 모두에 드래프트 예측을 조건화하는 드래프트 모델이에요. 아래 코드는 이런 MLP 기반 드래프트 모델이 제안을 생성하는 추측 디코딩을 vLLM에서 구성하는 예시예요. 자세한 내용은 "The Hitchhiker's Guide to Speculative Decoding"과 IBM Research의 기술 보고서를 참고하세요.

출처: 문서

본문

MLP Drafter 예시

from vllm import LLM, SamplingParams

prompts = ["The future of AI is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    tensor_parallel_size=1,
    speculative_config={
        "model": "ibm-ai-platform/llama3-8b-accelerator",
        "draft_tensor_parallel_size": 1,
        "method": "mlp_speculator",
    },
)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

알려진 문제: ibm-ai-platform/llama3-70b-accelerator는 AttributeError: 'MLPSpeculatorConfig' object has no attribute 'num_attention_heads' 오류로 실패할 수 있어요. 상태는 #34106과 #34163에서 추적 중이에요.

MLP 스펙큘레이터 동작 방식

MLP 스펙큘레이터는 컨텍스트 벡터(context vector)와 이미 샘플링된 토큰 모두를 드래프트 예측의 조건으로 사용하는 점이 일반 드래프트 모델과 달라요. 이 방식은 모델의 히든 상태에서 추출한 컨텍스트 벡터를 MLP가 받아 다음에 올 토큰 제안(드래프트)을 생성해요. vLLM에서는 method: "mlp_speculator"로 선택하며, draft_tensor_parallel_size로 드래프트 모델의 텐서 병렬 크기를 따로 지정할 수 있어요. 드래프트 모델은 기본 모델(target 모델)과 독립된 별도 가중치이므로, 추론 시 기본 모델과 드래프트 모델이 함께 로드돼요.

사전 학습된 MLP Drafter 모델

HF hub에서 이 유형의 다양한 스펙큘레이티브 모델을 사용할 수 있어요:

llama-13b-accelerator
llama3-8b-accelerator
codellama-34b-accelerator
llama2-70b-accelerator
llama3-70b-accelerator
granite-3b-code-instruct-accelerator
granite-8b-code-instruct-accelerator
granite-7b-instruct-accelerator
granite-20b-code-instruct-accelerator

주의 사항

MLP 스펙큘레이터를 사용할 때는 다음을 확인하세요:

드래프트 모델이 대상 모델과 호환되는지(동일 계열, 동일 토크나이저) 확인
tensor_parallel_size가 1보다 크면 메모리 사용량이 드래프트 모델 크기만큼 늘어남
추측 디코딩이 유익하려면 워크로드가 메모리 바운드여야 하며, 배치가 작고 생성 길이가 길수록 효과가 커요