Почему MoE (Mixture of Experts) быстрее dense модели при инференсе?

Краткий тезис

Mixture of Experts (MoE) быстрее dense модели при инференсе за счёт sparse activation (разреженной активации): на каждый токен активируется только top-k экспертов (обычно k=2 из 8–256). Это резко снижает FLOPs до dense_FLOPs * (k / n_experts), хотя все эксперты загружены в память. Таким образом, при одинаковом качестве MoE-модель (например, Mixtral 8x7B) использует столько же вычислительных операций на токен, сколько dense-модель в 2–3 раза меньшего размера, но требует большей memory bandwidth из-за загрузки всех параметров.

1. Термин: Mixture of Experts (MoE)

MoE — это архитектура нейронной сети, где несколько экспертов (отдельных подсетей, обычно Feed-Forward Network (FFN)) объединены gating network (сеть-шлюзом). Gating network для каждого входного токена вычисляет веса всех экспертов, затем выбирает top-k экспертов с наибольшими весами и активирует только их. Выход — взвешенная сумма результатов выбранных экспертов.

Почему это важно для инференса
В стандартной dense модели (например, LLaMA) на каждом слое обрабатываются все параметры всех подсетей. В MoE-модели на каждом слое активируется лишь малая часть параметров (экспертов), что экономит вычисления при сохранении общей ёмкости модели.

2. Сравнение MoE и dense модели

Важное уточнение ускорение MoE проявляется, когда узким местом являются именно вычисления (compute-bound), а не пропускная способность памяти]] (memory-bound). Для маленьких batch size или очень разреженных моделей загрузка всех экспертов может стать доминирующим фактором.

3. FLOPs vs Memory Bandwidth для MoE

FLOPs (Floating Point Operations) — количество операций с плавающей точкой, необходимых для одного токена. В dense модели FLOPs ≈ 2 * (число параметров) для одного forward pass (с учётом операций attention). В MoE — примерно 2 * (число активированных параметров).

Memory bandwidth — скорость, с которой модель может загружать веса из HBM (high-bandwidth memory) в вычислительные блоки. Даже если активируется только малая часть экспертов, для пересылки весов всех экспертов всё равно требуется ‘загрузка всех параметров’ (all-gather или заменяющие техники). Это делает инференс MoE memory-bound при малых batch size.

Типичный компромисс

• При batch size = 1 (как в онлайн-чате): время загрузки всех экспертов начинает доминировать, ускорение от снижения FLOPs сглаживается.
• При большом batch size (пакетная обработка): веса загружаются один раз и используются для многих токенов, тогда снижение FLOPs даёт чистое ускорение.

4. Пример: Mixtral 8x7B

Mixtral 8x7B — популярная открытая MoE-модель от Mistral. Её характеристики:

• Total parameters: 8 × 7B ≈ 47B (плюс общий attention — ~1B) → ~47B параметров.
• Активируемые параметры: только 2 эксперта из 8 на каждом слое → ~13B параметров на токен.
• FLOPs на токен: примерно как у dense 13B модели.
• Качество: сопоставимо с dense 70B моделью (например, LLaMA-2 70B).

Почему быстрее

• FLOPs для 13B dense модели меньше, чем для 70B dense. Mixtral достигает качества 70B при вычислительных затратах ~13B.
• На практике, при batch size ≥ 8, Mixtral на GPU типа A100 показывает задержку (latency) в 2–4 раза меньшую, чем dense 70B модель.

Минусы

• Нужно загрузить 47B параметров → требуется больше VRAM, чем для 13B dense.
• Для маленьких batch size (1–2) может быть медленнее, чем 13B dense, из-за memory bandwidth.

5. Почему MoE быстрее: механика top-k активации

Архитектура типичного MoE-слоя:

1. Входной вектор ( x ) проходит через gating network: [ w_i = [text](/wiki/text){softmax}(W_g \cdot x)i, \quad i=1..n{experts} ]
2. Выбираются top-k экспертов с наибольшими весами (обычно k=2).
3. Только выбранные эксперты вычисляют выход: [ y = \sum_{i \in top-k} w_i \cdot E_i(x) ]
4. Остальные эксперты не вычисляются (sparse computation).

