AI

شتاب‌دهی به مدل‌های زبانی بزرگ: رمزگشایی حدسی و بهینه‌سازی حافظه KV برای استنتاج زیر ۱۰۰ میلی‌ثانیه

مقدمه

در منظر سریعاً در حال تحول مدل‌های زبانی بزرگ (LLMs)، تأخیر استنتاج اغلب گلوگاهی است که از رسیدن کاربردهای بلادرنگ، مانند دستیاران کدنویسی تعاملی یا چت‌بات‌های با تأخیر کم، به پتانسیل کامل آن‌ها جلوگیری می‌کند. اگرچه کم‌عمق‌سازی مدل (Quantization) ردپای حافظه را کاهش می‌دهد، اما اغلب در ارائه شتاب‌های چشمگیر مورد نیاز برای زمان‌های پاسخگویی زیر ۱۰۰ میلی‌ثانیه با مشکل مواجه است. اینجاست که ترکیب رمزگشایی حدسی و بهینه‌سازی حافظه KV درخشش می‌کند. با پیش‌بینی هوشمندانه دنباله‌های توکن و مدیریت کارآمد وضعیت حافظه، توسعه‌دهندگان می‌توانند شتاب‌های خطی را بدون قربانی کردن کیفیت خروجی مدل به دست آورند.

حافظه KV: قهرمان خاموش عملکرد

برای درک نحوه دستیابی به استنتاج زیر ۱۰۰ میلی‌ثانیه، باید ابتدا حافظه کلید-مقدار (KV Cache) را بررسی کنیم. در طول تولید متن، مدل ترانسفورمر به توکن‌های قبلی توجه می‌کند. به جای محاسبه مجدد توجه برای کل دنباله در هر مرحله، مدل بردارهای کلید و مقدار توکن‌های قبلی را در حافظه پنهان (Cache) نگه می‌دارد. برای کاربردهای با زمینه طولانی، این حافظه پنهان می‌تواند به یک گلوگاه حافظه تبدیل شود. اگر به طور کارآمد مدیریت نشود، سربار تخصیص و مدیریت این حافظه پنهان می‌تواند مزایای پردازش موازی را خنثی کند.

استراتژی بهینه‌سازی کلیدی: پیاده‌سازی دسته‌بندی پیوسته (Continuous Batching) (که به عنوان دسته‌بندی آگاه از زمان‌بندی نیز شناخته می‌شود). برخلاف دسته‌بندی ایستا، دسته‌بندی پیوسته اجازه می‌دهد درخواست‌های جدید به محض اتمام یک درخواست قبلی به دسته اضافه شوند که بهره‌وری GPU را به حداکثر رسانده و حافظه KV را فشرده و مرتبط نگه می‌دارد.

# کد شبه برای نشان دادن مدیریت کارآمد حافظه KV
def generate_with_kv_cache(model, prompt, max_length):
    # مقداردهی اولیه حافظه KV با حافظه از پیش تخصیص یافته
    kv_cache = model.init_cache(max_length)
    
    # رمزگذاری درخواست و محاسبه حالت‌های اولیه KV
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model(input_ids=inputs["input_ids"], past_key_values=kv_cache)
    
    next_token = outputs.logits[:, -1, :].argmax(dim=-1)
    
    for _ in range(max_length):
        # به‌روزرسانی حافظه پنهان فقط با جفت‌های KV توکن جدید
        # این کار از محاسبه مجدد کل ماتریس توجه جلوگیری می‌کند
        kv_cache = update_kv_cache(kv_cache, outputs)
        
        # تولید توکن بعدی
        outputs = model(input_ids=next_token, past_key_values=kv_cache)
        next_token = outputs.logits[:, -1, :].argmax(dim=-1)
        
        if next_token == tokenizer.eos_token_id:
            break
            
    return tokenizer.decode(outputs)

رمزگشایی حدسی: موازی‌سازی توالی

رمزگشایی خودتوالی‌گر سنتی ذاتاً توالی‌وار است: توکن $T_n$ به $T_{n-1}$ وابسته است. این امر یک مسیر بحرانی ایجاد می‌کند که سرعت را محدود می‌سازد. رمزگشایی حدسی این مانع را با استفاده از یک مدل «پیش‌نویس» کوچک‌تر و سریع‌تر برای پیشنهاد چندین توکن می‌شکند که سپس توسط مدل «هدف» بزرگ‌تر و معتبرتر به صورت موازی تأیید می‌شوند. فرآیند به شرح زیر است: 1. مدل پیش‌نویس $N$ توکن کاندید را تولید می‌کند. 2. مدل هدف، درخواست به همراه این کاندیداها را همزمان پردازش می‌کند. 3. مدل هدف کاندیداها را با توزیع احتمالی خود تطبیق می‌دهد. 4. هرگونه عدم تطابق در توکن‌ها باعث بازگشت به عقب شده و تولید از آخرین توکن تأیید شده ادامه می‌یابد. وقتی مدل پیش‌نویس دقیق باشد، مدل هدف چندین توکن را در یک گذر رو به جلو تأیید می‌کند که منجر به شتاب‌های خطی می‌شود.

