ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить RWKV для инференса

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

Научиться запускать инференс модели RWKV на CPU, измерить её задержку (latency) для последовательностей разной длины и сравнить с классическим трансформером (например, GPT-2 или BERT). Задача направлена на понимание преимуществ RWKV в работе с длинными контекстами за счёт линейной сложности внимания.

Ключевой результат Наглядный отчёт (Jupyter Notebook или markdown-файл) с графиками и таблицами, демонстрирующий, что RWKV работает не менее чем в 10 раз быстрее трансформера на последовательностях длиной ≥ 4096 токенов при инференсе на CPU.

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

Если нет подходящего реального CPU — можно использовать colab.research.google.com с ручным ограничением на CPU runtime (среда выполнения → изменить → "none" GPU).

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

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

Этап 1: Установка окружения и загрузка моделей (~1 час)

Действия

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

   python -m venv rwkv_bench
   source rwkv_bench/bin/activate  # Linux/Mac
   rwkv_bench\Scripts\activate     # Windows
   

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

   pip install torch --index-url https://download.pytorch.org/whl/cpu
   pip install transformers rwkv tqdm matplotlib pandas psutil humanize
   

3. Загрузить модели через Hugging Face transformers:

   from transformers import GPT2Tokenizer, GPT2Model, AutoModelForCausalLM
   # Трансформер (GPT-2)
   gpt2_model = GPT2Model.from_pretrained("gpt2")
   gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
   

4. Загрузить RWKV через rwkv:

   from rwkv.model import RWKV
   model = RWKV("RWKV-4-1b5-pile", strategy="cpu fp32")
   # strategy можно менять: "cpu fp16" для ускорения
   

5. Сохранить модели локально для повторных запусков (кэш Hugging Face).

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

Этап 2: Подготовка тестовых последовательностей (~30 минут)

Действия

1. Определить длины: [128, 256, 512, 1024, 2048, 4096, 8192] (возможно добавить 16384, если хватает RAM).

2. Сгенерировать единый текстовый корпус, из которого брать фрагменты заданной длины — чтобы не зависеть от содержания.

3. Токенизировать корпус токенизатором трансформера. Для RWKV требуется специальный токенизатор (используйте rwkv.utils.RWKVTokenizer или общий 20B_tokenizer). Рекомендуется использовать RWKV_TOKENIZER:

   from rwkv.tokenizer import RWKV_TOKENIZER
   tokenizer = RWKV_TOKENIZER("20B_tokenizer.json")
   

4. Создать словарь sequences = {length: token_ids} — каждый закэшировать для повторяемости.

Ожидаемый результат этапа Готовые тензоры/списки токенов для каждой длины, идентичные при запуске на обеих моделях.

Этап 3: Написание бенчмарк-скрипта (1–1.5 часа)

Действия

1. Реализовать функцию measure_latency(model_func, input_ids, n_warmup=5, n_loops=20):

   • warmup: 5 прогонов без замера (для прогрева кэша)
   • n_loops: 20 замеров, из которых брать среднее и std
   • Для трансформера: передать input_ids.unsqueeze(0), замерить model.forward() без генерации (просто один forward pass)
   • Для RWKV: передать тензор int64, замерить model.forward(input_ids) (возвращает logits)

2. Важно трансформер и RWKV должны обрабатывать одинаковое количество токенов за вызов (без последующей генерации новых токенов) — только время одного forward pass.

3. Зафиксировать peak memory (через psutil.Process().memory_info().rss).

4. Сохранить результаты в pandas.DataFrame:

   results = []
   for length in lengths:
       data = sequences[length]
       # прогрев
       for _ in range(5): model_forward(data)
       # замер
       times = []
       for _ in range(20):
           t0 = time.perf_counter()
           model_forward(data)
           times.append((time.perf_counter() - t0) * 1000)  # ms
       results.append({
           "length": length,
           "model": "RWKV",
           "latency_ms_mean": np.mean(times),
           "latency_ms_std": np.std(times),
       })
   

5. Выполнить аналогично для GPT-2.

Ожидаемый результат этапа Полученная таблица results_all.csv с колонками: model, length, latency_ms_mean, latency_ms_std, peak_memory_mb.

