Когда технологические компании арендуют вычислительные кластеры для обучения передовых нейросетей, счет за электроэнергию и аренду серверных стоек идет на миллионы долларов в неделю. Тысячи графических процессоров Nvidia H100 и H200 гудят в серверных залах, решая колоссальные матричные уравнения. Но если заглянуть в профайлер производительности такого кластера, любого инженера охватит священный ужас: дорогущие тензорные ядра видеокарт простаивают почти половину рабочего времени!
Причина этих чудовищных потерь кроется в архитектуре Mixture of Experts (MoE, смесь экспертов), которая стала стандартом для современных моделей вроде Qwen 2.5, Qwen 3 и DeepSeek-V3. Чтобы токены попали к нужным «экспертам», видеокарты вынуждены непрерывно пересылать терабайты данных по межузловой сети, блокируя вычислительные блоки.
Команда платформы распределенного обучения Prime Intellect выкатила открытое решение этой проблемы — GPU-кернел Prime Mega MoE, опубликованный в открытом репозитории prime-kernels под лицензией Apache 2.0. За счет перекрытия сетевого обмена и математических вычислений на обратном проходе (Overlapping Backward Pass) инженерам удалось ускорить шаг обучения моделей MoE на феноменальные 55%. Давай разберем низкоуровневую механику этого кернела и посмотрим, как устроена эта магия в коде.
Анатомия простоя GPU: где спотыкается Mixture-of-Experts
Давай разберем устройство классического MoE-слоя простыми словами. В отличие от плотной (dense) нейросети, где каждый токен прогоняется через одну гигантскую полносвязную матрицу (FFN), в MoE этот слой разбит на десятки независимых блоков — «экспертов». Для каждого входящего слова динамический маршрутизатор (router) выбирает двух-четырех наиболее подходящих экспертов.
Пока эксперты живут на одном физическом GPU, все работает быстро. Но когда модель насчитывает сотни миллиардов параметров, разные эксперты неизбежно оказываются на разных серверах кластера. И здесь начинается кошмар:
- Операция All-to-All на прямом проходе (Forward Pass): видеокарты должны обменяться токенами, чтобы направить каждое слово на тот GPU, где сидит нужный эксперт;
- Вычисление матричных произведений (GEMM): эксперты обрабатывают токены на тензорных ядрах;
- Обратный проход (Backward Pass): когда вычисляются градиенты ошибок, токены нужно вернуть назад исходным серверам, попутно рассчитав градиенты по весам экспертов и градиенты по активациям.
В стандартных реализациях (включая популярный фреймворк Meta TorchTitan) эти операции выполняются последовательно: сначала видеокарта ждет, пока сетевой адаптер (InfiniBand или RoCE) завершит пересылку порции данных, и лишь затем начинает крутить тензорные ядра. В результате GPU простаивает до 40–50% времени каждого итерационного шага.
Overlapping Backward Pass: совмещение сети и матричных умножений
Инженеры Prime Intellect обратили внимание на фундаментальную особенность обратного прохода. Чтобы рассчитать градиенты активаций для передачи обратно по сети, требуются веса экспертов, но не требуются градиенты самих весов! Это означает, что вычисления можно разбить на два независимых потока выполнения.
В кернеле Prime Mega MoE, построенном на базе низкоуровневой библиотеки DeepGEMM, реализован механизм асинхронного совмещения (overlapping):
- Пока графический процессор рассчитывает градиенты по весам (weight gradients) для текущего эксперта на тензорных ядрах, коммуникационный модуль в фоновом потоке CUDA уже отправляет готовые градиенты активаций по сети на соседние узлы;
- И наоборот: получение входящих пакетов для следующего эксперта происходит параллельно с вычислением матричных блоков текущего шага.
Вот как выглядит концептуальная логика разделения потоков и управления тензорным ускорителем памяти (TMA) на уровне низкоуровневого кода CUDA C++:
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
// Ядро асинхронного обратного прохода с перекрытием вычислений и сетевого обмена
__global__ void prime_moe_overlapping_backward_kernel(
const half* __restrict__ grad_output,
const half* __restrict__ weights,
half* __restrict__ grad_input_buffer,
half* __restrict__ grad_weight_accumulator,
int num_tokens, int hidden_dim, int expert_dim
) {
cg::thread_block block = cg::this_thread_block();
// Асинхронное выделение разделяемой памяти под тайлы TMA
extern __shared__ half smem_tile[];
half* tile_a = smem_tile;
half* tile_b = smem_tile + (hidden_dim * 16);
// 1. Асинхронный расчет градиентов по входам (активациям)
#pragma unroll
for (int step = 0; step < expert_dim; step += 16) {
// Загрузка весов через асинхронный копировщик памяти Hopper TMA
asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1], [%2];"
: : "r"(tile_a), "l"(weights + step * hidden_dim), "r"(sizeof(half) * hidden_dim * 16));
asm volatile("mbarrier.arrive %0;" : : "r"(0));
// Синхронизация барьера транзакции и выполнение матричного умножения
asm volatile("mbarrier.try_wait %0, 0;" : : "r"(0));
// Вычисление градиентов входов для немедленной фоновой отправки в All-to-All
atomicAdd(&grad_input_buffer[block.thread_rank()], tile_a[block.thread_rank()]);
}
// 2. Фоновое накопление градиентов весов параллельно с сетевой передачей активаций
block.sync();
atomicAdd(&grad_weight_accumulator[block.thread_rank()], grad_output[block.thread_rank()]);
}
Этот низкоуровневый трюк полностью ликвидирует простои: пока сетевая карта перекачивает гигабайты градиентов, тензорные ядра загружены расчетом весовых матриц на 100%.
Замеры на 128 GPU: 55% ускорения шага для Qwen 3
Результаты тестирования кернела превзошли самые оптимистичные ожидания разработчиков:
- На уровне одиночного изолированного MoE-слоя ускорение составило в 3,11 раза на прямом проходе (Forward) и в 1,87 раза на обратном проходе (Backward);
- В масштабе полного сквозного цикла обучения (End-to-End step) открытой модели Qwen 3 235B-A22B на кластере из 128 GPU Nvidia H200 кернел Prime Mega MoE обеспечил суммарный прирост скорости в 55% по сравнению с эталонным стеком Meta TorchTitan!
Для интеграции в популярные фреймворки обучения Prime Intellect предоставила компактный слой-обертку на PyTorch с автоматической регистрацией пользовательских автодифф-операторов:
import torch
import torch.nn as nn
from prime_kernels import moe_overlapping_backward_op
class PrimeMoELayer(nn.Module):
def __init__(self, num_experts: int, hidden_dim: int, expert_dim: int):
super().__init__()
self.num_experts = num_experts
self.hidden_dim = hidden_dim
self.expert_dim = expert_dim
self.weights = nn.Parameter(torch.empty(num_experts, hidden_dim, expert_dim, dtype=torch.bfloat16))
nn.init.xavier_uniform_(self.weights)
def forward(self, hidden_states: torch.Tensor, router_indices: torch.Tensor):
# Вызов скомпилированного кернела с поддержкой асинхронного стриминга
return moe_overlapping_backward_op(
hidden_states,
self.weights,
router_indices,
overlap_comm=True
)
Как мы уже анализировали в разборе оптимизации инференса моделей со смесью экспертов, ускорение шага на 55% в реальном выражении экономит сотни тысяч долларов на аренде вычислительных мощностей и сокращает цикл дообучения моделей с двух месяцев до трех-четырех недель. Кернел уже интегрировали в открытую библиотеку обучения с подкреплением Prime RL.
Аппаратные требования и границы применимости
Однако у столь высокой производительности есть жесткая аппаратная привязка:
- Только Hopper и Blackwell: кернел критически завязан на аппаратные инструкции асинхронного копирования данных TMA (Tensor Memory Accelerator), которые впервые появились в чипах Nvidia архитектуры Hopper (H100, H200) и получили развитие в Blackwell (B200).
- Неприменимость на старом или потребительском железе: на видеокартах архитектуры Ampere (A100) или потребительских платах RTX 4090 кернел запустить не удастся — там физически нет аппаратных блоков TMA, и компилятор CUDA просто выдаст ошибку инструкций.
Prime Mega MoE — яркий пример того, как грамотная низкоуровневая оптимизация способна обогнать экстенсивное наращивание серверов. Открытый код кернела в очередной раз доказал мощь сообщества Open Source перед лицом закрытых корпоративных разработок.