ColossalAI: эффективные крупномасштабные решения для обучения моделей ИИ

Общее введение

ColossalAI - это платформа с открытым исходным кодом, разработанная HPC-AI Technologies для создания эффективного и экономичного решения для крупномасштабного обучения и вывода моделей ИИ. Благодаря поддержке множества параллельных стратегий, гетерогенному управлению памятью и обучению со смешанной точностью, ColossalAI позволяет значительно сократить время и затраты ресурсов на обучение и вывод моделей. Будь то параллелизм данных, тензорный параллелизм или конвейерный параллелизм, ColossalAI предоставляет мощные инструменты и библиотеки, чтобы помочь исследователям и разработчикам добиться эффективного обучения и вывода крупномасштабных моделей на кластерах с несколькими GPU.

Список функций

• Параллелизм данных, тензорный параллелизм, конвейерный параллелизм и поддержка других параллельных стратегий
• Смешанное прецизионное обучение и оптимизатор с нулевой избыточностью (ZeRO)
• Гетерогенное управление памятью для поддержки эффективного обучения больших моделей
• Поддержка множества моделей, специфичных для конкретной области, таких как Open-Sora, Colossal-LLaMA и др.
• Предоставление удобных инструментов для распределенного обучения и вывода
• Интеграция высокопроизводительного ядра, кэша KV, внимания к подкачке и последовательной пакетной обработки.
• Простая настройка параллельного обучения с помощью конфигурационных файлов
• Богатые примеры и документация помогут вам быстро начать работу.
• Предоставляет несколько вариантов установки образов Docker и сборки из исходного кода

Использование помощи

Руководство по установке

Установка из PyPI

Вы можете легко установить Colossal-AI с помощью следующей команды:

pip install colossalai

По умолчанию расширения PyTorch не собираются во время установки. Если вам необходимо собрать расширения PyTorch, вы можете установить параметрBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

Кроме того, каждую неделю мы выпускаем NIGHTLY-версии, позволяющие вам получить доступ к последним неизданным функциям и исправлениям ошибок. Установка осуществляется следующим образом:

pip install colossalai-nightly

Установка из источника

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install .

Ядра CUDA/C++ не компилируются по умолчанию. colossalAI будет собирать их во время выполнения. При необходимости включите слияние ядер CUDA:

BUILD_EXT=1 pip install .

Для пользователей CUDA 10.2 вы можете вручную загрузить библиотеку cub и скопировать ее в соответствующий каталог перед установкой.

Использование Docker

Извлечение образов из DockerHub

Вы можете получить информацию непосредственно изСтраница на DockerHubИзвлеките образ Docker.

Создайте свой собственный образ

cd ColossalAI
docker build -t colossalai ./docker

Запускает контейнер в интерактивном режиме:

docker run -ti --gpus all --rm --ipc=host colossalai bash

Функции Поток операций

параллелизм данных

Параллелизм данных - это процесс разделения набора данных на несколько подмножеств и параллельного обучения модели на нескольких графических процессорах. ColossalAI облегчает пользователям параллельное обучение с помощью упрощенного профиля параллелизма данных:

from colossalai.nn.parallel import DataParallel
model = DataParallel(model)

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

Тензорный параллелизм - это процесс разделения тензора параметров модели на множество субтензоров и их параллельного вычисления на нескольких графических процессорах. ColossalAI предоставляет реализацию тензорного параллелизма 1D, 2D, 2.5D и 3D:

from colossalai.nn.parallel import TensorParallel
model = TensorParallel(model, parallel_mode='1D')

параллельно сборочной линии

Конвейерный параллелизм - это разделение модели на несколько этапов, каждый из которых выполняется одним или несколькими графическими процессорами. ColossalAI обеспечивает простую настройку конвейерного параллелизма:

from colossalai.pipeline.parallel import PipelineParallel
model = PipelineParallel(model, num_stages=4)

Смешанная высокоточная подготовка

Обучение со смешанной точностью значительно сокращает расход памяти и ускоряет обучение за счет использования в процессе обучения комбинации 16-битных чисел с плавающей запятой (FP16) и 32-битных чисел с плавающей запятой (FP32):

from colossalai.amp import convert_to_amp
model, optimizer, criterion = convert_to_amp(model, optimizer, criterion)

Zero Redundancy Optimiser (ZeRO)

Оптимизатор ZeRO значительно сокращает объем графической памяти за счет распределения состояний, градиентов и параметров оптимизатора между несколькими графическими процессорами:

from colossalai.zero import ZeroOptimizer
optimizer = ZeroOptimizer(optimizer, model)

Применение в реальном мире

Open-Sora

Open-Sora - это комплексное решение ColossalAI для моделей генерации видео, включающее параметры модели, детали обучения и возможность генерировать 16-секундные видеоролики в формате 720p HD одним щелчком мыши:

# 训练
python train.py
# 推理
python infer.py

Для получения дополнительной информации см.Open-Sora.

Colossal-LLaMA

Colossal-LLaMA - это решение с открытым исходным кодом для доменно-специфических больших языковых моделей (LLM), которые могут достигать результатов, сравнимых с основными большими моделями, при небольшом количестве денег на обучение:

# 训练
python train_llama.py
# 推理
python infer_llama.py

Для получения дополнительной информации см.Colossal-LLaMA.