All-reduce: гиперкуб для коротких векторов, кольцо для длинных

Если ты знаешь, как устроен all-reduce гиперкубом и кольцом, почему кольцо выигрывает на длинных векторах и что такое reduce-scatter и all-gather, листай до «Теперь сам».

Как думаешь, сколько передаёт узел

Шестьдесят четыре видеокарты обучают одну нейросеть, у каждой после шага - вектор градиентов на 100 мегабайт. Всем нужна сумма. Сколько байт придётся передать каждой карте при лучшем алгоритме: около 100 МБ, около 600 МБ (по разу на каждый из шести уровней дерева) или 6,4 ГБ (каждой от каждой)?

Ответь себе до чтения. Кольцевой алгоритм передаёт от каждой карты 2 · 100 · 63 / 64 ≈ 197 мегабайт - чуть меньше двух векторов, и это почти не зависит от числа карт.

На что это похоже

Шестьдесят четыре человека считают голоса по 64 участкам, и у каждого свой протокол со всеми участками. Вместо того чтобы сдавать протоколы одному счётчику, каждый берёт на себя один участок: передаёт соседу свою строку по его участку, тот прибавляет свою и передаёт дальше. Через 63 передачи у каждого итог по «своему» участку. Ещё 63 передачи по тому же кругу - и итоги всех участков есть у всех. Каждый всё это время передавал по одной строчке, а не весь протокол.

Механизм: три алгоритма

Дерево плюс рассылка. Свести суммы деревом на узел 0 (log₂ p раундов) и разослать (log₂ p). Каждый раунд - весь вектор: 2 log₂ p (α + β · n).

Гиперкуб, рекурсивное удвоение. Пусть p - степень двойки. На шаге k узел i обменивается всем вектором с узлом, у которого переключён бит k номера: i + 2ᵏ, если этот бит у i нулевой, и i − 2ᵏ, если единичный, - и оба складывают. После шага k у каждого сумма по группе из 2ᵏ⁺¹ узлов. log₂ p (α + β · n) - вдвое лучше дерева, но всё ещё весь вектор на каждом шаге.

Кольцо (Патарасук и Юань, 2009; именно его популяризировали Baidu в 2017 году и Horovod для обучения нейросетей). Вектор режется на p кусков. Reduce-scatter: p − 1 шагов, на шаге s узел i отправляет соседу i + 1 кусок (i − s) mod p, сосед прибавляет к своему. В конце у каждого узла готова сумма одного куска: у узла i - куска (i + 1) mod p, который на последнем шаге пришёл к нему от соседа и вобрал вклады всех. All-gather: ещё p − 1 шагов по тому же кольцу, сосед не прибавляет, а запоминает. Время 2(p − 1)(α + β · n / p). Передача 2βn(p − 1)/p - нижняя граница all-reduce по объёму (Патарасук и Юань доказывают её для своей модели: каждый кусок суммы должен собрать вклады всех узлов и разойтись ко всем), зато задержек 2(p − 1) - линейно по p.

Две половины кольца - сами по себе коллективные операции: reduce-scatter (каждому - сумма своего куска) и all-gather (каждому - все куски). Их родственница all-to-all - каждый узел шлёт каждому свою порцию, как при перераспределении данных по ключу: (p − 1)(α + β · n) напрямую, где n - порция для одного адресата, или около log₂ p (α + β · p n) через гиперкуб - выбор опять по размеру.

Правило

All-reduce: дерево плюс рассылка - 2 log₂ p (α + β · n), гиперкуб - log₂ p (α + β · n), кольцо - 2(p − 1)(α + β · n / p). Гиперкуб выигрывает на коротких векторах (мало задержек), кольцо - на длинных (передача 2βn(p − 1)/p - это и есть нижняя граница). Кольцо - это reduce-scatter, за которым идёт all-gather.

Разбор: кольцо на четырёх узлах

Четыре узла, у каждого вектор из четырёх чисел - по куску на узел. Программа выполняет обе фазы по шагам: сначала все узлы одновременно формируют отправки, потом все применяют полученное.

