Параллельная библиотека моделей Amazon SageMaker теперь ускоряет рабочие нагрузки PyTorch FSDP до 20 % | Веб-сервисы Amazon

Обучение на больших языковых моделях (LLM) приобрело популярность за последний год с выпуском нескольких популярных моделей, таких как Llama 2, Falcon и Mistral. В настоящее время клиенты предварительно обучают и настраивают LLM с диапазоном от 1 до более чем 175 миллиардов параметров, чтобы оптимизировать производительность моделей для приложений в различных отраслях: от здравоохранения до финансов и маркетинга.

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

В этой статье мы расскажем о новых возможностях Создатель мудреца Амазонки Библиотека параллельных моделей (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.

## training_script.py
import torch.sagemaker as tsm
tsm.init() # Set up a PyTorch model
model = ... # Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
) optimizer = ...
...

Благодаря этим обновлениям 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 следующая:

{ "tensor_parallel_degree": 8, "tensor_parallel_seed": 0
}

В сценарии обучения используйте следующий код:

import torch.sagemaker as tsm
tsm.init() from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Дополнительные сведения об использовании тензорного параллелизма в SMP см. тензорный параллелизм раздел нашей документации.

Используйте расширенные функции для ускорения обучения модели до 20 %.

Помимо возможности распределенного обучения в кластерах с сотнями экземпляров, SMP также предлагает методы оптимизации, которые могут ускорить обучение модели до 20%. В этом разделе мы выделим некоторые из этих оптимизаций. Чтобы узнать больше, обратитесь к Основные особенности раздел нашей документации.

Гибридное шардинг

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

По умолчанию PyTorch FSDP сегментирует артефакты модели на всех устройствах-ускорителях в вашем кластере. В зависимости от вашей учебной работы этот метод сегментирования может увеличить накладные расходы на связь и создать узкое место. Чтобы помочь в этом, библиотека SMP предлагает настраиваемый гибридный параллелизм сегментированных данных поверх PyTorch FSDP. Эта функция позволяет вам установить степень шардинга, оптимальную для вашей тренировочной нагрузки. Просто укажите степень сегментирования в объекте конфигурации JSON и включите его в сценарий обучения SMP.

Конфигурация SMP следующая:

{ "hybrid_shard_degree": 16 }

Дополнительные сведения о преимуществах гибридного параллелизма сегментированных данных см. Почти линейное масштабирование обучения гигантской модели на AWS. Дополнительную информацию о реализации гибридного сегментирования с помощью существующего сценария обучения FSDP см. гибридный параллелизм общих данных в нашей документации.

Используйте операции коллективной связи SMDDP, оптимизированные для инфраструктуры AWS.

Вы можете использовать библиотеку SMP с Библиотека SageMaker для параллелизма распределенных данных (SMDDP) для ускорения рабочих нагрузок распределенного обучения. SMDDP включает оптимизированный AllGather операция коллективной связи, разработанная для обеспечения максимальной производительности на ускоренных экземплярах SageMaker p4d и p4de. В распределенном обучении операции коллективной связи используются для синхронизации информации между работниками графического процессора. AllGather — это одна из основных операций коллективной связи, обычно используемая при параллелизме сегментированных данных для материализации параметров уровня перед шагами прямых и обратных вычислений. Для задач по обучению, которые ограничены средствами связи, более быстрые коллективные операции могут сократить время и затраты на обучение без каких-либо побочных эффектов для конвергенции.

Чтобы использовать библиотеку SMDDP, вам нужно добавить в сценарий обучения всего две строки кода:

import torch.distributed as dist # Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl" # Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

Помимо SMP, SMDDP поддерживает PyTorch FSDP с открытым исходным кодом и DeepSpeed. Дополнительные сведения о библиотеке SMDDP см. Запускайте распределенное обучение с помощью библиотеки параллелизма распределенных данных SageMaker..

Разгрузка активации

Обычно прямой проход обучения модели вычисляет активации на каждом уровне и сохраняет их в памяти графического процессора до тех пор, пока не завершится обратный проход для соответствующего слоя. Эти сохраненные активации могут потреблять значительный объем памяти графического процессора во время обучения. Разгрузка активации — это метод, который вместо этого перемещает эти тензоры в память ЦП после прямого прохода, а затем возвращает их обратно в ГП, когда они необходимы. Этот подход может существенно сократить использование памяти графического процессора во время обучения.

