Нейросеть обучили с помощью распределенной вычислительной сети

Разработчики из России, США и Канады создали и протестировали новую платформу распределенного обучения нейросетей на множестве компьютеров, соединенных через интернет. Система сама распределяет задачи по компьютерам, учитывая при этом их характеристики, а также справляется с потерей части участников и их работой в разное время. Разработчики продемонстрировали работу платформы, обучив нейросеть для обработки бенгальского языка на компьютерах 40 добровольцев за девять дней. Статья о разработке опубликована на arXiv.org, также разработчики рассказали о ней в блоге Hugging Face.

Для некоторых научно-технических задач необходимо проводить огромный объем вычислений, выполнить которые в приемлемые сроки можно с использованием суперкомпьютеров. Поскольку не для всех задач удается найти достаточное финансирование, чтобы получить доступ к мощному суперкомпьютеру, в конце 1990-х годов исследователи начали создавать распределенные сети вычислений. Они состоят из огромного количества обычных людей, которые добровольно установили на свой компьютер программу, использующую часть его мощности для научных вычислений. Подход оказался успешным и некоторые проекты, такие как SETI@Home и Folding@Home достигли мощности, сопоставимой с лучшими суперкомпьютерами.

Поскольку обучение нейросетей тоже требует серьезных вычислительных мощностей, в последние годы разработчики стали предлагать использовать аналогичную схему и для машинного обучения. Подобные проекты уже существуют, самый известный из них — MLC@Home. Но у него, как и у аналогов, есть важное ограничение: каждый компьютер в сети обучает отдельные небольшие модели, а не одну общую гигантскую модель. Но в последние годы размеры успешных нейросетевых моделей растут с большой скоростью. Например, известную нейросеть для работы с текстом GPT-3 обучали на 570 гигабайт данных, а в самой модели используется 175 миллиардов параметров.

В прошлом году российские разработчики предложили архитектуру системы для распределенного обучения одной нейросети на множестве независимых компьютеров. В новой работе группа разработчиков из России, США и Канады, в которую входит и один из авторов той статьи Максим Рябинин (Max Ryabinin) из Яндекса и Высшей школы экономики, создала новую платформу и показала ее реальное применение на практике.

В начале обучения платформа определяет характеристики компьютеров в сети и параметры их сетевого подключения. Затем задачи распределяются по конкретным компьютером в зависимости от этих факторов: мощные компьютеры с высокоскоростным подключением к интернету и низкими задержками могут выполнять любые задачи, а компьютеры с тем или иным ограничением берут не себя только часть. Среди задач есть как непосредственно обучение и расчет функции потерь, так и другие, например, компьютеры могут проводить градиентный спуск, усреднять градиенты и оптимизировать веса модели после этапов обучения.

Поскольку соединение участников может быть нестабильным и слабым, разработчики оптимизировали архитектуру платформы, чтобы уменьшить частоту синхронизации. В частности, они решили использовать очень большой размер партий данных (batch — часть датасета, которая берется для одного этапа обучения). После того как партия обрабатывается, компьютеры в сети обмениваются градиентами и те из них, которые занимаются оптимизацией, подстраивают веса модели и рассылают по сети, чтобы компьютеры приступили к новой итерации обучения.

Платформа допускает изменение участников в сети прямо в процессе обучения: если к сети присоединяются новые компьютеры, обучение идет быстрее, а если старые отключаются и не передают свою часть данных, эти расчеты берут на себя остальные участники. Стоит отметить, что, как и предыдущие аналоги, метод не подходит для обучения огромных моделей наподобие GPT-3: хотя многие операции можно распределить по практически любым компьютерам, в сети все равно должны быть компьютеры, способные запускать всю модель и проводить ее обучение, пропуская данные из обучающей выборки.

Разработчики проверили работоспособность подхода на реальной задаче. Они собрали 40 участников с 91 компьютером, в которых были видеокарты разной мощности: от обычных и не очень мощных на текущий момент NVIDIA GTX 1060 до V100, предназначенного для мощных дата-центров. Задача участников была в том, чтобы предобучить языковую модель для обработки бенгальского языка, основанную на архитектуре ALBERT. В качестве обучающей выборки они использовали слепок бенгальской Википедии по состоянию на март 2021 года размером 657 мегабайт и бенгальскую часть датасета OSCAR размером 6,2 гигабайта. Для того чтобы участникам не приходилось загружать сразу такой объем данных перед началом работы разработчики создали стриминговый сервер, чтобы данные можно было подгружать по частям параллельно с обучением. При этом они отмечают, что в будущем отдельный сервер можно заменить p2p-протоколом наподобие BitTorrent.

Эксперимент занял девять дней. По результатам обучения разработчики оценили качество модели, сравнив ее с тремя аналогами по двум метрикам. Первая из них показывает качество классификации входящих токенов (слов или их частей) и отнесение их к одному из четырех типов: человек, организация, местоположение или другое. Вторая метрика показывает, насколько качественно модель определяет тип всего текста. В результате качество модели, обученной распределенным способом, оказалось сравнимо с обычными моделями, в том числе и имеющими гораздо большее количество параметров:

Нейросети адаптируют не только для распределенного обучения на обычных компьютеров, но и для работы с новыми платформами и новыми типами данных, например, на нейроморфных чипах, более близких к архитектуре естественных нейросетей, и с квантовыми данными.