Gemma 4 est disponible avec des entrées de texte, audio et image, et une fenêtre de contexte longue allant jusqu'à 256 000 jetons. En savoir plus

Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

Prédiction multi-jetons (MTP) Gemma 4 à l'aide de Hugging Face Transformers

Afficher sur ai.google.dev

Exécuter dans Google Colab

Exécuter dans Kaggle

Ouvrir dans Vertex AI

Afficher la source sur GitHub

Pour améliorer la vitesse d'inférence des modèles Gemma 4, une nouvelle série de modèles autorégressifs de "brouillon" a été publiée en même temps que la gamme principale. Au lieu de s'appuyer uniquement sur les principaux modèles Gemma 4 (appelés modèles "cibles"), le modèle brouillon prédit plusieurs jetons de manière autorégressive dans le temps nécessaire au modèle cible pour n'en traiter qu'un seul. Cette technique est également appelée décodage spéculatif.

Une fois que le rédacteur a prédit plusieurs jetons de brouillon, le modèle cible n'a plus qu'à vérifier ces jetons de brouillon suggérés. La validation est effectuée en parallèle, ce qui accélère considérablement l'inférence. Il réduit le nombre de passes directes que le modèle cible doit effectuer pour chaque jeton. Comme notre rédacteur génère une séquence de jetons pour la validation, nous l'appelons "tête de prédiction multi-jetons" (MTP, Multi-Token Prediction).

png

Les modèles brouillons publiés pour la famille Gemma 4 sont de petite taille et apportent plusieurs améliorations pour améliorer la qualité des jetons brouillons et accélérer davantage l'inférence, comme l'utilisation des activations du modèle cible et du cache KV pour obtenir de meilleures prédictions.

Ces améliorations permettent d'accélérer considérablement le décodage tout en garantissant une qualité similaire, ce qui rend ces points de contrôle parfaits pour les applications à faible latence et sur l'appareil.

Installer des packages Python

Installez les bibliothèques Hugging Face requises pour exécuter le modèle Gemma 4 et l'assistant Gemma 4.

# Install PyTorch & other libraries
pip install torch accelerate

# Install the transformers library
pip install transformers

Charger les modèles

Pour chaque modèle cible (l'un des principaux modèles du modèle Gemma 4), un assistant permet d'accélérer l'inférence. Vous allez donc charger deux modèles :

Cible (par exemple, google/gemma-4-E2B-it) : modèle cible Gemma 4 complet
Drafter (par exemple, google/gemma-4-E2B-it-assistant) : le drafteur MTP léger à quatre couches qui propose des jetons candidats

Notez que le rédacteur est souvent appelé assistant, car le modèle aide le modèle plus grand à choisir les jetons à prédire.

Utilisez les bibliothèques transformers pour créer une instance de processor et model à l'aide des classes AutoProcessor et AutoModelForCausalLM, comme indiqué dans l'exemple de code suivant :

TARGET_MODEL_ID = "google/gemma-4-E2B-it" # @param ["google/gemma-4-E2B-it", "google/gemma-4-E4B-it", "google/gemma-4-12B-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]

Gemma 4 avec l'Assistant

Heureusement, l'utilisation d'un assistant dans transformers est assez simple et nécessite de transmettre le modèle d'assistant à la fonction 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.

Voici le processus en détail :

Le rédacteur propose N jetons générés de manière autorégressive.
Le modèle cible vérifie tous les jetons N en une seule passe avant.
Les jetons rédigés avec une probabilité élevée sont acceptés.
Les jetons rédigés avec une faible probabilité sont rejetés.
Étant donné que le modèle cible effectue une propagation avant, il générera toujours un jeton par lui-même, quel que soit le nombre de jetons brouillons acceptés ou refusés.

Jetons de brouillon

L'auteur du brouillon peut générer n'importe quel nombre de jetons pour que le modèle cible puisse les valider. Toutefois, le modèle cible peut toujours choisir de refuser certains jetons. Dans ce cas, tous les jetons suivants sont ignorés.

png

Il est donc important de connaître le compromis lorsque vous utilisez différentes valeurs pour le nombre de jetons rédigés.

Plus de jetons de brouillon

Lorsque vous rédigez de nombreux jetons (par exemple, 15), il est fort probable qu'ils ne soient pas tous acceptés. Le risque de gaspillage de ressources de calcul est donc plus élevé. En revanche, il a tendance à accélérer l'inférence lorsque le taux d'acceptation est élevé.

png

Moins de jetons de brouillon

Lorsque vous rédigez moins de jetons, le taux d'acceptation a tendance à être plus élevé, car les jetons dont la position est plus proche de la requête initiale sont plus précis. Toutefois, comme seuls quelques jetons sont rédigés, le gain de vitesse que vous obtiendriez avec un modèle de rédaction plus rapide est réduit.

png

Heureusement, vous n'avez pas besoin de tester les meilleures valeurs pour votre cas d'utilisation dans transformers, car vous pouvez définir num_assistant_tokens_schedule sur "heuristic" (heuristique), ce qui adaptera automatiquement le nombre de jetons brouillons au moment de l'exécution :

Tous les jetons acceptés : augmentez de deux le nombre de jetons pour la version préliminaire, car le rédacteur est assez précis pour la requête. Augmenter le nombre de jetons rédigés peut accélérer le processus si ces jetons sont également acceptés.
Des jetons ont été refusés : si des jetons ont été refusés, réduisez de 1 le nombre de jetons à utiliser pour la rédaction. En réduisant le nombre de jetons, vous évitez de gaspiller trop de brouillons si le modèle cible continue de rejeter la plupart des jetons.

De même, vous pouvez mettre à jour le nombre de jetons brouillons en mettant à jour num_assistant_tokens dans le rédacteur comme suit :

# 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.