Экономия оборудования до 46 раз для AIGC и автоматического параллелизма в новом выпуске Colossal-AI

В последнее время крупномасштабные модели искусственного интеллекта привлекли внимание и получили широкое распространение. AIGC остается популярным, а чат-бот ChatGPT вызвал широкий интерес и привлек миллионы пользователей всего за две недели, даже вызвав «красный код» в Google, потому что он может разрушить бизнес поискового гиганта. Другие мощные приложения ИИ, такие как AlphaCode для разработки программного обеспечения и ESM2 для исследования лекарств, также получили распространение. Эти и другие приложения ИИ постоянно раздвигают границы возможного с большими моделями ИИ, исследуя новые области реализации и применения.

Новый выпуск Colossal-AI версии 0.2.0 включает в себя несколько ключевых функций, позволяющих значительно минимизировать затраты и ускорить вывод на рынок таких крупных моделей, а также повысить удобство использования для разработчиков.

Рецепт модели Stable Diffusion 2.0. Готовая конфигурация для Stable Diffusion 2.0 в Colossal-AI обеспечивает недорогое обучение, точную настройку и вывод, а также снижает потребление памяти графическим процессором до 5,6 раз. оборудование стоит до 46 раз. Всего этого можно достичь с помощью всего одной строки кода.
Рецепт модели BLOOM: готовое включение модели BLOOM со 175 миллиардами параметров обеспечивает автономный вывод с четырехкратным сокращением потребления памяти графическим процессором и снижением затрат на оборудование более чем в 10 раз.
Автоматический параллелизм. С помощью всего одной строки кода пользователи могут автоматически искать наилучшую стратегию параллелизма, упрощая распределенное обучение и поддерживая популярные библиотеки моделей ИИ, такие как Hugging Face и Timm.

Благодаря таким пользователям по всему миру, как AWS, Meta, BioMap, а также более чем 7000 звезд Github и доказанной непревзойденной производительности, Colossal-AI является надежным ресурсом для тех, кто хочет использовать возможности крупномасштабного ИИ.

Возможности оптимизации для Stable Diffusion 2.0

AIGC — очень востребованная тема в области искусственного интеллекта, которая была признана журналом Science «ПРОРЫВОМ ГОДА 2022 ГОДА». Стабильная диффузия также недавно была обновлена до версии 2.0. Некоторые пользователи отметили быстрые темпы развития этой технологии, отметив, что версия 2 выпускается еще до того, как версия 1 была полностью разработана.

Однако высокая стоимость внедрения AIGC в некоторой степени ограничила его широкое распространение. Например, Stability AI, лежащий в основе Stable Diffusion, поддерживает более 4000 кластеров графических процессоров NVIDIA A100 и понес операционные расходы в размере более 50 миллионов долларов США. Поскольку AIGC продолжает быстро развиваться с помощью итерационных моделей, алгоритмов и последующих задач, снижение затрат стало критически важным вопросом для его успешной реализации.

Stable Diffusion 2.0 построен на удобной для пользователя платформе PyTorch Lightning, и Colossal-AI оперативно выпустила всеобъемлющий рецепт, содержащий готовые решения с открытым исходным кодом для обучения, тонкой настройки и логических выводов, которые являются более эффективными и имеют более низкие требования к оборудованию в качестве официального крупномасштабного модельного решения для PyTorch Lightning. Рецепт включает:

Снижение потребления памяти графическим процессором для обучения в 5,6 раз и стоимости оборудования до 46 раз
Поддержка DreamBooth для быстрой персонализированной тонкой настройки на одном графическом процессоре
Снижение потребления памяти GPU для логических выводов в 2,5 раза

В ближайшем будущем это решение также будет интегрировано в Hugging Face, самое популярное сообщество моделей искусственного интеллекта, для дальнейшего улучшения пользовательского опыта.

Обучение

Увеличение размера партии является широко используемым методом ускорения обучения и снижения затрат. Однако ограниченный объем памяти графического процессора ограничивает размер пакета и повышает требования к оборудованию.

