Оглавление

Когда речь заходит о Mixture-of-Experts (MoE), большинство сразу представляет себе гигантские модели с триллионами параметров, доступные только Google и другим гигантам. Cerebras в своем практическом руководстве развенчивает этот миф, показывая, как можно обучать MoE-модели даже на обычном ноутбуке.

Демистификация MoE: не только для гигантов

MoE-модели действительно ассоциируются с колоссальными масштабами — триллионы параметров, как в случае с Kimi Team 2025. Но это похоже на отказ от нейросетей только потому, что существует GPT-4. Как отмечают авторы, можно создать перцептронную сеть с нуля менее чем в 20 строках кода.

Ключевой инсайт здесь в том, что MoE — это архитектурный паттерн, а не исключительно вопрос масштаба. Те же принципы работают и на скромных конфигурациях, просто с меньшим количеством экспертов и параметров.

Техническая основа: модернизированный GPT-2

В основе эксперимента — модифицированная версия nanoGPT от Andrej Karpathy. Модель представляет собой decoder-only трансформер с современными улучшениями:

  • ReLU2 вместо GELU (~1-2% прирост производительности)
  • RoPE для позиционного кодирования
  • RMSNorm вместо LayerNorm
  • QK-Norm для стабилизации обучения
  • Оптимизатор Muon

Такой подход позволяет обучить базовую модель на датасете FineWeb всего за 11 минут.

Архитектура экспертов: от теории к коду

Каждый эксперт в MoE-архитектуре — это буквально копия FFN-слоя из GPT-2:

class MLP(nn.Module):
 def __init__(self, config):
 super().__init__()
 self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
 self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
 self.c_proj.weight.data.zero_()

 def forward(self, x):
 x = self.c_fc(x)
 x = F.relu(x).square()
 x = self.c_proj(x)
 return x

Эксперты организуются в ModuleList, а их вывод комбинируется через функцию expert_mixing, которая обрабатывает батчи токенов. Авторы честно предупреждают: эта реализация не оптимизирована для скорости и тратит вычисления впустую, но идеально подходит для образовательных целей.

Learned Routing: промышленный стандарт с подводными камнями

Изучаемая маршрутизация — это то, что используют Google, OpenAI и Anthropic в своих MoE-моделях. Она решает две ключевые задачи:

  1. Находит лучшего эксперта для каждого токена через обучаемые веса
  2. Поддерживает балансировку нагрузки между экспертами

Реализация включает router на основе линейного слоя и вспомогательную функцию потерь для балансировки:

class LearnedRouter(nn.Module):
 def __init__(self, input_dim, num_experts, top_k):
 super().__init__()
 self.router = nn.Linear(input_dim, num_experts, bias=False)
 self.top_k = top_k

 def forward(self, x):
 logits = self.router(x)
 probs = logits.softmax(dim=-1)
 gate, topk_idx = torch.topk(probs, self.top_k, dim=-1)
 gate = gate / gate.sum(dim=-1, keepdim=True)

Вспомогательная потеря рассчитывает распределение нагрузки и сравнивает его с ожидаемым равномерным распределением.

MoE с 4 экспертами и маршрутизацией превосходит плотную модель GPT-2. Разрыв в потерях увеличивается в течение обучения, достигая 1% к целевому показателю в 2.4B токенов.

Утилизация экспертов (усредненная по слоям) быстро сходится к почти 25%, что представляет идеально сбалансированный случай для 4 экспертов.

Самая интересная проблема, которую обнаружили авторы — маршрутизатор практически не получает градиента от потерь перекрестной энтропии. Это означает, что он учится только на вспомогательной функции потерь, а не на основной задаче! Именно такие инсайты делают руководство ценным — оно показывает реальные проблемы, а не только успехи.

Материал с сайта Cerebras.