Оглавление

Модели Mixture-of-Experts (MoE) обещают революцию в эффективности обучения — больше параметров без пропорционального роста вычислительных затрат. Но когда дело доходит до практической реализации на реальном оборудовании, теоретические преимущества разбиваются о суровую реальность аппаратных ограничений, пишет Cerebras.

Проблема памяти: почему простые реализации не работают

Наивная реализация MoE, где все эксперты загружаются в память «на всякий случай», быстро упирается в ограничения GPU. При увеличении количества экспертов требования к памяти растут линейно, что делает невозможным обучение больших моделей на одном устройстве.

Пример смешения экспертов для пакетной обработки (см. train_gpt_moe.py):

def expert_mixing(x, experts, expert_ids, gate):
 # x: входные токены [B,T,C], experts: список экспертных сетей
 # expert_ids: выбранные эксперты на токен [B,T,top_k], gate: веса смешения экспертов [B,T,top_k]
 # B - размер батча, T - длина последовательности, C - размерность эмбеддинга
 B, T, C = x.shape
 x = x.reshape(B * T, C)
 gate = gate.reshape(B * T, -1)
 expert_ids = expert_ids.reshape(B * T, -1)
 output = torch.zeros_like(x)

 for i in range(len(experts)):
 mask = (expert_ids == i)
 rows, cols = torch.nonzero(mask, as_tuple=True)
 if len(rows) == 0:
 continue
 tokens = x[rows]
 expert_out = experts[i](tokens)
 weights = gate[rows, cols].unsqueeze(1)
 output.index_add_(0, rows, expert_out * weights)

 return output.reshape(B, T, C)

GPU решение и его фундаментальные ограничения

Стандартное решение для MoE, превышающего память одного GPU — экспертная параллелизация (EP). Однако этот подход создает серьезные компромиссы:

  • Необходимость балансировки нагрузки между экспертами
  • Коммуникационные накладные расходы при увеличении числа экспертов
  • Снижение общей утилизации оборудования при дисбалансе

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

Архитектурный прорыв: Cerebras WSE против GPU

В то время как GPU требуют сложных схем 3D-параллелизма, Cerebras Wafer Scale Engine предлагает принципиально иной подход:

  • В 1000 раз больше памяти на чипе по сравнению с GPU
  • Возможность хранить модели до 1B параметров непосредственно на чипе
  • Технология потокового измерения весов для моделей с триллионными параметрами

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

Проблема вычислений: когда эффективность не означает скорость

Даже решив проблему памяти, мы сталкиваемся с новыми проблемами — утилизацией вычислений. Обещанное 62% снижение FLOPs не превращается в пропорциональное ускорение обучения из-за:

  • Дробления батча между многочисленными экспертами
  • Простоя большей части экспертных сетей
  • Перехода в режим, ограниченный вводом-выводом, в экспертных слоях

Парадокс MoE: теоретическая эффективность не преобразуется в реальное ускорение из-за аппаратных ограничений и особенностей разреженных вычислений.