Технологии оптимизации памяти Colossal-AI и поддержка Stable Diffusion 2.0 снижают требования к памяти для использования Stable Diffusion с обучением большого размера партии 16 на каждом графическом процессоре с 64,5 ГБ до 11,6 ГБ, то есть в 5,6 раза. . Эти технологии также могут быть расширены на один или несколько графических процессоров параллельно. Это позволяет пользователям удовлетворять свои потребности, используя только графические карты потребительского уровня, такие как 3060, а не самую передовую A100 80 ГБ, что приводит к экономии затрат на оборудование до 46 раз. С Colossal-AI больше пользователей могут по доступной цене проводить исследования, связанные со стабильной диффузией, и внедрять их на графических процессорах потребительского уровня.

Стратегии оптимизации

Flash Attention
Компания Colossal-AI первой внедрила технологию Flash Attention для Stable Diffusion 1.0, которая увеличила скорость внимания на 104 % и уменьшила пиковую сквозную тренировочную память на 23 %. Flash Attention — это эффективная реализация внимания для длинных последовательностей задач, в которой используется сглаживание для сокращения операций чтения/записи памяти между памятью с высокой пропускной способностью (HBM) графического процессора и алгоритмом приблизительного внимания для блокировки разреженного внимания, что делает его более быстрым, чем существующие методы приблизительного внимания. В то время как в Stable Diffusion 1.0 было лишь небольшое количество слоев внимания, потенциал оптимизации памяти с помощью Flash Attention был дополнительно продемонстрирован в Stable Diffusion 2.0 путем замены многих сверточных слоев слоями внимания.

ZeRO + Gemini
Colossal-AI использует Zero Redundancy Optimizer (ZeRO) для устранения избыточности памяти, значительно повышая эффективность использования памяти по сравнению с классическим параллелизмом данных без ущерба для детализации вычислений и эффективности связи. Colossal-AI также включает управление памятью на основе фрагментов, что еще больше повышает производительность ZeRO. Управление памятью на основе фрагментов сохраняет последовательные наборы параметров в рабочем порядке в непрерывном пространстве памяти, называемом фрагментом, все одинакового размера для эффективного использования пропускной способности сети между PCI-e и GPU-GPU, сокращения обмена данными и предотвращения потенциальной фрагментации памяти. .

Диспетчер гетерогенной памяти Colossal-AI, Gemini, уменьшает объем памяти графического процессора, перекладывая состояния оптимизатора на ЦП, позволяя одновременно использовать память графического процессора и память ЦП (включая ЦП DRAM или память NVMe SSD) для увеличения масштаба доступных моделей за пределы памяти. одного графического процессора.

Одна строка кода для быстрого начала работы

Как официальный партнер PyTorch Lightning, оптимизация памяти Colossal-AI может быть легко доступна с помощью всего одной строки кода.

Тонкая настройка DreamBooth

DreamBooth – это метод, который использует всего 3–5 изображений нужного объекта для персонализации моделей преобразования текста в изображение, таких как Stable Diffusion, и создания серии изображений этого объекта. Доступ к оптимизации памяти Colossal-AI можно легко получить с помощью файла train_dreambooth_colossalai.py для быстрой персонализированной точной настройки модели и повышения простоты использования.

Вывод

Поскольку вывод модели нечувствителен к численной точности, это дает возможность реализовать недорогой вывод с низкой точностью. Модель Stable Diffusion 2.0 может быть квантована Int8 для вывода с помощью одной строки кода, что приводит к в 2,5 раза меньшему потреблению памяти (требуется 3,1 ГБ памяти) с минимальной потерей производительности.

model = replace_module(model)

Недорогой вывод для модели 175 миллиардов BLOOM

