Авторы: Николас Силфоун (GitHub), главный специалист по данным, Центр передового опыта в области искусственного интеллекта @ Fidelity Investments

В организации корпоративного уровня качество кода и вычислительные ресурсы могут быть совершенно разными. Мы оказались в ситуации, когда нам нужно было поддерживать как локальные вычислительные кластеры GPU (например, кластер k8s с несколькими подключенными узлами Nvidia DGX-1), так и облачные решения GPU (например, AWS SageMaker) при построении и обучении моделей глубокого обучения. . Учитывая архитектуру (ы) в то время (особенно зависимость AWS SageMaker от Horovod и сетевое взаимодействие нашего локального кластера), это требовало поддержки парадигм PyTorch DP и DDP, а также Horovod через OpenMPI. Вдобавок ко всему, начали появляться многие ускорители на основе программного обеспечения / оптимизации, построенные на основе PyTorch (например, DeepSpeed примерно в конце 2019 года).

Поддержка экспериментов в этом постоянно меняющемся ландшафте устройств (например, CPU, GPU, TPU и т. Д.), Распределенных методологий (например, PyTorch DDP, Horovod и т. Д.), Смешанной точности (например, Nvidia Apex, Pytorch AMP) и программного обеспечения / оптимизации «ускорители» (например, DeepSpeed ​​ZeRO 0–3) начали становиться довольно громоздкими. Казалось бы, каждый из них требует своего собственного немного другого синтаксиса, логики или контекста. Например, экспериментирование с двумя распространенными методами смешанной точности, APEX от Nvidia и собственным AMP PyTorch (который был выпущен с PyTorch 1.6.0 в июле 2020 года), потребовал значительного количества встроенной логической логики, такой как следующий пример (Примечание: self.config.amp может принимать amp или apex, который переключается между двумя методами):

Здесь, при создании экземпляра класса, который будет обрабатывать все обучающие операции, нам необходимо создать torch.cuda.amp.GradScaler, если мы хотим использовать PyTorch AMP. Однако, если мы хотели использовать APEX от Nvidia, нам потребовались как уже созданная модель типа torch.nn.Module, так и оптимизатор типа torch.optim.Optimizer, и нужно игнорировать torch.cuda.amp.GradScaler. Вдобавок, если мы хотим работать с полной точностью, нам нужно убедиться, что torch.cuda.amp.GradScaler по-прежнему создается, но enabled=False.

Затем при вычислении градиентов (с отсечением градиента для большей наглядности) для данного пакета код выглядит примерно так:

Приведенный выше код охватывает фундаментальную единицу обучения модели глубокого обучения с помощью PyTorch. Получение мини-пакета, вычисление градиентов, а затем выполнение шага с оптимизатором на основе этих градиентов. Надеюсь, ясно, насколько запутанным это становится, когда приходится иметь дело с логической логикой, необходимой для переключения между двумя методами смешанной точности. AMP нужен диспетчер контекста torch.cuda.amp.autocast, который обертывает вызов forward(), он требует прямого управления объектом torch.cuda.amp.GradScaler (особенно для отсечения градиента) и изменяет подпись для вызова step(). APEX отказывается от прямого управления масштабатором и диспетчером контекста, но требует вызовов собственных amp методов для получения параметров оптимизатора.

Кошмар комбинаторики

Глядя на этот простой пример, можно увидеть, что на самом деле между двумя библиотеками не так много точек соприкосновения, даже несмотря на то, что обе они используют одни и те же функции «ускорителя» со смешанной точностью. Теперь подключитесь ко всем другим библиотекам / фреймворкам «ускорителей» на разных устройствах (например, CPU, GPU, TPU и т. Д.), С которыми можно было бы поэкспериментировать: с распределенными методологиями (например, PyTorch DDP, Horovod, DeepSpeed ​​через DDP и т. Д.), смешанная точность (например, Nvidia Apex, Pytorch AMP, пользовательская реализация APEX DeepSpeed) и другие «расширения-ускорители» (например, DeepSpeed ​​ZeRO 0–3, Fairscale OSS, SDDP). Становится довольно очевидным, что комбинаторическая проблема возможности сконфигурировать каждый из этих «ускорителей» вместе чрезвычайно утомительна и трудоемка.