Эффект на вычисления

• Если dense FFN-слой имеет размерность ( d_{model} \times 4d_{model} ) (как в трансформерах), то один эксперт — это FFN тех же размеров. Для n_experts = 8, k=2, активируется 2/8 = 25% от полных FFN-параметров.
• FLOPs для одного MoE-слоя: ( [text](/wiki/text){FLOPs}{[text](/wiki/text){dense FFN}} \times \frac{k}{n{experts}} ).

6. Недостатки MoE, связанные с быстродействием

• Memory overhead: все эксперты должны быть загружены в VRAM. Для модели MoE 8×7B нужно ~47B × 2 bytes (float16) ≈ 94 GB, что требует нескольких GPU даже для инференса.
• Load balancing: gating network может перегружать одни эксперты и недоиспользовать другие; это снижает эффективность и может привести к дополнительному расчёту.
• Batch size зависимость: при инференсе с маленьким batch (онлайн-запросы) memory bandwidth часто bottleneck, и MoE может быть не быстрее dense при равном числе активированных параметров.
• Сложность параллелизации: требуется эффективная коммуникация между устройствами для обмена результатами экспертов (all-to-all).

7. Когда MoE даёт реальный выигрыш в скорости

Вывод на практике MoE применяют в серверных решениях, где можно собрать batch запросов или использовать модели с большим числом экспертов (например, Switch Transformer, Mixtral).

8. MoE в контексте RAG и AI-агентов

Хотя вопрос об архитектурной эффективности MoE, он прямо влияет на Agentic RAG:

• Multi-agent системы часто запускают несколько специализированных моделей (агентов) параллельно. MoE может быть внутренней оптимизацией каждого агента.
• Sparse активация позволяет иметь большой пул знаний (экспертов) с малым временем ответа, что критично для интерактивных агентов.
• Методика LoRA + MoE: для fine-tuning RAG-моделей может применяться разреженный апдейт экспертов — ускоряет дообучение и инференс.

Пет-проект для закрепления

Задача: Реализовать упрощённый MoE-слой и сравнить latency с эквивалентным dense-слоем при разных batch size.

Инструменты: Python, PyTorch, CUDA-профилировщик (torch.cuda.Event), библиотека matplotlib для графиков.

Шаги:

1. Напишите класс MoELayer с параметрами: d_model, d_ff, num_experts, k.
2. В forward: gate → top-k индексы → вычислить результаты только выбранных экспертов → взвешенная сумма.
3. Напишите класс DenseLayer с той же размерностью d_model, d_ff.
4. Для каждого batch size от 1 до 128 замеряйте среднее время 100 forward passов для обоих классов на одном и том же входном тензоре.
5. Постройте график batch_size vs время и FLOPs (расчётное) vs время.

Ожидаемый результат: Вы увидите, что при batch size > 32 MoE значительно быстрее, а при batch=1 dense может обгонять MoE.

Код (фрагмент):

import torch
import torch.nn as nn
import time

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, k=2):
        super().__init__()
        self.k = k
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.ReLU(),
                nn.Linear(d_ff, d_model)
            ) for _ in range(num_experts)
        ])
        self.gate = nn.Linear(d_model, num_experts)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        gate_logits = self.gate(x)  # (batch, seq_len, num_experts)
        weights, indices = torch.topk(gate_logits, self.k, dim=-1)
        weights = torch.softmax(weights, dim=-1)  # нормализация
        out = torch.zeros_like(x)
        for i in range(self.k):
            expert_idx = indices[..., i]
            expert_weight = weights[..., i].unsqueeze(-1)
            for b in range(x.shape[0]):
                # batch-wise вызов экспертов (упрощённо)
                expert = self.experts[expert_idx[b, 0].item()]
                out[b] += expert_weight[b] * expert(x[b])
        return out

Связь с другими вопросами

Навигация

• Предыдущий: 434
• Следующий: 436
• Индекс: 00. Индекс разборов