По мере того, как модели становятся больше, потребление памяти во время логического вывода становится ключевым фактором, который следует учитывать. Например, модель BLOOM со 175 миллиардами параметров с открытым исходным кодом, выпущенная Hugging Face, требует не менее 87,5 ГБ/43,8 ГБ памяти на каждый GPU при использовании обычных FP32/FP16 для логических выводов на 8 GPU в одном узле. Этот большой объем памяти не может поддерживаться даже самым передовым сервером 8*A100 (80 ГБ/40 ГБ), а многоузловой логический вывод требует дополнительных затрат и накладных расходов на связь.

Colossal-AI теперь предоставляет специальный рецепт BLOOM, который обеспечивает эффективный параллельный вывод квантованной модели Int8, что позволяет развертывать службы вывода больших моделей, такие как BLOOM (175 миллиардов параметров), на серверах с 8 графическими процессорами, использующих графические карты потребительского уровня, такие как 3090/4090. без значительного увеличения использования памяти ЦП или потери производительности. Это может снизить затраты на развертывание оборудования более чем в 10 раз по сравнению с использованием решений A100.

Квантуя модель с помощью Int8, общий объем памяти модели можно уменьшить с 352,3 ГБ (FP16) до 185,6 ГБ с помощью Colossal-AI, а его метод параллелизма модели снижает требования к памяти для каждого графического процессора до 23,2 ГБ. Параллелизм модели позволяет избежать увеличения объема памяти ЦП за счет использования lazy_init в основном процессе, тем самым получая метамодель практически без использования памяти. Colossal-AI квантифицирует и нарезает модель в основном процессе и использует lazy_init в каждом из оставшихся процессов для получения метамодели, которая почти не занимает памяти, а затем передает параметры модели между процессами через бэкенд gloo. Это приводит к тому, что пиковое использование памяти ЦП достигает теоретически оптимального уровня без загрузки параметров модели в сегменты, что повышает эффективность использования памяти при неинтенсивных запросах по сравнению с «конвейерным» распределением нарезки модели по слоям.

Автопараллелизм с одной строкой кода

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

Однострочная система автоматического параллелизма Colossal-AI упрощает процесс развертывания крупномасштабных моделей машинного обучения для разработчиков ИИ. По сравнению с другими решениями, которые требуют ручной настройки сложных параллельных политик и модификации модели, Colossal-AI требует от пользователя только одну строку кода, а также информацию о кластере и конфигурации модели, чтобы обеспечить распределенное обучение. Он легко интегрируется с популярными платформами моделирования ИИ, такими как Hugging Face и Timm.

# wrap the model using auto_engine
  model, optimizer = auto_engine(model, optimizer, cluster_info)
#  normal training loop
  ...

Новый выпуск Colossal-AI значительно снижает требования к разработчикам ИИ по использованию методов распределенного обучения для предварительной подготовки и тонкой настройки больших моделей. В то же время система автоматического параллелизма может искать оптимальные параллельные стратегии с большей степенью детализации.

Трассировка графика

Colossal-AI – это первая система автоматического параллелизма, использующая анализ статического графа на основе платформы PyTorch. Получение статического плана выполнения для PyTorch, среды динамического графа, уже давно является областью исследований в области систем машинного обучения. Система параллельных вычислений Colossal-AI использует ColoTracer, разветвленную версию torch.FX Tracer, для поиска оптимальной стратегии распараллеливания. Метаинформация каждого тензора, такая как форма тензора, размеры, dtype и т. д., вычисляется и записывается в процессе трассировки. Преимущество этого подхода заключается в лучшем обобщении, поскольку он не привязан к конкретным моделям или конфигурациям.

Мелкозернистый поиск параллелизма

Автоматический параллелизм Colossal-AI ищет стратегии в отношении каждого операнда с целью достижения самого быстрого времени выполнения при соблюдении ограничений бюджета памяти. В конечном итоге он определяет фактическую стратегию времени обучения, включая стратегию тензорного разделения для каждого тензора, тип операторов связи, которые должны быть вставлены между различными вычислительными узлами, следует ли заменять операторы и т. д. Тензор, данные и гибридный параллелизм, такой как столбец и разделение строк, используемое NVIDIA в Megatron-LM и других системах параллелизма, — все это подмножества стратегий, которые может искать Colossal-AI. В дополнение к этим параллелизмам, которые можно указать вручную, Colossal-AI может указать уникальный метод параллелизма для каждой операции и потенциально найти лучшую стратегию параллелизма, чем то, что могут предложить специалисты-люди.

