ColossalAI: soluciones eficaces de formación de modelos de IA a gran escala

Introducción general

ColossalAI es una plataforma de código abierto desarrollada por HPC-AI Technologies para proporcionar una solución eficiente y rentable para el entrenamiento y la inferencia de modelos de IA a gran escala. Al soportar múltiples estrategias paralelas, gestión de memoria heterogénea y entrenamiento de precisión mixta, ColossalAI es capaz de reducir significativamente el tiempo y el consumo de recursos para el entrenamiento y la inferencia de modelos. Ya sea en el paralelismo de datos, el paralelismo de tensor o el paralelismo de canalización, ColossalAI proporciona potentes herramientas y bibliotecas para ayudar a los investigadores y desarrolladores a lograr un entrenamiento y una inferencia eficientes de modelos a gran escala en clústeres multi-GPU.

Lista de funciones

• Paralelismo de datos, paralelismo tensorial, paralelismo de canalización y otras estrategias paralelas.
• Formación de precisión mixta y optimizador de redundancia cero (ZeRO)
• Gestión heterogénea de la memoria para un entrenamiento eficaz de modelos de gran tamaño
• Compatibilidad con múltiples modelos específicos de dominio, como Open-Sora, Colossal-LLaMA, etc.
• Herramientas de fácil uso para la formación y la inferencia distribuidas
• Integración de núcleo de alto rendimiento, caché KV, atención a la paginación y procesamiento secuencial por lotes.
• Fácil configuración del entrenamiento en paralelo mediante archivos de configuración
• Proporcionan numerosos ejemplos y documentación para ayudarle a empezar rápidamente
• Ofrece múltiples opciones de instalación para imágenes Docker y construcción a partir del código fuente.

Utilizar la ayuda

Guía de instalación

Instalación desde PyPI

Puedes instalar fácilmente Colossal-AI con el siguiente comando:

pip install colossalai

Por defecto, las extensiones de PyTorch no se construyen durante la instalación. Si necesitas construir extensiones PyTorch, puedes establecer la opciónBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

Además, publicamos versiones NOCTURNAS todas las semanas, lo que le permite acceder a las últimas funciones inéditas y correcciones de errores. La instalación es la siguiente:

pip install colossalai-nightly

Instalación desde el origen

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

Los kernels CUDA/C++ no se compilan por defecto. colossalAI los compilará en tiempo de ejecución. Activa la fusión de kernels CUDA si es necesario:

BUILD_EXT=1 pip install .

Los usuarios de CUDA 10.2 pueden descargar manualmente la biblioteca cub y copiarla en el directorio adecuado antes de instalarla.

Uso de Docker

Obtención de imágenes de DockerHub

Puede obtener la información directamente delPágina de DockerHubExtrae la imagen Docker.

Construye tu propia imagen

cd ColossalAI
docker build -t colossalai ./docker

Inicia el contenedor en modo interactivo:

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

Función Flujo de operaciones

paralelismo de datos

El paralelismo de datos es el proceso de dividir un conjunto de datos en varios subconjuntos y entrenar el modelo en paralelo en varias GPU.ColossalAI facilita a los usuarios el entrenamiento en paralelo de datos con un perfil de paralelismo de datos simplificado:

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

paralelismo tensorial

El paralelismo tensorial es el proceso de dividir el tensor de parámetros de un modelo en múltiples subtensores y calcularlos de forma paralela en varias GPU.ColossalAI proporciona implementaciones de paralelismo tensorial 1D, 2D, 2.5D y 3D:

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

en paralelo a la cadena de montaje

El paralelismo de canalización consiste en dividir un modelo en varias etapas, cada una de ellas ejecutada por una o varias GPU.ColossalAI proporciona una sencilla configuración del paralelismo de canalización:

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

Entrenamiento de precisión mixto

El entrenamiento de precisión mixta reduce significativamente el uso de memoria y acelera el entrenamiento mediante el uso de una combinación de números de coma flotante de 16 bits (FP16) y números de coma flotante de 32 bits (FP32) durante el entrenamiento:

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

Optimizador de redundancia cero (ZeRO)

El optimizador ZeRO reduce considerablemente el consumo de memoria gráfica al distribuir los estados, gradientes y parámetros del optimizador entre varias GPU:

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

Aplicaciones reales

Open-Sora

Open-Sora es la solución completa de ColossalAI para modelos de generación de vídeo, incluidos los parámetros del modelo, los detalles del entrenamiento y la capacidad de generar vídeos HD de 16 segundos y 720p con un solo clic:

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

Para más información, consulteOpen-Sora.

Colossal-LLaMA

Colossal-LLaMA ofrece una solución de código abierto para grandes modelos lingüísticos (LLM) específicos de un dominio que pueden lograr resultados comparables a los grandes modelos convencionales con una pequeña cantidad de dinero de formación:

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

Para más información, consulteColossal-LLaMA.