Hugging Face Transformers का इस्तेमाल करके, Gemma 4 के लिए मल्टी-टोकन प्रेडिक्शन (एमटीपी)

ai.google.dev पर देखें Google Colab में चलाएं Kaggle में चलाएं Vertex AI में खोलें GitHub पर सोर्स देखें

Gemma 4 मॉडल की अनुमान लगाने की स्पीड को बेहतर बनाने के लिए, मुख्य लाइनअप के साथ-साथ ऑटोरग्रेसिव “ड्राफ़्टर” मॉडल की एक नई सीरीज़ रिलीज़ की गई है. ड्राफ़्ट मॉडल, सिर्फ़ मुख्य Gemma 4 मॉडल (जिन्हें “टारगेट” मॉडल कहा जाता है) पर भरोसा करने के बजाय, कई टोकन का अनुमान लगाता है. इसमें उतना ही समय लगता है जितना टारगेट मॉडल को सिर्फ़ एक टोकन को प्रोसेस करने में लगता है. इस तकनीक को स्पेकुलेटिव डिकोडिंग भी कहा जाता है.

ड्राफ़्ट तैयार करने वाले व्यक्ति के कई ड्राफ़्ट टोकन का अनुमान लगाने के बाद, टारगेट मॉडल को अब सिर्फ़ सुझाए गए ड्राफ़्ट टोकन की पुष्टि करनी होगी. पुष्टि की प्रोसेस साथ-साथ चलती है. इससे अनुमान लगाने की प्रोसेस बहुत तेज़ी से पूरी होती है. इससे, टारगेट मॉडल को हर टोकन के लिए फ़ॉरवर्ड पास की संख्या कम करनी पड़ती है. हम इसे मल्टी-टोकन प्रेडिक्शन (एमटीपी) हेड कहते हैं, क्योंकि हमारा ड्राफ़्टर पुष्टि के लिए टोकन का एक क्रम जनरेट करता है.

png

Gemma 4 फ़ैमिली के लिए रिलीज़ किए गए ड्राफ़्ट मॉडल छोटे हैं. इनमें कई सुधार किए गए हैं, ताकि ड्राफ़्ट किए गए टोकन की क्वालिटी को बेहतर बनाया जा सके. साथ ही, अनुमान लगाने की प्रोसेस को और तेज़ किया जा सके. जैसे, बेहतर अनुमान पाने के लिए टारगेट मॉडल ऐक्टिवेशन और केवी-कैश का इस्तेमाल करना.

इन सुधारों की वजह से, डिकोडिंग की स्पीड काफ़ी बढ़ जाती है. साथ ही, क्वालिटी भी पहले जैसी ही बनी रहती है. इसलिए, ये चेकपॉइंट कम समय में प्रोसेस होने वाले और डिवाइस पर काम करने वाले ऐप्लिकेशन के लिए सबसे सही हैं.

Python पैकेज इंस्टॉल करना

Gemma 4 और Gemma 4 assistant मॉडल को चलाने के लिए, ज़रूरी Hugging Face लाइब्रेरी इंस्टॉल करें.

# Install PyTorch & other libraries
pip install torch accelerate

# Install the transformers library
pip install transformers

मॉडल लोड करना

हर टारगेट मॉडल (Gemma 4 मॉडल में मौजूद मुख्य मॉडल में से एक) के लिए, एक असिस्टेंट होता है. इससे अनुमान लगाने की प्रोसेस को तेज़ करने में मदद मिलती है. इसलिए, आपको दो मॉडल लोड करने होंगे:

  • टारगेट (जैसे, google/gemma-4-E2B-it): Gemma 4 का पूरा टारगेट मॉडल
  • Drafter (उदाहरण के लिए, google/gemma-4-E2B-it-assistant): यह चार लेयर वाला एमटीपी ड्राफ़्टर है. यह कैंडिडेट टोकन का सुझाव देता है

