Методы обучения больших нейронных сетей09.06.2022 07:00

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

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

Без распараллеливания

Обучение нейронной сети — это итеративный процесс. На каждой итерации мы выполняем прямой проход (forward pass) через слои⁠ модели, чтобы вычислить результат для каждого примера из пакета (батча) данных. Затем выполняется еще один проход в обратном направлении⁠ через слои (backward pass), распространяя информацию о том, насколько сильно каждый параметр влияет на конечный результат, путем вычисления градиента⁠ по отношению к каждому параметру. Средний градиент для батча, параметры и некоторое состояние оптимизации для каждого параметра передаются алгоритму оптимизации, например Adam⁠, который вычисляет параметры для следующей итерации (которые должны обеспечивать немного лучшую производительность на ваших данных) и новое состояние оптимизации. По мере того как обучение итеративно обрабатывает батчи данных, модель эволюционирует, выдавая всё более точные результаты.

Различные методы распараллеливания разделяют этот процесс обучения по разным измерениям,  включая:

• параллелизм данных (Data parallelism) — запуск различных подмножеств батча на разных GPU;
• конвейерный параллелизм (Pipeline parallelism) — запуск разных слоев модели на разных GPU;
• тензорный параллелизм (Tensor parallelism) — разбиение математических вычислений для отдельной операции, такой как умножение матриц, с распределением между GPU;
• смесь экспертов (Mixture-of-Experts) — обработка каждого примера только частью каждого слоя.

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

Параллелизм данных

Обучение с параллелизмом данных (Data Parallel) означает копирование одних и тех же параметров на несколько GPU (часто называемых «воркерами» или рабочими узлами) и назначение каждому из них разных примеров для одновременной обработки. Сам по себе параллелизм данных по-прежнему требует, чтобы ваша модель помещалась в память одного GPU, но позволяет задействовать вычислительную мощность множества GPU ценой хранения множества дублированных копий параметров. Тем не менее существуют стратегии увеличения эффективного объема оперативной памяти, доступной вашему GPU, например временная выгрузка параметров в память ЦП (CPU) между сеансами их использования.

Поскольку каждый воркер с параллелизмом данных обновляет свою копию параметров, им необходимо координировать свои действия, чтобы у каждого воркера оставались схожие параметры. Простейший подход заключается в использовании блокирующей связи между воркерами: (1) независимое вычисление градиента на каждом воркере; (2) усреднение градиентов по всем воркерам⁠; и (3) независимое вычисление одних и тех же новых параметров на каждом воркере. Шаг (2) представляет собой блокирующее усреднение, требующее передачи довольно большого объема данных (пропорционально количеству воркеров, умноженному на размер ваших параметров), что может снизить пропускную способность обучения. Существуют различные асинхронные схемы синхронизации⁠, позволяющие устранить эти накладные расходы, однако они ухудшают эффективность обучения; на практике люди обычно предпочитают синхронный подход.

Конвейерный параллелизм

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

Разбить большую модель на блоки последовательных слоев довольно просто. Тем не менее между входами и выходами слоев существует последовательная зависимость, поэтому наивная реализация может привести к большим периодам простоя, пока воркер ожидает выходных данных от предыдущей машины, чтобы использовать их в качестве своих входов. Эти периоды ожидания известки как «пузыри» (bubbles), что приводит к пустой трате вычислительных ресурсов, которые могли бы быть задействованы простаивающими машинами.

Мы можем использовать идеи параллелизма данных для снижения издержек от «пузырей», обязав каждого воркера обрабатывать только подмножество элементов данных за раз, что позволяет нам ловко совместить новые вычисления со временем ожидания. Основная идея заключается в разделении одного батча на несколько микробатчей; обработка каждого микробатча происходит пропорционально быстрее, и каждый воркер приступает к работе над следующим микробатчем сразу же, как только он становится доступным, ускоряя тем самым выполнение конвейера. При достаточном количестве микробатчей воркеры задействованы большую часть времени с минимальным «пузырем» в начале и конце шага. Градиенты усредняются по микробатчам, а обновление параметров происходит только после завершения обработки всех микробатчей.

Количество воркеров, между которыми разделена модель, обычно называют глубиной конвейера (pipeline depth).

Во время прямого прохода воркеры должны отправлять только выходные данные (называемые активациями) своего фрагмента слоев следующему воркеру; во время обратного прохода они отправляют только градиенты по этим активациям предыдущему воркеру. Существует широкое пространство для проектирования схем планирования этих проходов и агрегирования градиентов по микробатчам. GPipe⁠ организует последовательное выполнение прямых и обратных проходов каждым воркером с последующим синхронным агрегированием градиентов от множества микробатчей в самом конце. PipeDream⁠, напротив, планирует работу каждого воркера так, чтобы чередовать прямые и обратные проходы.

Тензорный параллелизм

