Speculative Decoding: как ускорить локальные LLM в 2-3 раза на одном GPU (полный гайд)

Почему ваши локальные LLM всё ещё тормозят (и как это исправить без покупки нового железа)

Вы загружаете Llama 3 70B в 4-битном формате на свою RTX 4090. Ждёте 15 секунд. Получаете первый токен. Потом ещё один. И ещё. Генерация одного ответа занимает минуту. Знакомо? Вы думаете: "Надо докупить ещё одну карту" или "Пора переходить на облако". Стойте. Прежде чем тратить тысячи долларов на железо, прочитайте это.

Проблема не в недостатке вычислительной мощности. Проблема в том, как мы её используем. Точнее — как не используем. Современные GPU при генерации текста работают на 10-30% своей пропускной способности. Остальное время они просто ждут. Ждут, пока завершится операция извлечения весов из памяти.

Speculative Decoding ломает эту парадигму. Не за счёт магии, а за счёт хитрости. Вместо того чтобы генерировать один токен и ждать, мы генерируем несколько предположений параллельно, а потом проверяем их все сразу. Результат? Скорость увеличивается в 2-3 раза на том же самом железе. Без потери качества. Без изменения модели. Просто другой алгоритм.

Как работает Speculative Decoding (простыми словами)

Представьте, что вы — супервайзер на заводе. У вас есть:

• Target модель — опытный мастер, который делает всё идеально, но медленно
• Drafter модель — стажёр, который работает быстро, но иногда ошибается

Вместо того чтобы ждать, пока мастер сделает одну деталь за 10 минут, вы даёте стажёру сделать 5 предположений за 2 минуты. Потом показываете мастеру все 5 вариантов и спрашиваете: "Какой правильный?" Мастер быстро проверяет — "первый верный, второй нет, третий верный..." — и вы принимаете все правильные варианты.

В терминах LLM:

1. Маленькая быстрая модель (Drafter) генерирует K токенов-предположений
2. Большая точная модель (Target) проверяет все K токенов за один проход
3. Принимаются все токены, где предположение совпало с проверкой
4. Следующий токен после первого несовпадения генерируется с нуля

Что нужно для реализации (проверьте свой стек)

Прежде чем переходить к коду, убедитесь, что у вас есть:

• Python 3.9+ (лучше 3.11)
• PyTorch 2.0+ с поддержкой CUDA
• Две модели: одна большая (Target), одна маленькая (Drafter)
• Хотя бы 8 ГБ VRAM свободно (для пары 7B+1B моделей)

Drafter модель должна быть в 3-10 раз меньше Target. Идеальные пары:

1 Устанавливаем зависимости

Начнём с чистого виртуального окружения. Не пытайтесь установить это в глобальное — сломаете другие проекты.

python -m venv spec_decode_env
source spec_decode_env/bin/activate  # для Linux/Mac
# или spec_decode_env\Scripts\activate для Windows

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers accelerate bitsandbytes
pip install vllm  # для оптимизированной реализации

2 Загружаем модели

Сначала базовый вариант — без speculative decoding. Чтобы было с чем сравнивать.

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Target модель (большая)
target_model_name = "meta-llama/Llama-3-8B-Instruct"
target_tokenizer = AutoTokenizer.from_pretrained(target_model_name)
target_model = AutoModelForCausalLM.from_pretrained(
    target_model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True  # квантуем до 4 бит для экономии памяти
)

# Drafter модель (маленькая)
drafter_model_name = "microsoft/Phi-3-mini-4k-instruct"
drafter_tokenizer = AutoTokenizer.from_pretrained(drafter_model_name)
drafter_model = AutoModelForCausalLM.from_pretrained(
    drafter_model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True
)

print(f"Target модель загружена: {target_model.config.model_type}")
print(f"Drafter модель загружена: {drafter_model.config.model_type}")

Обратите внимание: мы используем 4-битное квантование для обеих моделей. Это не обязательно, но сильно экономит память. Если у вас мощная карта (например, RTX 4090 с 24 ГБ), можете попробовать 8-битное или даже полную точность для Drafter модели.

3 Пишем ядро speculative decoding

Вот как НЕ надо делать (типичная ошибка новичков):

