Какие существуют подходы к обучению на нескольких GPU и в чем их сильные и слабые стороны?
Подходы к обучению на нескольких GPU можно разделить на две группы: разделение данных для параллельной обработки несколькими GPU и разделение модели по нескольким GPU для преодоления ограничений памяти, когда размер модели превышает возможности одной видеокарты. Параллелизм данных попадает в первую категорию, в то время как тензорный параллелизм и параллелизм моделей попадают во вторую. Такие подходы как конвейерный параллелизм, объединяют идеи из обеих категорий. В дополнение, современные программные реализации, такие как DeepSpeed, Colossal AI и другие смешивают различие техники в гибридные подходы.
В этой главе мы рассмотрим несколько подходов к обучению и поделимся рекомендациями по их эффективному применению на практике.
В этой главе мы в основном говорим о графических процессорах (GPU), чтобы описать аппаратное обеспечение, используемые для параллельной обработки. Однако те же концепции и техники, можно применять к другим специализированным устройствам, таким как тензорные процессоры (TPU) или другие ускорители, в зависимости от конкретной архитектуры и требований системы.
Подходы к обучению
В следующих разделах мы обсудим ключевые подходы к обучению на нескольких графических процессорах, такие как параллелизм моделей, параллелизм данных, тензорный параллелизм и параллелизм последовательностей
Параллелизм моделей
Параллелизм моделей, или межоперационный (inter-op) параллелизм, представляет собой метод, при котором разные части большой модели распределяются по разным GPU. Затем эти части последовательно вычисляются, передавая промежуточные результаты между устройствами. Такой подход позволяет обучать и запускать модели, которые не помещаются на одном GPU, но требует сложной координации для управления зависимостями между различными частями модели.
Параллелизм моделей, пожалуй, самый интуитивный способ распараллеливания по разным устройствам. Например, если у нас есть простая нейронная сеть, состоящая всего из двух слоев - скрытого и выходного, мы можем разместить первый слой на одном графическом процессоре, а второй - на другом. Само собой, количество слоев и GPU легко поддаются масштабированию.
Эта стратегия хорошо подходит для работы с ограниченным объемом памяти GPU, когда сеть не помещается на одном устройстве полностью. Однако, поскольку возникает структура, похожая на цепочку - слой 1 на GPU 1 → слой 2 на GPU 2, - создается узкое место. Существуют более эффективные способы использования нескольких GPU, такие как тензорный параллелизм. Другими словами, самый главный недостаток параллелизма моделей заключается в том, что GPU должны ждать друг друга. Они не могут эффективно работать одновременно, так как зависят от выходов друг друга.
Параллелизм данных
Параллелизм данных долгое время был стандартным способом обучения на нескольких видеокартах. В нем мы просто разделяем минибатч на несколько микробатчей. Затем каждый GPU обрабатывает свой микробатч отдельно, вычисляя функцию потерь и ее градиенты для весов модели. После того как все устройства завершат обработку своих микробатчей, градиенты объединяются, чтобы обновить веса для следующего раунда обучения.
Основное преимущество параллелизма данных над параллелизмом моделей состоит в том, что GPU могут работать одновременно. Во время обучения каждый графический процессор обрабатывает свою часть минибатча. Однако, подвох состоит в том, что для каждого GPU требуется полная копия модели. Очевидно, что мы не сможем использовать такой вид параллелизма, если у нас наша модель слишком большая и не помещается в видеопамять GPU.
Тензорный параллелизм
Тензорный параллелизм, или внутриоперационный (intra-op) параллелизм - более эффективная форма параллелизма моделей. Вместо того чтобы перемещать целые слои по различным устройствам, при тензорном параллелизме перемещаются только матрицы весов и активации. Эти отдельные матрицы разбиваются таким образом, что перемножение матриц распределяется между несколькими GPU.
Мы можем реализовать тензорный параллелизм используя базовые принципы линейной алгебры. Для этого мы можем разделить перемножение матриц на два GPU, распределяя их по строкам или столбцам, как показано на Рисунке для двух GPU. Аналогично, концепция масштабируется на любое количество устройств.
Подобно параллелизму моделей, тензорный параллелизм позволяет преодолеть ограничения памяти. При этом он обеспечивает параллельное выполнение операций, как и в случае с параллелизмом данных.
Небольшой недостаток тензорного параллелизма связан с высокими затратами на передачу данных между несколькими GPU, на которые разделяются матрицы. В частности, при использовании тензорного параллелизма параметры модели часто синхронизируются между различными устройствами, что может замедлить процесс обучения.
Рисунок 1.1 сравнивает параллелизм моделей, данных и тензорного параллелизма.
В параллелизме моделей мы распределяем различные слои по нескольким графическим процессорам, чтобы обойти ограничения памяти. В параллелизме данных мы разделяем батчи по разным GPU, чтобы обучить копии моделей одновременно, а затем усредняем их градиенты для последующего обновления параметров. В тензорном параллелизме мы делим матрицы по нескольким GPU для параллельной обработки, когда модели становятся слишком большими, чтобы поместиться в память GPU.
Конвейерный параллелизм
В рамках конвейерного параллелизма, как и в параллелизме моделей, активации передаются в процессе прямого распространения. Однако, в отличие от последнего, пока часть активаций передается вперед, градиенты в это же самое время распространяются назад, что позволяет избежать простоев в работе устройств. Таким образом, конвейерный параллелизм представляет собой своеобразную комбинацию параллелизма моделей и данных, обеспечивая более эффективное использование ресурсов.
Конвейерный параллелизм это своего рода форма параллелизма моделей, который старается преодолеть узкие места, связанные с последовательным выполнением вычислений. Однако, помимо этого, он также заимствует идеи из параллелизма данных, такие как разбиение минибатчей на микробатчи.
Конвейерный параллелизм, определенно, является значительным шагом вперед по сравнению с параллелизмом моделей, но, к сожалению, не лишен недостатков, и простои иногда будут возникать. Один из основных минусов конвейерного параллелизма - значительные затраты на разработку и реализацию шагов пайплайна и соответствующих моделей взаимодействия между ними. Кроме того, прирост производительности может оказаться не столь значимым, как при использовании других методов параллелизации, таких как чистый параллелизм данных, особенно для небольших моделей или в случаях, когда затраты на передачу данных между устройствами слишком высоки.
Для современных архитектур, которые слишком огромны для того, чтобы поместиться в память GPU, широкое распространение получил комбинированный подход, сочетающий параллелизм данных и тензорный параллелизм.
Параллелизм последовательностей
Параллелизм последовательностей направлен на устранение вычислительных трудностей, возникающих при работе с длинными последовательностями в LLM, основанных на трансформерах. Одним из недостатков архитектуры трансформеров является то, что механизм self-attention (исходный механизм на основе скалярного произведения) масштабируется квадратично по мере роста длины входной последовательности. Стоит отметить, что существуют более эффективные альтернативы реализации self-attention, которые масштабируются линейно.
Однако, на момент написания данной статьи. эти эффективные механизмы self-attention не пользуются такой популярностью, как исходный механизм на основе скалярного произведения. Параллелизм последовательностей, как показано на Рисунке 1.2, разделяет входную последовательность на меньшие части (чанки), чтобы распределить их по разным GPU. Это позволяет преодолеть ограничения памяти, с которыми сталкиваются механизмы self-attention во время вычислений.
Как параллелизм последовательностей соотносится с другими подходами, которые мы уже рассмотрели? Параллелизм последовательностей работает с последовательностями, тогда как тензорный параллелизм сосредоточен на внутренней структуре модели, а параллелизм данных регулирует способы разбиения тренировочных данных. Теоретически, каждый из этих видов параллелизма решает разные аспекты вычислительных задач, что позволяет комбинировать их различными способами для оптимизации процедуры обучения и инференса. Однако, в отличие от других подходов, параллелизм последовательностей до сих пор не так хорошо изучен.
В то время, как параллелизм последовательностей может быть полезен на практике, он также имеет свои недостатки, связанные с дополнительными затратами на передачу данных, похожие на те, которые отмечаются и у других видов параллелизма. Как и параллелизм данных, параллелизм последовательностей требует наличия полной копии модели на каждой GPU. Другой недостаток (в зависимости от реализации) обучения трансформеров на нескольких GPU состоит в том, что разбиение входной последовательности на несколько меньших подпоследовательностей может снизить точность модели, особенно, когда она применяется к длинным последовательностям.
Рекомендации
Практические рекомендации зависят от конкретных условий. Если мы обучаем небольшие модели, которые помещаются на одном графическом процессоре, то тогда параллелизм данных может оказаться наиболее эффективным подходом. Улучшения, достигаемые при использовании конвейерного параллелизма, могут оказаться несущественными, особенно для маленьких моделей или в случаях, когда затраты на передачу данных оказываются слишком велики.
Если ваши модели слишком большие, чтобы поместиться в память одной видеокарты, то необходимо подумать над использованием параллелизма моделей или тензорного параллелизма. Тензорный параллелизм, естественным образом, более эффективен, поскольку позволяет нескольким GPU работать одновременно в виду отсутствия последовательных зависимостей как в параллелизме моделей.
Современные подходы к обучению на нескольких GPU обычно сочетают параллелизм данных и тензорный параллелизм.
