DualPipe и EPLB: конвейер в обе стороны и балансировщик экспертов

От переводчика: что это и что повторить

Обе библиотеки DeepSeek открыла 27 февраля 2025 года, в четвёртый день «недели открытого кода» (Open Source Week). Это два небольших README, но за каждым стоит важная инженерная идея из техотчёта DeepSeek-V3:

• DualPipe — расписание конвейерный параллелизмконвейерный параллелизмEN pipeline parallelism, в котором микробатчи входят в конвейер с обоих концов, а пересылки между GPU прячутся за вычислениями. Нужно при обучении.
• EPLB — алгоритм, который решает, каких экспертов MoE скопировать и на какие GPU положить копии, чтобы ни одна GPU не простаивала и не была перегружена. Нужен в первую очередь на инференсе.

Перед чтением повторите:

• микробатчи, пузырь конвейерапузырь конвейераEN pipeline bubbleВремя, когда стадии конвейера простаивают в ожидании данных. У GPipe и 1F1B доля пузыря — (p − 1) / (m + p − 1) при p стадиях и m микробатчах. и расписание 1F1B (урок 5.4.3, раздел «Конвейерный параллелизм и пузыри»);
• параллелизм экспертовпараллелизм экспертовEN expert parallelism (EP)Разные эксперты MoE-слоя живут на разных GPU; токены пересылаются к своим экспертам и обратно через all-to-all-обмен. и балансировка нагрузкибалансировка нагрузкиEN load balancing в MoE (уроки 5.3.1 и 5.3.2);
• как DeepSeek раскладывает экспертов по GPU в боевом инференсе — обзор системы инференса V3/R1.

Рисунки оригинала (images/dualpipe.png, images/dualpipev.png, example.png) здесь описаны словами и таблицами. Комментарии переводчика вынесены в такие рамки, всё остальное — перевод.

Часть 1. DualPipe

Оригинал: README.md репозитория deepseek-ai/DualPipe.

DualPipe — новаторский алгоритм двунаправленного конвейерного параллелизма, представленный в техническом отчёте DeepSeek-V3. Он добивается полного наложения фаз вычислений и обмена данными при прямом и обратном проходах, а заодно уменьшает пузыри конвейера. Подробнее о наложении вычислений и обмена данными см. данные профилирования.

Расписания

Рисунок images/dualpipe.png.

Пример расписания DualPipe для 8 рангов конвейера (PP ranks) и 20 микробатчей, идущих в двух направлениях. Микробатчи обратного направления симметричны микробатчам прямого, поэтому для простоты их номера на рисунке не показаны. Две клетки в общей чёрной рамке — это вычисления и обмен данными, наложенные друг на друга.

От переводчика: как читать рисунок

По горизонтали — время, по вертикали — 8 устройств (Device 0…7). Цвета клеток: оранжевый — прямой проход (Forward), зелёный — обратный проход (Backward) или его часть «обратный по входу» (Backward for input), голубой — «обратный по весам» (Backward for weights), клетка из двух цветов — прямой и обратный проходы, наложенные друг на друга.

Что видно на рисунке:

1. Разогрев с двух сторон. Device 0 сразу запускает прямые проходы микробатчей 0, 1, 2… своего направления, а Device 7 — микробатчей другого направления. Поэтому в начале нет «лесенки» пустых клеток, как у 1F1B: пустует только небольшой треугольник у средних устройств.
2. Середина расписания — сплошные пары. Почти все клетки в центре двухцветные: пока одна стадия считает прямой проход одного микробатча, она же считает обратный проход другого, а пересылки (all-to-all и передача активаций соседней стадии) идут в это же время.
3. Голубые клетки в конце. Градиенты по весам (W) отложены и вставлены туда, где иначе был бы пузырь, — идея Zero Bubble, о которой говорится в уроке 5.4.3.

Каждое устройство хранит две стадии модели: Device 0 — первую и последнюю, Device 1 — вторую и предпоследнюю и т.д. Поэтому DualPipe держит две копии параметров (см. таблицу ниже).