def kolco(vektory):
    # p узлов, у каждого вектор длины n, n делится на p: куски по n / p
    # (программа этого требует; иначе куски длиной ceil(n / p)).
    # Фаза 1 (reduce-scatter): шаг s, узел i шлёт соседу i + 1 кусок
        # (i - s) mod p, сосед прибавляет. После неё у узла i готов кусок
    # (i + 1) mod p. Фаза 2 (all-gather): узел i шлёт кусок (i - s + 1)
    # mod p - начиная с готового, сосед не прибавляет, а заменяет.
    p = len(vektory)
    n = len(vektory[0])
    m = n // p
    v = []
    for x in vektory:
        v.append(x[:])
    otpravleno = [0] * p
    for faza in [1, 2]:
        for s in range(p - 1):
            poslaniya = []
            for i in range(p):
                nach = (i - s + faza - 1) % p * m
                poslaniya.append([i, nach, v[i][nach:nach + m]])
                otpravleno[i] = otpravleno[i] + m
            for i, nach, dannye in poslaniya:
                j = (i + 1) % p
                for t in range(m):
                    if faza == 1:
                        v[j][nach + t] = v[j][nach + t] + dannye[t]
                    else:
                        v[j][nach + t] = dannye[t]
    return [v, otpravleno]

vektory = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400],
           [1000, 2000, 3000, 4000]]
v, otpr = kolco(vektory)
for x in v:
    print(x)
print("каждый отправил элементов:", otpr, "при n = 4, p = 4")

[1111, 2222, 3333, 4444]
[1111, 2222, 3333, 4444]
[1111, 2222, 3333, 4444]
[1111, 2222, 3333, 4444]
каждый отправил элементов: [6, 6, 6, 6] при n = 4, p = 4

У всех четырёх узлов сумма 1111, 2222, 3333, 4444. Каждый отправил 6 элементов - ровно 2(p − 1) · n / p = 2 · 3 · 4 / 4. Гиперкуб на тех же данных переслал бы с каждого узла по два раза весь вектор - 8 элементов, а дерево с рассылкой - больше и неравномерно: корень отправил бы и принял больше всех.

Разбор: какой алгоритм когда

Та же сеть, 64 узла. Считаем время трёх алгоритмов на векторе в килобайт, мегабайт и сто мегабайт.

import math

alfa = 20.0
beta = 0.0001
p = 64
lg = math.ceil(math.log2(p))

print("байт       дерево+рассылка  гиперкуб   кольцо")
for n in [1000, 10 ** 6, 10 ** 8]:
    derevo = 2 * lg * (alfa + beta * n)
    giperkub = lg * (alfa + beta * n)
    kolco = 2 * (p - 1) * (alfa + beta * n / p)
    print(n, round(derevo), round(giperkub), round(kolco))
print("нижняя граница передачи для кольца:", round(2 * beta * 10 ** 8))

байт       дерево+рассылка  гиперкуб   кольцо
1000 241 121 2520
1000000 1440 720 2717
100000000 120240 60120 22208
нижняя граница передачи для кольца: 20000

Килобайт - гиперкуб в 20 раз быстрее кольца: 6 задержек против 126. Мегабайт - всё ещё гиперкуб, 720 против 2717. Сто мегабайт - кольцо: 22 миллисекунды против 60, при нижней границе передачи 20. Ровно поэтому NCCL и MPI переключаются: маленькие all-reduce - деревьями и удвоением, большие - кольцами (а на сотнях карт - двухуровневыми схемами, где кольцо внутри сервера, а дерево между серверами).

Что ломается без этого

All-reduce через параметр-сервер. Все карты шлют градиенты одному серверу и забирают у него сумму: сервер принимает p · n и отдаёт p · n, его канал - узкое место, и шаг обучения растёт линейно с числом карт.

Кольцо для мелких сообщений. Синхронизация скаляра (например, флаг «все закончили») кольцом - 2(p − 1) задержек вместо log₂ p. На 1024 узлах это 2046 против 10.

Неравные куски. Если n не делится на p, а куски берутся «как получится», самый длинный кусок задаёт время каждого шага, и выигрыш кольца тает.

Зачем это в работе

Обучение любой крупной модели на многих видеокартах - это all-reduce градиентов на каждом шаге, и его время напрямую входит во время обучения. Horovod, torch.distributed (бэкенд NCCL) и JAX опираются на кольцевые и древовидные all-reduce. В базах и движках обработки данных all-to-all - это перераспределение (shuffle) по ключу перед соединением или группировкой, и его цена часто и есть цена запроса.

Теперь сам

1. Сколько шагов и сколько элементов на узел передаёт кольцевой all-reduce на 8 узлах для вектора из 8000 элементов?
2. Почему гиперкубу нужно, чтобы p было степенью двойки, а кольцу - нет?
3. На каком размере вектора гиперкуб и кольцо из разбора сравняются по порядку величины, если p = 64?