ध्यान दें कि ड्राफ़्टर को अक्सर सहायक के तौर पर जाना जाता है, क्योंकि यह मॉडल, बड़े मॉडल को यह चुनने में मदद करता है कि किन टोकन का अनुमान लगाना है.

transformers लाइब्रेरी का इस्तेमाल करके, processor और model का एक इंस्टेंस बनाएं. इसके लिए, AutoProcessor और AutoModelForCausalLM क्लास का इस्तेमाल करें. जैसा कि यहां दिए गए कोड के उदाहरण में दिखाया गया है:

TARGET_MODEL_ID = "google/gemma-4-E2B-it" # @param ["google/gemma-4-E2B-it","google/gemma-4-E4B-it", "google/gemma-4-31B-it", "google/gemma-4-26B-A4B-it"]
ASSISTANT_MODEL_ID = TARGET_MODEL_ID + "-assistant"

import torch
from transformers import AutoProcessor, AutoModelForCausalLM

# Target Model
processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)
target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Assistant Model (the drafter)
assistant_model = AutoModelForCausalLM.from_pretrained(
    ASSISTANT_MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

[transformers] `torch_dtype` is deprecated! Use `dtype` instead!
Loading weights:   0%|          | 0/1951 [00:00<?, ?it/s]
Loading weights:   0%|          | 0/50 [00:00<?, ?it/s]

Assistant के साथ Gemma 4

transformers में किसी असिस्टेंट का इस्तेमाल करना बहुत आसान है. इसके लिए, आपको model.generate फ़ंक्शन में असिस्टेंट मॉडल पास करना होगा:

# Process inputs with the `target_model`
messages = [
    {
        "role": "user",
        "content": "Explain the concepts of speculative decoding and MTP in 3 sentences."
    }
]
input_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=input_text, return_tensors="pt").to(target_model.device)

# `assistant_model=assistant_model` is all you need to enable MTP!
outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
    do_sample=False,
)

# Decode the response into text
response = processor.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

**Speculative decoding** is a technique where a smaller, faster language model (the "draft model") generates several candidate tokens, which are then quickly verified by a larger, more accurate model to produce a final, high-quality output much faster than decoding the large model alone. **MTP (Multi-Task Prediction)** involves training a single model to perform multiple related tasks simultaneously, allowing it to leverage shared knowledge across different objectives. Together, these methods aim to significantly accelerate the inference speed of large language models while maintaining or improving output quality.

इस प्रोसेस के काम करने का तरीका यहां बताया गया है:

  • ड्राफ़्ट तैयार करने वाला व्यक्ति, N टोकन का सुझाव देता है. ये टोकन, ऑटोरेग्रेसिव तरीके से जनरेट किए जाते हैं
  • टारगेट मॉडल, फ़ॉरवर्ड पास के एक चरण में सभी N टोकन की पुष्टि करता है
  • ज़्यादा संभावना वाले ड्राफ़्ट किए गए टोकन स्वीकार किए जाते हैं
  • कम संभावना वाले ड्राफ़्ट किए गए टोकन अस्वीकार कर दिए जाते हैं
  • टारगेट मॉडल, फ़ॉरवर्ड पास करता है. इसलिए, यह हमेशा एक टोकन जनरेट करेगा. इससे कोई फ़र्क़ नहीं पड़ता कि ड्राफ़्ट किए गए कितने टोकन स्वीकार किए गए या अस्वीकार किए गए

ड्राफ़्ट टोकन

ड्राफ़्टर, टारगेट मॉडल की पुष्टि करने के लिए कितने भी टोकन जनरेट कर सकता है. हालांकि, टारगेट मॉडल अब भी कुछ टोकन को अस्वीकार कर सकता है. ऐसा होने पर, इसके बाद के सभी टोकन को अनदेखा कर दिया जाता है.

png

