ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить CUDA graphs для коротких запросов

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

Научиться применять CUDA graphs для ускорения инференса LLM при batch=1 и коротких промптах (до 50 токенов). Выполнить capture и replay графа вычислительных операций, измерить латентность до и после оптимизации. Добиться снижения медианной latency на 30% относительно baseline без потери качества генерации.

Ключевой результат Воспроизводимый скрипт, демонстрирующий ускорение инференса с помощью CUDA graphs на синтетических запросах.

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

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

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

1. Установите PyTorch в CPU-режиме (CUDA не нужна) — CUDA graph недоступен, но можно написать заглушку.
2. Напишите код, который вызывает ошибку при отсутствии CUDA, но логически моделирует capture/replay через хранение графа операций.
3. Для замера времени используйте time.perf_counter() на CPU — ускорение будет нулевым, но структура задачи сохранится.

Важно Реальное ускорение возможно только на GPU с CUDA. Если нет доступа, задачу можно выполнить как исследовательскую.

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

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

Этап 1: Подготовка окружения и базовая инференс-функция (1 час)

Действия

1. Установить зависимости:

pip install torch transformers matplotlib pandas tqdm

2. Написать функцию инференса без оптимизации (baseline):
   • Использовать модель из transformers (например prajjwal1/bert-tiny для быстрой отладки или gpt2 для реалистичного сценария).
   • Параметры: batch_size=1, max_length=50 (prompt 50 токенов, генерация 1 токена).
   • Режим no_grad, model.eval().
   • Вернуть latency одного forward pass (не включая токенизацию).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

device = torch.device("cuda")
model = AutoModelForCausalLM.from_pretrained("gpt2").to(device).eval()
tokenizer = AutoTokenizer.from_pretrained("gpt2")

def baseline_infer(input_ids: torch.Tensor) -> float:
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    with torch.no_grad():
        outputs = model(input_ids)
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end)  # ms

3. Проверить корректность: запустить инференс на синтетическом batch (1×50), убедиться в отсутствии ошибок.

Ожидаемый результат этапа Рабочая функция baseline_infer, которую можно замерить.

Этап 2: Профилирование baseline latency (30 минут)

Действия

1. Замерить latency на 100 прогонах, вычислить среднее, медиану, p95, min/max.
2. Зафиксировать baseline (например, медиана 12.5 ms).
3. Изучить профиль с помощью torch.cuda.profiler или nsys (опционально): посмотреть количество запусков CUDA kernel, их длительность.

Ожидаемый результат этапа Численные значения latency — цель для оптимизации.

Этап 3: Реализация CUDA graphs (2–3 часа)

Действия

1. Понимание ограничений CUDA graphs

   • Граф может быть записан только один раз, после чего воспроизводится с теми же размерами тензоров.
   • Не поддерживаются динамические управления потоком (if/for, зависящие от данных).
   • Тензоры должны быть фиксированного размера (static shapes).

2. Cоздаём «статический» инференс-пайплайн

   • Фиксируем input_ids размером (1, 50).
   • Используем torch.cuda.graph для захвата forward pass.

from torch.cuda import CUDAGraph

# Подготовка статических тензоров
static_input = torch.randint(0, 50256, (1, 50), device=device)

# Захват графа
graph = CUDAGraph()
with torch.cuda.graph(graph):
    static_output = model(static_input)

# Важно: static_output должен быть сохранён, иначе ссылка потеряется

3. Проверка, что граф захвачен:

   • Вызвать graph.replay() на том же static_input (или на скопированных данных через copy_).
   • Убедиться, что результат воспроизводится.

4. Обновление входных данных

   • При каждом новом запросе копируем новые токены в static_input:

     new_input = ... # (1, 50)
     static_input.copy_(new_input)
     graph.replay()
     # static_output теперь содержит результат для new_input
     

Ожидаемый результат этапа Функция cuda_graph_infer с использованием graph.replay(), которая выдаёт тот же логгит-массив, что и baseline.

Этап 4: Интеграция с генерацией (1 час)

Действия

1. Проблема При генерации следующего токена длина последовательности меняется (50→51). CUDA graph требует статического размера.
2. Решение Выделить максимальную длину (например, 1024) и делать padding справа. Для коротких запросов всегда подавать padding до max_len.
3. Реализовать пайплайн
   • Pad исходных 50 токенов до 1024 (слева или справа в зависимости от модели).
   • Использовать attention_mask с единицами на реальных токенах и нулями на паддинге.
   • Граф захватывать с attention_mask как статический тензор.
   • При каждом шаге копировать только актуальные токены и маску.

MAX_LEN = 1024

static_input = torch.zeros((1, MAX_LEN), dtype=torch.long, device=device)
static_mask = torch.zeros((1, MAX_LEN), dtype=torch.long, device=device)

with torch.cuda.graph(graph):
    static_output = model(static_input, attention_mask=static_mask)

Ожидаемый результат этапа Генерация одного токена с помощью CUDA graph. Latency замеряется на одном forward pass (без учёта первого раза на capture).

Этап 5: Бенчмаркинг и анализ (1 час)

Действия

1. Замерить latency с CUDA graphs (после capture, только replay) — 100 прогонов на новых случайных данных.
2. Сравнить с baseline
   • Вычислить процент улучшения: (baseline_median - graph_median) / baseline_median * 100.
   • Построить гистограмму latency.
3. Валидировать корректность
   • Для случайного примера сравнить выводы baseline и graph (torch.allclose(..., atol=1e-5)).
4. Документировать результаты в виде таблицы

Ожидаемый результат этапа Подтверждение уменьшения latency на ≥30%.

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

• Скрипт успешно захватывает CUDA graph для forward pass модели (batch=1, input length=50).
• После capture выполняется graph.replay() с копированием новых входных данных.
• Результат graph.replay() идентичен baseline (check allclose).
• Замеры latency проведены на 100 запусках каждого метода.
• Медианная latency с graph снижена минимум на 30% относительно baseline.
• Сценарий с padding до фиксированной длины (например, 1024) работает без ошибок.
• Код сопровождается README с инструкцией по запуску и результатами.

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

Основной артефакт — Python-скрипт cuda_graph_benchmark.py, который:

• Загружает модель (параметризуется через аргументы командной строки).
• Выполняет baseline замеры.
• Выполняет capture и замеры с CUDA graph.
• Выводит таблицу сравнения latency.
• Сохраняет график latency_comparison.png.

Опционально:

• Jupyter notebook с визуализацией profiling trace.
• Конфигурационный файл для разных моделей и длин.

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

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

Примечание Для первого раза рекомендуется заложить +2 часа на отладку capture. Если модель большая (≥7B), время может увеличиться.

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

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

• Я настроил окружение и убедился, что CUDA доступна (torch.cuda.is_available()).
• Я написал baseline замеры с использованием torch.cuda.Event и синхронизацией.
• Я создал статические тензоры для input и attention_mask фиксированного размера.
• Я успешно выполнил capture графа и проверил, что static_output живёт после выхода из контекстного менеджера.
• Я сравнил вывод baseline и graph на нескольких случайных примерах и убедился в идентичности.
• Я измерил latency 100 раз для каждого метода и вычислил процент ускорения.
• Я задокументировал результаты и подготовил README.