Войдите в PyTorch Lightning. Первоначально выпущенный в 2019 году, Lightning разделяет код PyTorch, чтобы отделить науку от инженерии. Это самоуверенная библиотека, которая помогает удалить общий шаблонный код, обычно необходимый для обучения модели глубокого обучения в PyTorch. Последовало широкомасштабное внедрение, и он стал бесценным инструментом в экосистеме PyTorch. Что наиболее важно, он предоставляет дополнительный API, называемый ускорителями, который помогает управлять переключением между устройствами (CPU, GPU, TPU), смешанной точностью (PyTorch AMP и Nvidia's APEX) и распределенными серверными модулями (PyTorch DDP, Horovod), если вы используете экосистема Lightning.

Однако большая часть наших кодовых баз была написана до выпуска Lightning или, даже после того, как это произошло, просто не использовала Lightning и была основана на чистом PyTorch. Поэтому мы начали с простого вопроса:

Можем ли мы взять лучшее из Lightning Accelerator и отключить его от остальной части PyTorch Lightning?

Путешествие по развитию Сток

Мы начали с использования в качестве вдохновения основ Lightning Accelerator (в то время он поддерживал «ускорители» PyTorch DDP, Horovod, AMP и APEX)… Позвольте пользователю смешивать и сопоставлять «ускорительные» фреймворки / библиотеки, просто изменяя несколько декларативных флагов. На основе декларативных флагов Stoke строит представление внутреннего состояния доступных «ускорителей» и использует классы стиля микширования для динамического создания объекта, который правильно поддерживает базовое объявленное состояние «ускорителей». Он также автоматически обрабатывает накопление градиента, отсечение градиента, ввод-вывод модели и размещение устройства.

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

Например, эти классы позволяют легко настроить базовый сервер PyTorch AMP или Nvidia APEX:

В конце 2020 - начале 2021 года мы достигли внутренней альфа-версии Stoke, которая по сути воспроизводила функциональность Lightning Accelerator. Примерно в то же время Fairscale начала набирать обороты благодаря поддержке Optimizer State Sharding (реализация, аналогичная тому, что DeepSpeed ​​установила со своими реализациями ZeRO), в то время как PyTorch Lightning добавил поддержку DeepSpeed ​​в качестве еще одного ускорителя. Мы поняли, что Stoke может иметь еще одну функциональную роль в качестве библиотеки с открытым исходным кодом, не только предоставляя функциональность, аналогичную Lightning Accelerator, которую легче настраивать, но также помогая справиться с эмпирической природой текущего состояния обучения моделей глубокого обучения. Давайте быстро объясним ...

Большинство методов «ускорителей» являются по сути «сокращениями» того, что можно описать как проблему «необходимы бесконечные вычисления» (спасибо Заку Семенову за эту идею). Фактически, наиболее распространенный метод, который мы используем для обучения моделей глубокого обучения, - это «ярлык» для решения именно этой проблемы. Мини-пакетный градиентный спуск разбивает очень большой набор обучающих данных на гораздо более мелкие пакеты и вычисляет градиенты по этому меньшему набору выборок вместо вычисления градиентов для всего обучающего набора данных. Почему мы это делаем…? Потому что весь набор обучающих данных не поместится в памяти на большинстве машин / устройств и потребует каких-то странных требований «бесконечных вычислений». Поэтому мы используем аппроксимацию градиентов методом выборки без каких-либо гарантированных границ. Смешанная точность - это просто «ярлык», который снижает точность с плавающей запятой и, следовательно, требования к памяти для числовых представлений, исходя из предположения, что «потерянная» точность не окажет значительного влияния на оптимизатор, проходящий через ландшафт оптимизации. Оптимизатор State Sharding - это просто «ярлык» связи для управления памятью за счет эффективного разделения градиентов по нескольким устройствам вместо сохранения избыточных копий одних и тех же данных на каждом устройстве.

Воздействие большинства этих «сокращенных» методов на обучение модели сводится в основном к одному простому месту, эффективному размеру пакета (эффективный размер пакета = per_device_batch_size * num_devices). Использование распределенного бэкэнда (например, DDP или Horovod) для горизонтального масштабирования обучения с несколькими устройствами GPU изменяет «эффективный размер пакета», поскольку общее количество устройств обычно увеличивается. Уменьшение объема памяти оптимизатора с помощью сегментирования состояния оптимизатора обычно позволяет практикующему увеличить размер пакета для каждого устройства, тем самым также увеличивая эффективный размер пакета. Каждый из «ярлыков» по-своему масштабируется в зависимости от эффективного размера партии. Однако любой, кто пытался обучить крупномасштабные модели глубокого обучения, знает, что увеличение эффективного размера пакета (скажем, с 32 на одном графическом процессоре до 96 на двух графических процессорах с использованием сегментирования состояния оптимизатора) означает изменение скорости обучения, расписаний скорости обучения и т. Д. ... и, скорее всего, изменяет сходимость модели и, следовательно, производительность. Итак, длинная касательная короткая:

довольно сложно заранее знать, какие «ярлыки» приведут к желаемому результату ... простое и быстрое экспериментирование имеет решающее значение для возможности использовать правые ярлыки "

Таким образом, мы осознали фундаментальное ценностное предложение Stoke (и в некоторой степени Lightning Accelerator) как библиотеки - предоставить легко доступную и настраиваемую «площадку для ускорителей» для PyTorch, где любой существующий или вновь появляющийся «ускоритель» может быстро и легко проверить на полезность. Поэтому мантрой для Стоука стала:

Никакого спагетти-кода, никакой перезаписи реализации модели, никакой сложной встроенной логической логики. Просто простой декларативный флаг, чтобы спросить «будет ли этот ускоритель работать с моей моделью»?

Что предлагает Сток

Stoke - это легкая оболочка для PyTorch, которая предоставляет простой декларативный API для переключения контекста между устройствами (например, CPU, GPU), распределенными режимами, смешанной точностью и другими расширениями «ускорителя» PyTorch. Он не накладывает ограничений на структуру / стиль кода для архитектуры модели, циклов обучения / вывода, функций потерь, алгоритма оптимизатора и т. Д. Он просто «обертывает» существующий код PyTorch для автоматической обработки необходимой базовой проводки для всех поддерживаемых «ускорителей». . Это позволяет переключаться с локального ЦП полной точности на распределенный мульти-ГП со смешанной точностью с сегментированием состояния оптимизатора путем простого изменения нескольких декларативных флагов. Вкратце, основные преимущества:

  • API декларативного стиля: объявите желаемое состояние (состояния) ускорителя, а остальное пусть stoke обработает.
  • Обернутый API отражает базовые вызовы стиля PyTorch model, loss, backward и step
  • Автоматическое размещение моделей и данных на устройстве
  • Универсальный интерфейс для сохранения и загрузки вне зависимости от серверной части или устройств
  • Автоматическая обработка накопления и отсечения градиента
  • Общий attrs интерфейс для всех параметров конфигурации серверной части (с полезными строками документации!)
  • Несколько дополнительных - Пользовательский torch.utils.data.distributed.Sampler: BucketedDistributedSampler, который объединяет данные в сегменты по отсортированному индексу, а затем произвольно производит выборку из определенных сегментов, чтобы предотвратить такие ситуации, как грубое несоответствие длины последовательности, ведущее к ненужным вычислительным накладным расходам ( т.е. лишнее заполнение). Вспомогательные методы для печати синхронизированных потерь, печати конкретного устройства, количества параметров модели и т. Д.

Помните код для переключения между AMP и APEX в разделе Мотивация выше? Давайте проведем простое сравнение того, как похожий код будет выглядеть в простом скрипте…

Без сток:

Со Стоуком:

Stoke поддерживает следующие «ускорители»:

  • Устройства: CPU, GPU
  • Распространяется: DDP, Horovod, deepspeed (через DDP).
  • Смешанная точность: AMP, Nvidia Apex, deepspeed (кастомный APEX как бэкэнд)
  • Расширения: fairscale (Optimizer State Sharding и Sharded DDP), deepspeed (ZeRO Stage 0–3 и т. Д.)

Определенные комбинации серверных программ / функций несовместимы друг с другом. В таблице ниже указано, какие комбинации были протестированы вместе:

Построение модели CIFAR10 с помощью Stoke

(Примечание: полный рабочий пример можно найти здесь)

Набор данных CIFAR-10 состоит из 60000 цветных изображений 32x32 в 10 классах, по 6K изображений на класс. Имеется 50К обучающих изображений и 10К тестовых изображений. Для простоты мы заимствуем код модели из Torchvision и используем встроенную модель ResNet-152 (arxiv). Давайте начнем с того, что запустим его и запустим без всяких наворотов, только на центральном процессоре (возможно, имитируя вашу локальную среду разработки). Для простоты мы жестко запрограммируем большинство параметров модели, оптимизатора и т. Д. (См. Spock, чтобы найти полезную библиотеку конфигурации параметров). Давайте создадим модель, оптимизатор и функцию потерь. Stoke требует немного другого способа определения оптимизатора (поскольку он обрабатывает создание экземпляров внутри) с помощью StokeOptimizer. Передайте неустановленный объект класса torch.optim.* и любые **kwargs, которые необходимо передать вызову __init__:

Теперь мы создаем объект Stoke, который будет основным интерфейсом для базовых функций API в стиле PyTorch: model, loss, backward и step. Здесь, поскольку мы просто используем ЦП без дополнительных «ускорителей», создание экземпляра довольно просто. Мы передаем модель, оптимизатор, функцию потерь и размер партии. Мы также устанавливаем степень детализации, чтобы получать полезные отладочные данные.

Затем создайте конвейеры для наборов данных изображений, а затем извлеките их, используя torchvision.datasets:

Наконец, нам нужно создать объект torch.utils.data.DataLoader. Подобно определению оптимизатора, это должно быть сделано в Stoke немного иначе, чтобы он правильно обрабатывал каждый из различных бэкэндов. Главный объект Stoke предоставляет зеркальную оболочку для собственного класса torch.utils.data.DataLoader (как метод DataLoader), который вернет правильно сконфигурированный объект torch.utils.data.DataLoader.

На данный момент мы успешно настроили Stoke для работы на одном процессоре. Следующий простой цикл обучения должен выглядеть довольно стандартным, за исключением того, что все вызовы модели forward, loss, backward и step вызываются для объекта Stoke, а не для каждого отдельного компонента (поскольку он внутренне поддерживает модель, потерю, оптимизатор и все необходимое код для всех бэкэндов / функций / расширений). Мы также используем некоторые встроенные функции печати Stoke:

Обязательно сохраним и модель. Удобно, что stoke предоставляет единый интерфейс для сохранения и загрузки контрольных точек модели независимо от серверной части / функциональности / расширений:

Добавление «ускорителей» в Stoke

Надеюсь, именно здесь простота настройки «ускорителей» со Споком сияет. В этом примере предположим, что у нас доступно 4 графических процессора. Давайте использовать собственный PyTorch AMP, сегментированный DDP FairScale и сегментирование состояния оптимизатора FairScale. Мы также собираемся настроить эти методы по своему вкусу. Кроме того, давайте также добавим накопление градиента и обрезку градиента. Для этого создание пользовательских конфигураций и создание экземпляра объекта Stoke становится следующим:

Когда мы делаем DataLoader, нам теперь нужно передать DistributedSampler, поскольку мы используем распределенный метод. Поскольку объект Stoke управляет бэкэндом (ами), мы можем фактически передавать world_size и rank непосредственно из объекта независимо от распределенного бэкэнда (т.е. Horovod и PyTorch DDP требуют разной семантики без Stoke):

Поскольку Stoke выполняет упаковку / сборку ваших torch.nn.Module и torch.utils.data.DataLoader, размещение устройств выполняется автоматически (в этом примере модель и данные перемещаются на правильные графические процессоры).

Цикл обучения и код сохранения модели остаются такими же, как и в случае с одним процессором. Обратите внимание, что не изменилось ... что-либо, связанное с кодом вашей модели. И только одно изменение в конвейере данных (которое можно удалить с помощью простого оператора if). Модель перешла от работы на одном процессоре к 4 графическим процессорам в настраиваемом сегментированном режиме DDP с использованием PyTorch AMP и сегментирования состояния оптимизатора с несколькими дополнительными декларативными конфигурациями / операторами.

А теперь давайте представим другой сценарий. Теперь у вас есть доступ к 8 графическим процессорам, но они разделены на 4 разных кластера. У вас есть доступ к Horovod через OpenMPI, и вы хотите протестировать реализацию APEX O1 от Nvidia вместо PyTorch AMP, но отключите все расширения Fairscale. Это так же просто, как изменить ваши декларативные параметры на следующие:

Вуаля! Совершенно другой набор «ускорителей» за счет простых изменений декларативных конфигураций / операторов!

Закрытие и взгляд в будущее

Потратив почти год на создание и тестирование Stoke внутри компании Fidelity, мы рады выпустить его как инструмент с открытым исходным кодом, который сможет использовать более широкое сообщество машинного обучения. Разработка будет продолжена в области открытого исходного кода. Мы представляем Stoke «площадкой для ускорителей» с его унифицированным интерфейсом, простым API конфигурации и постоянной поддержкой любых передовых функций «ускорителя» на основе PyTorch, выпущенных в домене с открытым исходным кодом (например, недавняя поддержка Full Model Sharding включены в Fairscale). Мы также нацелены на поддержание функционального паритета API Lightning Accelerators для тех, кто предпочитает просто использовать базовый PyTorch.

Кроме того, мы надеемся, что будет происходить постоянная передача функций «ускорителя», которые Stoke в настоящее время и в конечном итоге поддерживает, в версии PyTorch для GA. Например, PyTorch 1.8.0 представил бета-версию функции сегментирования состояния оптимизатора с добавлением torch.distributed.optim.ZeroRedundancyOptimizer (теперь стабильно с 1.10!), Которая перекрывается с некоторыми функциями, поддерживаемыми в Stoke (через Fairscale и DeepSpeed). Stoke всегда будет стараться быть более «передовыми» ускорителями и предоставлять больше «экспериментальных» функций, чем GA PyTorch.

Краткосрочное видение на оставшуюся часть 2021 года и начало 2022 года:

  • Bug squash, чтобы получить стабильную версию V1.0
  • Модульные тесты аппаратного обеспечения CPU / GPU для всего пространства комбинаторики «ускорителей»
  • Поддержка устройств TPU
  • ̶F̶a̶i̶r̶s̶c̶a̶l̶e̶ ̶F̶S̶D̶P̶ ̶s̶u̶p̶p̶o̶r̶t̶ (в настоящее время на мастере!)

Спасибо за прочтение!

Если вам интересно, мы приветствуем вклад сообщества! Любые взносы в Stoke должны поступать через отправленный пул реквест. Если вы столкнетесь с какими-либо ошибками / проблемами, пожалуйста, откройте проблему!

Благодарности

  1. Спасибо Джеффу Брауну, Ричу ДиБиасио, Заку Семенову и Амиту Шавиту за обсуждения, вычитку и обеспечение того, чтобы вся эта статья не была бессмыслицей!
  2. Несколько других библиотек с открытым исходным кодом от Fidelity, которые помогли создать Stoke: Spock - фреймворк, который помогает управлять сложными конфигурациями параметров, textwiser - унифицированный фреймворк для придания свойств тексту. Ознакомьтесь со всем доступным программным обеспечением с открытым исходным кодом Fidelity GitHub.
  3. Примечание. Было бы упущением не упомянуть, что HuggingFace Accelerate существует и пытается заполнить аналогичный пробел. Если вы предпочитаете оставаться в рамках вселенной HF, Accelerate должен предоставить возможности, аналогичные Stoke, однако ему не хватает поддержки Nvidia Apex, Horovod и Fairscale, а также DeepSpeed, все еще являющегося экспериментальной функцией. Вдобавок он имеет более ограниченный интерфейс конфигурации по сравнению со Stoke.