इसलिए, ड्राफ़्ट किए गए टोकन की संख्या के लिए अलग-अलग वैल्यू का इस्तेमाल करते समय, ट्रेडऑफ़ के बारे में जानना ज़रूरी है.

ज़्यादा ड्राफ़्ट टोकन

अगर आपने कई टोकन (जैसे कि 15) बनाए हैं, तो इस बात की संभावना ज़्यादा है कि सभी टोकन स्वीकार नहीं किए जाएंगे. इसलिए, कंप्यूटिंग के संसाधनों का ज़्यादा इस्तेमाल होने की संभावना होती है. इसके उलट, जब स्वीकार किए जाने की दर ज़्यादा होती है, तो यह अनुमान लगाने की प्रोसेस को तेज़ कर देता है.

png

कम ड्राफ़्ट टोकन

कम टोकन का ड्राफ़्ट बनाने पर, स्वीकार किए जाने की दर ज़्यादा होती है. ऐसा इसलिए, क्योंकि शुरुआती प्रॉम्प्ट के आस-पास की पोज़िशन वाले टोकन ज़्यादा सटीक होते हैं. हालांकि, सिर्फ़ कुछ टोकन ड्राफ़्ट किए जाते हैं. इसलिए, तेज़ी से ड्राफ़्ट करने वाले मॉडल से मिलने वाली स्पीड कम हो जाती है.

png

अच्छी बात यह है कि आपको transformers में अपने इस्तेमाल के उदाहरण के लिए सबसे सही वैल्यू के साथ एक्सपेरिमेंट करने की ज़रूरत नहीं है. ऐसा इसलिए, क्योंकि transformers को "अनुभव के आधार पर अनुमान लगाने की विधि" पर सेट किया जा सकता है. इससे, रनटाइम के दौरान ड्राफ़्ट किए गए टोकन की संख्या अपने-आप अडजस्ट हो जाएगी:num_assistant_tokens_schedule

  • सभी टोकन स्वीकार किए गए -- ड्राफ़्ट करने के लिए टोकन की संख्या को दो से बढ़ाएं, क्योंकि ड्राफ़्ट करने वाला टूल प्रॉम्प्ट के लिए काफ़ी सटीक है. ड्राफ़्ट किए गए टोकन की संख्या बढ़ाने से, प्रोसेस तेज़ी से पूरी हो सकती है. हालांकि, ऐसा तब होगा, जब उन टोकन को स्वीकार किया जाए.
  • अस्वीकार किए गए टोकन -- अगर कोई टोकन अस्वीकार किया जाता है, तो ड्राफ़्ट किए जाने वाले टोकन की संख्या में एक की कमी करें. टोकन की संख्या कम करने से, अगर टारगेट मॉडल ज़्यादातर टोकन को अस्वीकार करता है, तो बहुत ज़्यादा ड्राफ़्ट किए गए टोकन बर्बाद नहीं होते.

इसी तरह, ड्राफ़्ट किए गए टोकन की संख्या को अपडेट किया जा सकता है. इसके लिए, ड्राफ़्टर में num_assistant_tokens को इस तरह अपडेट करें:

# Update how many draft tokens are generated at the start of inference
assistant_model.generation_config.num_assistant_tokens = 4

# Update how the number of draft tokens are updated ("heuristic" for a dynamic schedule and "constant" for a constant schedule)
assistant_model.generation_config.num_assistant_tokens_schedule = "heuristic"

# Run with MTP
outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
    do_sample=False,
)

# Decode the response into text
response = processor.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

**Speculative decoding** is a technique where a smaller, faster language model (the "draft model") generates several candidate tokens, which are then verified by a larger, more accurate model to quickly produce a high-quality output. **MTP (Multi-Task Prediction)** involves training a single model to perform multiple related tasks simultaneously, allowing it to leverage shared knowledge across different objectives. Together, these methods aim to significantly speed up the inference process of large language models while maintaining or improving output quality.