Распределенная система тензора и согласованности формы

Система Colossal-AI использует сетку устройств, аналогичную последней версии PyTorch DTensor, для управления своим кластером. Colossal-AI использует спецификацию сегментирования, чтобы аннотировать состояние хранения каждого тензора и облегчить их распределение по кластеру. В системе также используется диспетчер согласованности формы для автоматического преобразования тензоров между различными спецификациями сегментирования, что позволяет плавно нарезать и нарезать тензоры, в то время как диспетчер согласованности формы гарантирует, что выходные данные восходящих операндов последовательно сохраняются в кластере, независимо от как хранится ввод нижестоящих операндов. Это делает Colossal-AI очень универсальным и простым в использовании, и пользователи не беспокоятся о состоянии хранения тензоров при выполнении операций с ними.

Вот некоторые ключевые преимущества Colossal-AI по сравнению с PyTorch DTensor:

Сетка устройств Colossal-AI использует показатели производительности кластера и результаты профилирования для оценки времени, затрачиваемого различными операторами связи. Это помогает Colossal-AI оптимизировать связь между узлами и повысить общую эффективность системы.
Менеджер согласованности формы Colossal-AI использует жадный алгоритм поиска, чтобы найти относительно эффективные способы преобразования тензоров между различными спецификациями сегментирования, а не просто преобразовывать измерения одно за другим. Это может привести к более эффективным и действенным преобразованиям.
Интеграция операций «все ко всем» в Colossal-AI повышает масштабируемость системы, обеспечивая более эффективную связь между узлами. Это особенно полезно для крупномасштабных задач машинного обучения, требующих передачи больших объемов данных между узлами.

Интеграция с контрольной точкой активации

Colossal-AI использует контрольные точки активации для сокращения использования памяти во время обучения больших моделей машинного обучения. Вместо сохранения всей модели сохраняются активации промежуточных слоев. Функция автоматического поиска Colossal-AI для контрольной точки активации находит наиболее эффективную контрольную точку в рамках заданного бюджета памяти, а не просто стремится к максимальному сжатию памяти. Чтобы избежать длительного процесса поиска оптимальной контрольной точки активации, Colossal-AI реализовал двухэтапный процесс поиска. Это позволяет системе найти подходящее решение для распределенного обучения за разумное время, в то же время используя контрольные точки активации для управления памятью. Интеграция контрольных точек активации в Colossal-AI повышает эффективность и результативность обучения больших моделей.

О Колоссальном-ИИ

Colossal-AI — это унифицированная система глубокого обучения для эпохи больших моделей, которая поддерживает эффективное и быстрое развертывание обучения и вывода больших моделей ИИ, а также снижает стоимость приложений для больших моделей ИИ. С тех пор, как Colossal-AI стал открытым исходным кодом, он несколько раз занимал первое место в рейтинге GitHub Trending с более чем 7000 звезд. Он был успешно принят в качестве официального руководства для ведущих международных конференций по искусственному интеллекту и высокопроизводительным вычислениям, таких как SC, AAAI и PPoPP.

Соответствующие решения были успешно применены известными технологическими гигантами для автономного вождения, облачных вычислений, розничной торговли, медицины и чипов и получили широкое признание. Например, недавний горячий ChatGPT еще не имеет открытого исходного кода и не имеет доступа в Интернет. Colossal-AI успешно помог предприятию из списка Fortune 500 разработать модель чат-бота с расширенными возможностями онлайн-поиска.

Найдите исходный код и дополнительную информацию о Colossal-AI на GitHub по адресу: https://github.com/hpcaitech/ColossalAI.