flowchart LR
subgraph D0["GPU 0"]
  S0a["стадия 1"]
  S0b["стадия 8"]
end
subgraph D1["GPU 1"]
  S1a["стадия 2"]
  S1b["стадия 7"]
end
subgraph D3["GPU 7"]
  S3a["стадия 8"]
  S3b["стадия 1"]
end
A["микробатчи направления А"] --> S0a
S0a --> S1a
S1a -.-> S3a
B["микробатчи направления Б"] --> S3b
S3b -.-> S1b
S1b --> S0b

Идея DualPipe при 8 GPU (схема переводчика): у каждой GPU две стадии модели, микробатчи входят с обоих концов конвейера

DualPipeV

DualPipeV — компактное V-образное расписание, полученное из DualPipe процедурой «разрезания пополам». Его предложила Sea AI Lab под названием «Cut-in-half» в своей заметке в блоге. Спасибо им за это эффективное расписание!

Расписания

Рисунок images/dualpipev.png.

Пример расписания DualPipeV для 4 рангов конвейера (8 стадий конвейера) и 10 микробатчей.

От переводчика: в чём «V»

В DualPipe каждая GPU хранит стадии $i$ и $PP - 1 - i$, а микробатчи двух направлений — это две половины батча. DualPipeV оставляет одно направление, но «складывает» конвейер пополам: микробатч проходит стадии 1→2→3→4 по GPU 0→1→2→3, а затем стадии 5→6→7→8 по GPU 3→2→1→0 — путь похож на букву V. На каждой GPU по-прежнему две стадии, но GPU вдвое меньше: на рисунке 4 устройства и 8 стадий. Плюс в том, что одна и та же стадия не хранится дважды: нужна половина устройств при той же глубине конвейера.

Сравнение пузырей конвейера и расхода памяти (при одинаковом числе стадий PP)

Метод	Пузырь	Параметры на устройство	Активации на устройство	Число устройств
1F1B	$(PP-1)(F+B)$	1×	$PP$	$PP$
ZB1P	$(PP-1)(F+B-2W)$	1×	$PP$	$PP$
DualPipe	$(PP/2-1)(F\&B+B-3W)$	2×	$PP+1$	$PP$
DualPipeV	$(PP/2-1)(F\&B+B-3W)$	2×	$PP+1$	$PP/2$

$PP$ — число стадий конвейера (чётное). $F$ — время выполнения прямого чанка, $B$ — время выполнения полного обратного чанка, $W$ — время выполнения чанка «обратный проход по весам», $F\&B$ — время выполнения двух наложенных друг на друга прямого и обратного чанков.

От переводчика: считаем на числах

Как читать формулы: «пузырь» — это суммарное время, которое каждая GPU простаивает за один шаг обучения. Чанк (chunk) — кусок работы одной стадии над одним микробатчем.

Возьмём условные числа: $PP = 8$, прямой проход $F = 1$, полный обратный $B = 2$ (обратный обычно вдвое дольше прямого), из них по весам $W = 1$; пусть наложенная пара длится $F\&B = 2{,}5$ — меньше, чем $F + B = 3$, потому что пересылки одного чанка спрятаны за счётом другого.

• 1F1B: $7 \cdot (1 + 2) = 21$ единица простоя.
• ZB1P (Zero Bubble, Sea AI Lab): $7 \cdot (1 + 2 - 2) = 7$ — отложенные градиенты по весам заполняют пузыри.
• DualPipe: $3 \cdot (2{,}5 + 2 - 3) = 4{,}5$ — множитель $PP/2 - 1 = 3$ вместо $PP - 1 = 7$: разогрев идёт с двух сторон одновременно.

Цена — вдвое больше памяти под параметры и одна лишняя порция активаций. Заметьте, что ни пузырь, ни память не зависят от числа микробатчей. DeepSeek-V3 могла позволить себе двойные параметры: при 64-путевом параллелизме экспертов на каждой GPU лежит малая доля весов (урок 5.4.3).

