 Посмотреть на ai.google.dev |  Запустить в Google Colab |  Запуск в Kaggle |  Открыть в Vertex AI |  Посмотреть исходный код на GitHub |
Для повышения скорости вывода в моделях Gemma 4, наряду с основной линейкой, была выпущена новая серия авторегрессивных «черновых» моделей. Вместо того чтобы полагаться исключительно на основные модели Gemma 4 (называемые «целевыми» моделями), черновая модель авторегрессивно предсказывает несколько токенов за то время, пока целевая модель обрабатывает только один. Этот метод также известен как спекулятивное декодирование.
После того, как составитель спрогнозировал несколько вариантов токенов, целевой модели остается только проверить эти предложенные варианты. Проверка выполняется параллельно, что значительно ускоряет вывод. Это уменьшает количество проходов, которые целевой модели приходится выполнять для каждого токена. Поскольку наш составитель генерирует последовательность токенов для проверки, мы называем его блоком прогнозирования нескольких токенов (MTP).
Выпущенные черновые модели для семейства Gemma 4 невелики по размеру и содержат ряд улучшений, направленных на повышение качества созданных токенов и дальнейшее ускорение вывода, например, использование активаций целевой модели и KV-кэша для получения более точных прогнозов.
Эти усовершенствования обеспечивают значительное ускорение декодирования, гарантируя при этом аналогичное качество, что делает эти контрольные точки идеальными для приложений с низкой задержкой и для приложений, работающих непосредственно на устройстве.
Установите пакеты Python.
Установите библиотеки Hugging Face, необходимые для работы Gemma 4 и модели-ассистента Gemma 4.
# Install PyTorch & other librariespip install torch accelerate# Install the transformers librarypip install transformers
Загрузите модели
Для каждой целевой модели (одной из основных моделей в модели Gemma 4) существует помощник, который помогает ускорить вывод. Таким образом, вам нужно будет загрузить две модели:
- Цель (например,
google/gemma-4-E2B-it): Полная целевая модель Gemma 4. - Drafter (например,
google/gemma-4-E2B-it-assistant): Легковесный 4-слойный инструмент для создания MTP-токенов, предлагающий варианты токенов-кандидатов.
Следует отметить, что составителя часто называют помощником , поскольку модель помогает более крупной модели в выборе токенов для прогнозирования.
Используйте библиотеки 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]
Джемма 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 токенов за один прямой проход.
- Принимаются созданные токены с высокой вероятностью успеха.
- Созданные токены с низкой вероятностью отклоняются.
- Поскольку целевая модель выполняет прямой проход, она всегда будет генерировать 1 токен самостоятельно, независимо от того, сколько созданных токенов было принято или отклонено.
Черновые токены
Разработчик может сгенерировать любое количество токенов для проверки целевой моделью. Однако целевая модель по-прежнему может отклонить определенные токены. В этом случае все последующие токены игнорируются.
Поэтому важно понимать компромисс, возникающий при использовании различных значений количества созданных токенов.
Больше жетонов для драфта
При создании большого количества токенов (например, 15) высока вероятность того, что не все токены будут приняты. Таким образом, возрастает вероятность неэффективного использования вычислительных ресурсов. В то же время, высокая вероятность принятия токенов, как правило, ускоряет процесс вывода.
Меньше жетонов для драфта
При выборе меньшего количества токенов процент принятия, как правило, выше, поскольку токены, расположенные ближе к исходному предложению, более точны. Однако, поскольку выбирается лишь небольшое количество токенов, ускорение, которое вы получили бы от более быстрой модели выбора, снижается.
К счастью, вам не нужно экспериментировать с оптимальными значениями для вашего конкретного случая в transformers поскольку вы можете установить параметр num_assistant_tokens_schedule в значение "heuristic", которое автоматически адаптирует количество созданных токенов во время выполнения:
- Все жетоны приняты — Увеличьте количество жетонов для розыгрыша на 2, поскольку составитель достаточно точно соответствует подсказке. Увеличение количества разыгрываемых жетонов может привести к ускорению игры, если эти жетоны также будут приняты.
- Отклоненные токены — Если какие-либо токены отклонены, уменьшите количество токенов для драфта на 1. Уменьшение количества токенов позволяет избежать чрезмерной потери большого количества драфтованных токенов, если целевая модель продолжит отклонять большинство из них.
Аналогичным образом, вы можете обновить количество токенов черновика, изменив значение 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.