ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать speculative decoding с draft моделью

1. Цель задачи

Освоить технику speculative decoding — ускорение инференса большой языковой модели (target) с помощью маленькой быстрой модели (draft). Реализовать pipeline, в котором draft-модель (1B параметров) генерирует последовательность кандидатов, а target-модель (8B) проверяет их параллельно. Добиться практического ускорения в 1.5–2x по сравнению с последовательной генерацией только target-моделью.

Ключевой результат Рабочий скрипт или ноутбук, демонстрирующий speculative decoding на заданном промпте с замером скорости (токенов/сек) и измеренным speedup в диапазоне 1.5–2x.

2. Исходные данные

Перед началом необходимо иметь:

Если нет реального GPU (только CPU) — симулируем:

1. Используем stub для target-модели: симулируем задержку (~0.1 с на токен) и возвращаем детерминированный ответ.
2. Draft-модель — любая маленькая (например, distilgpt2), работает на CPU.
3. Измеряем логическое ускорение на синтетическом трафике без реального вычисления.

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка окружения и загрузка моделей (30-45 мин)

Действия

1. Создать окружение

   conda create -n spec_dec python=3.10
   conda activate spec_dec
   pip install torch transformers accelerate sentencepiece tqdm matplotlib pandas
   

2. Загрузить draft-модель (1B)
   Использовать AutoModelForCausalLM.from_pretrained с torch_dtype=torch.float16 (если GPU).
   Пример:

   from transformers import AutoModelForCausalLM, AutoTokenizer
   draft_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
   draft_model = AutoModelForCausalLM.from_pretrained(draft_name, torch_dtype=torch.float16, device_map="auto")
   draft_tokenizer = AutoTokenizer.from_pretrained(draft_name)
   

3. Загрузить target-модель (8B)
   Аналогично, но для экономии VRAM использовать bitsandbytes 4-bit квантизацию:

   target_name = "meta-llama/Llama-3.1-8B-Instruct"
   target_model = AutoModelForCausalLM.from_pretrained(
       target_name,
       torch_dtype=torch.float16,
       load_in_4bit=True,            # 4-bit экономит память
       device_map="auto"
   )
   target_tokenizer = AutoTokenizer.from_pretrained(target_name)
   

4. Проверить работоспособность каждой модели по отдельности
   Сгенерировать один промпт (например, "What is the capital of France?") и вывести ответ.

Ожидаемый результат этапа
Обе модели загружены, успешно генерируют текст при прямом вызове model.generate().

Этап 2: Реализация базового speculative decoding (2-3 часа)

Действия

