NCCL: как несколько GPU NVIDIA работают над одной задачей
Введение
Современные нейросети обучают не на одной видеокарте, а на десятках, сотнях и даже тысячах сразу. Кажется, что всё просто: чем больше ускорителей, тем быстрее идёт работа. На деле восемь видеокарт в одном сервере почти никогда не работают в восемь раз быстрее одной. Каждая из них решает свою часть задачи, но время от времени им нужно сверять результаты, и на эту сверку уходит заметная часть времени.
Один из основных инструментов обмена данными между видеокартами NVIDIA — библиотека NCCL. Многие разработчики используют её через PyTorch и другие фреймворки, не обращаясь к ней напрямую. Однако NCCL не является единственным способом организации обмена между GPU: распределённые вычисления могут использовать и другие коммуникационные библиотеки. Разберёмся, что она делает, как устроена и почему даже с ней ускорение никогда не бывает идеальным.
Почему одной видеокарты мало
Причины две. Первая связана с памятью. Возьмём модель на 70 млрд параметров, это обычный размер для современных открытых языковых моделей. Одни только её веса в 16-битном формате занимают около 140 ГБ. Но при обучении для каждого параметра нужно хранить ещё градиент, 32-битную копию веса и две вспомогательные величины оптимизатора, и в сумме выходит около 16 байт на параметр, то есть примерно 1,1 ТБ. Даже у самых ёмких ускорителей NVIDIA, Blackwell Ultra и Rubin, по 288 ГБ памяти, так что такой модели нужно минимум четыре видеокарты, и это без учёта промежуточных результатов. Для моделей на сотни миллиардов и триллионы параметров счёт идёт на десятки и сотни ускорителей.
Вторая причина связана со временем. Обучение большой модели на одном ускорителе заняло бы годы. Поэтому работу делят между многими видеокартами, и делят её по-разному.
Самый понятный способ называется параллелизмом по данным. Каждая видеокарта получает полную копию модели и свою порцию обучающих примеров. Посчитав, как нужно изменить веса, видеокарты должны договориться: усреднить свои поправки и применить к модели одно общее изменение, иначе копии со временем перестанут совпадать. Такая сверка происходит после каждого шага обучения, а её объём сопоставим с размером всей модели.
Если модель слишком велика, её режут на части. Можно разделить каждый слой между несколькими видеокартами, и тогда им придётся обмениваться промежуточными результатами прямо в процессе вычислений. Можно раздать разным GPU разные слои, и данные будут переходить от одной видеокарты к следующей, как деталь по сборочной линии. В моделях со «смесью экспертов» каждый фрагмент текста отправляется на ту видеокарту, где находится подходящий для него «эксперт». Способы разные, но у всех есть общее: без постоянного обмена данными совместная работа невозможна.
Что такое NCCL
NCCL (NVIDIA Collective Communications Library) произносится как “никель”. Это библиотека, которая организует обмен данными между видеокартами NVIDIA. Она не обучает нейросети и не считает их слои. Её задача уже: как можно быстрее доставить данные от одних GPU к другим и при необходимости сложить их по дороге.
Первая версия появилась в конце 2015 года и умела работать только с видеокартами внутри одного компьютера. В 2017 году вторая версия научилась связывать между собой целые серверы по сети. С тех пор библиотека стала стандартом: через неё обмениваются данными PyTorch, системы обучения крупных языковых моделей и программы, которые запускают готовые нейросети на нескольких видеокартах. У AMD есть похожая библиотека RCCL, созданная по тому же образцу.
Язык коллективных операций
Обмены, которые нужны при совместной работе, повторяются раз за разом, поэтому их удобно свести к небольшому набору стандартных действий. Они называются коллективными операциями, потому что в каждой участвуют сразу все видеокарты группы.
Самая важная из них называется AllReduce. Каждая видеокарта приносит свой массив чисел, а в итоге все получают одинаковый результат, обычно поэлементную сумму. Именно так усредняются поправки к весам при параллелизме по данным. Есть и другие операции. Broadcast рассылает данные одной видеокарты всем остальным, например начальные веса модели перед стартом обучения. AllGather собирает кусочки со всех участников и выдаёт каждому полную картину. ReduceScatter, наоборот, складывает данные и раздаёт каждой видеокарте только её долю итога. AlltoAll позволяет каждому участнику отправить каждому свой отдельный пакет, как раз это нужно моделям со смесью экспертов.
Операция AllReduce: четыре видеокарты (ранги) приносят свои данные, и каждая получает их общую сумму. Источник:
Программисту достаточно сказать "выполни AllReduce", а всё остальное NCCL берёт на себя: по каким проводам пойдут данные, в каком порядке и какими порциями. От этих решений зависит, займёт операция миллисекунды или секунды.
Кольцо: как сложить данные без давки
Самый очевидный способ выполнить AllReduce выглядит так: все видеокарты отправляют данные одной, она всё складывает и рассылает результат обратно. Но тогда к одной видеокарте сходится весь поток, она становится бутылочным горлышком, а линии связи остальных большую часть времени простаивают.
NCCL с самого начала использует другой подход. Видеокарты мысленно выстраиваются в кольцо, и каждая передаёт данные только соседу. Массив делится на столько частей, сколько участников. Каждая часть проходит по кругу, и на каждом шаге очередная видеокарта добавляет к ней свои числа. Когда часть обойдёт всех, в ней окажется полная сумма. Затем готовые части совершают ещё один круг, чтобы попасть ко всем участникам.
Четыре видеокарты, подключённые через два коммутатора, образуют кольцо. Порядок обхода выбран так, чтобы как можно меньше данных проходило через процессор. Источник: .
В таком кольце все линии связи работают одновременно, и ни одна видеокарта не перегружена. Самое приятное свойство кольца: объём данных, который передаёт каждый участник, почти не зависит от их числа. Для восьми видеокарт он составляет 1,75 размера массива, для ста приближается к двум. Поэтому кольцо отлично подходит для пересылки больших объёмов.
У кольца есть и слабое место. Каждой части нужно обойти всех участников, а значит, число шагов растёт вместе с количеством видеокарт. Для восьми GPU это незаметно, но когда их тысячи, даже маленькое сообщение долго идёт по кругу.
Деревья: когда участников тысячи
Для очень больших систем в 2019 году в NCCL добавили алгоритм на основе дерева. Видеокарты выстраиваются в иерархию, похожую на турнирную сетку: пары складывают свои данные, победители объединяются в новые пары, и так до самого верха. Затем итог спускается обратно по тем же ветвям. Число уровней в таком дереве растёт очень медленно: для тысячи участников их нужно около десяти, а не тысяча, как в кольце.
У простого дерева есть недостаток: видеокарты на нижнем уровне только отправляют данные и почти ничего не получают, то есть часть линий связи простаивает. Поэтому NCCL строит сразу два дерева, перевёрнутых относительно друг друга. Видеокарта, которая в одном дереве находится внизу, в другом оказывается в середине, и каждое дерево обрабатывает половину данных. Нагрузка распределяется равномерно.
Два дерева на 32 видеокарты дополняют друг друга: участник, стоящий внизу одного дерева, оказывается в середине другого. Источник: .
Новый подход проверили на суперкомпьютере Summit с 24 576 видеокартами, и задержка оказалась до 180 раз меньше, чем у кольца. При этом выбирать вручную не нужно: NCCL сама решает, какой алгоритм выгоднее для конкретного размера данных и числа участников. В последние годы появились и способы складывать данные прямо в сетевых коммутаторах, чтобы видеокартам не приходилось пересылать числа туда и обратно.
Дороги между видеокартами
Алгоритм определяет порядок обмена, но скорость в первую очередь зависит от того, по каким «дорогам» идут данные. Внутри одного сервера самая быстрая дорога называется NVLink. Это специальное соединение NVIDIA между видеокартами: у ускорителя A100 оно передаёт до 600 ГБ/с, почти в десять раз больше обычного разъёма PCIe четвёртого поколения. В крупных серверах видеокарты соединяются через коммутатор NVSwitch, и каждая может обмениваться данными с любой другой на полной скорости.
Если NVLink нет, данные идут по шине PCIe, той самой, через которую видеокарта подключается к материнской плате. Это заметно медленнее, а в худшем случае данные приходится копировать через оперативную память компьютера. Между разными серверами используются быстрые сети вроде InfiniBand, причём сетевая карта умеет обращаться к памяти видеокарты напрямую, без участия процессора.
При запуске NCCL изучает устройство компьютера: какие видеокарты соединены NVLink, какие подключены к одному процессору, где находятся сетевые карты. На основе этой карты она прокладывает кольца и деревья так, чтобы данные шли по самым быстрым дорогам. Поэтому одна и та же программа без изменений работает и на рабочей станции с двумя видеокартами, и в кластере из тысяч серверов.
Почему восемь видеокарт не дают восьмикратного ускорения
Теперь можно ответить на главный вопрос. Вычисления при добавлении видеокарт действительно делятся между ними, а обмен данными нет. Каждый участник на каждом шаге должен отправить и получить почти двойной объём градиентов, сколько бы видеокарт ни было.
Посмотрим на цифры. Градиенты модели с семью миллиардами параметров занимают около 14 ГБ. Чтобы сложить их на восьми видеокартах кольцом, каждой нужно передать около 24,5 ГБ. По шине PCIe 4.0 со скоростью около 32 ГБ/с в одну сторону это займёт не меньше 0,8 секунды, а на практике больше. По NVLink у A100 тот же обмен займёт меньше десятой доли секунды. Если сами вычисления на шаге длятся секунду, в первом случае почти половина времени уходит на ожидание данных, а во втором лишь несколько процентов.
Есть и другие потери. Кольцо работает со скоростью самого медленного своего участка, поэтому одна плохо подключённая видеокарта тормозит всех. Маленькие сообщения страдают от задержек: им не хватает объёма, чтобы разогнать линию. Отчасти это компенсируется тем, что программы начинают пересылать уже готовые градиенты, пока вычисления для остальных слоёв ещё продолжаются. Но полностью скрыть обмен удаётся редко.
Куда движется NCCL
Чтобы пересылать данные, NCCL использует часть вычислительных блоков видеокарты, и эти блоки на время обмена недоступны для расчётов. Поэтому разработчики последних версий стремятся уменьшить использование вычислительных ресурсов GPU при передаче данных. В частности, NCCL поддерживает выполнение некоторых коллективных операций с использованием специализированных копирующих движков, что позволяет снизить нагрузку на вычислительные блоки видеокарты. Кроме того, появилась возможность запускать обмен прямо из программы, выполняющей расчёт на видеокарте, без помощи центрального процессора. Особенно это полезно моделям со смесью экспертов и быстрым ответам чат-ботов, где сообщения маленькие и каждая микросекунда задержки заметна.
Заключение
Когда нейросеть обучается на многих видеокартах, её скорость определяется не только мощностью ускорителей, но и тем, насколько быстро они договариваются между собой. NCCL организует этот обмен: учитывает особенности соединений между видеокартами, выбирает подходящий алгоритм из нескольких поддерживаемых вариантов и определяет маршруты передачи данных с учётом топологии системы. Отменить законы физики она не может, и обмен данными всегда отнимает часть времени. Поэтому самые мощные системы для обучения ИИ строят не только из быстрых видеокарт, вроде тех же A100 или H100, но и из быстрых соединений между ними, а библиотека, которая управляет этими соединениями, остаётся незаметной, но обязательной частью любой большой нейросети.