Хотя PyTorch поддерживает разгрузку активации, ее реализация неэффективна и может привести к простою графических процессоров, в то время как активации извлекаются обратно из ЦП во время обратного прохода. Это может привести к значительному снижению производительности при использовании разгрузки активации.

SMP v2 предлагает оптимизированный алгоритм разгрузки активации, который может повысить эффективность обучения. Реализация SMP предварительно выбирает активации до того, как они потребуются на графическом процессоре, сокращая время простоя.

Поскольку SMP построен на основе API-интерфейсов PyTorch, для включения оптимизированной разгрузки активации требуется всего лишь несколько строк кода. Просто добавьте связанные конфигурации (sm_activation_offloading и activation_loading_horizon параметры) и включите их в свой сценарий обучения.

Конфигурация SMP следующая:

{ "activation_loading_horizon": 2, "sm_activation_offloading": True
}

В скрипте обучения используйте следующий код:

import torch.sagemaker as tsm
tsm.init() # Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
) model = FSDP(...) # Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
) model = offload_wrapper(model)

Чтобы узнать больше об инструментах контрольных точек PyTorch с открытым исходным кодом для разгрузки активации, см. checkpoint_wrapper.py скрипт в репозитории PyTorch GitHub и Контрольная точка активации в сообщении блога PyTorch Масштабирование мультимодальных моделей фундамента в TorchMultimodal с Pytorch Distributed. Чтобы узнать больше об оптимизированной реализации разгрузки активации в SMP, см. разгрузка активации раздел нашей документации.

Помимо гибридного сегментирования, SMDDP и разгрузки активации, SMP предлагает дополнительные оптимизации, которые могут ускорить рабочую нагрузку по обучению больших моделей. Сюда входит оптимизированная контрольная точка активации, отложенная инициализация параметров и другие. Чтобы узнать больше, обратитесь к Основные особенности раздел нашей документации.

Заключение

Поскольку наборы данных, размеры моделей и обучающие кластеры продолжают расти, эффективное распределенное обучение становится все более важным для своевременной и доступной доставки моделей и продуктов. Последняя версия параллельной библиотеки моделей SageMaker поможет вам добиться этого за счет сокращения изменений кода и согласования с API-интерфейсами PyTorch FSDP, обеспечивая обучение на огромных кластерах с помощью тензорного параллелизма и оптимизации, которые могут сократить время обучения до 20%.

Чтобы начать работу с SMP v2, обратитесь к нашему документации и наш образцы записных книжек.

Об авторах

Роберт Ван Дусен — старший менеджер по продукту в Amazon SageMaker. Он руководит фреймворками, компиляторами и методами оптимизации для обучения глубокому обучению.

Луис Кинтела является менеджером по разработке программного обеспечения для параллельной библиотеки моделей AWS SageMaker. В свободное время его можно встретить за рулем своего Harley в Сан-Франциско.

Гаутам Кумар — инженер-программист в AWS AI Deep Learning. Он увлечен созданием инструментов и систем для искусственного интеллекта. В свободное время он любит кататься на велосипеде и читать книги.

Рахул Хуилгол является старшим инженером-разработчиком программного обеспечения в области распределенного глубокого обучения в Amazon Web Services.

• SEO-контент и PR-распределение. Получите усиление сегодня.
• PlatoData.Network Вертикальный генеративный ИИ. Расширьте возможности себя. Доступ здесь.
• ПлатонАйСтрим. Интеллект Web3. Расширение знаний. Доступ здесь.
• ПлатонЭСГ. Углерод, чистые технологии, Энергия, Окружающая среда, Солнечная, Управление отходами. Доступ здесь.
• ПлатонЗдоровье. Биотехнологии и клинические исследования. Доступ здесь.
• Источник: https://aws.amazon.com/blogs/machine-learning/amazon-sagemaker-model-parallel-library-now-accelerates-pytorch-fsdp-workloads-by-up-to-20/