Быстрый старт

Использование показано в примерах:

python examples/example_dualpipe.py
python examples/example_dualpipev.py

Примечание: в реальных приложениях вам нужно реализовать собственный метод overlapped_forward_backward, подогнанный под ваш модуль.

Требования

• PyTorch 2.0 и выше

Разработчики

DualPipe создали и разработали Цзяши Ли (Jiashi Li), Чэнци Дэн (Chengqi Deng) и Вэньфэн Лян (Wenfeng Liang).

Цитирование

@misc{deepseekai2025deepseekv3technicalreport,
      title={DeepSeek-V3 Technical Report},
      author={DeepSeek-AI},
      year={2025},
      eprint={2412.19437},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2412.19437},
}

От переводчика: что внутри примера и как запустить на вашей машине

examples/example_dualpipe.py строит игрушечную модель: каждая стадия — два линейных слоя с GELU между ними (hidden_size = 512, 20 микробатчей по 3 последовательности длиной 256). Главное в нём — два приёма:

1. Свой autograd.Function для линейного слоя. В backward градиент по входу считается сразу, а вычисление градиента по весам заворачивается в функцию grad_weight_fn и кладётся в очередь WeightGradStore, если та включена. Так обратный проход делится на $B$ (по входу, нужен предыдущей стадии срочно) и $W$ (по весам, можно отложить).
2. Метод overlapped_forward_backward. Он получает модуль и входы прямого прохода одного микробатча и модуль, loss или градиенты для обратного прохода другого. В примере он просто выполняет их по очереди — README подчёркивает, что настоящую стратегию наложения (например, запуск all-to-all в отдельном CUDA-потоке, пока идёт счёт) надо писать под свою модель.

Пример сам проверяет себя: сравнивает loss и градиенты DualPipe с обычным проходом по всей модели на одной GPU (расхождение градиентов меньше $10^{-13}$ по косинусной мере).

