Мы хотели знать, сколько времени (и денег) будет стоить обучение модели Stable Diffusion с нуля с использованием наших наборов потоковых данных, Composer и облачной платформы MosaicML. Наши результаты: нам потребуется 79 000 часов A100 за 13 дней, а общая стоимость обучения составит менее 160 000 долларов. Наши инструменты не только сокращают время и затраты в 2,5 раза, но также являются расширяемыми и простыми в использовании.

Попробуйте наш код Stable Diffusion здесь!

Мир ИИ гудит благодаря мощи крупных генеративных нейронных сетей, таких как ChatGPT, Stable Diffusion и других. Эти модели способны обеспечить впечатляющую производительность при решении широкого круга задач, но из-за их размера и сложности лишь немногие организации могут их обучать. Как следствие, доступ к этим моделям может быть ограничен организацией, которой они принадлежат, и пользователи не могут контролировать данные, которые модель видела во время обучения.

Вот как мы можем помочь: в MosaicML мы упрощаем эффективное обучение больших моделей, позволяя большему количеству организаций обучать свои собственные модели на своих данных. Как показано в предыдущем сообщении в блоге, наша библиотека StreamingDataset, наша обучающая среда Composer и наша платформа MosaicML Cloud значительно упрощают процесс обучения больших языковых моделей (LLM). Для этого сообщения в блоге мы использовали тот же процесс для измерения времени и стоимости обучения Модели стабильной диффузии с нуля. Мы оценили верхний предел в 79 000 A100-часов для обучения базового уровня Stable Diffusion v2 за 13 дней на нашей облачной платформе MosaicML, что соответствует общей стоимости обучения менее 160 000 долларов США. Это сокращение в 2,5 раза. время и стоимость указаны в карточке модели от Stability AI. Помимо экономии времени и денег, наши инструменты Streaming, Composer и MosaicML Cloud упрощают настройку и масштабирование обучения Stable Diffusion на сотнях графических процессоров без каких-либо дополнительных усилий. Код, который мы использовали для этого эксперимента, имеет открытый исходный код и готов к запуску; загляните здесь! И если вы заинтересованы в самостоятельном обучении моделей распространения в облаке MosaicML, свяжитесь с нами для демонстрации.

Оценка времени и стоимости

Таблица 1 и Рисунок 1 ниже иллюстрируют, как оценки базового времени и стоимости обучения Stable Diffusion V2 зависят от количества используемых графических процессоров. Наша окончательная оценка для 256 A100 составляет 12,83 дня для обучения со стоимостью 160 000 долларов США, что в 2,5 раза сокращает время и стоимость, указанные в карточке модели Stable Diffusion. Эти оценки были рассчитаны с использованием измеренной пропускной способности и предполагали обучение на 2,9 миллиардах образцов. Пропускная способность была измерена путем обучения на изображениях и подписях с разрешением 512x512 с максимальной длиной токенов 77. Мы масштабировали графические процессоры с 8 до 128 NVIDIA A100 40 ГБ, а затем экстраполировали пропускную способность до 256 A100 на основе этих измерений.

Настройка эталона

Как мы получили эти результаты? Мы воспользовались преимуществами набора данных MosaicML Streaming, нашей обучающей среды Composer и облака MosaicML для измерения пропускной способности при обучении модели стабильной диффузии. Вы можете воспроизвести наши результаты, используя код в этом репозитории. Читайте дальше для более подробной информации:

Потоковое

Основной проблемой при обучении Stable Diffusion является работа с огромными наборами данных, такими как LAION-5B. Библиотека MosaicML StreamingDataset значительно упрощает управление и использование этих массивных наборов данных. Он работает путем преобразования целевого набора данных в формат потоковой передачи, а затем сохранения преобразованного набора данных в нужном облачном хранилище (например, в корзине AWS S3). Чтобы использовать сохраненный набор данных, мы просто определяем класс StreamingDataset, который извлекает и преобразует образцы из сохраненного набора данных.

Для наших результатов мы транслировали подмножество набора данных LAION-400M с изображениями 256x256 и соответствующими подписями. Размер изображений был изменен до 512x512 для окончательных измерений пропускной способности. Мы оценили пропускную способность загрузчика потоковых данных как минимум в 30 раз выше пропускной способности модели во всех протестированных нами конфигурациях, поэтому вряд ли в ближайшее время мы увидим узкие места, связанные с загрузчиком данных. Ознакомьтесь с нашим скриптом данных, чтобы узнать, как мы определили наш набор данных потоковой передачи, и следите за более подробной информацией в нашей публикации в блоге о потоковой передаче, которая скоро будет выпущена!

Композитор

Наша библиотека Composer с открытым исходным кодом содержит множество современных методов для ускорения обучения нейронных сетей и улучшения обобщения. Для этого проекта мы сначала определили ComposerModel для Stable Diffusion, используя модели из HuggingFace’s Diffusers library и конфиги из stabilityai/stable-diffusion-2-base. Набор данных ComposerModel и Streaming был предоставлен Тренеру Composer с оптимизатором AdamW, алгоритмом EMA 2, обратным вызовом для измерения пропускной способности и регистратором весов и смещений. Наконец, мы вызвали fit() для объекта Trainer, чтобы начать обучение.

Облако MosaicML

MosaicML Cloud организует и контролирует вычислительную инфраструктуру для крупномасштабных учебных заданий. Наш планировщик заданий упрощает запуск и масштабирование заданий.На рис. 2 показана конфигурация обучения MosaicML Cloud (слева) и интерфейс командной строки, используемый для запуска тренировочного прогона (справа). Из конфигурации мы можем легко масштабировать количество графических процессоров с помощью параметра gpu_num. Один и тот же код, который мы пишем для одного узла, может автоматически использовать десятки узлов и сотни графических процессоров.

Что дальше?

В этом сообщении блога мы оценили сокращение времени и затрат на обучение Stable Diffusion в 2,5 раза при использовании наших наборов потоковых данных, библиотеки Composer и облака MosaicML. Это отличный предварительный результат, но мы еще не закончили. В следующем сообщении в блоге мы проверим, как мы можем тренироваться для конвергенции на такой скорости. Чтобы быть в курсе наших последних работ, присоединяйтесь к нашему Community Slack или следите за нами в Twitter. Если ваша организация хочет начать обучение диффузионным моделям сегодня, вы можете запланировать демонстрацию онлайн или написать нам по адресу [email protected].

Примечания:

  1. Размер микропакета для каждого устройства связан с накоплением градиента, но размер микропакета проще использовать при увеличении количества графических процессоров. Размер микропакета на устройство и накопление градиента связаны следующим образом: grad_accum = глобальный_размер_пакета / (n_devices * per_device_microbatch_size). Дополнительные сведения см. в переменной device_train_microbatch_size в Composer Справочник по API Trainer.
  2. В Composer есть встроенный алгоритм EMA, но нам пришлось повысить его эффективность использования памяти для нашего теста, поэтому реализация EMA в нашем репозитории тестов. Вскоре мы обновим алгоритм EMA в Composer, чтобы он стал более эффективным с точки зрения использования памяти.

Первоначально опубликовано на https://www.mosaicml.com.