# ПЛОХО: последовательная генерация и проверка
def bad_speculative_decode(prompt, k=5):
    generated = []
    for i in range(k):
        # Drafter генерирует токен
        drafter_out = drafter_model.generate(...)
        # Target проверяет этот один токен
        target_out = target_model.generate(...)
        # Сравниваем
        if drafter_out == target_out:
            generated.append(drafter_out)
        else:
            break
    return generated

Почему это плохо? Потому что мы теряем всё преимущество параллельной проверки. Target модель всё равно загружается k раз. Правильный подход — проверять все k токенов за один forward pass.

Вот рабочая реализация:

def speculative_decode(prompt, max_new_tokens=100, k=5):
    """
    Speculative decoding с параллельной верификацией.
    
    Args:
        prompt: входной текст
        max_new_tokens: максимальная длина генерации
        k: сколько токенов предсказывать заранее
    """
    
    # Токенизируем промпт
    input_ids = target_tokenizer.encode(prompt, return_tensors="pt").cuda()
    
    generated = []
    current_input = input_ids
    
    with torch.no_grad():
        for step in range(max_new_tokens // k + 1):
            # 1. Drafter генерирует k предположений
            drafter_output = drafter_model.generate(
                current_input,
                max_new_tokens=k,
                do_sample=False,  # greedy decoding для consistency
                pad_token_id=target_tokenizer.eos_token_id
            )
            
            # Берем только новые токены (предположения)
            draft_tokens = drafter_output[:, current_input.shape[1]:]
            
            if draft_tokens.shape[1] == 0:
                break
            
            # 2. Target проверяет все предположения за один проход
            # Конкатенируем промпт с предположениями
            verification_input = torch.cat([current_input, draft_tokens], dim=1)
            
            # Получаем логиты от Target модели
            target_logits = target_model(verification_input).logits
            
            # 3. Параллельная верификация (rejection sampling)
            accepted_tokens = []
            
            # Сравниваем предсказания Drafter с распределением Target
            for i in range(draft_tokens.shape[1]):
                # Индекс текущего токена для проверки
                token_idx = current_input.shape[1] + i
                
                # Получаем распределение вероятностей от Target
                target_probs = torch.softmax(target_logits[:, token_idx-1, :], dim=-1)
                
                # Вероятность, которую Target назначил токену от Drafter
                draft_token_prob = target_probs[0, draft_tokens[0, i]].item()
                
                # Сэмплируем для принятия решения
                if torch.rand(1).item() < draft_token_prob:
                    accepted_tokens.append(draft_tokens[0, i].item())
                else:
                    # Если токен отвергнут, сэмплируем новый из Target
                    new_token = torch.multinomial(target_probs, 1).item()
                    accepted_tokens.append(new_token)
                    break  # прерываем цепочку после первого отвергнутого
            
            # 4. Добавляем принятые токены к результату
            if accepted_tokens:
                generated.extend(accepted_tokens)
                # Обновляем current_input для следующей итерации
                current_input = torch.cat([
                    current_input,
                    torch.tensor([accepted_tokens], device=current_input.device)
                ], dim=1)
            else:
                # Если ни один токен не принят, генерируем один от Target
                target_next = target_model.generate(
                    current_input,
                    max_new_tokens=1,
                    do_sample=True,
                    temperature=0.7
                )
                next_token = target_next[:, current_input.shape[1]:].item()
                generated.append(next_token)
                current_input = torch.cat([
                    current_input,
                    torch.tensor([[next_token]], device=current_input.device)
                ], dim=1)
    
    # Декодируем результат
    return target_tokenizer.decode(generated)

4 Тестируем и сравниваем скорость

Теперь сравним обычную генерацию со speculative decoding:

import time

def benchmark_generation(method, prompt, num_runs=5):
    """Бенчмарк скорости генерации."""
    times = []
    
    for _ in range(num_runs):
        torch.cuda.synchronize()  # ждем завершения всех GPU операций
        start = time.time()
        
        result = method(prompt)
        
        torch.cuda.synchronize()
        end = time.time()
        times.append(end - start)
    
    avg_time = sum(times) / len(times)
    tokens_per_second = len(target_tokenizer.encode(result)) / avg_time
    
    return avg_time, tokens_per_second, result[:100]

# Промпт для теста
test_prompt = "Объясни, как работает speculative decoding, в трёх предложениях."

print("=== Бенчмарк скорости ===")
print(f"Промпт: {test_prompt}")
print()

# 1. Обычная генерация (Target модель)
def regular_generation(prompt):
    inputs = target_tokenizer(prompt, return_tensors="pt").cuda()
    outputs = target_model.generate(
        inputs.input_ids,
        max_new_tokens=100,
        do_sample=True,
        temperature=0.7
    )
    return target_tokenizer.decode(outputs[0])

reg_time, reg_tps, reg_preview = benchmark_generation(regular_generation, test_prompt)
print(f"Обычная генерация:")
print(f"  Время: {reg_time:.2f} сек")
print(f"  Токенов в секунду: {reg_tps:.1f}")
print(f"  Превью: {reg_preview}...")
print()

# 2. Speculative decoding
spec_time, spec_tps, spec_preview = benchmark_generation(
    lambda p: speculative_decode(p, max_new_tokens=100, k=5),
    test_prompt
)
print(f"Speculative decoding (k=5):")
print(f"  Время: {spec_time:.2f} сек")
print(f"  Токенов в секунду: {spec_tps:.1f}")
print(f"  Ускорение: {reg_time/spec_time:.1f}x")
print(f"  Превью: {spec_preview}...")

Оптимизации и подводные камни

Базовая реализация работает, но есть нюансы. Вот что чаще всего ломается:

Проблема 1: Разные токенизаторы

Если Target и Drafter модели используют разные токенизаторы (например, Llama и GPT-2), speculative decoding не сработает. Токены будут несовместимы. Решение:

1. Используйте модели из одного семейства (Llama 3 + Llama 2, Qwen + Qwen)
2. Или используйте универсальный токенизатор (tiktoken), но тогда придется переобучать Drafter

Проблема 2: Слишком большое K

Если установить K=20, Drafter будет часто ошибаться, и Target придется проверять длинные цепочки, большинство из которых будут отвергнуты. Оптимальное K зависит от схожести моделей:

• Для очень похожих моделей (например, 70B и 13B из одного семейства): K=5-8
• Для умеренно похожих: K=3-5
• Для разных семейств: K=2-3

Проблема 3: Качество Drafter модели

Drafter не должен быть совсем уж тупым. Если его accuracy ниже 70%, speculative decoding замедлит, а не ускорит генерацию. Как проверить:

def measure_drafter_accuracy(prompts_dataset, k=5):
    """Измеряет accuracy Drafter модели относительно Target."""
    correct = 0
    total = 0
    
    for prompt in prompts_dataset[:100]:  # первые 100 промптов
        # Получаем ground truth от Target
        inputs = target_tokenizer(prompt, return_tensors="pt").cuda()
        target_output = target_model.generate(
            inputs.input_ids,
            max_new_tokens=k,
            do_sample=False
        )
        target_tokens = target_output[:, inputs.input_ids.shape[1]:]
        
        # Получаем предсказания Drafter
        drafter_output = drafter_model.generate(
            inputs.input_ids,
            max_new_tokens=k,
            do_sample=False
        )
        drafter_tokens = drafter_output[:, inputs.input_ids.shape[1]:]
        
        # Сравниваем
        for i in range(min(k, target_tokens.shape[1])):
            if target_tokens[0, i] == drafter_tokens[0, i]:
                correct += 1
            total += 1
    
    return correct / total if total > 0 else 0

accuracy = measure_drafter_accuracy(["test prompt 1", "test prompt 2"], k=5)
print(f"Accuracy Drafter модели: {accuracy:.1%}")
print(f"Рекомендуемое K: {max(2, int(accuracy * 5))}")  # эвристика

Продвинутые техники

Когда освоите базовый вариант, попробуйте эти оптимизации:

1. Lookahead decoding

Вместо одного Drafter используйте несколько маленьких моделей, которые предсказывают токены на разных шагах вперед. Это увеличивает accuracy предсказаний, но требует больше памяти.

2. Adaptive K

Динамически меняйте K в зависимости от confidence Drafter модели. Если модель уверена в своих предсказаниях — увеличивайте K. Если нет — уменьшайте.

3. Пакетная обработка

Генерируйте несколько промптов параллельно. Это особенно эффективно в продакшн-средах, где запросы приходят пачками. В статье про масштабирование LLM я подробно разбирал техники пакетной обработки.

Когда speculative decoding НЕ работает

Не ждите чуда в этих сценариях:

• Очень короткие ответы (1-2 токена). Нет смысла предсказывать вперед
• Креативные задачи (поэзия, код). Drafter не может угадать творческий подход
• Модели с разными архитектурами (например, трансформер + Mamba)
• Когда VRAM на пределе. Две модели занимают больше памяти

Готовые решения и фреймворки

Не хотите писать код с нуля? Используйте эти библиотеки:

1. vLLM с speculative decoding

Самый простой способ. Установите vLLM 0.3.0+ и используйте встроенную поддержку:

pip install vllm

from vllm import LLM, SamplingParams
from vllm.model_executor.layers.spec_decode import MultiStepSpecDecodeWorker

# Инициализация с speculative decoding
llm = LLM(
    model="meta-llama/Llama-3-8B-Instruct",
    speculative_model="microsoft/Phi-3-mini-4k-instruct",
    speculative_draft_length=5,
    enforce_eager=True  # для отладки
)

# Генерация как обычно
sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
outputs = llm.generate(["Your prompt here"], sampling_params)
print(outputs[0].outputs[0].text)

2. Hugging Face TGI

Text Generation Inference от Hugging Face тоже поддерживает speculative decoding:

docker run --gpus all \
  -p 8080:80 \
  -v ~/.cache/huggingface:/data \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model-id meta-llama/Llama-3-8B-Instruct \
  --draft-model-id microsoft/Phi-3-mini-4k-instruct \
  --max-batch-total-tokens 4096 \
  --speculative-draft-tokens 5

3. Custom implementation в llama.cpp

Если вы используете GGUF-модели, в llama.cpp есть экспериментальная поддержка:

# Собираем llama.cpp с поддержкой speculative decoding
cd llama.cpp
make LLAMA_CUDA=1

# Запускаем с двумя моделями
./main -m ./models/llama-3-8b.Q4_K_M.gguf \
  --draft-model ./models/phi-3-mini.Q4_K_M.gguf \
  -p "Your prompt" \
  -n 100 \
  --speculative 5

Что делать, если нет подходящей Drafter модели?

Частая проблема: у вас есть Llama 3 70B, но нет маленькой Llama для Drafter. Варианты:

1. Дистиллируйте большую модель в маленькую. Обучите 1B-модель предсказывать выходы 70B-модели. Долго, но эффективно.
2. Используйте универсальную маленькую модель (Phi-3, Gemma, Qwen). Accuracy будет ниже, но работать будет.
3. Сгенерируйте synthetic dataset и дообучите маленькую модель. В гайде про синтетические данные я показывал, как это делать безопасно.

Мой выбор? Phi-3-mini (4B) как универсальный Drafter для любых 7B+ моделей. Работает в 80% случаев.

Итоговый чеклист перед продакшеном

• ☑️ Измерили accuracy Drafter (должно быть > 70%)
• ☑️ Подобрали оптимальное K (accuracy * 5)
• ☑️ Проверили совместимость токенизаторов
• ☑️ Протестировали на разных типах промптов
• ☑️ Настроили мониторинг скорости и качества
• ☑️ Добавили fallback на обычную генерацию при ошибках

Speculative decoding — не серебряная пуля. Это инструмент, который требует настройки. Но когда он настроен, результат впечатляет: ваши локальные LLM начинают работать в 2-3 раза быстрее без апгрейда железа. Вы экономите тысячи долларов на видеокартах и получаете отзывчивый интерфейс.

Самый частый вопрос, который мне задают после внедрения: "Почему эта техника не используется везде по умолчанию?" Ответ прост: потому что нужно думать. Выбрать Drafter модель, настроить K, проверить совместимость. Большинство разработчиков предпочитают "просто добавить ещё одну карту". Но теперь вы не из большинства.

P.S. Если speculative decoding дал вам ускорение 1.1x вместо ожидаемых 2.5x — не расстраивайтесь. Скорее всего, проблема в неоптимальном K или слабой Drafter модели. Уменьшите K до 2-3, попробуйте другую Drafter модель. Работает в 95% случаев.

P.P.S. Хотите увидеть, как speculative decoding работает в реальном времени на разных железах? В следующей статье разберу бенчмарки на RTX 3090, 4090 и даже на старом железе вроде майнинг-ригов. Спойлер: на 4x RTX 3090 speculative decoding даёт ускорение 4-5x — потому что memory-bound bottleneck ещё более выражен в многокарточных системах.