ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить expert parallelism для Mixtral

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

Научиться распределять эксперты модели Mixtral 8x7B (MoE) между несколькими GPU, используя техники expert parallelism. Требуется сконфигурировать инференс (и/или дообучение) так, чтобы 8 экспертов были размещены на 4 GPU (по 2 эксперта на GPU), избежав Out‑of‑Memory (OOM) при полной загрузке модели (~47B параметров в FP16). Ключевой результат модель Mixtral‑8x7B успешно выполняет forward pass и генерацию текста на 4×A100‑40GB (или эквивалентных GPU) с использованием expert parallelism, измерены утилизация GPU и скорость инференса.

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

Если нет реальных 4×40GB GPU — симулируем:

1. Использовать 2×GPU с CPU offloading (accelerate device_map="auto") — но это не даст expert parallelism. Для симуляции можно использовать accelerate с num_processes=4 на одной машине с меньшим числом GPU, используя cpu для части экспертов (но цель — настоящий parallel). Рекомендуется арендовать кластер на $1‑2/час (RunPod, Vast.ai) — это дешевле, чем покупать оборудование.
2. Валидировать конфигурацию на малой версии модели Mixtral‑8x7B (Minitron 8x2B) если нет доступа к полной модели.

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

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

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

Действия

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

   conda create -n mixtral_ep python=3.10 -y
   conda activate mixtral_ep
   

2. Установить библиотеки

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   pip install transformers accelerate bitsandbytes sentencepiece
   

3. Запросить доступ к модели на Hugging Face (mistralai/Mixtral-8x7B-Instruct-v0.1), создать токен и сохранить в переменную окружения:

   export HF_TOKEN="hf_yourtoken"
   

4. Проверить доступность GPU

   import torch
   print(torch.cuda.device_count())  # должно быть ≥4
   

5. Загрузить модель без шардинга (только для проверки OOM — не запускать на 1 GPU, если памяти <80GB; используем CPU‑offload):

   from transformers import AutoModelForCausalLM, AutoTokenizer
   model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", device_map="cpu", offload_folder="./offload")
   tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
   print("Модель загружена (CPU).")
   

Ожидаемый результат этапа Рабочее окружение с установленными зависимостями и загруженной моделью (хотя бы на CPU).

Этап 2: Анализ архитектуры и шардинг с помощью Accelerate (1 час)

Действия

1. Изучить архитектуру MoE определить, какие модули являются экспертами (внимание: в Mixtral 8x7B каждый блок имеет 8 экспертов в MLP).

   for name, module in model.named_modules():
       if 'block_sparse_moe' in name and 'experts' in name:
           print(name, module)
   

2. Настроить accelerate config для expert parallelism (размещение каждого эксперта на отдельной GPU):

   • Создать файл deepseed_config.json (или inference_config.json) с tensor_parallel:

   {
     "tensor_parallel_size": 4,
     "dtype": "float16",
     "replace_with_kernel_inject": true
   }
   

   • Загрузить модель с device_map="auto" и max_memory:

   max_memory = {i: "38GB" for i in range(4)}
   model = AutoModelForCausalLM.from_pretrained(
       "mistralai/Mixtral-8x7B-Instruct-v0.1",
       device_map="auto",
       max_memory=max_memory,
       torch_dtype=torch.float16,
       load_in_8bit=False,
       trust_remote_code=True
   )
   

   Если device_map="auto" не распределяет эксперты равномерно, применить accelerate.infer_auto_device_map с пользовательской схемой.

3. Использовать accelerate с dispatch_model для ручного размещения:

   • Получить карту устройства для каждого параметра, принудительно назначив эксперты на разные GPU:

   from accelerate import infer_auto_device_map, dispatch_model
   device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["MixtralSparseMoeBlock"])
   # Ручная корректировка: все эксперты блока X -> GPU (X % 4)
   for name, param in model.named_parameters():
       if 'expert' in name:
           # определить номер эксперта
           ... 
   model = dispatch_model(model, device_map=device_map)
   

   (Упрощённый вариант: полагаться на Accelerate, но проверить распределение экспертов.)

4. Верифицировать распределение

   for name, _ in model.named_parameters():
       print(name, _.device)
   

   Убедиться, что каждый эксперт лежит на своём GPU (4 группы по 2 эксперта).

Ожидаемый результат этапа Модель загружена на 4 GPU с равномерным распределением экспертов (каждый GPU — ~12GB экспертов + shared layers).

