Распределенное обучение и вычисления в сети. Доклад Яндекса
Все эксперты по сетям знакомы с принципом end-to-end, когда специфичные для конкретной задачи фичи реализовываются на конечных нодах, а промежуточные ноды должны только передавать данные, не взаимодействуя с ними. Но есть случаи, когда полезны вычисления внутри сети — с использованием устройств сети, занятых передачей трафика. Один из примеров таких задач — распределенный ML. В докладе Дмитрий Афанасьев дал краткое введение в особенности вычислений для распределенного ML, паттерны обмена данными и коллективные операции. Вторая половина доклада — о том, как редукция увеличивает производительность при обучении, и о некоторых реализациях.— Меня зовут Дмитрий Афанасьев, я сетевой архитектор Яндекса. И я сегодня расскажу про достаточно экзотические — по крайней мере пока — технологии. Но думаю, что они будут становиться менее экзотическими, и шансы с ними встретиться возрастают.
Итак, распределенное обучение для machine learning и вычисления в сети. Что я хочу рассказать? Сначала очень быстрое введение про то, что вообще такое нейросети, как они устроены, какие у них есть режимы функционирования и как они обучаются. Дальше специфика распределенного обучения, как устроен трафик при распределенном обучении и как он может ложиться на топологию сети Compute Implementations.
И, наконец, расскажу про две реализации, где вычисление в сетях используется, чтобы оптимизировать распределенное обучение.
В пакетных и IP-сетях мы привыкли руководствоваться e2e-принципом. И говорим, что всю функциональность желательно реализовывать на endpoint, а сеть, которая посередине, должна быть быстрой, простой, дешевой и достаточно тупой, и заниматься только передачей данных, не вмешиваясь в процесс их обработки.
Но из любых хороших правил есть исключения. И есть приложения, у которых специфика, структура вычислений и структура обмена данными таковы, что, реализовав некоторые функции (вычисления) в сети, можно очень хорошо оптимизировать эти вычисления.
Дальше мы посмотрим, как устроены сами задачи и что в них такое происходит с вычислениями, что можно внести много пользы, добавив вычисления в сети.
Итак, что такое нейросети? Слева на картинке у нас биологический нейрон с кучей своих запчастей. Справа у нас его математический аналог. Устроен просто. У него есть входы, на которые поступают какие-то сигналы. У этих входов есть веса. И еще есть смещение.
Ссылка со слайда
Внутри происходит следующее — вектор входных активаций умножается на вектор весов и добавляется смещение, и все это подается на вход функции активации, которая выдает некоторый сигнал на выход. Собственно, все.
Функция активации для всех наших последующих упражнений с этой конструкцией должна быть монотонной, и очень желательно, чтобы она была как-то нормально дифференцируема, потому что на этом построено обучение. 
Ссылка со слайда
Вот то же самое, здесь чуть более подробно про веса и так далее. Видно, что здесь вычисляется классическое скалярное возведение весов на входы. Добавляем смещение, отправляем результат дальше.
Из таких нейронов мы собираем собственно нейронную сеть. У нее есть входной уровень, который ничего не делает, это точки, в которые поступают какие-то входы. Дальше идут один или несколько слоев, которые производят обработку. Слои, которые не являются ни входными, ни выходными, мы называем скрытыми.
Слоев может быть много. Нетрудно заметить, что нейросети имеют такую слоистую структуру, то есть соединения внутри не произвольные, а мы производим обработку в одном слое, то, что получилось, подается на следующий и так далее. Это будет важно дальше, когда мы посмотрим на то, как они обучаются. 
Ссылка со слайда
Наконец, у нейросетей есть два основных режима функционирования или два способа с ними работать.
Во-первых, есть режим inference. Это когда мы на обученную сеть со всеми параметрами нейронов, связей и так далее подаем входные данные, а она выдает ответ. Например, что-то классифицирует.
Во-вторых, чтобы получить такую сеть, которая будет осуществлять классификацию и другие полезные вещи, ее сначала нужно обучить. Это режим обучения.
Режимы характеризуются разным количеством вычислений и данных. Inference мы делаем на небольшом наборе входных данных, возможно, даже на одном семпле каждый раз. Обучение мы делаем на очень большом наборе данных, и в нем данные проходят через сеть не только в прямом, но и в обратном направлении.
То есть мы прогоняем прямой путь, получаем результат. Смотрим на ошибку, на отклонение от того, что должно было получиться. И распространяем эту ошибку в обратном направлении по сети, чтобы скорректировать коэффициенты и уменьшить ошибку.
Таким образом, задача обучения — это задача минимизации. Вся нейросеть в итоге выдает нам результирующую функцию, и у нас есть результаты, которые мы хотели бы получить в идеале. Функция разницы между тем, что выдала нейросеть, и тем, что хотелось получить — это некоторая функция потерь, и мы ее хотим минимизировать. 
Для минимизации обычно используется метод стохастического градиентного спуска. Но если представить, как выглядит нейросеть целиком, в ней много нейронов, много связей. И когда мы ее минимизируем, то пытаемся посчитать градиент функции потерь относительно всех весов и всех смещений всех нейронов в сети.
Это задача очень большой размерности. А собственно градиент мы считаем на какой-то большой обучающей выборке. И выборка реально должна быть большой, потому что иначе сеть может научиться не тому, чему хочется, и плохо обобщать данные.
Есть эвристика, состоящая в том, что не обязательно за каждый проход делать вычисления на полном наборе входных данных. Можно взять их подмножество, так называемый minibatch, провести расчеты на нем, и это будет неплохой аппроксимацией градиента. В результате на полном наборе входных данных получается не одна, а несколько итераций.
А когда мы проходим через весь набор входных семплов, это называется эпохой обучения, и дальше начинаем прогонять весь набор по новой. 
Чтобы вычислить градиент после того, как мы прогнали вычисления в прямом направлении, мы идем обратно по сети. Не будем смотреть уравнения детально, но это действительно обратное распространение, то есть эффективный метод расчета учитывает структуру сети.
Этот эффективный метод расчета называется алгоритмом обратного распространения или backpropagation algorithm.
Он значительно эффективнее, чем если бы мы пытались считать градиент грубой силой за счет того, что он учитывает структуру сети и использует расчеты, которые были сделаны на разных этапах. 
Процесс можно представить как состоящий из следующих потоков распространения данных и прохождения вычислений. Здесь — прямое распространение вычислений через сеть. 
А здесь — обратное. То есть, действительно, распространение данных и вычислений проходит в обратном порядке по сети и вычисления должны достичь тех точек в этом графе, тех весов и смещений, которые мы оптимизируем.
Если суммировать, алгоритм обратного распространения — это действительно оптимизированный способ посчитать градиент. Он использует информацию о том, как устроена нейросеть. Вычисления распространяются в обратном направлении, и поскольку вычисляется градиент, естественно, мы вычисляем производные. Он использует правила вычисления сложных производных, потому что практически следующий уровень нейросети — это функция от выходов предыдущих уровней.
Основано на туториале «Fundamentals of Scaling Out DL Training» Паулиуса Мицикевичуса (Nvidia) на симпозиуме HotChips
Теперь, когда мы примерно разобрались с тем, что происходит и как вообще устроена нейросеть внутри, мы можем посмотреть, что происходит с вычислениями.
Мы стартуем со случайно инициализированными весами и итерируем наши данные по minibatch за проход в прямом направлении, затем в обратном направлении и апдейтим веса.
Как я уже говорил, сейчас нас интересует обучение, потому что inference — более простая задача, в ней меньше вычислений, и не возникает таких больших объемов передачи данных. 
Проход в прямом направлении выглядит как распространение вектора данных. То есть у нас есть входной вектор, каждый слой преобразует его в выходной вектор этого слоя, и он как-то проходит по сети. В результате мы получаем некоторый выход.
Если minibatch имеет размер больше единицы, то входом становится уже не вектор, а матрица. И каждый промежуточный результат тоже является матрицей. На практике так обычно и происходит, и, более того, это эффективно, потому что GPU заточен на работу с матрицами. 
В этом случае внутри каждого слоя происходит перемножение уже не вектора на матрицу весов, а матрицы входных активаций на матрицу весов. Получается матрица выходных активаций, просто поменялась размерность данных. 
В итоге, когда мы прогнали вычисления в прямом направлении по сети, у нас получается какой-то результат. И у нас есть значения, которые мы хотели бы получить. Взяв и то, и другое, считаем функцию потерь от них. Получается результат потерь, который мы и хотим минимизировать в итоге. 
Когда мы вычислили потери, то начинаем проходить в обратном направлении по сети, и задача этого прохода — посчитать модификаторы, апдейт для весов и смещения всех нейронов сети.
Мы распространяем ошибку в обратном направлении по сети. Каждый уровень считает градиент весов и использует его, чтобы потом проапдейтить веса. Каждый уровень считает градиент активации, чтобы его отдать на предыдущий уровень.
По структуре вычисления очень похожи на то, что было в прямом направлении, но матрица оказывается транспонированной. 
Мы идем обратно по сети и вычисляем апдейты весов. Это шаг оптимизации. На входе у него текущие веса и градиенты весов. На выходе — новые модифицированные веса. Операция состоит в том, что мы обновляем каждый вес с учетом соответствующего значения градиента. На практике вычисления могут быть несколько более сложными, но сама структура вычислений и структура обмена данными меняется не сильно.
Важно, что оптимизации могут использовать значительно больше памяти, чем сама модель. Особенно методы, которые хранят дополнительное состояние, например момент. Это связано с желанием улучшить входимость и уменьшить, например, застревание градиентного спуска на каких-то участках.
Как выглядит стадия вычислений на каждом уровне? На прямом пути это матричные вычисления, где мы входные активации умножаем на матрицу весов и получаем выходные активации. На обратном — градиенты весов и активаций.
Наконец, мы берем матрицу весов, добавляем к ней матрицу апдейтов весов и получаем новые веса.
Важно, что эти вычисления можно распараллеливать, потому что, вообще говоря, вычислительная задача обучения нейросети — сложная, в ней очень много вычислений. Один из способов как-то ускорить процесс — распараллеливать, потому что иначе обучение некоторых сетей может занимать недели или месяцы, что очень печально. Хочется быстрее. Как можно параллелить? По-разному.
Один из методов — Data Parallel, другой — Model Parallel. Внутри модельного параллелизма можно еще делать параллелизм внутри уровней и между уровнями.
Data Parallel — это когда мы разным нодам в сети отдаем вычисления на разных входных данных. То есть мы действительно разбиваем данные и считаем на разных нодах результат для разных сэмплов.
Если же говорить про модельный параллелизм, то можно на разных нодах или акселераторах считать разные участки, буквально разные нейроны внутри одного уровня, обсчитываются разными акселераторами и даже на разных машинах. Либо разные уровни могут считаться на разных нодах сети.
Каждая из моделей параллелизма подразумевает разную структуру потоков данных между вычислительными нодами. Мы будем фокусироваться на параллелизме данных, который устроен проще и весьма популярен. 
Что происходит при распараллеливании по данным? Каждый worker получает полную копию всей нейросети и отвечает за вычисления только на части данных. На прямом проходе он считает входные активации для своей части minibatch, и никаких коммуникаций не происходит.
А вот на обратном проходе появляются коммуникации, потому что мы считаем градиенты и активации для своей части minibatch, и мы считаем вклад для градиентов весов на основании своей порции minibatch, после чего все workers должны просуммировать свои вклады в апдейт весов. То есть каждый worker считает только свою долю. Чтобы посчитать истинную модификацию весов на этом minibatch, нужно просуммировать или редуцировать результат всех workers.
Как это будет выглядеть в разбивке на ноды по разным workers? Вот иллюстрация. Они все работают с одной и той же матрицей весов, но с разными входными данными. 
На обратном проходе они тоже работают с разными данными, потом это все нужно собрать. 
Здесь мы подходим к самому интересному — к коммуникациям. Основная операция, которую мы здесь видим, — так называемый Allreduce. Это очень типовая операция в суперкомпьютерных вычислениях. В мире суперкомпьютеров так называемые групповые или коллективные операции очень типовые, с ними все знакомы. Но за пределами этого мира они достаточно малоизвестны.
Allreduce — одна из самых популярных из них. Ее смысл в том, что мы сначала редуцируем или собираем вместе, в данном случае суммируем промежуточные результаты с каждой ноды, а потом на все ноды распространяем получившийся результат. 
Есть довольно много способов реализовать Allreduce. Эти способы имеют разные затраты по трафику, разные структуры обмена данными. Хуже всего наивный способ. Это когда каждый worker просто отдает свои результаты всем остальным, и каждый сам выполняет все суммирование. Никто так не делает.
К эффективным методам относятся кольцевые редукции — чуть позже покажу, как они выглядят — и так называемые One-shot reduction, то есть редукции за один проход, которые обычно используются в типологиях с коммутаторами, где все ноды видят других — достижимых, хоть и не обязательно через один ход. Такие типы редукции могут использовать иерархические оптимизации. 
Есть довольно большой коммуникационный overhead, который растет вместе со степенью параллелизма. И есть способы его оптимизации, на которые мы дальше посмотрим. 
Что касается модельного параллелизма внутри уровня и между уровнями, я просто привожу картинки о том, как это устроено с точки зрения разделения задач. Но мы не будем в них углубляться. 
Почему важно оптимизировать коммуникации между нодами? Это графики, снятые с тестового кластера с тестовой нагрузкой. Синее — GPU power, то есть энергопотребление GPU, которое является довольно хорошим индикатором того, насколько занят GPU и насколько он активно считает.
Другой график — обмен данными. Здесь видно, что наш дорогущий GPU простаивает, пока идет обмен и синхронизация данных. В жизни очень хотелось бы минимизировать этот простой. 
Как устроены потоки трафика и как они связаны с интерконнектами? Небольшое напоминание про коллективные или групповые операции, какие они бывают. 
Ссылка со слайда
Бывают операции broadcast, когда один набор данных раздается всем. Бывает scatter, когда есть набор данных, и мы по кусочку раздаем его остальным нодам. Gather — строго обратное, собираем данные вместе по кусочкам.
И reduction, когда мы собираем куски и вычисляем из них что-то новое. Например, суммируем. Хотя редукция не обязана быть суммированием, в случае нейросети она является суммированием. Важная свойство редукции — то, что размер блока данных после редукции не меняется. То есть мы собрали по вектору какой-то размерности с каждой ноды, просуммировали, получилось не в n раз больше, а столько же данных, но в которых уже учтены промежуточных результаты.
Ссылка со слайда
Как я уже говорил, наивная реализация One-shot reduction за один проход не эффективна. 
Ссылка со слайда
Есть другие реализации, например так называемый кольцевой Allreduce. 
Он сводится к тому, что на первой фазе ноды по замкнутому кольцу коммуницируют промежуточный результат друг другу. Причем со смещением. 
В итоге каждая нода получает свою завершенную часть результата. 
И потом реплицирует на все остальные. 
Ссылка со слайда
Если представить, что кольца могут быть ориентированы по разным осям, то этот способ также может быть иерархическим. И такая конструкция очень хорошо ложится на типологию классических интерконнектов. 
Потому что с чем мы обычно сталкиваемся в IP-сетях? IP сети, которые мы строим, — всё это коммутируемые или indirect-сети. Ноды в них подключены не напрямую к другим нодам, если мы говорим про вычислительные ноды, а к коммутаторам. И только через них видят другие ноды.
А есть еще интерконнекты/сети, которые называются прямыми, direct, или же point-to-point. Это когда в одной ноде совмещены и линки для коммуникации с соседями, и вычислительные элементы. Многие суперкомпьютерные сети и многие кластеры для high performance-компьютинга исторически так и выглядят. Это неспроста. 
Ссылка со слайда
Дело в том, что многие коллективные операции очень хорошо ложатся на такую топологию интерконнекта. Если мы говорим конкретно про машинное обучение, то например, Google, как известно, делает свои тензорные процессоры для машинного обучения. А еще он делает на них кластеры. И в этих кластерах топология — как раз 2-D torus, что позволяет кольцевой Allreduce непосредственно отобразить на топологию этого кластера. 
Ссылка со слайда
Но конечно, не все могут делать свои акселераторы и специализированные интерконнекты для них. Что же можно сделать? А если бы свитчи тоже умели что-то считать?
Итак, вычисление в сетях на сетевых элементах. 
Ссылка со слайда
Мы посмотрим на две реализации. Первая — SHARP, Scalable Hierarchical Aggregation and Reduction Protocol. Это реализация вычислений в сети от компании Mellanox, теперь Nvidia.
И это production ready-реализация. Она оптимизирует все ту же задачу расчета редукций. Один из способов такого расчета — построение дерева, по которому могут идти редукции. Это логическое дерево. Оно отображается на топологию сети. Если мы говорим про индиректную или switched-сеть, то отображаться оно может довольно неудачно. Например, вызывая перегрузку на определенных линках.
Пример одной из реализации Allreduce — Recursive Doubling. Он происходит в несколько этапов. 
Что делает SHARP? Во-первых, SHARP основан на том, что внутри коммутаторов, в данном случае это InfiniBand-коммутаторы, реально есть логические устройства, которые могут что-то считать. Считать им нужно не так много, в данном случае достаточно суммировать. Но они могут работать со всеми популярными типами данных, которые нужны в HPC и в машинном обучении.
В случае с SHARP у нас также строится дерево агрегации, это логическое дерево, наложенное на физическую топологию сети. 
Что происходит? У нод появляются промежуточные расчеты данных. Они отсылают эти данные по дереву, и при распространении вверх по этому дереву происходит редукция. У нас есть switch снизу. Ему несколько нод присылают свои данные для редукции. Он редуцирует эти наборы данных до одного.
Мы получили пачку векторов от разных нод, например, при расчетах для машинного обучения. Просуммировали эти вектора, наверх ушел один вектор. То есть данные уменьшились в размере.
Все это доходит до корня дерева. Корень дерева получает итоговый результат, потому что у него есть все необходимые для расчета данные. После чего идем в обратном направлении — сверху вниз дерева. У нас начинается распространение итогового результата на все ноды, участвующие в вычислении.
И поскольку здесь речь идет об InfiniBand, у которого надежный транспорт, и есть реализация надежного мультикаста, это очень эффективная операция. Этот номер не прошел бы так легко в произвольной пакетной сети, где есть потери. Нам пришлось бы подставлять некоторое количество подпорок. Но в случае с InfiniBand это работает очень неплохо. 
Дело даже не в том, что свитчи очень эффективно считают. Они делают простую операцию — суммируют, просто находятся в правильном месте, чтобы это считать, потому что раз мы просуммировали векторы внутри свитча, дальше уходит меньше данных.
То есть мы не просто сделали полезную работу, выполнив операцию редукции. Мы еще и значительно уменьшили объем передаваемых дальше данных. 
На графике видно, как использование SHARP влияет на использование операции Allreduce. Влияет не только на полосу, а еще и на Latency. Видно, что Latency остается примерно постоянной. Конечно, на самом деле нет. Latency является логарифмической. Потому что если посмотреть на то, что было здесь, это на самом деле подмножество внутри сети. И Latency при использовании SHARP зависит в основном от количества уровней сети, а не от количества нод, принимающих участие в вычислениях.
С учетом того, что сейчас доступны свитчи довольно большого радикса, это позволяет иметь довольно хорошо ограниченную сверху Latency даже в очень больших кластерах.
Но улучшается не только полоса и задержка.
Есть данные тестов с обучением различных нейросетей. В зависимости от того, какие модели мы обучаем, сколько нод в кластере, и от другой специфики, есть выигрыш от единиц процентов до весьма заметных 10–16% на некоторых моделях.
Это очень много с учетом того, какими цифровыми молотилками являются современные GPU и как соотносится стоимость самой вычислительной части нод GPU со стоимостью интерконнекта. 10% роста производительности могут пару раз окупить весь интерконнект.
sands.kaust.edu.sa/project/switchml
И еще одна реализация. Она не production ready, но куда более доступна, потому что подразумевает менее экзотическое оборудование.
То есть в исследовательских целях она более доступна, и, вообще, больше шансов до нее добраться, даже если у вас нет большого кластера. Это SwitchML.
Та же проблема — распределенное обучение, Data Parallel, worker посчитали свои промежуточные результаты. 
Их нужно собрать вместе, и мы хотим сделать это эффективнее. 
И хотим сделать это на switch.
На чем это можно было бы сделать? Тут мы вспоминаем, что у нас есть программируемые чипы от компании Barefoot, теперь Intel. У них нет поддержки операций с плавающей запятой и продвинутой арифметики, но в целом с числами они работать умеют. Выясняется, что на этом тоже можно собрать конструкцию для редукции в сети. 
Поработать придется немножко больше, и у нас будет разделение обязанностей между worker и switch. Worker будет нарубать данные на векторы ограниченного размера, делать масштабирование и нормализацию. Потому что нам нужно данные с плавающей точкой перевести в целочисленные значения: switch ничего другого не умеет. И, наконец, поскольку это нормальная, привычная пакетная сеть, в которой есть потери, то worker нужно будет обнаруживать и восстанавливаться при потерях.
Switch же будет выполнять суммирование векторов, а также отслеживать, как заполняются слоты для вычислений. 
Подробно вдаваться не будем. У нас есть данные с плавающей точкой. У них разные порядки. Мы хотим это все привести к целочисленному представлению, и желательно, чтобы они по масштабу были более-менее вместе, чтобы не было сильно различающихся масштабов данных. Это делается на worker.
Важная вещь — устойчивость к потере пакетов. Пакеты могут теряться в двух направлениях — от worker к switch и в обратном направлении. worker определяют потери, используя таймеры. Пакеты перепосылаются. Если switch получает один и тот же пакет больше одного раза, он его просто игнорирует.
Switch при помощи bitmap отслеживают, какие worker уже прислали свои данные для данного слота вычислений. 
Собрана вся эта конструкция с использованием такого знатного зоопарка технологий. Но в то же время реализована поддержка для различных фреймворков машинного обучения.
Это тестовый кластер, на котором производились эксперименты. 
Здесь видно, что этот кластер — даже такой прототип, собранный не всегда на самых эффективных технологиях, а на том, что мне было доступно, — очень неплохо себя показывает относительно варианта, не использующего вычисление на switch. И как видно, это в любом случае больше единицы, а в некоторых случаях выигрыш достигает трешки. 
Из экспериментов видно, что производительность не очень зависит от количества workers, в отличие от реализаций, не использующих вычисления на switch.
SwithML достаточно неплохо себя ведет при разумных значениях количества потерь. 1% — для типовой сети это уже как-то очень много. Меньше 0,1% — это то, на что можно ориентироваться. Даже 0,1% — много в контролируемой среде.
И, наконец, чуть подробнее о том, как реализована устойчивость относительно потерь пакетов. На этапе агрегации мы игнорируем дублирующиеся данные, ждем того, что не пришло. И на этапе распространения результатов мы отсылаем результат. Если worker не получил результат, он его перезапрашивает. Система устойчива к потерям в обоих направлениях. 
Ссылка со слайда
И про SHARP, и про SwitchML есть статьи и презентации. Они в открытом доступе, их можно посмотреть.