Конвейерный параллелизм разделяет модель «вертикально» по слоям. Также возможно «горизонтальное» разделение определенных операций внутри слоя, что обычно называется тензорным параллельным обучением (Tensor Parallel). Для многих современных моделей (таких как Transformer⁠) узким местом вычислений является умножение матрицы батча активаций на большую весовую матрицу. Умножение матриц⁠ можно рассматривать как скалярные произведения между парами строк и столбцов; можно вычислять независимые скалярные произведения на разных GPU или вычислять части каждого скалярного произведения на разных GPU с последующим суммированием результатов. При любом из этих подходов мы можем разделить весовую матрицу на «осколки» (shards) равного размера, разместить каждый осколок на отдельном GPU и использовать его для вычисления соответствующей части общего матричного произведения перед последующим обменом данными для объединения результатов.

Одним из примеров является Megatron-LM⁠, который распараллеливает умножение матриц внутри слоев самовнимания (self-attention) и многослойного перцептрона (MLP) трансформера. PTD-P⁠ использует тензорный, дата- и конвейерный параллелизм; расписание его конвейера назначает несколько неконсекутивных (непоследовательных) слоев каждому устройству, уменьшая накладные расходы на «пузыри» ценой увеличения сетевого взаимодействия.

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

Смесь экспертов (Mixture-of-Experts, MoE)

При использовании подхода «смеси экспертов» (Mixture-of-Experts, MoE)⁠ для вычисления ответа на любой конкретный входной сигнал задействуется лишь малая доля сети. Один из вариантов подхода заключается в создании множества наборов весов, причем сеть может выбирать, какой набор использовать, с помощью механизма гейтинга (gating) во время инференса. Это позволяет использовать гораздо большее количество параметров без увеличения вычислительных затрат. Каждый набор весов называется «экспертом» в надежде, что сеть научится распределять специализированные вычисления и навыки между различными экспертами. Разные эксперты могут размещаться на разных графических процессорах, что обеспечивает понятный способ масштабирования количества GPU, используемых для модели.

GShard⁠ масштабирует модель MoE Transformer до 600 миллиардов параметров с помощью схемы, при которой только слои MoE разделяются между несколькими устройствами TPU, а остальные слои полностью дублируются. Switch Transformer⁠ масштабирует размер модели до триллионов параметров с еще более высокой разреженностью (sparsity) за счет маршрутизации каждого входного сигнала к единственному эксперту.

Другие решения для экономии памяти

Существует множество других вычислительных стратегий, делающих обучение всё более крупных нейронных сетей более управляемым. Например:

• Для вычисления градиента необходимо сохранить исходные активации, что может потреблять большой объем оперативной памяти устройства. Чекпоинтинг (Checkpointing)⁠ (также известный как ревычисление активаций или activation recomputation) сохраняет произвольное подмножество активаций и пересчитывает промежуточные значения «на лету» (just-in-time) во время обратного прохода. Это существенно экономит память ценой дополнительных вычислительных затрат, составляющих не более одного полного прямого прохода. Также можно постоянно балансировать между затратами вычислительных ресурсов и памяти с помощью выборочного ревычисления активаций⁠ (selective activation recomputation), которое представляет собой чекпоинтинг подмножеств активаций, относительно более дорогих для сохранения в памяти, но более дешевых для вычисления.
• Обучение со смешанной точностью (Mixed Precision Training)⁠ предполагает обучение моделей с использованием чисел пониженной точности (чаще всего FP16⁠). Современные ускорители способны достигать гораздо более высоких показателей FLOP при использовании чисел пониженной точности, что также экономит оперативную память устройства. При правильном подходе результирующая модель практически не теряет в точности.
• Выгрузка (Offloading) заключается во временной выгрузке неиспользуемых данных на ЦП (CPU) или между различными устройствами с последующим чтением обратно по мере необходимости. Наивные реализации сильно замедляют обучение, однако изощренные подходы осуществляют предварительную выборку данных (pre-fetch), благодаря чему устройство никогда не простаивает в ожидании. Одной из реализаций этой идеи является ZeRO⁠, которая распределяет параметры, градиенты и состояния оптимизатора по всему доступному оборудованию и материализует их по мере необходимости.
• Были предложены эффективные с точки зрения памяти оптимизаторы (Memory Efficient Optimizers) для уменьшения объема памяти, занимаемого рабочим состоянием, поддерживаемым оптимизатором, например Adafactor⁠.
• Сжатие также может использоваться для хранения промежуточных результатов в сети. Например,  Gist⁠ сжимает активации, которые сохраняются для обратного прохода;  DALL·E⁠ сжимает градиенты перед их синхронизацией.

В OpenAI мы занимаемся обучением и совершенствованием больших моделей — начиная с базовой инфраструктуры и вплоть до их развертывания для решения реальных задач. Если вы хотите применить идеи из этой статьи на практике (что особенно актуально для наших команд по масштабированию и прикладным исследованиям), мы нанимаем сотрудников⁠!

Спасибо Николасу Тезаку (Nikolas Tezak), Сэму Альтману (Sam Altman), Дэниелу Гэклзу (Daniel Gackle), Илье Суцкевичу (Ilya Sutskever) и Стивену Адлеру (Steven Adler) за отзывы о черновиках статьи. Спасибо Джастину Джею Вангу (Justin Jay Wang), Бьянке Мартин (Bianca Martin) и Стиву Даулингу (Steve Dowling) за коммуникацию и дизайн.