مثال پیاده‌سازی عملی

با استفاده از کتابخانه‌های مدرن مانند Hugging Face Transformers و vLLM، پیاده‌سازی رمزگشایی حدسی در حال در دسترس‌تر شدن است. در زیر یک پیاده‌سازی مفهومی با استفاده از یک مدل پیش‌نویس (مثلاً یک مدل ۷ میلیارد پارامتری فشرده شده) برای شتاب‌دهی به یک مدل هدف (مثلاً یک مدل ۷۰ میلیارد پارامتری) آورده شده است.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# بارگذاری مدل سبک پیش‌نویس و مدل سنگین هدف
draft_model = AutoModelForCausalLM.from_pretrained("draft-model-7b")
target_model = AutoModelForCausalLM.from_pretrained("target-model-70b")

draft_tokenizer = AutoTokenizer.from_pretrained("draft-model-7b")

def speculative_decode(prompt, draft_model, target_model, num_drafts=4):
    # 1. فاز پیش‌نویس
    draft_inputs = draft_tokenizer(prompt, return_tensors="pt")
    draft_outputs = draft_model.generate(**draft_inputs, max_new_tokens=num_drafts)
    draft_tokens = draft_outputs[0]
    
    # 2. فاز تأیید
    # مدل هدف درخواست + توکن‌های پیش‌نویس را یکجا پردازش می‌کند
    target_inputs = target_tokenizer(prompt + draft_tokenizer.decode(draft_tokens), 
                                     return_tensors="pt")
    
    with torch.no_grad():
        target_outputs = target_model(**target_inputs)
    
    # 3. منطق پذیرش/رد
    # مقایسه احتمالات پیش‌نویس با لگیت‌های مدل هدف
    # در صورت پذیرش، به دنباله اضافه می‌شود. در صورت رد، کوتاه می‌شود.
    accepted_tokens = verify_acceptance(draft_tokens, target_outputs)
    
    return accepted_tokens

ترکیب تکنیک‌ها برای دستیابی به اهداف زیر ۱۰۰ میلی‌ثانیه

برای رسیدن ثابت به زمان‌های استنتاج زیر ۱۰۰ میلی‌ثانیه، نمی‌توان تنها به یک تکنیک تکیه کرد. باید آن‌ها را ترکیب کنید:
  • شتاب‌دهنده سخت‌افزاری: از هسته‌های CUDA بهینه‌شده برای حافظه KV (مانند FlashAttention) برای کاهش استفاده از پهنای باند حافظه استفاده کنید.
  • استحصال مدل (Distillation): یک مدل پیش‌نویس کوچک‌تر را به طور خاص برای تقلید از توزیع خروجی مدل بزرگ‌تر آموزش دهید تا نرخ پذیرش در رمزگشایی حدسی افزایش یابد.
  • کارایی دسته‌بندی: همانطور که ذکر شد، از دسته‌بندی پیوسته استفاده کنید تا مطمئن شوید GPU همیشه در حال پردازش داده است و زمان بیکاری به حداقل می‌رسد.

نتیجه‌گیری

دستیابی به استنتاج LLM زیر ۱۰۰ میلی‌ثانیه دیگر یک تمرین نظری نیست، بلکه یک چالش مهندسی عملی است. با بهره‌گیری از کارایی حافظه بهینه‌سازی حافظه KV و قدرت پردازش موازی رمزگشایی حدسی، توسعه‌دهندگان می‌توانند مدل‌های قدرتمندی را مستقر کنند که با سرعتی مشابه انسان پاسخ می‌دهند. با بلوغ ادامه‌دار اکوسیستم‌های سخت‌افزاری و نرم‌افزاری، این تکنیک‌ها به شیوه‌های استاندارد برای هر برنامه هوش مصنوعی در سطح تولید تبدیل خواهند شد. با پروفایل‌سازی گلوگاه‌های تأخیر فعلی خود شروع کنید، حافظه پنهان KV کارآمد را پیاده‌سازی کنید و با مدل‌های پیش‌نویس سبک آزمایش کنید تا پتانسیل واقعی مدل‌های زبانی بزرگ خود را آزاد نمایید.
Share: