Обучение на больших языковых моделях (LLM) приобрело популярность за последний год с выпуском нескольких популярных моделей, таких как Llama 2, Falcon и Mistral. В настоящее время клиенты предварительно обучают и настраивают LLM с диапазоном от 1 до более чем 175 миллиардов параметров, чтобы оптимизировать производительность моделей для приложений в различных отраслях: от здравоохранения до финансов и маркетинга.
Обучение производительных моделей такого масштаба может оказаться непростой задачей. Высокоточные LLM могут потребовать терабайты обучающих данных и тысячи или даже миллионы часов вычислительного времени ускорителя для достижения целевой точности. Чтобы своевременно завершить обучение и запустить продукты, клиенты полагаются на методы параллелизма, позволяющие распределить огромную рабочую нагрузку между тысячами ускорительных устройств. Однако эти методы параллелизма могут быть трудными в использовании: различные методы и библиотеки совместимы только с определенными рабочими нагрузками или ограничены определенными архитектурами моделей, производительность обучения может быть очень чувствительной к неясным конфигурациям, а уровень техники быстро развивается. В результате специалистам по машинному обучению приходится тратить недели на подготовку, чтобы масштабировать свои рабочие нагрузки LLM на большие кластеры графических процессоров.
В этом посте мы освещаем новые функции библиотеки параллельных моделей Amazon SageMaker (SMP), которые упрощают процесс обучения больших моделей и помогают быстрее обучать LLM. В частности, мы рассмотрим новый упрощенный пользовательский интерфейс библиотеки SMP, основанный на API-интерфейсах PyTorch Fully Sharded Data Parallel (FSDP) с открытым исходным кодом, расширенную функциональность тензорного параллельного параллелизма, которая позволяет обучать модели с сотнями миллиардов параметров, а также оптимизацию производительности, которая сокращает время обучения модели. и стоимость до 20%.
Чтобы узнать больше о библиотеке параллельного моделирования моделей SageMaker, обратитесь к документации библиотеки параллелизма моделей SageMaker v2. Вы также можете обратиться к нашему примеры блокнотов для начала.
Новые функции, которые упрощают и ускоряют обучение больших моделей.
В этом посте обсуждаются новейшие функции, включенные в версию 2.0 параллельной библиотеки моделей SageMaker. Эти функции повышают удобство использования библиотеки, расширяют функциональные возможности и ускоряют обучение. В следующих разделах мы суммируем новые функции и обсуждаем, как вы можете использовать библиотеку для ускорения обучения больших моделей.
Согласование SMP с PyTorch с открытым исходным кодом
С момента своего запуска в 2020 году SMP обеспечил высокопроизводительное крупномасштабное обучение на вычислительных экземплярах SageMaker. В последней основной версии SMP библиотека упрощает работу пользователя за счет согласования своих API с PyTorch с открытым исходным кодом.
PyTorch предлагает Полностью сегментированный параллелизм данных (FSDP) в качестве основного метода поддержки больших обучающих нагрузок на многих вычислительных устройствах. Как показано в следующем фрагменте кода, обновленные API-интерфейсы SMP для таких методов, как параллелизм сегментированных данных, отражают API-интерфейсы PyTorch. Вы можете просто запустить import torch.sagemaker
и использовать его вместо torch
.
Благодаря этим обновлениям API-интерфейсов SMP вы теперь можете реализовать преимущества производительности SageMaker и библиотеки SMP, не пересматривая существующие сценарии обучения PyTorch FSDP. Эта парадигма также позволяет использовать ту же базу кода при локальном обучении, что и в SageMaker, упрощая взаимодействие с пользователем для клиентов, которые обучаются в нескольких средах.
Дополнительную информацию о том, как включить SMP с помощью существующих сценариев обучения PyTorch FSDP, см. в разделе Начало работы с SMP.
Интеграция тензорного параллелизма для обучения на массивных кластерах
Этот выпуск SMP также расширяет возможности PyTorch FSDP, включая методы тензорного параллелизма. Одна из проблем при использовании только параллелизма сегментированных данных заключается в том, что вы можете столкнуться с проблемами конвергенции при увеличении размера кластера. Это связано с тем, что параметры сегментирования, градиенты и состояние оптимизатора для параллельных рангов данных также увеличивают размер глобального пакета; в больших кластерах этот глобальный размер пакета может выйти за пределы порога, ниже которого модель сходится. Вам необходимо внедрить дополнительный метод параллелизма, который не требует увеличения размера глобального пакета при масштабировании кластера.
Чтобы смягчить эту проблему, в SMP v2.0 реализована возможность объединения параллелизма сегментированных данных с тензорным параллелизмом. Тензорный параллелизм позволяет увеличивать размер кластера, не изменяя глобальный размер пакета и не влияя на сходимость модели. С помощью этой функции вы можете безопасно увеличить производительность обучения, предоставив кластерам 256 или более узлов.
Сегодня тензорный параллелизм с PyTorch FSDP доступен только в SMP v2. SMP v2 позволяет включить эту технику с помощью нескольких изменений кода и разблокировать стабильное обучение даже на больших кластерах. SMP v2 интегрируется с Трансформаторный двигатель за реализацию тензорного параллелизма и делает его совместимым с API-интерфейсами PyTorch FSDP. Вы можете одновременно включить тензорный параллелизм PyTorch FSDP и SMP, не внося никаких изменений в модель PyTorch или конфигурацию PyTorch FSDP. В следующих фрагментах кода показано, как настроить словарь конфигурации SMP в формате JSON и добавить модуль инициализации SMP. torch.sagemaker.init()
который принимает словарь конфигурации в серверной части при запуске задания обучения, в ваш сценарий обучения.
Конфигурация SMP следующая:
В сценарии обучения используйте следующий код:
Чтобы узнать больше об использовании тензорного параллелизма в SMP, обратитесь к разделу тензорного параллелизма нашей документации.
Используйте расширенные функции для ускорения обучения модели до 20 %.
Помимо возможности распределенного обучения в кластерах с сотнями экземпляров, SMP также предлагает методы оптимизации, которые могут ускорить обучение модели до 20%. В этом разделе мы выделим некоторые из этих оптимизаций. Чтобы узнать больше, обратитесь к разделу основных функций нашей документации.
Гибридное шардинг
Параллелизм сегментированных данных — это метод распределенного обучения, позволяющий экономить память, который распределяет состояние модели (параметры модели, градиенты и состояния оптимизатора) между устройствами. Меньший объем памяти позволяет разместить в кластере более крупную модель или увеличить размер пакета. Однако параллелизм сегментированных данных также увеличивает требования к коммуникации в вашей работе по обучению, поскольку артефакты сегментированной модели часто собираются с разных устройств во время обучения. Таким образом, степень сегментирования является важной конфигурацией, которая обеспечивает компромисс между потреблением памяти и накладными расходами на связь.
По умолчанию PyTorch FSDP сегментирует артефакты модели на всех устройствах-ускорителях в вашем кластере. В зависимости от вашей учебной работы этот метод сегментирования может увеличить накладные расходы на связь и создать узкое место. Чтобы помочь в этом, библиотека SMP предлагает настраиваемый гибридный параллелизм сегментированных данных поверх PyTorch FSDP. Эта функция позволяет вам установить степень шардинга, оптимальную для вашей тренировочной нагрузки. Просто укажите степень сегментирования в объекте конфигурации JSON и включите его в сценарий обучения SMP.
Конфигурация SMP следующая:
Дополнительные сведения о преимуществах гибридного параллелизма сегментированных данных см. Почти линейное масштабирование обучения гигантских моделей на AWS. Дополнительные сведения о реализации гибридного сегментирования с помощью существующего сценария обучения FSDP см. в разделе «Гибридный параллелизм общих данных» в нашей документации.
Используйте операции коллективной связи SMDDP, оптимизированные для инфраструктуры AWS.
Вы можете использовать библиотеку SMP с библиотекой параллелизма распределенных данных SageMaker (SMDDP), чтобы ускорить рабочие нагрузки распределенного обучения. SMDDP включает оптимизированный AllGather
операция коллективной связи, разработанная для обеспечения максимальной производительности на ускоренных экземплярах SageMaker p4d и p4de. В распределенном обучении операции коллективной связи используются для синхронизации информации между работниками графического процессора. AllGather
— это одна из основных операций коллективной связи, обычно используемая при параллелизме сегментированных данных для материализации параметров уровня перед шагами прямых и обратных вычислений. Для задач по обучению, которые ограничены средствами связи, более быстрые коллективные операции могут сократить время и стоимость обучения без каких-либо побочных эффектов для конвергенции.
Чтобы использовать библиотеку SMDDP, вам нужно добавить в сценарий обучения всего две строки кода:
Помимо SMP, SMDDP поддерживает PyTorch FSDP с открытым исходным кодом и DeepSpeed. Дополнительные сведения о библиотеке SMDDP см. в разделе Запуск распределенного обучения с помощью библиотеки параллелизма распределенных данных SageMaker.
Выгрузка активации
Обычно прямой проход обучения модели вычисляет активации на каждом уровне и сохраняет их в памяти графического процессора до тех пор, пока не завершится обратный проход для соответствующего слоя. Эти сохраненные активации могут потреблять значительный объем памяти графического процессора во время обучения. Разгрузка активации — это метод, который вместо этого перемещает эти тензоры в память ЦП после прямого прохода, а затем возвращает их обратно в ГП, когда они необходимы. Этот подход может существенно сократить использование памяти графического процессора во время обучения.
Хотя PyTorch поддерживает разгрузку активации, ее реализация неэффективна и может привести к простою графических процессоров, в то время как активации извлекаются обратно из ЦП во время обратного прохода. Это может привести к значительному снижению производительности при использовании разгрузки активации.
SMP v2 предлагает оптимизированный алгоритм разгрузки активации, который может повысить эффективность обучения. Реализация SMP предварительно выбирает активации до того, как они потребуются на графическом процессоре, сокращая время простоя.
Поскольку SMP построен на основе API PyTorch, для включения оптимизированной разгрузки активации требуется всего лишь несколько строк кода. Просто добавьте связанные конфигурации (sm_activation_offloading
и activation_loading_horizon
параметры) и включите их в свой сценарий обучения.
Конфигурация SMP следующая:
В скрипте обучения используйте следующий код:
Чтобы узнать больше об инструментах контрольных точек PyTorch с открытым исходным кодом для разгрузки активации, см. checkpoint_wrapper.py скрипт в репозитории PyTorch GitHub и Контрольная точка активации в сообщении блога PyTorch Масштабирование мультимодальных моделей фундамента в TorchMultimodal с Pytorch Distributed. Дополнительные сведения об оптимизированной реализации разгрузки активации в SMP см. в разделе «Разгрузка активации» нашей документации.
Помимо гибридного сегментирования, SMDDP и разгрузки активации, SMP предлагает дополнительные оптимизации, которые могут ускорить рабочую нагрузку по обучению больших моделей. Сюда входит оптимизированная контрольная точка активации, отложенная инициализация параметров и другие. Чтобы узнать больше, обратитесь к разделу основных функций нашей документации.
Заключение
Поскольку наборы данных, размеры моделей и обучающие кластеры продолжают расти, эффективное распределенное обучение становится все более важным для своевременной и доступной доставки моделей и продуктов. Последняя версия параллельной библиотеки моделей SageMaker поможет вам добиться этого за счет сокращения изменений кода и согласования с API-интерфейсами PyTorch FSDP, обеспечивая обучение на больших кластерах с помощью тензорного параллелизма и оптимизации, которые могут сократить время обучения до 20%.
Чтобы начать работу с SMP v2, обратитесь к нашей документации и нашему образцы тетрадей.
Об авторах
Роберт Ван Дюзен — старший менеджер по продукту в Amazon SageMaker. Он руководит фреймворками, компиляторами и методами оптимизации для обучения глубокому обучению.
Луис Кинтела является менеджером разработчиков программного обеспечения для параллельной библиотеки моделей AWS SageMaker. В свободное время его можно встретить катающимся на своем Харлее в районе залива Сан-Франциско.
Гаутам Кумар — инженер-программист в AWS AI Deep Learning. Он увлечен созданием инструментов и систем для искусственного интеллекта. В свободное время он любит кататься на велосипеде и читать книги.
Рахул Уильгол — старший инженер по разработке программного обеспечения в области распределенного глубокого обучения в Amazon Web Services.