Этап 3: Инференс и бенчмарк (45 мин)

Действия

1. Создать тестовый промпт

   tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
   prompt = "[INST] Explain the concept of Mixture of Experts in 3 sentences. [/INST]"
   inputs = tokenizer(prompt, return_tensors="pt").to(0)  # input_ids на GPU 0
   

2. Запустить генерацию через model.generate с torch.no_grad():

   import time
   start = time.time()
   outputs = model.generate(
       **inputs,
       max_new_tokens=128,
       do_sample=False
   )
   latency = time.time() - start
   tokens = outputs.shape[1] - inputs['input_ids'].shape[1]
   print(f"Generated {tokens} tokens in {latency:.2f}s -> {tokens/latency:.2f} tok/s")
   

3. Измерить пиковое использование памяти на каждом GPU (до и после):

   import subprocess
   def get_gpu_mem(gpu_id):
       result = subprocess.run(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'], capture_output=True, text=True)
       lines = result.stdout.strip().split('\n')
       return int(lines[gpu_id])
   for i in range(4):
       print(f"GPU {i}: {get_gpu_mem(i)} MB")
   

4. Сравнить с baseline на 1 GPU (если возможно с offloading) — должно быть значительное снижение пиковой памяти (с ~80GB+ до <38GB на GPU).

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

Этап 4: Оптимизация и повторное тестирование (45 мин)

Действия

1. Использовать FlashAttention v2 (если поддерживается):

   model.to_bettertransformer()  # из optimum
   

   Или attn_implementation="flash_attention_2" при загрузке.
2. Попробовать снизить точность до torch.bfloat16 (если поддерживается GPU).
3. Изменить количество экспертов на GPU (например, 3+3+1+1) и сравнить производительность.
4. Использовать deepseed / inference для expert parallelism:
   • Установить deepseed (deepseed‑moe)
   • Запуск через deepspeed --num_gpus 4 inference.py
5. Провести sweep с разными max_memory (35GB, 38GB, 40GB) и замерить, как меняется скорость.

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

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

Действия

1. Написать скрипт run_mixtral_ep.py, который автоматически:
   • Загружает модель с expert parallelism;
   • Запускает тестовый промпт;
   • Выводит метрики (tokens/s, память).
2. Сохранить конфигурацию accelerate в файл (accelerate config) и приложить.
3. Оформить краткий отчёт (Markdown): архитектура, схема распределения экспертов, результаты бенчмарка, выводы.

Ожидаемый результат этапа Воспроизводимый артефакт с инструкцией и результатами.

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

• Модель Mixtral‑8x7B загружается на 4 GPU без OOM (максимум 38GB/GPU).
• Каждый эксперт размещён на отдельной группе GPU (не более 3 экспертов на один GPU).
• Успешная генерация текста длиной 128 токенов (output не пустой).
• Скорость инференса ≥ 5 tokens/s (в сумме по GPU).
• Пиковое использование GPU memory на каждом устройстве < 40GB.
• В логах nvidia-smi зафиксировано, что все 4 GPU активны (utilization >50%).
• Конфигурация device_map зафиксирована в коде и может быть повторно использована.
• Написан скрипт для однокомандного запуска (+ инструкция).

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

Основной артефакт

• Файл mixtral_ep_demo.py — Python-скрипт, который выполняет:
  1. Загрузку модели с expert parallelism на 4 GPU;
  2. Прогон тестового промпта;
  3. Вывод latency, tokens/s и памяти GPU.

Дополнительно

• Файл benchmark_results.md — таблица результатов (попытки с разными настройками).
• Конфиг accelerate.yaml (если использовался accelerate config).

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

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

Примечание для первого раза: возможно потребуется дополнительное время на отладку device_map (20–40 мин) и на скачивание модели (зависит от скорости интернета — заложите 15 мин на скачивание 90GB).

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

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

• Я загрузил модель Mixtral‑8x7B на 4 GPU и проверил, что каждый эксперт находится на отдельной группе GPU.
• Я измерил пиковое использование памяти – ни один GPU не превысил 38GB.
• Я получил осмысленный вывод генерации (текст, а не только пустые токены).
• Я записал скорость (tokens/s) и сравнил с теоретической производительностью GPU (приблизительно 300 TFLOPS FP16 на A100).
• Я оформил код в виде воспроизводимого скрипта с комментариями по требуемым зависимостям.
• Я проверил, что при повторном запуске результаты стабильны (не падает с OOM из-за утечек памяти).