1. Написать функцию принятия / отклонения (rejection sampling)
   Алгоритм (стандартный speculative decoding):

   • Draft-модель генерирует K токенов за один шаг (каждый токен выбирается из её распределения ( q )).
   • Target-модель вычисляет логиты для последнего токена и получает распределение ( p ).
   • Для каждого токена ( t_i ) принимаем его с вероятностью ( \min(1, p(t_i)/q(t_i)) ).
   • При первом отклонении генерируем токен из скорректированного распределения ( \max(0, p - q) ) и останавливаем цикл.

   def speculative_generate(draft_model, target_model, input_ids, max_new_tokens=50, K=5, temperature=1.0):
       # input_ids: (1, seq_len)
       accepted = 0
       for _ in range(max_new_tokens // K + 1):
           # 1. Draft: generate K candidates
           draft_output = draft_model.generate(input_ids, max_new_tokens=K, do_sample=True, temperature=temperature, pad_token_id=draft_tokenizer.eos_token_id)
           candidates = draft_output[0][input_ids.shape[1]:]   # (K,)
           # 2. Get target logits for last position
           with torch.no_grad():
               target_logits = target_model(input_ids).logits[0, -1, :]   # vocab
           # 3. Rejection sampling loop
           for i in range(K):
               q = draft_model output distribution for token i
               p = target_logits.softmax(-1)
               ratio = p[candidates[i]] / q[candidates[i]]
               if torch.rand(1).item() < min(1, ratio):
                   accepted += 1
                   input_ids = torch.cat([input_ids, candidates[i].unsqueeze(0).unsqueeze(0)], dim=-1)
               else:
                   # resample from (p - q)+
                   new_probs = torch.clamp(p - q, min=0)
                   new_probs /= new_probs.sum()
                   new_token = torch.multinomial(new_probs, 1)
                   input_ids = torch.cat([input_ids, new_token], dim=-1)
                   break
           if len(candidates) == accepted:
               # all K accepted, continue
               pass
       return input_ids
   

   Примечание: на практике нужно получать q для каждого шага; проще сохранять логиты draft-модели во время её генерации.

2. Оптимизация: за раз передавать target-модели весь блок принятых токенов
   Вместо пошагового forward'a, можно передать всю последовательность input_ids + accepted_tokens и получить логиты сразу для всех позиций. Это повышает эффективность.

3. Настроить параметры: ( K ) (число кандидатов)
   Обычно ( K = 5..10 ). Экспериментально подобрать под задачу.

Ожидаемый результат этапа
Функция speculative_generate корректно выводит текст (визуально похожий на target-модельный) и не падает с ошибками.

Этап 3: Интеграция и тестирование на простом промпте (30-45 мин)

Действия

1. Написать бенчмарк-функцию
   Сравнить три режима:

   • Только target-модель (baseline).
   • Только draft-модель (для справки).
   • Speculative decoding (draft + target).

   Для каждого режима замерить:

   • Время генерации (time.perf_counter).
   • Количество сгенерированных токенов.
   • Скорость (токенов/сек).
   • Speedup = (время baseline) / (время spec).

2. Запустить на 3-5 промптах длиной 20-50 токенов
   Выбрать разнообразные темы (факты, рассуждения, код).

3. Вывести таблицу результатов
   Пример:

   Промпт	Baseline (т/с)	Draft (т/с)	Spec (т/с)	Speedup
   "Explain quantum computing"	12.5	85.2	19.8	1.58x
   "Write fibonacci in Python"	11.3	79.1	18.1	1.60x

Ожидаемый результат этапа
Чёткое численное подтверждение speedup >1.5x хотя бы на одном промпте.

Этап 4: Оптимизация и тонкая настройка (1-2 часа)

Действия

1. Варьировать ( K )
   Попробовать ( K = 3, 5, 7, 10, 15 ). Построить график speedup vs K. Определить оптимальное.

2. Включить greedy decoding vs sampling
   Проверить, как работает speculative при temperature=0 (greedy). Желательно использовать одинаковый seed для сравнения.

3. Оптимизация batch-обработки
   Если позволяет память, запустить несколько промптов в батче.

4. Проверить влияние промпта на акцепт-рейт (acceptance rate)
   Вычислить средний % принятых токенов. Если <60%, возможно, draft-модель плохо аппроксимирует target.

Ожидаемый результат этапа
Оптимальное ( K ), acceptance rate 60-80%, стабильный speedup 1.5-2x.

Этап 5: Документирование и демонстрация (30-45 мин)

Действия

1. Оформить ноутбук или скрипт с чёткими ячейками: загрузка, speculative-функция, бенчмарк, графики.
2. Добавить комментарии к ключевым шагам алгоритма.
3. Подготовить краткое описание (в README или Markdown-ячейке) с таблицей результатов и выводом.

Ожидаемый результат этапа
Готовый артефакт (.ipynb или .py), который можно запустить и воспроизвести ускорение.

5. Критерии приемки (Definition of Done)

• Реализована функция speculative decoding, принимающая две модели и промпт.
• Код корректно работает на GPU (или CPU-симуляции).
• Проведено не менее 3 замеров для каждого режима (baseline, draft, spec).
• Speedup >= 1.5x при ( K = 5 ) на хотя бы одном промпте.
• График зависимости speedup от ( K ) построен и интерпретирован.
• Acceptance rate указан в результатах.
• Ноутбук/скрипт воспроизводим (указаны версии библиотек, seed).
• Написаны комментарии к алгоритму rejection sampling.

6. Ожидаемый результат

Основной артефакт
speculative_decoding.ipynb (Jupyter Notebook) или speculative_decoding.py (Python-скрипт), содержащий:

• Загрузку draft (1B) и target (8B) моделей.
• Функцию speculative_generate.
• Бенчмарк-функцию с измерением времени.
• Таблицу и график speedup.
• Выводы о параметре ( K ) и acceptance rate.

Дополнительно (опционально):

• requirements.txt с фиксированными версиями.
• README.md с краткой инструкцией.

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: для первого раза может потребоваться больше времени на отладку rejection sampling. Рекомендуется начинать с симуляции на CPU.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я загрузил обе модели и проверил их отдельную работоспособность.
• Моя функция speculative decoding корректно обрабатывает случай отсутствия принятых токенов.
• Я сравнил вывод speculative decoding с выводом только target-модели — тексты не расходятся катастрофически.
• Я зафиксировал seed для воспроизводимости результатов.
• Я построил график speedup от K и выбрал оптимальное значение на основе acceptance rate.