🖥 На вашей машине (4×GPU). Пример запускает по процессу на GPU через torch.multiprocessing.spawn с бэкендом NCCL и берёт чётное число карт (torch.cuda.device_count() // 2 * 2), так что на 4× RTX 4080 он пойдёт на 4 GPU (по очереди прогоняет конфигурации на 4 и на 2 картах). Связь между картами будет через PCIe — для игрушечной модели это не важно. Установите пакет (pip install -e . в корне репозитория) и запустите обе команды выше; ожидайте секунды–минуту на конфигурацию. Это хороший способ увидеть, что двойное хранение стадий и обмен с двух сторон дают тот же результат, что и обучение без конвейера.

Часть 2. EPLB — балансировщик нагрузки для параллелизма экспертов

Оригинал: README.md репозитория deepseek-ai/EPLB.

При параллелизме экспертов (expert parallelism, EP) разные эксперты назначаются на разные GPU. Нагрузка на экспертов меняется в зависимости от текущих запросов, поэтому важно поддерживать нагрузку на GPU сбалансированной. Как описано в статье о DeepSeek-V3, мы применяем стратегию избыточных экспертов: дублируем сильно нагруженных экспертов. Затем эвристически раскладываем дубликаты по GPU, чтобы нагрузка между GPU была сбалансирована. Кроме того, благодаря маршрутизации экспертов с ограничением по группам (group-limited expert routing), которая используется в DeepSeek-V3, мы по возможности стараемся класть экспертов одной группы на один узел, чтобы уменьшить трафик между узлами.

Чтобы облегчить воспроизведение и развёртывание, мы открываем код нашего боевого алгоритма балансировки EP в файле eplb.py. Алгоритм вычисляет сбалансированный план копирования и размещения экспертов на основе оценок их нагрузки. Заметим, что точный способ предсказывать нагрузку экспертов выходит за рамки этого репозитория. Распространённый способ — скользящее среднее по историческим статистикам.

От переводчика: термины

• Логический эксперт — эксперт модели как таковой (в DeepSeek-V3 их 256 на слой). Физический эксперт — его копия, реально лежащая на какой-то GPU. Избыточные эксперты — это «лишние» физические копии самых популярных логических экспертов.
• Нагрузка эксперта — сколько токенов роутер отправил ему за период наблюдения. Если у эксперта две копии, каждая получает примерно половину.
• Группы экспертов. В V3 256 экспертов разбиты на 8 групп, и токен может пойти к экспертам не более чем из 4 групп (а значит, не более чем на 4 узла). Если группа целиком лежит на одном узле, токены реже пересекают медленную межузловую сеть.

Алгоритм

У алгоритма балансировки две политики для разных случаев.

Иерархическая балансировка

Когда число серверных узлов делит число групп экспертов, мы используем иерархическую политику балансировки, чтобы воспользоваться маршрутизацией с ограничением по группам. Сначала мы равномерно раскладываем группы экспертов по узлам, так чтобы нагрузки узлов были сбалансированы. Затем копируем экспертов внутри каждого узла. Наконец, раскладываем скопированных экспертов по отдельным GPU, чтобы нагрузка на разные GPU была сбалансирована. Иерархическую политику можно использовать на стадии prefill с меньшим размером параллелизма экспертов.

Глобальная балансировка

В остальных случаях мы используем глобальную политику балансировки: она копирует экспертов глобально, не обращая внимания на группы, и раскладывает копии по отдельным GPU. Эту политику можно применять на стадии decode с большим размером параллелизма экспертов.

flowchart TD
L["Нагрузка логических экспертов (скользящее среднее)"] --> Q{"Число групп делится на число узлов?"}
Q -- "да" --> H1["1. Группы → узлы (поровну по нагрузке)"]
H1 --> H2["2. Копии экспертов внутри узла"]
H2 --> H3["3. Копии → GPU узла (поровну по нагрузке)"]
Q -- "нет" --> G1["Все эксперты = одна группа на одном «узле»"]
G1 --> H2
H3 --> R["phy2log, log2phy, logcnt"]

Две политики EPLB (схема переводчика по коду eplb.py): глобальная — тот же алгоритм с одной группой и одним узлом

Интерфейс и пример

Главная функция балансировщика — eplb.rebalance_experts.

Код ниже — пример для двухслойной MoE-модели, в каждом слое которой 12 экспертов. Мы добавляем по 4 избыточных эксперта на слой, и все 16 копий размещаются на 2 узлах по 4 GPU в каждом.

import torch
import eplb

# Нагрузка экспертов: 2 слоя × 12 экспертов (сколько токенов досталось каждому)
weight = torch.tensor([[ 90, 132,  40,  61, 104, 165,  39,   4,  73,  56, 183,  86],
                       [ 20, 107, 104,  64,  19, 197, 187, 157, 172,  86,  16,  27]])

num_replicas = 16   # физических экспертов на слой: 12 логических + 4 избыточных
num_groups = 4      # группы экспертов (по 3 эксперта)
num_nodes = 2       # серверные узлы
num_gpus = 8        # всего GPU (по 4 на узел)

phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
print(phy2log)

# Вывод:
# tensor([[ 5,  6,  5,  7,  8,  4,  3,  4, 10,  9, 10,  2,  0,  1, 11,  1],
#         [ 7, 10,  6,  8,  6, 11,  8,  9,  2,  4,  5,  1,  5,  0,  3,  1]])

Этот вывод получен иерархической политикой балансировки и задаёт следующий план копирования и размещения экспертов.

Рисунок example.png.

От переводчика: что на рисунке

На рисунке два узла (Node 0 и Node 1), в каждом по 4 GPU, на каждой GPU по два физических эксперта в каждом из двух слоёв. Это тот же phy2log, нарезанный по 2 элемента на GPU:

	GPU 0	GPU 1	GPU 2	GPU 3	GPU 4	GPU 5	GPU 6	GPU 7
Узел	0	0	0	0	1	1	1	1
Слой 0	5, 6	5, 7	8, 4	3, 4	10, 9	10, 2	0, 1	11, 1
Слой 1	7, 10	6, 8	6, 11	8, 9	2, 4	5, 1	5, 0	3, 1

(На рисунке GPU второго узла пронумерованы заново, 0–3.)

Проверим слой 0. Группы по 3 эксперта: $\{0,1,2\}$ с нагрузкой 262, $\{3,4,5\}$ — 330, $\{6,7,8\}$ — 116, $\{9,10,11\}$ — 325. Жадная раскладка по узлам (самую тяжёлую группу — на самый лёгкий узел, по 2 группы на узел) даёт узлу 0 группы $\{3,4,5\}$ и $\{6,7,8\}$ (446 токенов), узлу 1 — остальные (587). Внутри узла 0 копируются самые нагруженные эксперты 5 (165) и 4 (104), внутри узла 1 — эксперты 10 (183) и 1 (132). Каждый раз копируется эксперт с наибольшей нагрузкой на одну копию, поэтому logcnt для слоя 0 — [1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1].

Нагрузка на GPU слоя 0 получается от 86,5 до 156 токенов при среднем 129. Без копий так не выйдет: один эксперт 10 уже даёт 183 токена своей GPU. Если вызвать функцию с num_groups = 3 (3 не делится на 2 узла — сработает глобальная политика), разброс по GPU слоя 0 станет меньше: от 95,5 до 138,5. Это и есть компромисс двух политик: иерархическая экономит межузловой трафик, глобальная лучше выравнивает GPU.

Лицензия

Этот репозиторий кода выпущен под лицензией MIT.

Приложение: код eplb.py с переведёнными комментариями

От переводчика

Ниже — файл eplb.py из репозитория без изменений в коде; переведены только докстринги и комментарии. Он работает на CPU и требует лишь PyTorch: сохраните его как eplb.py рядом с примером выше и запустите пример — вывод совпадёт с README (проверено на PyTorch для CPU).

Как устроен алгоритм в трёх функциях:

• balanced_packing — жадная раскладка: объекты сортируются по убыванию веса, и каждый кладётся в самый лёгкий пакет, где ещё есть место (в каждом пакете ровно $n/m$ объектов);
• replicate_experts — каждая новая копия достаётся эксперту с наибольшей нагрузкой на одну копию (weight / logcnt);
• rebalance_experts_hierarchical — три шага: группы → узлы, копии внутри узла, копии → GPU. Индексы mlog («логический эксперт после перестановки по узлам») и pphy («физический эксперт после раскладки по GPU») нужны, чтобы собрать итоговые таблицы соответствия.

from typing import Tuple

import torch

def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Разложить n объектов с весами по m пакетам так, чтобы в каждом пакете было ровно n/m объектов,
    а суммарные веса пакетов были как можно ближе друг к другу.

    Параметры:
        weight: [X, n], вес каждого объекта
        num_packs: число пакетов

    Возвращает:
        pack_index: [X, n], номер пакета для каждого объекта
        rank_in_pack: [X, n], номер объекта внутри его пакета
    """
    num_layers, num_groups = weight.shape
    assert num_groups % num_packs == 0
    groups_per_pack = num_groups // num_packs

    if groups_per_pack == 1:
        pack_index = torch.arange(weight.size(-1), dtype=torch.int64, device=weight.device).expand(weight.shape)
        rank_in_pack = torch.zeros_like(weight, dtype=torch.int64)
        return pack_index, rank_in_pack

    indices = weight.float().sort(-1, descending=True).indices.cpu()
    pack_index = torch.full_like(weight, fill_value=-1, dtype=torch.int64, device='cpu')
    rank_in_pack = torch.full_like(pack_index, fill_value=-1)
    for i in range(num_layers):
        pack_weights = [0] * num_packs
        pack_items = [0] * num_packs
        for group in indices[i]:
            pack = min((i for i in range(num_packs) if pack_items[i] < groups_per_pack),
                       key=pack_weights.__getitem__)
            assert pack_items[pack] < groups_per_pack
            pack_index[i, group] = pack
            rank_in_pack[i, group] = pack_items[pack]
            pack_weights[pack] += weight[i, group]
            pack_items[pack] += 1
    return pack_index, rank_in_pack


def replicate_experts(weight: torch.Tensor, num_phy: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Размножить `num_log` экспертов до `num_phy` копий так, чтобы максимальная нагрузка
    среди всех копий была как можно меньше.

    Параметры:
        weight: [X, num_log]
        num_phy: общее число экспертов после копирования

    Возвращает:
        phy2log: [X, num_phy], номер логического эксперта для каждого физического
        rank: [X, num_phy], номер копии
        logcnt: [X, num_log], число копий каждого логического эксперта
    """
    n, num_log = weight.shape
    num_redundant = num_phy - num_log
    assert num_redundant >= 0
    device = weight.device
    phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
    rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
    logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
    arangen = torch.arange(n, dtype=torch.int64, device=device)
    for i in range(num_log, num_phy):
        redundant_indices = (weight / logcnt).max(dim=-1).indices
        phy2log[:, i] = redundant_indices
        rank[:, i] = logcnt[arangen, redundant_indices]
        logcnt[arangen, redundant_indices] += 1
    return phy2log, rank, logcnt


def rebalance_experts_hierarchical(weight: torch.Tensor, num_physical_experts: int,
                      num_groups: int, num_nodes: int, num_gpus: int):
    """
    Параметры:
        weight: [num_moe_layers, num_logical_experts]
        num_physical_experts: число физических экспертов после копирования
        num_groups: число групп экспертов
        num_nodes: число серверных узлов; сеть внутри узла (например, NVLink) быстрее, чем между узлами
        num_gpus: число GPU, должно быть кратно `num_nodes`

    Возвращает:
        physical_to_logical_map: [num_moe_layers, num_physical_experts]
        logical_to_physical_map: [num_moe_layers, num_logical_experts, X]
        logical_count: [num_moe_layers, num_logical_experts]
    """
    num_layers, num_logical_experts = weight.shape
    assert num_logical_experts % num_groups == 0
    group_size = num_logical_experts // num_groups
    assert num_groups % num_nodes == 0
    groups_per_node = num_groups // num_nodes
    assert num_gpus % num_nodes == 0
    assert num_physical_experts % num_gpus == 0
    phy_experts_per_gpu = num_physical_experts // num_gpus

    def inverse(perm: torch.Tensor) -> torch.Tensor:
        inv = torch.empty_like(perm)
        inv.scatter_(1, perm, torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand(perm.shape))
        return inv

    # Шаг 1: разложить группы по узлам
    tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
    group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes)
    log2mlog = (((group_pack_index * groups_per_node + group_rank_in_pack) * group_size).unsqueeze(-1) +
                torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)).flatten(-2)
    mlog2log = inverse(log2mlog)

    # Шаг 2: создать избыточных экспертов внутри узлов
    # [num_layers * num_nodes, num_logical_experts // num_nodes]
    tokens_per_mlog = weight.gather(-1, mlog2log).view(-1, num_logical_experts // num_nodes)
    phy2mlog, phyrank, mlogcnt = replicate_experts(tokens_per_mlog, num_physical_experts // num_nodes)

    # Шаг 3: разложить физических экспертов по GPU
    # [num_layers * num_nodes, num_physical_experts // num_nodes]
    tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
    pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
    phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
    pphy2phy = inverse(phy2pphy)

    pphy2mlog = phy2mlog.gather(-1, pphy2phy) # [num_layers * num_nodes, num_log_per_nodes]
    pphy2mlog = (pphy2mlog.view(num_layers, num_nodes, -1) +
                 torch.arange(0, num_logical_experts, num_logical_experts // num_nodes,
                              device=group_pack_index.device).view(1, -1, 1)).flatten(-2)
    pphy2log = mlog2log.gather(-1, pphy2mlog)
    pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
    logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
    return pphy2log, pphyrank, logcnt

def rebalance_experts(weight: torch.Tensor, num_replicas: int, num_groups: int,
                      num_nodes: int, num_gpus: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Точка входа балансировщика нагрузки для параллелизма экспертов.

    Параметры:
        weight: [layers, num_logical_experts], статистика нагрузки всех логических экспертов
        num_replicas: число физических экспертов, должно быть кратно `num_gpus`
        num_groups: число групп экспертов
        num_nodes: число серверных узлов; сеть внутри узла (например, NVLink) быстрее, чем между узлами
        num_gpus: число GPU, должно быть кратно `num_nodes`

    Возвращает:
        physical_to_logical_map: [layers, num_replicas], номер эксперта для каждой копии
        logical_to_physical_map: [layers, num_logical_experts, X], номера копий каждого эксперта
        expert_count: [layers, num_logical_experts], число физических копий каждого логического эксперта
    """
    num_layers, num_logical_experts = weight.shape
    weight = weight.float().cpu()
    if num_groups % num_nodes == 0:
        # иерархическая политика балансировки
        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(weight, num_replicas,
                                                                  num_groups, num_nodes, num_gpus)
    else:
        # глобальная политика балансировки
        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(weight, num_replicas, 1, 1, num_gpus)
    maxlogcnt = logcnt.max().item()
    log2phy: torch.Tensor = torch.full((num_layers, num_logical_experts, maxlogcnt),
                                       -1, dtype=torch.int64, device=logcnt.device)
    log2phy.view(num_layers, -1).scatter_(-1, phy2log * maxlogcnt + phyrank,
            torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(num_layers, -1))
    return phy2log, log2phy, logcnt

__all__ = ['rebalance_experts']

От переводчика: что сделать руками

1. Запустите пример из README и выведите ещё logcnt и log2phy[0]. Для слоя 0 вы увидите, например, что у эксперта 1 копии с номерами 15 и 13, а у эксперта 0 — одна копия номер 12 (вторая ячейка заполнена -1).
2. Посчитайте нагрузку на каждую GPU: сложите weight[l, e] / logcnt[l, e] для двух экспертов каждой GPU. Сравните отношение «максимум / среднее» у иерархической и глобальной политики (для глобальной передайте num_groups = 3).
3. Проверьте масштаб DeepSeek-V3 из обзора системы инференса: 256 логических экспертов, 8 групп, 32 избыточных копии. Prefill — 4 узла по 8 GPU (num_replicas = 288, num_nodes = 4, num_gpus = 32, сработает иерархическая политика), decode — 18 узлов (num_gpus = 144; 8 групп не делятся на 18 узлов, поэтому сработает глобальная политика). Нагрузку сгенерируйте случайно, например по закону Ципфа.

Словарик

EN	RU
pipeline parallelism (PP), PP rank / stage	конвейерный параллелизм, ранг / стадия конвейера
micro-batch	микробатч
pipeline bubble	пузырь конвейера (простой GPU)
forward / backward chunk	прямой / обратный чанк (кусок работы стадии над микробатчем)
backward for input / backward for weights	обратный проход по входу / по весам
computation-communication overlap	наложение вычислений и обмена данными
1F1B, ZB1P (Zero Bubble)	расписания «один прямой — один обратный» и «без пузыря»
expert parallelism (EP)	параллелизм экспертов
redundant experts	избыточные (дублированные) эксперты
logical / physical expert	логический эксперт / физическая копия
group-limited expert routing	маршрутизация с ограничением по группам
hierarchical / global load balancing	иерархическая / глобальная балансировка
prefilling / decoding stage	стадия prefill / стадия decode