Этап 4: Анализ и визуализация (~30 минут)

Действия

1. Построить график latency (ms) vs sequence length с двумя линиями и доверительными интервалами (std):

   import matplotlib.pyplot as plt
   plt.figure(figsize=(10,6))
   for model in ["RWKV", "GPT-2"]:
       subset = results[results["model"]==model]
       plt.errorbar(subset["length"], subset["latency_ms_mean"],
                    yerr=subset["latency_ms_std"], label=model, marker='o')
       # регрессия O(n) vs O(n^2)
   plt.xscale('log')
   plt.yscale('log')
   plt.legend()
   plt.xlabel("Sequence length (tokens)")
   plt.ylabel("Latency (ms)")
   plt.grid(True)
   plt.savefig("latency_comparison.png")
   

2. Вычислить ускорение RWKV = latency_gpt2 / latency_rwkv для каждой длины. Построить гистограмму.

3. Проверить критическое условие задачи: для длины 4096 ускорение >= 10? Если нет — исследовать причину (например, неверная стратегия CPU, не та модель).

Ожидаемый результат этапа Графики latency_vs_length.png, speedup_over_gpt2.png и финальная таблица с выводами.

Этап 5: Формирование отчёта (30 минут)

Действия

1. Создать markdown-документ benchmark_report.md со структурой:

   • Введение (цель, модели)
   • Методология (настройки, количество прогонов, среда)
   • Результаты (таблица средних задержек и memory)
   • Графики (встроить base64 или ссылками)
   • Выводы (подтверждение/опровержение гипотезы 10x)
   • Возможные улучшения (fp16, batching, better CPU BLAS)

2. Если ускорение не достигло 10x — написать анализ причин и предложить альтернативы (более крупная модель RWKV, использование strategy='cpu fp16', увеличение batch_size до 1).

Ожидаемый результат этапа Финальный отчёт, готовый к отправке.

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

• Установлена и работоспособна модель RWKV (версия 4.x) на CPU.
• Установлена и работоспособна модель GPT-2 (или аналогичный трансформер) на CPU.
• Проведены замеры latency для как минимум 6 различных длин (128, 512, 1024, 2048, 4096, 8192).
• Каждый замер выполнен с 20 повторениями, вычислены среднее и стандартное отклонение.
• Построены графики «latency vs sequence length» для обеих моделей (с логарифмической шкалой).
• Вычислено ускорение RWKV относительно трансформера для каждой длины.
• Для длины 4096 (или более) ускорение ≥ 10 раз.
• Создан отчёт (Jupyter Notebook или markdown) с таблицами, графиками и выводами.
• В отчёте указаны технические детали: версии библиотек, параметры strategy, количество ядер CPU, объем RAM.

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

Основной артефакт Файл benchmark_report.md (или .ipynb), содержащий:

• описание методологии,
• таблицу результатов (latency, std, memory),
• два графика,
• вывод с подтверждением/опровержением 10x-ускорения.

Дополнительные артефакты

• results_all.csv — сырые данные замеров.
• latency_comparison.png и speedup_over_gpt2.png — отдельные изображения.
• benchmark_script.py — воспроизводимый Python-скрипт.

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

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

Примечание При первом выполнении время может увеличиться на 1–2 часа за счёт отладки и настройки strategy для RWKV.

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

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

• Я скачал(а) и загрузил(а) обе модели без ошибок.
• Я замерил время только одного forward pass (без генерации новых токенов).
• Прогревочные проходы (warmup) были сделаны перед замерами.
• Количество прогонов для каждой длины ≥ 20.
• Результаты сохранены в CSV, графики построены с доверительными интервалами.
• Я построил график ускорения и проверил, >=10x на длине 4096.
• В отчёте указаны версии библиотек, количество потоков CPU (OMP_NUM_THREADS), модель RWKV (точное название), модель трансформера.
• Отчёт воспроизводим: другой инженер может повторить бенчмарк по моим шагам.
• Я описал причины, если ускорение не достигнуто, и что можно попробовать (например, larger RWKV model, fp16, batch=1 vs batch на трансформере).