ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить Mixture of Experts (Mixtral)

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

Научиться загружать и анализировать механизм маршрутизации (routing) в архитектуре Mixture of Experts (MoE) на примере модели Mixtral 8x7B. Практически исследовать, как gating network распределяет токены между экспертами, и визуализировать специализацию экспертов на синтетическом наборе текстов разных жанров (технический, художественный, разговорный).

Ключевой результат Набор визуализаций (тепловая карта распределения токенов по экспертам, график доминирующих экспертов для каждого жанра) и краткий аналитический отчёт о том, какие эксперты активируются на разных типах данных.

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

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

1. Ограничение ресурсов вместо полной Mixtral-8x7B используем allenai/tiny-MoE (0.5B), которая тоже имеет routing и 4 эксперта. Код будет идентичен.
2. Недостаток данных генерируем 3 текста (tech, fiction, dialogue) с помощью ChatGPT или пишем вручную по 500 токенов каждый.
3. Отсутствие GPU используем CPU с квантованием 8-bit (BitsAndBytes) — скорость упадёт, но анализ выполним на 100–200 токенах.

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

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

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

Действия

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

   pip install torch transformers accelerate bitsandbytes matplotlib seaborn tqdm
   

2. Создать скрипт analyze_moe.py

   • Загрузить модель с опциями:

     from transformers import AutoModelForCausalLM, AutoTokenizer
     import torch
     
     model_name = "mistralai/Mixtral-8x7B-v0.1"  # или "allenai/tiny-MoE"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         device_map="auto",           # auto-распределение по GPU/CPU
         torch_dtype=torch.bfloat16,  # экономия памяти
         load_in_8bit=True,           # при нехватке VRAM
         trust_remote_code=True       # для tiny-MoE
     )
     model.eval()
     

3. Проверить архитектуру

   • Вывести model.config.num_local_experts и model.config.num_experts_per_tok.
   • Убедиться, что в каждом слое есть block_sparse_moe.gate (для Mixtral).

Ожидаемый результат этапа Модель загружена, получено количество экспертов (8) и количество экспертов на токен (2). Скормить один тестовый токен и проверить, что gate.forward возвращает тензор формы [batch, seq_len, num_experts].

Этап 2: Подготовка тестовых данных (15 минут)

Действия

1. Написать 3–5 текстов разного жанра

   • Технический "The gradient of the loss with respect to the weights is computed using backpropagation. Batch normalization reduces internal covariate shift."
   • Художественный "The old man stared at the sea, his weathered hands gripping the oars. The waves whispered secrets only he could understand."
   • Разговорный "Hey, what's up? I was thinking we could grab a coffee later. Did you finish the report?"
   • (опционально: научная статья / новости)

2. Токенизировать

   tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(model.device)
   

Ожидаемый результат этапа Список из 3–5 тензоров input_ids разной длины, готовых к прогону через модель.

Этап 3: Прямой проход и сбор routing-информации (1–2 часа)

Действия

1. Написать хук на слой MoE
   Зарегистрировать forward_hook для каждого слоя, чтобы перехватывать routing_weights (gating probabilities).

   routing_data = []  # список: (layer_idx, token_idx, expert_probs)
   
   def hook_fn(module, input, output):
       # внутри MoE: gate возвращает routing_weights
       gate_output = module.gate(input[0])  # [batch, seq, num_experts]
       routing_data.append(gate_output.detach().cpu())
       return output
   
   for layer in model.model.layers:
       if hasattr(layer, 'block_sparse_moe'):
           layer.block_sparse_moe.gate.register_forward_hook(hook_fn)
   

2. Прогнать каждый текст через модель в режиме with torch.no_grad():

   model(tokens.input_ids)
   

3. После прохода обработать routing_data

   • Для каждого слоя получить матрицу [seq_len, num_experts] вероятностей.
   • Определить, какой эксперт получил максимальную вероятность для каждого токена (argmax).
   • Сохранить результаты в pandas DataFrame: колонки layer, token_index, expert_id, prob, text_id.

Ожидаемый результат этапа Собранные gating probabilities для каждого слоя и каждого токена. Данные сохранены в CSV или pickle.

Этап 4: Визуализация специализации экспертов (1 час)

Действия

1. Построить тепловую карту распределения экспертов по слоям:

   • Для каждого слоя подсчитать долю токенов, отнесённых к каждому эксперту.
   • sns.heatmap с осью X — эксперты (0–7), Y — номера слоёв.

2. Построить столбчатую диаграмму доминирующего эксперта по жанрам:

   • Для каждого текста определить, какой эксперт активируется чаще всего (по всем слоям).
   • Сгруппировать по text_id и посчитать counts. Построить plt.bar.

3. Опционально: анализ top-2 экспертов (Mixtral использует 2 активных эксперта на токен):

   • Для каждого токена взять два эксперта с наибольшими вероятностями. Оценить, насколько часто пары (i,j) встречаются в разных слоях/жанрах.

4. Сохранить все графики в PNG (300 dpi).

Ожидаемый результат этапа Две визуализации в файлах heatmap.png и expert_specialization_by_genre.png.

Этап 5: Анализ и выводы (30 минут)

Действия

1. Написать краткий отчёт (markdown), включающий:

   • Гистограмму загрузки экспертов (сколько токенов обработал каждый эксперт).
   • Выводы: какие эксперты специализируются на каких типах данных (например, эксперт 3 — технический текст, эксперт 5 — диалоги).
   • Заметки, если специализация неочевидна (возможно, для маленьких моделей специализация слабая).

2. Сравнить с ожидаемым поведением в Mixtral эксперты не обучаются на конкретные жанры, но анализ может показать определённые склонности.

Ожидаемый результат этапа Файл analysis_report.md с графиками и выводами.

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

• Модель загружена и корректно выполнен прямой проход без ошибок.
• Собраны gating probabilities минимум для 10 слоёв (или всех слоёв модели).
• Получено распределение токенов по экспертам для каждого текста.
• Построена тепловая карта “Layer vs Expert activation frequency”.
• Построена столбчатая диаграмма “Top expert per genre”.
• В отчёте указаны минимум две жанровые специализации (или обоснование их отсутствия).
• Все визуализации сохранены в PNG и приложены к результату.
• Код воспроизводим: единственный скрипт или Colab notebook с комментариями.

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

Опционально: анимация потока токенов по слоям (только при наличии времени).

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

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

Примечание: Время может увеличиться при первом выполнении из-за загрузки модели и отладки хуков. Для маленького MoE (tiny-MoE) время сокращается до 2–3 часов.

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

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

• Я загрузил модель и убедился, что она содержит block_sparse_moe.
• Я зарегистрировал forward_hook на gate каждого MoE-слоя.
• Я прогнал минимум 3 текста разных жанров и собрал gating probabilities.
• Я построил тепловую карту и столбчатую диаграмму.
• Я написал отчёт с конкретными наблюдениями (например, “эксперт 4 чаще активируется на технических текстах”).
• Код сохранён в формате .ipynb или .py с комментариями.
• Все артефакты (графики, CSV, отчёт) лежат в одной папке.