Сегодня мы дадим ответ на простой вопрос: "Как работает распределённое обучение (в контексте MXNet)?" Все примеры кода протестированные на MXNet v0.10.0 и могут не работать (или работать по-другому) в других версиях, однако полагаю, что общие концепции будут неизменимы еще долго. Ну и последнее перед тем, как мы перейдем к основной части, я хочу выразить благодарность за помощь в написании статьи моим коллегам, без которых эта статья не была бы возможной: Еще хотел бы порекомендовать поднять машинку с DLAMI и выполнить все примеры из статьи самостоятельно, тем более, что они достаточно простые. Для выполнения кода вполне себе подойдет бесплатная машинка на AWS.  С преамбулой окончено, лезем под кат... В MXNet все участники процесса обучения поделены на 3 логические группы: Это чисто логическое распределение, так что все участники могут работать на одной машине. Для начала посмотрим на поверхностное объяснение, что каждый из участников собой представляет: Планировщик является центральным узлом кластера, отвечает за начальную настройку кластера, обеспечение нужной информацией каждого участника процесса обучения и… ничего более. Мы еще увидим, как он впадает в анабиоз сразу, как только кластер готов начать обучение. И даже когда кластер закончит свое обучение его задачей будет лишь выключить себя. Думаю, все уже догадались, что в кластере может быть только один планировщик. Сервер выступает в качестве хранилища параметров модели обучения. То есть, если обучается модель в стиле: Y = AX + B, сервер хранит вектора A и B. Еще он отвечает за их корректное обновление. Серверов может быть более чем один, а соответственно есть правило, по которому модель распределяется на несколько серверов. Но это тема отдельной статьи. Это собственно те участиники кластера, которые непосредственно выполняют обучение модели. Каждый рабочий получает свою часть данных, на которых нужно обучится, считает шаг градиента и отправляет его серверам для обновления модели. Давайте возьмем бутафорский пример кластера с: Сам кластер будет выглядеть вот так:  Эта картинка, как и описанная конфигурация, будет использована только для визуализации потока данных.  Мы не будем на практике создавать такой большой кластер как было описано выше, а обойдемся намного более маленьким кластером с 3-мя нодами на одной физической машине. Есть несколько тому причин: Перед тем как продолжить, нужно уточнить одну деталь. Для MXNet, распределённое обучение по сути означает, что необходимо использовать KVStore. Имя это — акроним от "Key Value Storage". И по существу — это распределенное хранилище, которое выполняется на серверах и имеет некоторую дополнительную функциональность (например знает, как именно нужно обновлять модель, получив градиентный шаг от рабочего). Настало время начать создавать логических участников кластера. Чтобы создать участника достаточно лишь создать некоторые переменные среды и после заимпортировать модуль mxnet. Первым делом запустим планировщик: Давайте остановимся тут на секунду, чтобы получить представление о том, что происходит. Первые 4ре линии кода не должны вызывать много вопросов у Python программистов: просто импорт зависимостей и создание окружения ОС. Что интересно здесь, так это то, какие именно обновления в переменные окружающей среды будут внесены: Начнем с рассмотрения DMLC_ROLE. Давайте посмотрим где именно она используется, а именно в пакете ps-lite. В соответствии с официальным README(в вольном переводе): Думаю, не стоит быть С++ гуру, чтобы понять, что тут происходит. Логичиская роль нода определяется по строке в этой самой переменной "DMLC_ROLE". Забавно, но похоже тут нет проверки на то, что данная переменная содержит одно из разрешенных значений. Это, потенциально, может привести к интерестным проблемам. Второе, что нас интересует, это не только где переменная читается, но и где она используется. Что бы рассказать об этом, нужно обратится к файлу van.cc, который будет встречаться нам не раз, вот конкретная линия, где переменная используется и создается переменная "is_scheduler": Если быстро пробежаться далее по коду, чтобы посмотреть, что там происходит можно увидеть следующее интересное место: В этом конкретном примере переменная "role" никогда не будет равна Node::SCHEDULER. Так что у вас есть шанс создать pull-request, чтобы это поправить (если еще никто этого не сделал). Так же глядя на это место понимаешь, что для планировщика не так уж и много работы. Это потому что, в отличии от рабочего и сервера — планировщик использует IP адрес и порт, которые ему были переданы, а не ищет свободный порт в системе. Идем далее, параметр: DMLC_PS_ROOT_PORT. С этим мы быстро разберемся с учетом уже имеющихся знаний. Вот код, который уже видели: Вновь, это из van.cc. Как не трудно догадаться, это порт, на котором планировщик должен слушать сообщения. Надеюсь на этом этапе понятно, что DMLC_PS_ROOT_URI это просто ip адрес планировщика. Так что давайте сразу прыгнем к обсуждению DMLC_NUM_SERVER и DMLC_NUM_WORKER. Так сложилось, что каждый логический нод MXNet в кластере должен знать о всех других нодах. Так что для каждого нода, перед тем как он запустился, в переменных среды записывается как много в кластере рабочих и серверов (число планировщиков ненужно, ибо это всегда 1). К слову эта информация хранится в классе Postoffice (вместе с другой информацией о кластере). Ну и последний параметр, но пожалуй один из архи-главнейших — PS_VERBOSE. Это заставит наш новосозданный процесс выводить отладочную информацию, что жизненно важно для нас сейчас. С точки зрения нашей бутафорской диаграммы наш кластер выглядит сейчас как-то так:  Теперь, когда у нас есть планировщик, давайте поднимем сервер. Так как мы поднимаем все логические ноды на одной машине, то нам придётся создать копию параметров окружения и вновь внести туда нужные изменения для того, что бы запустить сервер: Надеюсь теперь происходящее в коде не вызывает вопросов, но на всякий случай: Тут кто-то может спросить: погодите, я думал что DMLC_PS_ROOT_PORT и DMLC_PS_ROOT_URI для указания IP и порта логического нода, который мы запускаем? Ответом будет — нет, это адрес и порт планировщика, а вот все остальные должны сами разобраться какой у них адрес и найти себе доступный порт в системе. Информация о планировщике им нужна, чтобы постучаться к нему и попросить, чтобы он добавил их в кластер. После запуска серверов наша диаграмма выглядит вот так:  Настало время запустить, собственно, самого рабочего и создать KVStore: К слову, KVStore может работать в двух режимах: Я оставлю за пытливым читателем вопрос о том, чем эти режимы отличаются, об этом можно почитать вот тут.  После запуска рабочих, наша диаграмма будет выглядеть вот так:  Перед тем как бросится обсуждать, что происходит в момент создания KVStore, нужно рассказать о том, что у каждой ноды есть жизненный цикл, у которого есть следующие события: Так же, тот самый класс (Van), что отвечает за обработку этих событий, имеет несколько других, не менее важных методов. О части из них мы поговорим детально позже в других статьях, а сейчас просто перечислим: Вот что выполняет каждая нода в момент, когда приходит сигнал Start: Как только все команды, приведенные выше, будут выполнены, на экране должно показаться много отладочной информации, которая приходит от трех ранее запущенных процессов одновременно. Пройдемся теперь по каждой линии, чтобы детально обсудить, что происходит и как будет выглядеть наша диаграмма на каждом этапе. Это запускается процесс рабочего. В данном случае это метод Start, который сообщает нам, что его адрес 1.1.1.1, роль "worker" и порт, который он нашел 37350. Теперь он моментально попробует уведомить планировщик, что он готов быть добавлен в кластер, указав свой адрес и порт: Это конкретное сообщение сгенерировано в методе Send, вот тут. В нем нужно обратить внимание на несколько вещей: На нашей диаграмме этот обмен сообщениями выглядит следующим образом:  Идем далее Это проснулся наш сервер. Нашел себе порт (54160) и тут же пытается уведомить об этом планировщик: На диаграмме это выглядит вот так:  Так же, как и в случае с рабочим, наш сервер отправляет команду "ADD_NODE", чтобы его зарегистрировали в кластере. Так, как сервер еще не зарегистрирован в кластере и не имеет ранга, то мы видим: "? => 1". Наконец то планировщик запущен. Он использует локальный IP и порт 9000 (все ноды в кластере должны уже знать об его адресе и порте). Так как планировщик поднят, то логично ожидать, что в этот момент он получит все входящие сообщения, что были ему отправлены и… вуаля: Сообщение от сервера. Эта часть логов сгенерирована методом Receive, если быть еще более точным то вот тут. Тут же планировщик получает второе сообщение, на этот раз от рабочего: Первым делом планировщик берется назначать ранги, вначале рабочему (9): Теперь серверу (8): После идет довольно важная часть: Сообщения вроде этих показывают, что планировщик получил команды "ADD_NODE" от всех, нод кластера (в нашем случае от 1го рабочего и 1го сервера) и теперь начал уведомлять все ноды обратно о их рангах и об информации о всех других нодах в кластере. То есть планировщик отправляет ВСЮ информацию о КАЖДОМ ноде кластера КАЖДОМУ ноду кластера. В данном конкретном сообщении мы видим все данные о кластере и это сообщение отправлено ноде с рангом 9 (это работник). Информация о кластере жизненно важна, так как она нужна рабочему, например, чтобы понять, на какой сервер отправлять обновление модели. На диаграмме этот процесс выглядит вот так:  Следующий вывод: Такое же подтверждение планировщик отправляет ноде с рангом 8 (сервер). На диаграмме выглядит так:  Планировщик радостно сообщил, что он подключён к одному рабочему и одному серверу (ко всем нодам кластера). Напоминание — при запуске на реальном кластере все эти логи находятся на разных машинах, посему сейчас может показаться, что тут информации больше чем нужно. Это рабочий получил сообщения от планировщика и сообщает, что он подключен к кластеру. Можно спросить, а что такое "2147483647". Ответ — понятия не имею =) скорее всего бага, Я бы ожидал увидеть: "1 =>9". Так, как рабочий корректно видит свой ранг: "W[9]", баг скорее всего где-то в процессе логирования, так что можете его пофиксить и стать контрибьютором проекта. Тоже самое для сервера: он получил сообщение и довольный поведал об этому миру. Еще одна важная часть. До сих пор мы видели только одну команду "ADD_NODE". Тут мы наблюдаем новую: "BARRIER". Если кратко, то эта концепция барьеров, которая, надеюсь, знакома читателю по многопоточному программированию и означает: "остановитесь до тех пор, пока все не дойдут до этого барьера". Планировщик отвечает за то, чтобы сообщить, когда именно все достигнут барьера и могут продолжать выполнение. Первый барьер расположен сразу после того, как кластер стартовал, но перед началом обучения. Все три ноды (включая сам планировщик) отправили сообщения, которое по существу значит: "я достиг барьера, дай мне знать, когда можно двигаться далее". Так же, как видно из сообщения тут есть понятие барьерной группы(barrier_group). Барьерная группа эту группа нодов, которые участвуют в том или ином барьере. Эти группы: 1 — планировщик
2 — сервера
4 — рабочие Как не трудно догадаться, это степень двойки, так что наша группа 7 это: 4 + 2 + 1. По существу данный барьер распространяется на всех. Ну и само собой, так как в наших логах, мы увидели три отправки сообщения, логично ожидать три строки о получении планировщиком этих сообщений: Происходящее на нашей диаграмме выглядит вот так:  Теперь настало время обсудить, что делает планировщик, когда получает новое сообщение о том, что нода достигла барьера в определенной группе: В логах выше можно увидеть, как счетчик увеличивался по мере получения каждого нового сообщения. Ну а в момент, когда он достиг ожидаемого размера (3), планировщик начал отправлять подтверждения: На нашей диаграмме это выглядит вот так:  Как можно заметить, планировщик даже отправляет подтверждение самому себе. Ну и само собой, раз было отправлено сообщение от планировщика (аж 3), то мы должны увидеть логи о том, что эти сообщения получены: Ну и последние прикосновения. В данный момент планировщик достиг второго барьера, который будет достигнут всеми нодами после окончания обучения, однако, так как планировщик не принимает участие в обучении, то он уже достиг этот самый барьер. Так что он отправляет группе barrier_group=7 что он достиг барьер, с мгновенным подтверждением получения сообщения и установкой счетчика барьерной группы 7 в 1. На этом этапе инициализация кластера закончена, можно начать обучение... Выполнив весь код, мы имеем инициализированный KVstore. Что же теперь? Давайте используем его для непосредственного обучения. Я воспользуюсь очень простым примером линейного регрессора взятого вот от сюда. Только прошу, перед тем, как продолжить, пройдитесь по примеру, чтобы понять происходящее. Что бы сделать тренировку в описанном примере распределенной, нужно поменять всего 1 линию в коде. Вместо: Нужно написать: Так просто? если коротко — да. Надеюсь у читателя теперь есть более детальное понимание того, что происходит с кластером MXNet в момент его старта. Также, я надеюсь, эта статья поможет с отладкой кластера в случае каких либо проблем. Ну и плюс, имея эти знания, можно сделать некоторые выводы о характеристике сети для кластера, а именно: Буду очень благодарен за рекомендации оригинальной статьи на Medium. Так же, если вдруг вы занимаетесь построением распределённых систем машинного обучения на базе AWS с использованием MXNet и у вас есть какие либо вопросы, то я с радостью готов помочь и ответить (viacheslav@kovalevskyi.com).  Google Cloud / DeepLearning